CookLLM文档系统工程分布式训练
数据并行
会员专享理解通信原语和 DDP 的梯度同步机制
在权益中心获取代码训练大模型的第一步，是让多块 GPU 一起干活。数据并行（Data Parallelism）是最直觉的方式：每块 GPU 持有完整的模型副本，各自处理不同的数据，最后把梯度汇总。本章我们先理解单卡的内存瓶颈，再学习多卡通信的基础，最后深入 DDP 的实现。
单卡训练的内存组成
一个 7B 参数的模型，fp16 权重只需要 14 GB（7B × 2 bytes），但在 80GB 的 A100 上却训练不了。内存都去哪了？
答案是：训练时 GPU 上不只存参数，还要存梯度和优化器状态。
混合精度训练的内存需求
现代训练普遍采用混合精度（Mixed Precision）：前向和反向用 fp16 计算（快），但参数更新用 fp32（精确）。以 Adam 优化器为例，假设模型有 个参数：
fp16 部分（前向/反向）：
- 参数： bytes
- 梯度： bytes
fp32 部分（优化器）：
- 参数副本： bytes（用于精确更新）
- 一阶矩 ： bytes（梯度的指数移动平均）
- 二阶矩 ： bytes（梯度平方的移动平均）
总计： bytes
| 组件 | 精度 | 内存（bytes） |
|---|---|---|
| 参数 | fp16 | |
| 梯度 | fp16 | |
| 参数副本 | fp32 | |
| 一阶矩 | fp32 | |
| 二阶矩 | fp32 | |
| 总计 |
为什么需要 fp32 参数副本？
训练时前向和反向都用 fp16 计算（速度快、省显存），但 fp16 精度只有约 3 位有效数字。当学习率很小时，参数 += 学习率 × 梯度 这个更新量可能小到被 fp16 直接舍入为零，模型就学不动了。所以 Adam 在内部维护一份 fp32 精度的参数副本，用 fp32 做更新，再把结果转回 fp16 给下一轮前向使用。
为什么需要 和 ？
Adam 不是简单地"沿着梯度走一步"。它需要维护两个跨步骤的历史统计量：（梯度的指数移动平均，相当于动量）和 （梯度平方的移动平均，用来自适应调节每个参数的学习率）。梯度是"这一步往哪走"， 和 是"过去所有步的经验"。
7B 模型的实际内存：
- 参数：
14 GB - 梯度：
14 GB - 优化器状态：
84 GB（参数副本 28GB + 28GB + 28GB） - 总计：约 112 GB（还没算激活值）
一张 A100 只有 80 GB 显存，连模型的训练状态都放不下。
Training Memory per GPU
Mixed Precision + AdamTotal = 16Φ bytes
12Φ
2Φ
2Φ
Optimizer States(fp32 params + m + v)
Gradients(fp16)
Parameters(fp16)
Example: 7B Model (Φ = 7×10⁹)
Optimizer States84 GB
Gradients14 GB
Parameters14 GB
Total112 GB
A100: 80 GB
登录以继续阅读
这是一篇付费内容，请登录您的账户以访问完整内容。