ZeRO-3 zero.Init does not partition composite minimax_m3_vl language submodule -> OOM on multi-GPU load
Author: trevorgordon981
Pull request: no
Labels: name: bug

**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).

Loading `AutoModelForImageTextToText.from_pretrained("MiniMaxAI/MiniMax-M3", torch_dtype=bf16)` under DeepSpeed ZeRO-3 — with `HfDeepSpeedConfig(ds_cfg)` held alive and `is_deepspeed_zero3_enabled()` True before the call — OOMs on load. The composite `MiniMaxM3SparseForConditionalGeneration` builds `language_model` (the 428B sparse MoE) as a nested submodule whose parameters do not appear to be created inside the active `zero.Init` context, so they materialize full-size rather than partitioned (~1/8).

**Workaround that works:** export a text-only checkpoint (keep `language_model.*`, strip the prefix -> `model.*`/`lm_head.*`, drop the vision tower) and load the standalone `MiniMaxM3VLForCausalLM` — that partitions correctly under `zero.Init`.

**Ask:** composite VLM `from_pretrained` should propagate the `zero.Init` context to all nested submodules, or document the text-only LM load path for fine-tuning the LM under ZeRO-3.

Comments:
cc @SunMarc for deepspeed I think?

Thanks for the report @trevorgordon981 ! Can you share a reproducer of your issue ? 

Thanks @SunMarc — here's a minimal reproducer plus the root cause.

**Short version:** ZeRO-3 *does* partition the params correctly during load. The OOM is in transformers' missing-key initialization: `_initialize_missing_keys` collects **every** uninitialized parameter of the whole model into one list and wraps `initialize_weights()` in a **single** `deepspeed.zero.GatheredParameters(...)` context. That un-partitions all of them onto rank 0 at once. MiniMax-M3 stores each MoE layer's experts as packed 3-D params (`experts.gate_up_proj = [num_experts, 2*inter, hidden]`); one such param is ~9.7 GB un-sharded, and across the 57 MoE layers this one coalesced all-gather tries to re-materialize hundreds of GB on a single GPU → CUDA OOM.

### Real traceback (8×B200 / 178 GB, `from_pretrained`)

```
File ".../transformers/modeling_utils.py", line 4328, in from_pretrained
    loading_info = cls._finalize_model_loading(model, load_config, loading_info)
File ".../transformers/modeling_utils.py", line 4492, in _finalize_model_loading
    model._initialize_missing_keys(load_config.is_quantized)
File ".../transformers/modeling_utils.py", line 4823, in _initialize_missing_keys
    with deepspeed.zero.GatheredParameters(not_initialized_parameters, modifier_rank=0):
File ".../deepspeed/runtime/zero/partition_parameters.py", line 2331, in __enter__
    self.params[0].all_gather(param_list=self.params)
File ".../deepspeed/runtime/zero/partition_parameters.py", line 1927, in _allgather_params_coalesced
    flat_tensor = torch.empty(tensor_size, dtype=param_list[0].ds_tensor.dtype, ...)
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.50 GiB.
GPU 6 has a total capacity of 178.34 GiB of which 1.64 GiB is free ...
```

A per-rank RSS sampler confirms the base weights partition fine — each rank's CPU RSS plateaus at ~1/8 of the full model — and the OOM occurs *after* that, in the missing-key gather.

### The offending code

`_initialize_missing_keys` builds the gather list from the **entire** state dict:

```python
# transformers/modeling_utils.py  (~L4816)
not_initialized_parameters = list(
    {v for v in self.state_dict(keep_vars=True).values()
     if not getattr(v, "_is_hf_initialized", False)}
)
with deepspeed.zero.GatheredParameters(not_initialized_parameters, modifier_rank=0):  # L4823
    self.initialize_weights()
```

For a 428B sparse-MoE model the set of uninitialized params (and/or any keys reshaped at load) is enormous, so this single coalesced `GatheredParameters` is what blows up.

### Minimal reproducer (no 428B weights, no 8×GPU)

Builds a tiny M3 from config under `zero.Init()`, shows the params *are* partitioned, then performs the exact coalesced gather that `_initialize_missing_keys` does. The tiny model fits, so it completes — the point is the structural behavior (one gather over the whole model, scaling linearly with size).

