ZeRO-3: stream partitioning of oversized parameters in zero.Init

[Achyuthan-S]

Problem

Under zero.Init (ZeRO-3), every parameter is moved to the accelerator, broadcast in full, and then sliced into per-rank partitions. A single very large fused parameter — e.g. a 128-expert MoE weight — must be fully materialized on one device during this step, which can OOM that device during a from_pretrained load even when the sharded model fits. offload_param: {device: cpu} does not help: it only controls where the resulting partition is stored, not where the full tensor is staged.

Closes #8085.

Change

Adds an opt-in ZeRO-3 config stage3_partition_stream_chunk_size (default 0 = disabled). When set, a parameter with more elements than the threshold that is not already on the accelerator (the host-staged from_pretrained / low_cpu_mem_usage path) is partitioned by streaming its flattened data through fixed-size chunks: stage a chunk on the accelerator → broadcast from the owner rank → copy only this rank's slice into ds_tensor. The full tensor is never materialized on a single device, bounding the partition-time peak to roughly the chunk size.

With the default (0) the standard broadcast-then-partition path runs unchanged. Streaming is skipped for the nvme / quantized / ZeRO++ secondary-partition paths, which stage parameters differently.

Validation

Correctness — new unit test covers the chunk/partition overlap math (incl. padding, single-rank). End-to-end, the streamed partition reconstructs bit-for-bit identically to the standard path across world sizes 1–3, with padding, all_gather round-trip, and offload_param: cpu.

NCCL + peak memory (2× NVIDIA L40S):
[A] NCCL correctness (gathered streamed == standard): True
[B] peak GPU memory during zero.Init (world=2, dim=22528, fp32)
full param : 2.03 GB partition/rank: 1.02 GB chunk: 40 MB
streaming OFF peak : 3.05 GB
streaming ON peak : 1.10 GB
peak reduction : 1.95 GB (64% lower)

Scope

Applies to parameters that reach partitioning off-GPU (the from_pretrained / low_cpu_mem_usage path this issue targets). Parameters constructed directly on the accelerator inside zero.Init are unaffected — the spike there happens at construction time, which can be addressed as a follow-up.

cc @tohtana @tjruwase @loadams

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: cc35972244

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Avoid checking streaming before ZeRO metadata exists

When stage3_partition_stream_chunk_size is set and zero.Init(module=prebuilt_model, ...) is used, this new pre-check runs on ordinary torch.nn.Parameters before _zero_init_param() calls _convert_to_deepspeed_param(). _should_stream_partition() immediately asks for _partition_world_size(param), which dereferences param.ds_process_group; that attribute is only installed later in _convert_to_deepspeed_param(), so the module-conversion path raises AttributeError even for parameters smaller than the chunk size. Move the stream decision until after conversion, or make the pre-check use the default process group without requiring ZeRO metadata.

Useful? React with 👍 / 👎.

[Achyuthan-S]

Fixed. _should_stream_partition now gates on the global num_partitions instead of _partition_world_size(param), so the zero.Init(module=...) path no longer dereferences param.ds_process_group before _convert_to_deepspeed_param attaches it. The per-parameter group is still used in the actual partitioning (_partition_param_streaming), which runs after conversion. Added a DistributedTest that exercises the module= path with streaming enabled to guard this.

[Achyuthan-S]

Hey @tohtana , I have been working on this issue and opened a PR with the solution.
It would be great if you review this and let me know if this works.
Thank you !

[tohtana]

Hi @Achyuthan-S,
Thank you for submitting a very useful PR!

The code looks correct to to me. But I think most users use zero.Init in the following pattern (including HuggingFace integration). However, this PR currently doesn't cover this.

with deepspeed.zero.Init(config_dict_or_path=config):
    model = Model()

Can you extend this PR to support this pattern?

[Achyuthan-S]

Hi @Achyuthan-S, Thank you for submitting a very useful PR!

The code looks correct to to me. But I think most users use zero.Init in the following pattern (including HuggingFace integration). However, this PR currently doesn't cover this.

with deepspeed.zero.Init(config_dict_or_path=config):
    model = Model()

Can you extend this PR to support this pattern?

