二十二：分布式训练：DeepSpeed 与 ZeRO¶

来源:http://mp.weixin.qq.com/s?__biz=MzYyNTk3Njg1NA==&mid=2247484211&idx=1&sn=02c50dadc80d6909a906a4d367024d50&chksm=f01eb04ac769395ce7cc9c183c2b9f2336571893010a6ba4336d0899736f564f92e9116e0feb#rd

1. 学习定位¶

第 21 天学习了 DDP。DDP 每张卡保存完整参数、梯度和优化器状态，因此显存冗余大。DeepSpeed 的核心价值之一是 ZeRO（Zero Redundancy Optimizer），通过分片优化器状态、梯度和参数，让多卡共同承载更大的模型。本日知识链路：

DDP 显存冗余
-&gt; optimizer states / gradients / parameters
-&gt; ZeRO Stage 1/2/3
-&gt; CPU/NVMe offload
-&gt; activation checkpointing
-&gt; DeepSpeed config
-&gt; fp16/bf16、gradient accumulation、micro batch
-&gt; Hugging Face Trainer/Accelerate 集成
-&gt; checkpoint、性能排错

2. 大模型训练的显存组成¶

训练一个模型时，显存通常包括： - 参数 parameters。

梯度 gradients。
优化器状态 optimizer states，Adam 通常包含 momentum 和 variance。
激活 activations。
临时 buffer、通信 buffer、fragmentation。

以 Adam + FP16/BF16 训练为例，优化器状态可能比参数本身更占显存。DDP 每张卡都保存完整副本，因此模型越大，冗余越严重。

3. ZeRO 的核心思想¶

ZeRO 的目标是消除数据并行中的冗余状态。DDP 中每张卡保存完整：

parameters + gradients + optimizer states

问题在于：所有卡存了一模一样的完整数据，显存是高度冗余的，卡数越多、模型越大，这种浪费就越明显。 ZeRO 把这些状态按 data parallel ranks 分片，让每张卡只保存其中一部分。核心直觉：

需要计算时再通信聚合
不需要每张卡都常驻完整训练状态
用通信换显存

接下来是ZeRO的三级分片，逐级解放显卡。

4. ZeRO Stage 1---只分片优化器状态¶

ZeRO Stage 1 分片 optimizer states。

parameters: replicated（全卡复制）
gradients: replicated（全卡复制）
optimizer states: sharded（按DP卡数分片）

适合优化器状态显存占比高，但模型参数还能放下的场景。 - 前向、反向和标准 DDP 一致，参数、梯度全程每张卡都持有完整副本；

反向得到全局梯度后，通过reduce-scatter将梯度分片分发到对应 rank；
各 rank 仅用本地分片的 Adam 动量m、v，完成局部参数更新。

优点： - 改动相对小。

显存节省明显。
通信复杂度相对低。

5. ZeRO Stage 2---分片优化器状态+梯度¶

ZeRO Stage 2 在 Stage 1 基础上进一步分片 gradients。

parameters: replicated
gradients: sharded
optimizer states: sharded

相比 Stage 1，进一步减少梯度显存。训练时需要 reduce-scatter 等通信，把梯度分片分发给对应 rank。 - 前向传播同 DDP，参数全局复制；

反向计算梯度过程中，执行reduce-scatter，梯度直接按分片归属分发到对应 rank，不再全局保留完整梯度；
各 rank 结合本地梯度分片 + 优化器状态分片，完成参数更新。

Stage 2 常用于中大型模型训练，显存效率和通信成本之间较平衡。

6. ZeRO Stage 3---分片优化器状态+梯度+模型参数（全分片）¶

ZeRO Stage 3 进一步分片 parameters。

parameters: sharded
gradients: sharded
optimizer states: sharded

每张卡只常驻一部分参数。forward/backward 需要临时 all-gather 当前层参数，用完后释放。 - 前向计算某一层时，执行all-gather从其他 rank 拉取当前层完整参数；

计算完成后，立即释放临时聚合的参数，只保留本地分片；
反向传播同理，梯度经reduce-scatter归位到对应 rank；
各 rank 仅更新自己管辖的参数、梯度、优化器状态分片。

Stage 3 显存最省，但通信最复杂。它适合单卡无法保存完整模型参数或 optimizer states 的超大模型。

7. ZeRO 与 DDP/FSDP 对比¶

DDP：

