LlamaFactory 全参微调 + FSDP2：`init_adapter` 的 pre-shard fp32 upcast 使 `cpu_ram_efficient_loading` 失效，载入权重时 host 峰值达 N×M（每卡一份完整模型）

现象

• 环境:transformers 4.57.1 + accelerate 1.11.0 + LlamaFactory(hyper-parallel FSDP2 集成),全参微调(full),bf16 混合精度,cpu_ram_efficient_loading=True。
• 根据实验:开启 cpu_ram_efficient_loading 后,模型加载阶段 host 内存(free -h 的 used)仍随卡数线性增长——每个 rank 各持一份完整模型,总量约为 N×M(N=每节点 rank 数,M=单份模型大小),而非预期的 ~1×M。规模大时直接把节点 OOM。
• 逐点打点定位:非 rank0 在 from_pretrained 后仍是空占位(host≈baseline),真正被实体化发生在其后的 init_adapter。

背景:cpu_ram_efficient_loading 的设计意图

FSDP 下开启该选项后,只有每节点的 local rank0 真正从磁盘读取权重到 CPU(占 1×M),其余 rank 的参数保持 meta / 空占位(物理内存≈0);真实权重由 FSDP 在设备侧从 rank0 broadcast 到各 rank 的 shard。因此 host 峰值应为 ~1×M。该机制由 is_fsdp_enabled() 门控。

根因:init_adapter 在分片前于 host 上 upcast 到 fp32

加载后状态(正常)

from_pretrained 时门控正常触发,非 rank0 的参数落到 meta / 以 torch.empty_like(..., device="cpu") 建惰性空占位(不写内容、物理 RSS≈0):

https://github.com/huggingface/transformers/blob/v4.57.1/src/transformers/modeling_utils.py#L5826-L5851

此时:rank0 = 完整 bf16(1×M);非 rank0 = 空占位(≈0)。

正常路径:fp32 upcast 应在 分片之后

bf16 混合精度需要可训练参数的 fp32 主权重,但正确时机是 FSDP 分片之后,每个 rank 只 upcast 自己 1/N 的 shard(设备侧),不触碰 host 全量。accelerate / hyper-parallel 都在分片后执行:

https://github.com/huggingface/accelerate/blob/v1.11.0/src/accelerate/utils/fsdp_utils.py#L701C1-L713C17

问题路径:init_adapter 在 分片前、于 host upcast

全参微调时,LlamaFactory init_adapter 会在 FSDP 分片之前,把每个可训练参数在 host 上 upcast 到 fp32:

https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/model/adapter.py#L299-L344

https://github.com/hiyouga/LlamaFactory/blob/main/src/llamafactory/model/adapter.py#L292-L294

冲突点:此时模型还没分片,非 rank0 的参数正是上面那批 empty_like 空占位。param.data.to(torch.float32) 对每个参数做写入,把非 rank0 的空占位全部 fault 进物理 host → 每个 rank 都实体化一份完整模型 → host = N×M(fp32 比 bf16 还翻倍,趋向 N×2M)。这与 cpu_ram_efficient_loading 的设计直接冲突:本应保持空占位到设备侧 broadcast,却被这步 pre-shard host upcast 提前实体化。

注:hyper-parallel 的权重分发是设备侧(.to(device) + dist.broadcast),不增加 host;host 的 N×M 完全来自这步 host fp32 upcast。

device / dtype 变化对照

阶段	rank0 host	非 rank0 host	dtype
from_pretrained 后	完整(1×M)	空占位 empty/lazy(≈0)	bf16
init_adapter 带 host upcast(问题)	完整 fp32	完整 fp32(占位被 fault)	bf16→fp32
init_adapter 跳过 host upcast(修复)	完整 bf16	仍为空占位(≈0)	bf16
分片 + post-shard upcast	1/N shard fp32	1/N shard fp32	fp32(仅 1/N,设备侧)

修复

在 LlamaFactory 集成里 monkeypatch adapter._setup_full_tuning / _setup_freeze_tuning:当 is_fsdp_enabled() 且模型为 fp16/bf16 时,强制 cast_trainable_params_to_fp32=False,跳过 host 上的 pre-shard upcast,把 upcast 推迟到分片之后(由 _maybe_upcast_trainable_params 完成)。patch 在 load_model 之前安装,幂等。

正确性:跳过的只是冗余的 pre-shard host upcast;fp32 主权重仍由 post-shard 的 _maybe_upcast_trainable_params 保证(分片后每 rank 仅 1/N),混合精度语义不变。patch 仅在 FSDP + bf16/fp16 时生效,不影响纯 bf16 / zero3 / 量化等其它路径。

根据实验:修复后非 rank0 全程保持空占位,只有 local rank0 持一份完整模型,host 峰值从 N×M 降到 ~1×M,符合 cpu_ram_efficient_loading 的设计目标。

注意:与 torch.use_deterministic_algorithms(True) 的交互(另一条独立的 N×M 来源)

即使上面已修复,只要训练脚本里开了 torch.use_deterministic_algorithms(True),host 仍会回到 N×M。

机理:它会联动开启 torch.utils.deterministic.fill_uninitialized_memory(默认 True),使 torch.empty()/empty_like() 在分配后立即用已知值(NaN)写一遍。于是非 rank0 那批惰性 empty_like 占位被全部写实 → 物理页 fault → 整模型实体化 → N×M。(只有 use_deterministic_algorithms 有此影响;各类 seed、HCCL_DETERMINISTIC 不影响 host 分配。)

安全:这些占位会被设备侧 broadcast 的真值覆盖、从不被读,故不填充不影响数值或确定性。(NPU 上建议实测确认该 flag 被 torch_npu 的 empty 继承。)

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 222
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/222

Start at the LlamaFactory integration's load_model path and inspect adapter.py's _setup_full_tuning and _setup_freeze_tuning around the cited lines. Verify how is_fsdp_enabled and the fp16/bf16 settings reach cast_trainable_params_to_fp32, then confirm post-shard _maybe_upcast_trainable_params still handles the conversion. Done means non-local-rank0 hosts retain lazy placeholders during loading without changing non-FSDP paths.