ZeroRedundancyOptimizer仅支持ZeRO-1（分片优化器状态），不处理梯度和参数，显存节省有限；而DeepSpeed ZeRO-2/3分别分片梯度和参数，通信与内存优化更彻底，支持offload、overlap等高级特性，适用于大模型训练。
PyTorch 原生 ZeroRedundancyOptimizer 只支持 ZeRO-1（切分优化器状态），不能替代 DeepSpeed 的 ZeRO-2/3；它适用于轻量级多卡训练，但对大模型显存节省有限。
为什么用 ZeroRedundancyOptimizer 而不是直接写 DDP + Adam？
传统 DistributedDataParallel 每个 rank 都保存一份完整的 Adam 状态（exp_avg 和 exp_avg_sq），共占 8 字节/参数。10B 模型光这部分就吃掉 80GB 显存（8 卡 × 10GB）。ZeroRedundancyOptimizer 把这部分切开：每张卡只存 1/N 的动量和方差，显存直降约 50% —— 但它不碰梯度、不碰参数，所以省不了激活值或前向缓存。
常见错误现象：RuntimeError: Expected all tensors to be on the same device —— 因为它默认要求所有参数在同一个设备上初始化，但你若手动把模型分到多卡（如用 model.to(f'cuda:{rank}')），就会崩。
- 必须先用
torch.nn.parallel.DistributedDataParallel包裹模型，再把 DDP module 传给ZeroRedundancyOptimizer - 不能和
torch.compile或torch._dynamo混用（会触发未实现的 hook） - 不支持
load_state_dict直接恢复 optimizer 状态，需用state_dict+scatter_state_dict手动分发
ZeroRedundancyOptimizer 初始化时的关键参数
它本质是 wrapper，构造函数接收一个普通 optimizer（如 torch.optim.Adam）和一个 process_group：
立即学习“Python免费学习笔记（深入）”；
from torch.distributed.optim import ZeroRedundancyOptimizer from torch.distributed import init_process_group <h1>假设已初始化 process group</h1><p>optimizer = ZeroRedundancyOptimizer( params=model.parameters(), # 注意：必须是 model.parameters()，不能是 DDP.module.parameters() optimizer_class=torch.optim.Adam, lr=3e-5, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, process_group=process_group # 必须显式传入，不能依赖默认 )
容易踩的坑：
-
params必须来自原始模型（非 DDP 包裹后），否则报错Parameter is not registered in the optimizer -
process_group若不传，会 fallback 到全局 default group，但在多 process group 场景下（如混合并行）极易出错 - 不支持
foreach=True（PyTorch 2.0+ 的批量更新优化），启用会触发NotImplementedError
它和 DeepSpeed ZeRO 的根本区别在哪？
DeepSpeed 的 zero_optimization.stage 是端到端控制流，从梯度归约、参数分片、通信调度全链路接管；而 ZeroRedundancyOptimizer 只在 optimizer.step() 阶段介入，其余仍走 PyTorch 原生 DDP 流程。
这意味着：
- 梯度仍是 AllReduce 后全量存在每张卡上（ZeRO-2 的
reduce_scatter它不做） - 参数仍是全量拷贝（ZeRO-3 的
allgather_partitions它不提供） - 没有 offload 支持（无法把优化器状态卸载到 CPU）
- 通信模式固定为 blocking，不支持
overlap_comm这类重叠计算与通信的优化
性能影响明显：7B 模型在 8 卡 A100 上，ZeroRedundancyOptimizer 比原生 DDP 省约 4.2GB 显存/卡；DeepSpeed ZeRO-2 同配置可省约 9.6GB/卡。
什么时候该选它？什么时候该换 DeepSpeed？
如果你的场景满足以下全部条件，ZeroRedundancyOptimizer 是轻量可行解：
- 模型参数 ≤ 3B，单卡显存 ≥ 40GB（如 A100-40G）
- 仅需多卡加速，不追求极致吞吐或最大 batch size
- 不想引入 DeepSpeed 的 JSON 配置、
deepspeed.initialize初始化流程 - 训练脚本已重度依赖 PyTorch 原生 API（如自定义
grad_scaler、checkpoint）
一旦遇到这些信号，就得切 DeepSpeed：
- 训练时报
OutOfMemoryError即使开了ZeroRedundancyOptimizer - 想用
offload_optimizer.device: cpu把动量卸到内存 - 需要 pipeline parallel 或 tensor parallel 配合 ZeRO
- 模型 checkpoint 加载失败，或 optimizer state 恢复后 loss 突增（说明状态没对齐）
最易被忽略的一点：它不处理 torch.nn.functional.cross_entropy 中隐式创建的临时 buffer，这些 buffer 在多卡间不共享也不分片，仍按原大小驻留每张卡——这在长序列生成任务中可能吃掉额外 5–10GB 显存，得靠 torch.utils.checkpoint 或 torch.compile(fullgraph=True) 手动干预。
