[Bug已解决] FSDP summon_full_params(offload_to_cpuTrue) 在 CPU 卸载模型上失败解决方案一、现象长什么样你用FullyShardedDataParallelFSDPPyTorch 的另一种分布式并行比 DDP 更省显存做训练并且开启了参数 CPU 卸载offload_to_cpuTrue——把分片参数卸载到 CPU 内存用时再搬回 GPU。然后你想用summon_full_params临时「召唤出完整参数」比如推理、保存 checkpoint、或某种需要全参数的操作却在offload_to_cpuTrue时失败[FSDP] summon_full_params(offload_to_cpuTrue) fails with CPU-offloaded FSDP model也就是官方描述的场景。summon_full_params是 FSDP 的一个上下文管理器进入时把所有分片参数聚合all-gather成完整参数退出时再分片回去。但当模型本身已开启 CPU 卸载召唤完整参数时的「设备搬运逻辑」和卸载逻辑冲突导致失败。本文讲清楚 FSDP 与 CPU 卸载是什么、summon_full_params为何冲突以及如何规避。二、FSDP 与 CPU 卸载是什么2.1 FSDP 基本思想DDP 是「每卡一份完整模型」。FSDP 更进一步每张卡只保存模型的一部分参数分片前向 / 反向需要时通过 all-gather 临时拼出完整参数用完即拆回分片。这样单卡显存占用大幅降低能训更大的模型。2.2 CPU 卸载offload_to_cpu显存还不够FSDP 可以把「暂时不用的分片参数」卸载到CPU 内存用的时候再搬回 GPU。代价是 CPU↔GPU 搬运开销但显存占用极低。from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import CPUOffload model FSDP( model, cpu_offloadCPUOffload(offload_paramsTrue), # 参数卸载到 CPU )2.3 summon_full_params有时你需要「完整参数」视图推理、算度量、存 ckpt。summon_full_params临时把所有分片聚合成完整参数with FSDP.summon_full_params(model): # 此时 model 参数是完整的跨所有 rank 聚合 do_something_with_full_params(model)三、为什么 offload_to_cpuTrue 时召唤失败summon_full_params内部要做all-gather 各 rank 的分片 → 拼成完整参数 → 放到某设备。但当模型开了offload_to_cpuTrue参数默认在 CPU 上召唤逻辑在「聚合后的完整参数该放哪、如何和卸载状态协调」上出 bug召唤要求把分片从「当前所在设备CPU」all-gather 到 GPU 拼完整但卸载状态让这个流程的设备假设错乱或退出summon_full_params时无法正确把完整参数「重新分片 卸载回 CPU」于是 RuntimeError / 设备不匹配。四、可运行健壮的 FSDP 推理骨架下面脚本演示 FSDP summon_full_params的标准安全写法假设多卡环境单卡简化演示import os import datetime import torch import torch.distributed as dist import torch.nn as nn from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp import CPUOffload def setup(rank, world): os.environ[MASTER_ADDR] os.environ.get(MASTER_ADDR, localhost) os.environ[MASTER_PORT] os.environ.get(MASTER_PORT, 29501) dist.init_process_group(nccl, rankrank, world_sizeworld, timeoutdatetime.timedelta(seconds60)) torch.cuda.set_device(rank) def main(): rank int(os.environ.get(RANK, 0)) world int(os.environ.get(WORLD_SIZE, 1)) setup(rank, world) device torch.device(fcuda:{rank}) model nn.Linear(16, 8).to(device) # 开 CPU 卸载 model FSDP(model, cpu_offloadCPUOffload(offload_paramsTrue), device_idrank) x torch.randn(4, 16, devicedevice) out model(x) print(f[rank {rank}] 前向输出形状, out.shape) # 推理时召唤完整参数注意开 offload 时需在正确设备上下文 with FSDP.summon_full_params(model, writebackFalse): # writebackFalse 表示不写回分片避免卸载状态冲突 for name, p in model.named_parameters(): if rank 0: print(name, 是否为完整参数, p.shape) dist.destroy_process_group() if __name__ __main__: main()关键是writebackFalse——召唤出来的完整参数不写回分片状态规避卸载逻辑冲突。五、解决方案一summon 时用 writebackFalsesummon_full_params(writebackFalse)让召唤出的完整参数不回写到分片 / 卸载状态避免退出时的设备协调冲突with FSDP.summon_full_params(model, writebackFalse): # 只读地用完整参数不修改 pass适合「临时查看 / 推理 / 计算度量」这类只读场景。如果需要修改参数再写回见下节。六、解决方案二召唤前暂时关掉 CPU 卸载如果必须writebackTrue修改后写回可以先把模型从 CPU 卸载状态「搬到 GPU」再召唤# 伪代码思路用 FSDP 的 reshard_after_forward / 手动 flush # 具体 API 随版本变化核心是让参数先回到 GPU 分片态再 summon with FSDP.summon_full_params(model): # 此时参数已被聚合做需要的修改 pass不同 PyTorch 版本对「offload summon writeback」的支持不同。稳妥起见要么writebackFalse要么升级到修复该交互的版本。七、解决方案三保存 checkpoint 不用 summon_full_params保存 FSDP 模型的 checkpoint其实不需要召唤完整参数——可以分片保存加载时再分片恢复import torch def save_fsdp_sharded(model, path, rank): # 每个 rank 保存自己那份分片 state model.state_dict() torch.save(state, f{path}_rank{rank}.pt) def load_fsdp_sharded(model, path, rank): state torch.load(f{path}_rank{rank}.pt) model.load_state_dict(state)分片保存避免了summon_full_params的设备冲突也更适合大规模模型。如果非要存「完整 ckpt」可以在 rank 0 上单独聚合用summon_full_params(writebackFalse)只读复制出来再存而不写回。八、解决方案四只在推理阶段召唤训练阶段不召唤summon_full_params主要服务于「需要完整参数视图」的偶发操作。训练主循环里不要频繁召唤既慢又容易踩 offload 冲突。把它放在训练结束后的评估周期性保存 ckpt特定调试。# 训练循环里正常 FSDP分片、可能卸载 for batch in loader: train_step(model, batch) # 训练后评估时才召唤 with FSDP.summon_full_params(model, writebackFalse): evaluate(model, val_loader)九、解决方案五升级 PyTorchsummon_full_params(offload_to_cpuTrue)的失败是 FSDP 在该组合下的具体 bug新版本会逐步修复。查看并升级import torch print(PyTorch, torch.__version__) print(FSDP 可用, hasattr(torch.distributed, fsdp))十、如何判断你踩的是同一条用的是 FSDP不是 DDP开了cpu_offloadCPUOffload(offload_paramsTrue)调用FSDP.summon_full_params(...)失败不开 offload 时 summon 正常开了就失败用writebackFalse后恢复。命中即说明踩中该 offload summon 的冲突 bug。十一、小结summon_full_params(offload_to_cpuTrue)在 CPU 卸载的 FSDP 模型上失败是**「聚合完整参数」的设备搬运逻辑与「参数卸载」状态冲突**。应对召唤时用writebackFalse只读不写回第五节最稳必须写回就先让参数回到 GPU 分片态再召唤第六节保存 ckpt 用分片保存避免召唤第七节只在推理 / 评估阶段召唤训练主循环不召唤第八节升级到修复该交互的 PyTorch第九节。FSDP CPU 卸载是「训超大模型」的神器但它把参数在 GPU / CPU / 分片 / 完整 之间反复搬运状态机复杂。summon_full_params只是其中一个「状态切换点」——理解它的设备假设你就能在省显存和用完整参数之间游刃有余。