完整复制参数、梯度、优化器状态
通信主要是 gradient all-reduce

ZeRO：

按 stage 分片训练状态
通过 reduce-scatter / all-gather 等通信恢复计算所需状态

FSDP：

PyTorch 原生参数分片方案
和 ZeRO Stage 3 思想接近

实际选型要看框架生态、模型结构、checkpoint、集成成本和团队熟悉度。

8. Offload¶

Offload = 把 GPU 常驻的训练状态，转移到 CPU 内存 / NVMe 硬盘存放，本质是用低速大容量存储，置换高速 GPU 显存。 DeepSpeed 支持把部分状态 offload 到 CPU 或 NVMe：

optimizer offload:
&nbsp; optimizer states 放 CPU/NVMe

parameter offload:
&nbsp; 参数分片放 CPU/NVMe

优点： - 显著降低 GPU 显存。

可训练更大模型。

代价： - 带宽瓶颈：PCIe/NVMe 带宽远低于 GPU 显存。

训练速度可能明显下降。
需要调 prefetch、pin memory、buffer 等参数。

Offload 是“用更慢的存储换更大容量”。

9. ZeRO-Infinity 与 ZeRO-Offload¶

ZeRO-Offload 主要把 optimizer states 和计算转移到 CPU，降低 GPU 显存压力。 ZeRO-Infinity 扩展到 NVMe offload，面向更大规模模型，利用 CPU/NVMe 层级存储承载模型状态。利用「GPU 显存 → CPU 内存 → NVMe 硬盘」三级存储架构，彻底突破硬件显存 / 内存上限。面试中重点不是背所有参数，而是说明：

Offload 可以突破 GPU 显存限制
但会引入主机内存、PCIe、NVMe IO 瓶颈

10. Activation Checkpointing 激活重计算¶

激活显存随 batch size、sequence length、hidden size、layers 增长。activation checkpointing 的思想是：

forward 时不保存所有中间激活
backward 时重新计算部分 forward
用额外计算量换取显存开销下降

DeepSpeed 支持 activation checkpointing，可与 ZeRO 组合。代价： - backward 计算量增加。

训练时间变长。

优点： - 显著降低长序列和大模型训练显存。

11. DeepSpeed 配置文件¶

DeepSpeed 通常通过 JSON 配置：

{
&nbsp; "train_micro_batch_size_per_gpu": 1,
&nbsp; "gradient_accumulation_steps": 16,
&nbsp; "bf16": { "enabled": true },
&nbsp; "zero_optimization": {
&nbsp; &nbsp; "stage": 2
&nbsp; }
}

常见字段： - train_micro_batch_size_per_gpu。

gradient_accumulation_steps。
train_batch_size。
fp16 / bf16。
zero_optimization.stage。
optimizer。
scheduler。
gradient_clipping。
activation_checkpointing。
zero_allow_untested_optimizer。

配置错误是 DeepSpeed 使用中最常见问题之一。

12. Batch 语义¶

DeepSpeed 中有三个容易混淆的 batch 概念：

micro batch per GPU:
&nbsp; 每张卡每次 forward 的样本数

gradient accumulation steps:
&nbsp; 累积多少 micro step 后做 optimizer step

train batch size:
&nbsp; global batch

关系：

train_batch_size
= train_micro_batch_size_per_gpu
&nbsp; * gradient_accumulation_steps
&nbsp; * data_parallel_world_size

如果配置不一致，DeepSpeed 可能报错或自动推断，建议明确计算。

13. DeepSpeed 初始化¶

典型代码：

model_engine, optimizer, train_loader, scheduler = deepspeed.initialize(
&nbsp; &nbsp; model=model,
&nbsp; &nbsp; model_parameters=model.parameters(),
&nbsp; &nbsp; training_data=dataset,
&nbsp; &nbsp; config=ds_config
)

loss = model_engine(batch)
model_engine.backward(loss)
model_engine.step()

DeepSpeed 接管 backward、optimizer step、梯度累积、ZeRO 通信和状态管理。使用 Hugging Face Trainer 或 Accelerate 时，DeepSpeed 可通过配置文件集成，用户不一定直接调用 deepspeed.initialize。

14. Mixed Precision¶

DeepSpeed 支持 FP16、BF16 和部分低精度训练配置。 FP16：

需要 loss scaling
可能有 overflow

BF16：

动态范围大
通常更稳定
需要硬件支持

