CUDA out of memory

快速结论：该报错发生在 DeepSpeed ZeRO-3 加载大型稀疏 MoE 多模态模型（如 MiniMax-M3 428B）时，问题根源并非 ZeRO-3 分区失败，而是 transformers 的 _initialize_missing_keys 将整个模型所有未初始化参数放进单个 GatheredParameters 上下文，一次性全部反分区到 rank 0 导致显存溢出。优先排查路径是确认是否因 MoE 层超大参数（单个约 9.7GB）在合并 all-gather 时被同时物化。

适用环境：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；8×B200（178GB）GPU、2TB RAM；模型 MiniMaxAI/MiniMax-M3（428B 稀疏 MoE VLM，minimax_m3_vl）。

最快修复方案：Issue 中已确认有效的方案是导出仅文本的 checkpoint（保留 language_model.*，去掉前缀转为 model.*/lm_head.*，移除 vision tower），然后加载独立的 MiniMaxM3VLForCausalLM，该路径能在 zero.Init 下正确分区。

注意事项：此方案为绕过问题而非修复 transformers 内部逻辑；复合 VLM 完整加载路径下该问题尚未修复，若需直接加载完整模型需等待官方补丁或自行修改 _initialize_missing_keys 行为。

问题场景

用户在使用 AutoModelForImageTextToText.from_pretrained("MiniMaxAI/MiniMax-M3", torch_dtype=bf16) 加载 428B 稀疏 MoE 多模态模型时，DeepSpeed ZeRO-3 已启用（HfDeepSpeedConfig 持有且 is_deepspeed_zero3_enabled() 返回 True），加载过程中触发 torch.OutOfMemoryError: CUDA out of memory。复合模型 MiniMaxM3SparseForConditionalGeneration 的 language_model 子模块参数本应在 zero.Init 上下文中创建并分区，但实际在加载后的缺失键初始化阶段发生 OOM。

报错原文

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

原因分析

根据 Issue 中用户提供的 root cause 分析，ZeRO-3 实际在加载时正确分区了参数（每个 rank 的 CPU RSS 稳定在完整模型的约 1/8），OOM 发生在 transformers 的缺失键初始化阶段。具体原因如下：

_initialize_missing_keys 从整个 state dict 收集所有未初始化参数到一个列表，然后用单个 deepspeed.zero.GatheredParameters(...) 上下文包裹 initialize_weights()。这个单个合并的 all-gather 会把所有未初始化参数一次性反分区到 rank 0 上。对于 MiniMax-M3 这种稀疏 MoE 模型，每个 MoE 层的 expert 参数被打包为 3-D 参数（如 experts.gate_up_proj = [num_experts, 2*inter, hidden]），单个参数在未分区时约 9.7GB；57 个 MoE 层全部合在一起时，一次 coalesced all-gather 试图在单张 GPU 上重新物化数百 GB 数据，直接触发 CUDA OOM。

因此，这不是 ZeRO-3 分区本身的问题，而是 transformers 在缺失键初始化时的实现缺陷——它把整个模型的未初始化参数（数量庞大）放进单个 gather 上下文，导致内存峰值远超单卡容量。

环境排查

确认 transformers 版本是否为 5.12.1（或检查 _initialize_missing_keys 是否存在相同逻辑）
确认 deepspeed 版本是否为 0.18.9，并检查 GatheredParameters 行为是否一致
确认 torch 版本为 2.12.0+cu130（CUDA 13.0 配套）
确认 GPU 型号为 8×B200（每卡 178GB 显存）
检查模型是否为 MiniMaxAI/MiniMax-M3（428B 稀疏 MoE VLM），并确认 MoE 层参数打包格式（3-D experts.* 参数）
确认 HfDeepSpeedConfig 已持有且 is_deepspeed_zero3_enabled() 在加载前返回 True

解决步骤

（已验证）采用文本-only 加载路径：导出仅文本 checkpoint，保留 language_model.* 键，去除前缀转换为 model.*/lm_head.*，丢弃 vision tower 相关权重；然后加载独立的 MiniMaxM3VLForCausalLM。该路径已验证能在 zero.Init 下正确分区。
（可优先尝试）自行修补 transformers 源码：修改 modeling_utils.py 中 _initialize_missing_keys 的实现，避免将所有未初始化参数放入单个 GatheredParameters 上下文——可考虑分批 gather 或逐层初始化，但需自行验证性能与正确性。
（可能可行）升级/回退 transformers 版本：检查是否有更新版本修复了 _initialize_missing_keys 的 gather 逻辑，但 Issue 中没有明确证据，需自行验证。
（可作为绕过手段）使用 CPU offload 或 NVMe offload：在 ZeRO-3 配置中启用 offload_param 将参数卸载到 CPU/磁盘，可能降低单卡内存峰值，但 Issue 中未验证，需谨慎尝试。

验证方法

加载完成后，确认模型能在 ZeRO-3 下正常执行前向传播（如跑一次小型推理或训练 step），同时观察各 rank 的 CPU/GPU 内存峰值：如果使用文本-only 路径，应看到各 rank 的 RSS 稳定在完整模型的约 1/8；若直接加载完整 VLM，需确认不再触发 torch.OutOfMemoryError 且所有参数均被分区（检查 hasattr(param, "ds_id") 的属性存在性）。