[FSDP2] cpu_ram_efficient_loading materializes full model on CPU for every non-rank-0, defeating rank-0 broadcast (host OOM)

快速结论：这个报错通常出现在多卡、单机多进程的 FSDP2 训练场景中：用户启用了 cpu_ram_efficient_loading（FSDP_CPU_RAM_EFFICIENT_LOADING=1），本以为只有 rank-0 在 CPU 上加载权重、再广播给其他 rank，但 from_pretrained 阶段每个非 rank-0 进程也会在 CPU 上完整实例化一份模型，导致主机内存被成倍占用并触发 OOM。优先确认当前 transformers 版本是否已包含针对该路径的修复（见 #45649）。

适用环境：transformers 5.12.0、accelerate 1.15.0、torch 2.10；1 进程/GPU，单节点 16 rank；FSDP2（FullyShardedDataParallelPlugin(fsdp_version=2, cpu_ram_efficient_loading=True)）；模型约 27B 参数，torch_dtype=bfloat16，完整 checkpoint 约 54 GB。Issue 未给出操作系统、Python、CUDA 和具体显卡型号。

最快修复方案：升级到已合入 #45649 修复的 transformers 版本后重试。Issue 讨论中维护者确认“I think this was fixed by #45649!”，这是目前唯一有明确指向的修复来源；但该结论未在评论中附带验证记录，升级后仍需按“验证方法”实测。

注意事项：该修复能否覆盖你使用的 transformers 版本需要自行核对；Issue 本身被自动标记为 stale 后关闭，没有贴出逐个 rank 的内存对照数据。若升级后非 rank-0 仍出现整模型 CPU 分配，应回到 #45649 的改动范围继续排查，而不是只调整节点内存。

问题场景

用户在单节点 16 卡（1 进程/GPU，world=16）上，通过 accelerate 的 FSDP2 插件加载一个约 27B 参数、bfloat16（完整 checkpoint 约 54 GB）的因果语言模型。启动时设置了 ACCELERATE_USE_FSDP=1 与 FSDP_CPU_RAM_EFFICIENT_LOADING=1，并调用 AutoModelForCausalLM.from_pretrained(ckpt, torch_dtype=torch.bfloat16)。

由于 FSDP2 的包装（accelerator.prepare）发生在 from_pretrained 之后，主机内存实际消耗集中在 from_pretrained 阶段。用户观察到每个 rank 的 host RSS 在这个调用中都会上涨，16 个 rank 各自持有约 54 GB 主机内存，最终在内存小于约 1 TB 的节点上被 SIGKILL。用户还指出非 rank-0 会重复从磁盘读取 checkpoint 分片，产生多余 I/O。

报错原文

[FSDP2] cpu_ram_efficient_loading materializes full model on CPU for every non-rank-0, defeating rank-0 broadcast (host OOM)

对应的关键代码路径（Issue 中引用 modeling_utils.py v5.12.0 约 L4762）：

if is_fsdp_enabled() and not is_local_dist_rank_0() and not is_quantized:
    for key, param in self.named_parameters():
        value = torch.zeros_like(param, device="cpu")   # full ~54 GB per rank
        _load_parameter_into_model(self, key, value)
    ...

原因分析

最可能的原因是：在 from_pretrained 期间，_move_missing_keys_from_meta_to_device 对所有非 rank-0 进程无条件执行了“在 CPU 上物化完整模型”的逻辑，即对每个参数调用 torch.zeros_like(param, device="cpu") 再写回模型。这样一来，cpu_ram_efficient_loading 的设计意图被抵消：原本应只由 rank-0 在 CPU 上加载权重，其余 rank 保持 meta 张量，等 FSDP2 加载器广播时再填充。

而 accelerate 的 FSDP2 加载路径（fsdp2_load_full_state_dict，由 parallelize_model_fsdp2 调用，后者会执行 model.to("meta")）本身已经通过 dist.broadcast(..., src=0) 再 load_state_dict(assign=True) 把权重从 rank-0 广播给其他 rank，其 docstring 也写明非主进程应持有空的 meta 模型。因此非 rank-0 在 CPU 上物化权重属于多余行为，是主机内存 OOM 的直接来源。

环境排查

确认 transformers 版本是否为 5.12.0，以及是否已包含 #45649 的修复。
确认 accelerate 版本（Issue 中为 1.15.0）与 torch 版本（Issue 中为 2.10）是否匹配。
确认启动环境变量 ACCELERATE_USE_FSDP=1、FSDP_CPU_RAM_EFFICIENT_LOADING=1 是否生效。
确认 FSDP 插件配置为 fsdp_version=2 且 cpu_ram_efficient_loading=True，并注意 FSDP 包装发生在 from_pretrained 之后。
确认节点主机内存总量与模型规模（约 27B 参数、bfloat16、约 54 GB checkpoint）的关系，判断是否有余量容纳意外的多份副本。
确认 Python、CUDA、显卡型号等未在 Issue 中提供的环境项，便于与官方 CI 环境对比。

解决步骤

记录问题复现时的环境：transformers / accelerate / torch 版本、rank 数、模型大小与 dtype，以及各 rank 的 host RSS 峰值，作为基线。
对照 #45649 确认该修复是否已进入你当前使用的 transformers 版本；若未包含，升级到包含该修复的版本。
用与复现相同的启动方式（ACCELERATE_USE_FSDP=1 FSDP_CPU_RAM_EFFICIENT_LOADING=1，world=16）和相同的 from_pretrained 调用重新运行。
在 from_pretrained 调用前后按 rank 采样 host RSS，重点观察非 rank-0 是否仍出现约 54 GB 级别的 CPU 分配。
若升级后问题依旧，回到 #45649 讨论帖补充版本信息与 per-rank 内存数据，说明修复未覆盖你的场景。

验证方法

在 AutoModelForCausalLM.from_pretrained 执行期间逐 rank 观察主机 RSS：修复生效时，非 rank-0 不应再出现与完整模型体积（约 54 GB）相当的 CPU 内存分配，只有 rank-0 承担加载开销，其余 rank 在广播前保持 meta 模型。继续完成 accelerator.prepare 与后续前向/反向，确认训练正常启动，且节点不再因 host OOM 被 SIGKILL。同时可观察非 rank-0 是否还重复读取 checkpoint 分片。