2、DeepSpeed基础：分布式训练原理、ZeRO优化阶段、混合精度、通信机制、常见应用场景
在排查 DeepSpeed 训练故障之前，必须先建立一张清晰的“系统地图”：一次训练步里，数据、梯度、优化器状态和通信到底在哪里流动，哪些环节由 DeepSpeed 接管。只有理解这些基础机制，后续遇到 OOM、卡死、loss 不收敛、吞吐骤降时，才能快速定位根因，而不是盲目调参。
一、分布式训练原理：从数据并行到分片并行
DeepSpeed 建立在 PyTorch 分布式之上，最常用的起点是数据并行（DP）：每张 GPU 持有完整模型副本，输入不同 mini-batch，反向传播后通过 AllReduce 汇总梯度，再各自更新参数。它的优点是实现简单、扩展性好；缺点是显存冗余——模型参数、梯度、优化器状态在每个 rank 上都完整保存一份。
当模型规模超过单卡容量时，就需要引入模型并行（张量并行 TP、流水线并行 PP）与 ZeRO 分片。DeepSpeed 的核心价值在于：在不修改模型代码的前提下，把原本冗余存储的状态切分到各卡，并在需要时通过通信临时聚合。理解“哪些状态被切分、何时通信”，是判断故障类型的关键。
二、ZeRO 优化阶段：显存换通信的阶梯
ZeRO（Zero Redundancy Optimizer）把训练状态分为三块：优化器状态、梯度、参数。阶段越高，切分越彻底，单卡显存占用越低，但通信量也越大。
| 阶段 | 切分对象 | 显存收益 | 通信开销 | 典型适用 |
|---|---|---|---|---|
| ZeRO-1 | 优化器状态 | 约 4 倍 | 低 | 中小模型，追求吞吐 |
| ZeRO-2 | 优化器状态 + 梯度 | 约 8 倍 | 中 | 主流大模型训练 |
| ZeRO-3 | 优化器状态 + 梯度 + 参数 | 与卡数线性 | 高 | 超大模型，显存极限 |
| ZeRO-Infinity | 以上 + CPU/NVMe offload | 进一步扩展 | 很高 | 单机多卡跑超大模型 |
配置示例：
{
"zero_optimization": {
"stage": 2,
"allgather_partitions": true,
"reduce_scatter": true,
"overlap_comm": true,
"contiguous_gradients": true
},
"fp16": { "enabled": true, "loss_scale": 0 }
}三、混合精度：速度与稳定性的平衡
混合精度让前向/反向用 FP16 或 BF16 计算，同时保留 FP32 主权重和优化器状态。FP16 动态范围小，需要 loss scaling 防止梯度下溢；BF16 动态范围接近 FP32，通常无需 loss scaling，但精度略低。DeepSpeed 中通过 fp16 或 bf16 配置启用。
常见故障信号：loss 出现 NaN/Inf、梯度范数异常、loss scale 频繁回退。排查时应先确认硬件是否支持 BF16，再检查 loss scale 初始值与窗口参数。
四、通信机制：NCCL 与通信桶
DeepSpeed 的集合通信由 NCCL 承载，核心操作包括 AllReduce、ReduceScatter、AllGather。ZeRO-2 中梯度先 ReduceScatter 到各卡分片，再在更新时 AllGather；ZeRO-3 额外对参数做 AllGather。通信桶（bucket）大小、overlap_comm 是否开启，直接影响通信与计算的重叠效率。
五、常见应用场景与误区复盘
- 单机多卡微调：ZeRO-2 + FP16，重点调 batch size 与 gradient accumulation。
- 多机大模型预训练：ZeRO-3 + BF16 + TP/PP，重点排查跨机通信与负载均衡。
- 显存受限推理/训练：ZeRO-Infinity + offload，注意 CPU 内存与 NVMe 带宽瓶颈。
一个典型误区是“显存不够就无脑升 ZeRO-3”。实际上，若瓶颈在通信而非显存，升阶段只会让 step time 更长。正确做法是先看显存峰值与通信占比，再决定是调 bucket、开 overlap，还是升阶段。另一个误区是忽略 loss scale 与 BF16 的差异，导致 loss 曲线抖动却误判为数据问题。