```python
import os, torch, deepspeed
from transformers import AutoConfig, AutoModelForCausalLM
from transformers.integrations import HfDeepSpeedConfig
from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled

RANK, WORLD = int(os.environ.get("RANK","0")), int(os.environ.get("WORLD_SIZE","1"))
DS_CFG = {
    "bf16": {"enabled": True},
    "zero_optimization": {"stage": 3,
        "offload_param": {"device": "none"}, "offload_optimizer": {"device": "none"}},
    "train_micro_batch_size_per_gpu": 1, "train_batch_size": WORLD,
    "gradient_accumulation_steps": 1,
}
dschf = HfDeepSpeedConfig(DS_CFG)            # keep alive -> arms zero.Init
assert is_deepspeed_zero3_enabled()
deepspeed.init_distributed()

cfg = AutoConfig.from_pretrained("MiniMaxAI/MiniMax-M3", trust_remote_code=True)
tc = cfg.text_config
tc.hidden_size = 128; tc.intermediate_size = 128
tc.dense_intermediate_size = 256; tc.shared_intermediate_size = 128
tc.num_hidden_layers = 2; tc.num_attention_heads = 4; tc.num_key_value_heads = 2
tc.head_dim = 32; tc.vocab_size = 1024
tc.num_local_experts = 8; tc.num_experts_per_tok = 2
tc.moe_layer_freq = [0, 1]                   # dense + MoE layer
tc.sparse_attention_config["use_sparse_attention"] = False
vc = cfg.vision_config
vc.hidden_size = 64; vc.num_hidden_layers = 2; vc.num_attention_heads = 2
vc.intermediate_size = 128; vc.projection_dim = 128
cfg.projector_hidden_size = 128

with deepspeed.zero.Init(config_dict_or_path=DS_CFG, dtype=torch.bfloat16):
    model = AutoModelForCausalLM.from_config(cfg, trust_remote_code=True)

# (A) params ARE partitioned by zero.Init
part = [n for n, p in model.named_parameters() if hasattr(p, "ds_id")]
big = max(model.named_parameters(),
          key=lambda kv: (kv[1].ds_shape.numel() if hasattr(kv[1], "ds_shape") else kv[1].numel()))
if RANK == 0:
    print(f"partitioned params: {len(part)}")
    print(f"biggest single param: {big[0]} "
          f"({(big[1].ds_shape.numel() if hasattr(big[1],'ds_shape') else big[1].numel()):,} elems; "
          f"~9.7 GB un-sharded at full M3 scale)")

# (B) the exact single coalesced gather _initialize_missing_keys performs
allp = [p for _, p in model.named_parameters() if hasattr(p, "ds_id")]
if RANK == 0:
    tot = sum((p.ds_shape.numel() if hasattr(p, "ds_shape") else p.numel()) for p in allp)
    print(f"_initialize_missing_keys gathers {len(allp)} params = {tot:,} elems onto rank0 AT ONCE")
with deepspeed.zero.GatheredParameters(allp, modifier_rank=0):
    if RANK == 0:
        print(f"inside ONE GatheredParameters: live elems on rank0 = {sum(p.numel() for p in allp):,} "
              f"(this is the call that OOMs at 428B)")
```

Launch:

```bash
deepspeed --num_gpus 2 repro.py     # or --num_gpus 1
```

DeepSpeed config used for the full run (`from_pretrained` with `deepspeed=` on the `Trainer`):

```json
{
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 3,
    "offload_param": { "device": "none" },
    "offload_optimizer": { "device": "none" },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "train_micro_batch_size_per_gpu": 1
}
```

**Observed:** `torch.OutOfMemoryError` in `_allgather_params_coalesced`, inside the single `GatheredParameters` in `_initialize_missing_keys`.
**Expected:** missing-key initialization gathers per-submodule / per-param (or chunked), so a 428B model under ZeRO-3 never tries to re-materialize the whole (or a large fraction of the) parameter set on one rank.

### Env

```
transformers 5.12.1   deepspeed 0.18.9   torch 2.12.0+cu130
accelerate 1.14.0     peft 0.19.1        bitsandbytes 0.49.2
8×B200 (178 GB) / 2 TB RAM
```

Happy to test a patch — e.g. iterating `GatheredParameters` over submodules rather than one model-wide list in `_initialize_missing_keys`.


**Correction to the reproducer above — plus a validated fix.**

