不能混用ddp与deepspeed,因后者已封装分布式通信和设备管理;混用会导致张量分散在不同device而报runtimeerror;正确做法是模型定义后直传deepspeed.initialize(),禁用所有.to()/.cuda()及ddp包装。

能用,但必须绕开 PyTorch 原生 DDP、禁掉所有 .to() 和 .cuda()、且 ZeRO-3 下推理要手动干预参数拉取逻辑——否则大概率 OOM 或报 RuntimeError: Expected all tensors to be on the same device。
为什么不能和 torch.nn.parallel.DistributedDataParallel 混用
DeepSpeed 自己封装了完整的分布式通信、参数同步和设备调度。一旦你手动 wrap 一层 DDP,或在模型初始化后调用 model.cuda(),就会导致:
- DeepSpeed engine 管理的参数在 GPU 上,而 DDP 又试图把副本放到各卡上,张量分散在不同 device
-
engine.module和原始model指向不同设备,forward()时输入数据和模型不在同一卡 - 触发
RuntimeError: Expected all tensors to be on the same device(这是最常见报错)
实操唯一正确路径:模型定义完直接传给 deepspeed.initialize(),中间不加任何设备转移或并行包装。
zero_optimization.stage 选 2 还是 3?
别一上来就上 stage 3。ZeRO-2 已切分优化器状态 + 梯度,显存节省约 8 倍;ZeRO-3 额外切分参数,显存更省,但每次 forward/backward 都隐式触发 all-gather 拉取全量参数——这在 batch 小、层数多、激活大时极易拖慢训练,甚至比 stage 2 还慢。
建议步骤:
- 先用
--zero-stage 2跑 baseline,看显存是否仍超限 - 若显存 OK,就停在 stage 2;若仍 OOM,再升 stage 3
- stage 3 必须配
offload_param或offload_optimizer,否则通信开销无补偿 - 启用
stage3_gather_16bit_weights_on_model_save: true防保存时 OOM,但会拖慢save_checkpoint()
ZeRO-3 下 torch.no_grad() 推理仍 OOM 怎么办
DeepSpeed 默认在 forward() 前仍执行参数 all-gather,哪怕进了 torch.no_grad()——这是为兼容某些 hook 行为,但对纯推理完全冗余。
不能只靠 model.eval(),必须显式干预:
- 调用
engine._allgather_params(reload=False)手动预热/拉取一次,之后 forward 不再重复拉 - 或者在 eval loop 开头加
engine.zero_grad()清 optimizer 状态(避免梯度相关 gather) - 慎用
engine.disable_untested_features(),它可能破坏内部状态一致性 - 绝对不要在 forward 中写
input.to('cuda'),统一用input.to(engine.device)
torch.compile() 和 ZeRO-3 冲突怎么解
PyTorch 2.0+ 的 torch.compile() 会构建 CUDA graph,而 ZeRO-3 的 stage3_prefetch_bucket_size(默认非零)会在编译期触发动态内存操作,导致 graph capture 失败。
解决办法只有一条:
- 必须设
stage3_prefetch_bucket_size: 0或直接删掉该字段(让它 fallback 到默认 0) - 其他 ZeRO-3 参数如
reduce_bucket_size、contiguous_gradients可照常使用 - 开启
torch.compile()后务必验证 loss 是否收敛,因某些自定义 op 可能未被 graph 正确覆盖
真正麻烦的从来不是配置写几行,而是 ZeRO-3 把“参数在哪”这个本该静态的事变成了运行时按需调度——你得时刻清楚哪一步在拉、哪一步在释放、哪一步被 compile 锁死了。漏掉一个 .to(device),或少清一次 gather 状态,OOM 就在下一 epoch。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











