Python开发中PyTorch如何使用ZeroRedundancyOptimizer节省显存？-Python教程-PHP中文网

ZeroRedundancyOptimizer仅支持ZeRO-1（分片优化器状态），不处理梯度和参数，显存节省有限；而DeepSpeed ZeRO-2/3分别分片梯度和参数，通信与内存优化更彻底，支持offload、overlap等高级特性，适用于大模型训练。

python开发中pytorch如何使用zeroredundancyoptimizer节省显存？

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：

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  # 必须显式传入，不能依赖默认
)</p>

容易踩的坑：

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 流程。

Python Hardener

通过完整的安全审查、代码整洁重构、健壮的错误处理、正确的日志记录以及完整的文档字符串覆盖，强化 Python 脚本。

下载

这意味着：

梯度仍是 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) 手动干预。