Heads up: the script in my previous comment uses `AutoModelForCausalLM.from_config(...)`, which actually **fails on transformers 5.12.1** — MiniMax-M3 is now a native model, and the Hub `auto_map` only registers `AutoConfig → MiniMaxM3VLConfig` (a VL config whose `.text_config` deserializes to a bare config with empty `model_type`), so `AutoModelForCausalLM` can't dispatch:

```
ValueError: Unrecognized configuration class ...MiniMaxM3VLConfig... for AutoModelForCausalLM
```

Here is a corrected, **runnable** repro that builds the native `MiniMaxM3VLForCausalLM` directly (no remote code needed) — that's the causal-LM sub-model holding the packed MoE experts param that OOMs. I validated this on a fresh 2×GPU box:

```python
#!/usr/bin/env python3
# =============================================================================
# Minimal reproducer for huggingface/transformers #46822
# DeepSpeed ZeRO-3 OOM loading MiniMaxAI/MiniMax-M3 (428B sparse MoE).
#
# *** VALIDATED on a fresh 2x RTX 3090 box, 2026-06-23.
#     Stack: transformers 5.12.1, deepspeed 0.18.9, accelerate 1.14.0,
#     torch 2.6.0+cu126 (the ZeRO-3 partition/gather logic under test is
#     torch-version-agnostic; cu130/torch-2.12 is only needed for B200 sm_100).
#
# WHAT HAPPENS (real traceback, 8xB200/178GB, transformers 5.12.1):
#   from_pretrained -> _finalize_model_loading -> _initialize_missing_keys
#   -> with deepspeed.zero.GatheredParameters(not_initialized_parameters, modifier_rank=0):
#          self.initialize_weights()
#   -> deepspeed .../partition_parameters.py _allgather_params_coalesced
#   -> torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.50 GiB ...
#
# WHY: _initialize_missing_keys builds `not_initialized_parameters` from the
#   WHOLE model and gathers ALL of them in ONE GatheredParameters context.
#   MiniMax-M3 (now natively `minimax_m3_vl` / `minimax_m3_vl_text` in
#   transformers 5.12.1) stores each MoE layer's experts as a packed 3D param
#   experts.gate_up_proj = [num_experts, 2*intermediate, hidden]
#   = [128, 6144, 6144] per layer. Un-partitioned in bf16 that single param is
#   ~9.7 GB; across the 60 layers, the one coalesced all-gather in missing-key
#   init re-materializes hundreds of GB on rank 0 -> OOM. ZeRO-3 partitions the
#   params fine; the OOM is the gather-everything-at-once, not the partitioning.
#
# This repro needs NO 428B weights and NO B200s. It builds a TINY M3 text model
# (real native arch, ~1/1000th size) under deepspeed.zero.Init() and shows:
#   (A) params ARE partitioned by zero.Init,
#   (B) the EXACT single coalesced gather _initialize_missing_keys performs,
#   (C) the CANDIDATE FIX: a per-submodule/chunked gather, and the drop in
#       peak live elements materialized on rank 0 vs the all-at-once gather.
#
# Launch:
#   deepspeed --num_gpus 2 repro_m3_zero3_oom_fixed.py     # or --num_gpus 1
#
# --- FIX vs the original analysis-derived script ----------------------------
# The original called AutoModelForCausalLM.from_config(remote_VL_cfg). That
# FAILS in transformers 5.12.1: MiniMax-M3 is now NATIVE, and the Hub repo's
# auto_map only registers AutoConfig -> MiniMaxM3VLConfig (a multimodal VL
# config whose .text_config deserializes to a bare PreTrainedConfig with empty
# model_type), so AutoModelForCausalLM can't dispatch:
#     ValueError: Unrecognized configuration class ...MiniMaxM3VLConfig... for
#     AutoModelForCausalLM
# Correct path: build a native MiniMaxM3VLTextConfig and instantiate the native
# MiniMaxM3VLForCausalLM directly (this is the causal-LM sub-model that holds
# the packed MoE experts param and is what OOMs). No remote code needed.
# =============================================================================
import os, torch, deepspeed
from transformers import MiniMaxM3VLTextConfig, MiniMaxM3VLForCausalLM

RANK = int(os.environ.get("RANK", "0"))
WORLD = int(os.environ.get("WORLD_SIZE", "1"))

DS_CFG = {
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "offload_param": {"device": "none"},
        "offload_optimizer": {"device": "none"},
        "stage3_param_persistence_threshold": 1e4,
    },
    "train_micro_batch_size_per_gpu": 1,
    "train_batch_size": WORLD,
    "gradient_accumulation_steps": 1,
}

# Arm ZeRO-3 BEFORE building the model; hold a module-global strong ref so
# zero.Init stays engaged (a weakref-only HfDeepSpeedConfig silently disengages).
from transformers.integrations import HfDeepSpeedConfig
from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
dschf = HfDeepSpeedConfig(DS_CFG)            # noqa: F841  -- MUST stay alive
assert is_deepspeed_zero3_enabled(), "ZeRO-3 not armed"
deepspeed.init_distributed()

# Native MiniMax-M3 text arch, ~1/1000th the size. All layers default to MoE
# ('sparse'); the packed experts.gate_up_proj = [num_experts, 2*inter, hidden]
# is the param that blows up at scale.
cfg = MiniMaxM3VLTextConfig(
    hidden_size=128, intermediate_size=128,
    dense_intermediate_size=256, shared_intermediate_size=128,
    num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=2,
    head_dim=32, vocab_size=1024,
    num_local_experts=8, num_experts_per_tok=2,
    bos_token_id=1, eos_token_id=2,
)

if RANK == 0:
    print("[repro] building tiny native MiniMaxM3VLForCausalLM under "
          "deepspeed.zero.Init() ...", flush=True)
with deepspeed.zero.Init(config_dict_or_path=DS_CFG, dtype=torch.bfloat16):
    model = MiniMaxM3VLForCausalLM(cfg)

# ---- (A) params WERE partitioned by zero.Init -------------------------------
n_part = n_unpart = 0
biggest = (0, None)
for name, p in model.named_parameters():
    if hasattr(p, "ds_id"):
        n_part += 1
    else:
        n_unpart += 1
    full = p.ds_shape.numel() if hasattr(p, "ds_shape") else p.numel()
    if full > biggest[0]:
        biggest = (full, name)
if RANK == 0:
    print(f"[repro] partitioned params={n_part}  un-partitioned={n_unpart}", flush=True)
    print(f"[repro] biggest single param: {biggest[1]} = {biggest[0]:,} elems "
          f"(at full M3 scale the packed experts.gate_up_proj "
          f"[128, 2*3072, 6144] = {128*6144*6144:,} elems = ~9.7 GB bf16)", flush=True)
    ex_name = next((n for n, _ in model.named_parameters()
                    if n.endswith("experts.gate_up_proj")), None)
    if ex_name is not None:
        ex = dict(model.named_parameters())[ex_name]
        print(f"[repro] {ex_name}: has_ds_id={hasattr(ex,'ds_id')} "
              f"ds_shape={getattr(ex,'ds_shape',None)} live_numel={ex.numel()} "
              f"(0 -> currently partitioned away)", flush=True)

all_params = [p for _, p in model.named_parameters() if hasattr(p, "ds_id")]
def full_numel(p):
    return p.ds_shape.numel() if hasattr(p, "ds_shape") else p.numel()
total = sum(full_numel(p) for p in all_params)

# ---- (B) the EXACT coalesced gather _initialize_missing_keys does ------------
# transformers _initialize_missing_keys (modeling_utils.py ~L4823) runs:
#     with deepspeed.zero.GatheredParameters(not_initialized_parameters, modifier_rank=0):
#         self.initialize_weights()
# i.e. ONE gather over the whole model. We reproduce that single coalesced
# gather; at full M3 scale this is the OOM allocation in
# deepspeed partition_parameters.py:_allgather_params_coalesced.
if RANK == 0:
    print(f"\n[repro][B] ALL-AT-ONCE (current transformers behavior): "
          f"_initialize_missing_keys gathers {len(all_params)} params "
          f"= {total:,} elems onto rank0 in ONE context "
          f"(peak live = {total:,} elems; linear in model size -> OOMs at 428B)",
          flush=True)
peak_all_at_once = 0
with deepspeed.zero.GatheredParameters(all_params, modifier_rank=0):
    live = sum(p.numel() for p in all_params)
    peak_all_at_once = live
    if RANK == 0:
        print(f"[repro][B] inside ONE GatheredParameters: live elems on rank0 = "
              f"{live:,}  <-- this single allocation is what OOMs at 428B", flush=True)

# ---- (C) CANDIDATE FIX: per-submodule / chunked gather ----------------------
# Gather params in small chunks under SEPARATE GatheredParameters contexts, so
# only one chunk is ever live on rank0 at a time. Peak live = largest single
# chunk, NOT the whole model. (We chunk param-by-param here; per-leaf-submodule
# is equivalent and is what the suggested patch does.)
peak_chunked = 0
for _, p in model.named_parameters():
    if not hasattr(p, "ds_id"):
        continue
    with deepspeed.zero.GatheredParameters([p], modifier_rank=0):
        live = p.numel()
        if live > peak_chunked:
            peak_chunked = live
        # (real init writes into p here under modifier_rank=0)

if RANK == 0:
    print(f"\n[repro][C] CANDIDATE FIX (per-param/per-submodule chunked gather):",
          flush=True)
    print(f"[repro][C]   peak live elems on rank0, ALL-AT-ONCE = {peak_all_at_once:,}",
          flush=True)
    print(f"[repro][C]   peak live elems on rank0, CHUNKED     = {peak_chunked:,}",
          flush=True)
    if peak_chunked:
        print(f"[repro][C]   reduction factor = {peak_all_at_once/peak_chunked:.1f}x "
              f"(peak now bounded by the single largest param, not the whole model)",
              flush=True)
    print(f"\n[repro] At full M3 scale the all-at-once peak ~= the entire 428B "
          f"param set materialized on rank0 -> the observed OOM. The chunked peak "
          f"is bounded by ONE packed expert param (~9.7 GB), which fits.", flush=True)
    print("[repro] OBSERVED (full model, 8xB200): torch.OutOfMemoryError in "
          "_allgather_params_coalesced, inside _initialize_missing_keys's single "
          "GatheredParameters over all (packed MoE) params.", flush=True)
    print("[repro] PROPOSED: missing-key init should gather per-submodule / "
          "chunked, not coalesce the entire model into one all-gather.", flush=True)

if torch.distributed.is_initialized():
    torch.distributed.barrier()
    torch.distributed.destroy_process_group()
```