大模型训练中，BF16 通常是更稳的默认选择。使用 FP16 时要关注 loss scale、overflow 次数和梯度异常。

15. DeepSpeed Checkpoint¶

ZeRO 下 checkpoint 不再只是普通 state_dict。尤其 ZeRO-3 中参数被分片保存。常见操作：

model_engine.save_checkpoint(path)
model_engine.load_checkpoint(path)

需要注意： - 每个 rank 可能保存自己的 shard。

从 ZeRO checkpoint 合并成普通 FP32 state_dict 需要专门工具。
保存/加载时 world size 变化可能涉及转换。
Hugging Face 格式导出要确认是否已 gather 完整权重。

16. DeepSpeed 与 Hugging Face¶

Hugging Face Trainer 可通过参数启用 DeepSpeed：

--deepspeed ds_config.json

Accelerate 也支持 DeepSpeed plugin。常见组合： - SFT + DeepSpeed ZeRO-2/3。

LoRA/QLoRA + DeepSpeed。
RLHF/PPO + DeepSpeed。
多机训练 + DeepSpeed launcher。

需要注意 PEFT、量化、ZeRO-3 参数 gather 和保存之间的交互。

17. 性能与通信¶

ZeRO 节省显存，但通信更多。常见通信： - all-gather 参数。

reduce-scatter 梯度。
all-reduce 部分统计。
offload prefetch/evict。

性能优化方向： - 合理选择 ZeRO stage。

增大 micro batch 或 gradient accumulation。
减少 offload。
使用 BF16。
开启 activation checkpointing。
调整 bucket size。
避免频繁 checkpoint。

18. 常见 OOM 排查¶

OOM 时先判断显存来源： - 参数太大：考虑 ZeRO-3/FSDP。

优化器状态太大：ZeRO-1/2。
梯度太大：ZeRO-2。
activation 太大：activation checkpointing、减 sequence length、减 micro batch。
临时 buffer：减 bucket size、检查 kernel。

常见策略：

降低 micro batch
增加 gradient accumulation
启用 activation checkpointing
提高 ZeRO stage
使用 offload
使用 LoRA/QLoRA
缩短 sequence length

19. 常见误区¶

误区一：DeepSpeed 一定比 DDP 更快。 DeepSpeed 主要优势是显存和大模型能力。ZeRO stage 越高通信越多，不一定更快。误区二：ZeRO-3 永远最好。 ZeRO-3 最省显存但通信重。如果模型 DDP/ZeRO-2 能放下，ZeRO-3 可能更慢。误区三：Offload 是免费显存扩展。 Offload 会消耗 CPU/NVMe 带宽，速度可能明显下降。误区四：DeepSpeed batch 配置可以随便写。 micro batch、accumulation、world size、global batch 必须一致。误区五：ZeRO checkpoint 就是普通 PyTorch checkpoint。 ZeRO checkpoint 可能是分片格式，导出和恢复要按 DeepSpeed 规则处理。

20. 核心总结¶

第 22 天需要掌握的最小闭环：

DeepSpeed:
&nbsp; distributed training system for large models

ZeRO:
&nbsp; remove data parallel redundancy

Stage 1:
&nbsp; shard optimizer states

Stage 2:
&nbsp; shard optimizer states + gradients

Stage 3:
&nbsp; shard optimizer states + gradients + parameters

Offload:
&nbsp; move states to CPU/NVMe

Activation checkpointing:
&nbsp; recompute activations to save memory

Batch:
&nbsp; global = micro_per_gpu * grad_accum * dp_world_size

Tradeoff:
&nbsp; memory saving usually costs communication or recomputation

21. 参考资料¶

DeepSpeed 官方文档：https://www.deepspeed.ai/
DeepSpeed ZeRO 文档：https://www.deepspeed.ai/tutorials/zero/
ZeRO 论文：https://arxiv.org/abs/1910.02054
ZeRO-Offload 论文：https://arxiv.org/abs/2101.06840
ZeRO-Infinity 论文：https://arxiv.org/abs/2104.07857
DeepSpeed 实战视频：https://www.bilibili.com/video/BV1hb421E7WY/
DeepSpeed 详解：https://blog.csdn.net/zwqjoy/article/details/130732601
分布式训练方法汇总：https://zhuanlan.zhihu.com/p/598714869
```
        预览时标签不可点
```
<div class="