zero.Init partitioning of large fused MoE-expert tensors spikes a single GPU (transient full materialization) -> OOM during load even when the sharded model fits

[trevorgordon981]

Env: transformers 5.12.1, deepspeed 0.18.9, torch 2.12.0+cu130, peft 0.19.1, bitsandbytes 0.49.2, accelerate 1.14.0; 8x B200 (178GB) / 2TB RAM; model MiniMaxAI/MiniMax-M3 (428B sparse MoE VL, minimax_m3_vl).

During the correctly-sharded (world_size=8) load, a single GPU spikes to ~180GB vs the ~94GB steady partition and OOMs. The giant fused MoE-expert parameter (128 experts) appears to be materialized in full on a GPU before being scattered. This is a load-time transient — independent of sequence length — and offload_param: {device: cpu} does not prevent the GPU spike during the from_pretrained init path.

Ask: stream/scatter very large parameters during partitioning without a full single-GPU materialization, and honor remote_device='cpu' during the from_pretrained init path so the load can stage through CPU.

[Achyuthan-S]

Opened a PR for this -> #8103 . It adds an opt-in ZeRO-3 setting stage3_partition_stream_chunk_size that partitions oversized parameters by streaming them through fixed-size chunks, so a giant fused tensor is never fully materialized on one GPU. On 2× L40S, the zero.Init peak for a ~2 GB param dropped 3.05 → 1.10 GB (−64%), and the gathered result is bit-identical to the default path. Default is 0 (disabled), so existing behavior is unchanged.

@trevorgordon981 if you can test on MiniMax-M3, set "stage3_partition_stream_chunk_size": 10000000 in your ZeRO-3 config — would be great to confirm it clears the OOM.

[trevorgordon981]

Thanks @Achyuthan-S for the quick PR! I tested #8103 (stage3_partition_stream_chunk_size) on MiniMaxAI/MiniMax-M3 under ZeRO-3, model-load only, on 8× A100-80GB / 2TB RAM (transformers 5.12.1, deepspeed 0.19.3+cc0fb6ab, torch 2.5.1+cu124, offload_param: cpu).

The new config key is recognized (zero/config.py:280). But I could not reproduce the original single-GPU OOM with this checkpoint, so the result is inconclusive — and I think the premise needs a second look:

• MiniMax-M3 on the Hub ships experts as ~38 MB split per-expert tensors (...block_sparse_moe.experts.N.{w1,w2,w3}, each [3072, 6144] ≈ 0.038 GB). The largest single tensor in the entire 856 GB checkpoint is the 2.46 GB embed/lm_head. There is no fused ~180 GB expert tensor.
• Consequently zero.Init never materializes a >80 GB (or even >2.5 GB) tensor on any one GPU — GPU memory stayed flat at ~18.5 GB/rank the whole time, no CUDA OOM.
• Since the patch only chunks oversized params during partitioning and nothing here exceeds ~2.5 GB, the knob has nothing to chunk → chunk_size=10000000 behaved identically to chunk_size=0 (no-op, not a fix or a regression).
• With offload_param: cpu, both settings instead died via the host OOM-killer (SIGKILL/-9, exit 137) while building the ~856 GB CPU offload across 8 ranks — RAM peaked at ~539 GB used / 45 GB avail before the kernel reaped it. That's a host-RAM/offload-sizing issue, orthogonal to the GPU transient the PR targets.

So on this build the patch is correctly wired but can't be exercised by this checkpoint — the >180 GB single-GPU expert transient simply doesn't arise here.

Question: was the original OOM seen on a fused-expert M3 variant — a merged checkpoint or a custom modeling file that concatenates the 128 experts into one parameter (or kv_b/fuse_qkv-style fusion) — rather than the stock split-expert Hub weights? If you can share the exact checkpoint/revision + ZeRO config (with/without offload) that triggered the >80 GB single-GPU transient, I'll re-run against the structure that actually reproduces it and confirm whether stage3_partition_stream_chunk_size clears it.

[Achyuthan-S]

Thanks for the really thorough test, @trevorgordon981 — agreed on all of it.

On the streaming PR: you're right that stock M3 ships split per-expert tensors (~38 MB each), with the largest single tensor being the 2.46 GB embed/lm_head — so there's no oversized fused param for the patch to bound; it's correct but unexercised by these weights. (Minor note: the threshold is in elements, so at chunk_size=1e7 the experts (18.9M) and embed do cross it and stream — it just has no measurable effect because those tensors are small in bytes. So it's a no-op here, not broken.) I built the PR to the structure described in the issue (a fused 128-expert parameter materialized in full). Since the original >80 GB single-GPU transient was on your 8× B200 run, you'd have much better visibility than me into whether that came from a merged checkpoint or a custom modeling file that concatenates the experts — if you can pin that down, I'm glad to confirm the knob clears it on that structure.

On the offload OOM: I chased this down with a controlled repro (small Llama, from_pretrained + ZeRO-3 + offload_param: cpu, pin_memory=False to match the default, 2× A100):

offload ratio = 1.00   (sharded 1/N, not duplicated across ranks)
dtype = bf16           (no fp32 upcast on the offloaded partitions)
RSS grew ~1.48 GB/rank ≈ the 1.36 GB sharded partition + overhead

So the CPU offload is structurally correct — each rank holds 1/world_size of the params, in bf16, no duplication. That points the 539 GB OOM at host-RAM sizing rather than a DeepSpeed bug: ZeRO-3's CPU offload total is always the full model (~856 GB), and with all 8 ranks on one node that whole thing sits on a single node plus load transients. Your "~539 GB used / 45 GB avail" suggests the run's effective host-RAM budget was ~584 GB — below the ~856 GB needed.

What should clear it:

• Spread the 8 ranks across 2+ nodes (offload total = model size, so fewer ranks/node = less host RAM/node — e.g. 4 × 2 ≈ 428 GB/node).
• NVMe offload: offload_param: {device: nvme, nvme_path: ...} instead of cpu.
• Or ensure the job's --mem covers ≥ ~856 GB if keeping all 8 ranks on one node.

If free -g / your job's --mem during the run confirms the budget, that'd close it out — and if per-rank RSS ever exceeds model/world_size on the real load (it didn't in my repro), that points back at the loader and I'm happy to dig further.

[NovusEdge]

Interested in looking at this — the transient full materialization during partitioning is a nasty footgun. Will dig into the zero.Init path and see where the scatter could be streamed instead of staged. Might take a bit to ramp up on the partitioning logic.