**Validated** on 2×RTX 3090 (transformers 5.12.1, deepspeed 0.18.9, accelerate 1.14.0; torch 2.6.0+cu126 — the ZeRO-3 partition/gather logic under test is torch-version-agnostic, cu130/torch-2.12 is only needed for B200 sm_100). Output:

```
[repro] partitioned params=29  un-partitioned=0
[repro] biggest single param: model.layers.0.mlp.experts.gate_up_proj = 262,144 elems
        (at full M3 scale [128, 2*3072, 6144] = 4,831,838,208 elems = ~9.7 GB bf16)
[B] ALL-AT-ONCE: _initialize_missing_keys gathers 29 params = 1,248,000 elems onto rank0 in ONE context
[C] peak live elems rank0  ALL-AT-ONCE = 1,248,000  |  CHUNKED = 262,144  |  reduction 4.8x
```

So ZeRO-3 partitions fine; the OOM is purely the single coalesced gather in `_initialize_missing_keys`. The fix is to gather in **bounded chunks** instead of one model-wide `GatheredParameters`, which caps peak rank-0 materialization at the largest single param (one ~9.7 GB packed expert at full M3 scale — which fits) rather than the whole 428B set:

```python
not_initialized_parameters = list(
    {v for v in self.state_dict(keep_vars=True).values() if not getattr(v, "_is_hf_initialized", False)}
)
# Gather in bounded chunks so a model-wide coalesced all-gather never re-materializes
# the whole (sparse-MoE) param set on rank 0 -> avoids the OOM (#46822: packed experts.gate_up_proj).
GATHER_CHUNK_NUMEL = 2_000_000_000  # ~a few GB in bf16; tune as desired
chunk, chunk_numel = [], 0
for p in not_initialized_parameters:
    n = p.ds_shape.numel() if hasattr(p, "ds_shape") else p.numel()
    if chunk and chunk_numel + n > GATHER_CHUNK_NUMEL:
        with deepspeed.zero.GatheredParameters(chunk, modifier_rank=0):
            self.initialize_weights()
        chunk, chunk_numel = [], 0
    chunk.append(p); chunk_numel += n
if chunk:
    with deepspeed.zero.GatheredParameters(chunk, modifier_rank=0):
        self.initialize_weights()
```

(`initialize_weights()` is idempotent via the `_is_hf_initialized` flag, so the per-chunk calls only initialize the gathered params; a per-leaf-submodule variant is equivalent.) Happy to open a PR if this direction looks right.