理解通信原语和 DDP 的梯度同步机制

配套代码

训练大模型的第一步，是让多块 GPU 一起干活。数据并行（Data Parallelism）是最直觉的方式：每块 GPU 持有完整的模型副本，各自处理不同的数据，最后把梯度汇总。本章我们先理解单卡的内存瓶颈，再学习多卡通信的基础，最后深入 DDP 的实现。

单卡训练的内存组成

一个 7B 参数的模型，fp16 权重只需要 14 GB（7B × 2 bytes），但在 80GB 的 A100 上却训练不了。内存都去哪了？

答案是：训练时 GPU 上不只存参数，还要存梯度和优化器状态。

混合精度训练的内存需求

现代训练普遍采用混合精度（Mixed Precision）：前向和反向用 fp16 计算（快），但参数更新用 fp32（精确）。以 Adam 优化器为例，假设模型有 Φ Φ 个参数：

fp16 部分（前向/反向）：

• 参数： 2 Φ 2Φ bytes
• 梯度： 2 Φ 2Φ bytes

fp32 部分（优化器）：

• 参数副本： 4 Φ 4Φ bytes（用于精确更新）
• 一阶矩 𝑚 m： 4 Φ 4Φ bytes（梯度的指数移动平均）
• 二阶矩 𝑣 v： 4 Φ 4Φ bytes（梯度平方的移动平均）

总计：2 Φ + 2 Φ + 4 Φ + 4 Φ + 4 Φ

16 Φ 2Φ+2Φ+4Φ+4Φ+4Φ=16Φ bytes

Φ 2Φ | | 梯度 | fp16 | 2 Φ 2Φ | | 参数副本 | fp32 | 4 Φ 4Φ | | 一阶矩 𝑚 m | fp32 | 4 Φ 4Φ | | 二阶矩 𝑣 v | fp32 | 4 Φ 4Φ | | 总计 | | 16 Φ 16Φ |

为什么需要 fp32 参数副本？