Thanks @tohtana! Done — extended to cover with deepspeed.zero.Init(): model = Model() (and HF's from-scratch construction).

The gap was that under the context manager zero.Init patches the tensor constructors to allocate on the accelerator, so a large parameter is fully materialized at construction time, before partitioning runs. Now, when streaming is enabled, the shape-based constructors (empty/zeros/ones/randn/new) build tensors above the threshold on the host instead; the existing streaming partition then stages only one chunk at a time onto the device. Small params and all other constructors are unchanged, and with the default (0) nothing changes at all.

Measured on 2× A100 for exactly this pattern (with zero.Init(): Linear(22528, 22528), fp32): peak GPU memory during init drops 3.05 GB → 1.10 GB (−1.95 GB, ~64%) with stage3_partition_stream_chunk_size set — about one full-parameter's worth.

Added a unit test for the construction-device decision, and the context-manager correctness test now uses deterministic weights (random init draws from a device-specific RNG, so a CPU-constructed streamed weight vs a GPU-constructed reference wouldn't be directly comparable otherwise).

[tohtana]

Hi @Achyuthan-S,
Thank you for updating this PR!

But I have concerns about the new approach:

• A large implicit-device registered buffer is now redirected to CPU and remain there after initialization. Actually, arbitrary large temporaries are subject to the same size-only policy.
• Valid keyword size= also bypass the host-construction decision because only positional shape arguments are inspected.
  Can you address these?

[Achyuthan-S]

Thanks @tohtana — both good catches. Addressed: (1) the host redirect now only applies to floating-point tensors (large integer masks/ids stay on the accelerator), and _post_init_method moves any buffer that landed on the host back to the accelerator, with a test asserting it. (2) the decision now reads the shape from the size= keyword too, with a test. One inherent limitation flagged: an unregistered large float temporary used cross-device inside init would still build on the host — standard modules don't do this, but happy to add a stricter gate if you'd prefer…

[tohtana]

Hi @Achyuthan-S,

Thank you for the update! The size= issue and implicit floating-buffer restoration seem to be addressed.

However, I still see several challenges. Consider the following pattern:

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        payload = torch.zeros(
            1024,
            1025,
            dtype=torch.float32,
            device="cpu",
        )
        self.register_buffer("cache", payload)

with deepspeed.zero.Init(config_dict_or_path=config):
    model = Model()

The current implementation cannot distinguish a tensor that was implicitly redirected to CPU for streaming from one that was explicitly placed on CPU by the user. As a result, the registered buffer above is automatically moved to the GPU, changing the module's intended device semantics.

Another issue can occur after streaming a parameter:

for chunk in parameter:
    chunk.to(cuda)
    copy_chunk_into_local_ds_tensor()

free_param(param)

free_param() replaces param.data with an empty tensor on the parameter's current device:

param.data = torch.empty(
    0,
    dtype=param.dtype,
    device=param.device,
)

Because the streamed source parameter is on CPU, this leaves param.data as an empty CPU placeholder while param.ds_tensor is on the GPU. This inconsistency can affect later code that uses param.device as the destination for the materialized parameter.

For example, the sequential path used by all_gather_coalesced() when the group contains a single parameter or when stage3_allgather_sequential=true does the following:

DeepSpeed/deepspeed/runtime/zero/partition_parameters.py

Lines 1433 to 1504 in 3a8eb21

	def _all_gather_sequential(params, world_size, use_secondary_tensor, ds_process_group, quantize):
	handles = []
	for param in params:
	buffer_size = math.ceil(param.ds_numel / world_size) * world_size
	if use_secondary_tensor:
	buffer_size = param.ds_secondary_tensor.shape[0] * world_size #make sure out is appropriately sized
	param_ds_tensor = param.ds_secondary_tensor if use_secondary_tensor else param.ds_tensor
	original_dtype = param_ds_tensor.dtype
	if quantize:
	allgather_dtype = torch.int8
	else:
	allgather_dtype = get_allgather_dtype(param, param_ds_tensor)
	param_buffer = torch.empty(
	buffer_size,
	dtype=allgather_dtype,
	device=get_accelerator().current_device_name(),
	requires_grad=False,
	)
	if not quantize:
	handle = _dist_allgather_fn(
	param_ds_tensor.to(get_accelerator().current_device_name()).to(allgather_dtype),
	param_buffer,
	ds_process_group,
	)
	if original_dtype == allgather_dtype:
	param.data = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to(param.device)
	handles.append(AllGatherHandle(handle, param))
	else:
	# This case is complicated:
	# We use `register_post_accumulate_grad_hook` to set allgather hooks. Normally, the hook is
	# called once per parameter, even if that parameter is tied to multiple layers.
	# However, when the dtype changes, the hook may be triggered multiple times.
	# If we directly do:
	# param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to(param.device)
	# as above, the dtype may differ, causing the gradient-reduce hook
	# to be invoked multiple times.
	# To avoid this, we leave `param.data` in a partitioned state.
	# This prevents duplicate gradient-reduce hook calls.
	# In theory, this path could be consolidated with the case where
	# (original_dtype == allgather_dtype), but because it changes the
	# state transition of DeepSpeed parameters, we keep it separate for safety.
	handles.append(
	AllGatherHandle(handle, param, param_buffer=param_buffer, original_dtype=original_dtype))
	else:
	if hasattr(param_ds_tensor, "ds_quant_scale"):
	scales = param_ds_tensor.ds_quant_scale
	quantized_param = param_ds_tensor.data
	else:
	quantized_param, scales = self.quantizer_module.quantize(param_ds_tensor)
	handle = _dist_allgather_fn(quantized_param.to(get_accelerator().current_device_name()),
	param_buffer, ds_process_group)
	quant_scale_buffer = torch.empty(
	scales.numel() * world_size,
	dtype=scales.dtype,
	device=get_accelerator().current_device_name(),
	requires_grad=False,
	)
	quant_handle = _dist_allgather_fn(scales.to(get_accelerator().current_device_name()),
	quant_scale_buffer, ds_process_group)
	quant_info = QuantizationInfo()
	quant_info.quantized_param = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to(
	param.device)
	quant_info.backend = self.quantizer_module
	quant_info.quant_handle = quant_handle
	quant_info.scale_buffer = quant_scale_buffer
	handles.append(AllGatherHandle(handle, param, quantization=quant_info))
	return MultipleAllGatherHandles(handles)

param.data = param_buffer.narrow(0, 0, param.ds_numel).view(param.ds_shape).to(param.device)

In this case, param.device is derived from the CPU placeholder, so the full gathered parameter may be copied back to CPU before it is used by a GPU forward. The current GatheredParameters probe does not cover this behavior because it uses a different gather implementation rather than the normal coordinator-driven forward path.

More generally, many of these issues appear to come from the fact that global wrappers around constructors such as torch.empty() and torch.zeros() do not know how the resulting tensor will eventually be used. At construction time, the tensor may later become a parameter, a registered buffer, or an arbitrary temporary consumed inside __init__. Post-init restoration cannot reliably recover that intent, especially for explicit CPU placement or temporaries used before post-init runs.

I was initially open to exploring this approach, but given the issues above, I now think we need to clarify its practical benefit. Even if streaming reduces peak memory during initialization, ZeRO-3 still materializes full parameters during the forward all-gather. I would like to understand how common it is for initialization memory to be the blocker while the subsequent forward still fits.

[Achyuthan-S]

Thanks @tohtana — I've addressed the concrete issues, and I want to give you a straight answer on the value question rather than talk around it.

Fixes:

• Reverted the constructor-level host redirection entirely. You're right that global constructor wrappers can't tell a parameter from a user-placed CPU buffer or a temporary, and post-init restoration can't recover that intent — so the mechanism is removed, not patched. The tensor-creation wrappers are back to their state on master.
• Fixed the free_param inconsistency: a streamed parameter now leaves its empty param.data placeholder on the accelerator, identical to the standard path, so downstream code that targets param.device (the sequential all-gather you pointed at) behaves correctly. Added a regression test over that exact _all_gather_sequential path.

On the value question — you're right, and here's the honest analysis.

Since ZeRO-3 materializes each parameter fully on one device during the forward all-gather, a parameter that can't be built at init generally can't be gathered at forward either — the init peak exceeds the forward peak only by about one partition. So for a parameter that's gathered whole, there is no general regime where initialization is the blocker while the forward still fits. I couldn't construct one, and on the model that motivated the issue (MiniMax-M3) the experts ship as split ~38 MB tensors, so the spike doesn't even arise on the stock weights.

The only places it genuinely helps are narrow: a non-contiguous large parameter (the standard path's .contiguous() doubles it on-device, whereas streaming does that copy on the host), and the semantic of making offload_param: cpu actually stage the load through host memory (today it only controls where the partition lands). Neither is the broad win the issue implied.

Given that, I don't think this earns a place in the hot path on its own. The real answer for a parameter too large for a single device is tiled parameters (TiledLinear), which never gather it whole and therefore bound the forward too — something init-side streaming fundamentally can't do.

Given that, my recommendation is to keep this strictly as the scoped, default-off safeguard for the host-staged (CPU-offload) load path, together with the free_param fix — it's opt-in and changes nothing at all when disabled. I don't want to overstate the benefit, so if it'd help, I'm glad to put together an end-to-end measurement of the load-time peak vs a forward step under offload_param: cpu, so the decision is on real numbers rather than argument. Let me know how you'd like to proceed.