【LLM技术全景】混合精度与分布式训练：训练大模型的工程奥秘

摘要

• 核心问题：单卡显存与算力无法承载百亿乃至万亿参数模型的训练，训练效率受限于显存墙、计算墙与通信墙。
• 主要贡献：系统性拆解混合精度训练、激活重计算与四类并行策略（数据/张量/流水线/专家并行），并给出工程选型依据。
• 阅读收获：① 理解 FP16 与 BF16 的本质差异与损失缩放机制；② 掌握激活重计算以时间换空间的权衡；③ 理解三维并行（DP×TP×PP）如何协同突破单卡极限。

---

一、引言：为什么训练大模型是工程奇迹

训练一个 175B 参数的 GPT-3，若用 FP32 单卡（按第14篇 KV Cache 的显存估算方法类比），仅模型参数就需要约 700GB 显存，而单张 A100 仅 80GB。更现实地看：

训练 175B 模型的最小显存需求（粗略下界）：
  参数 (FP32)        : 175B × 4B  ≈ 700 GB
  梯度 (FP32)        : 175B × 4B  ≈ 700 GB
  优化器状态 (Adam)   : 175B × 12B ≈ 2.1 TB   ← 动量+方差+FP32副本
  激活值 (per micro-batch): 随序列长度与batch线性增长，可达数百GB
  ─────────────────────────────────────────
  合计保守估计        : > 3.5 TB 显存

结论：单机单卡（80GB）连参数的 1/40 都放不下。
      必须依赖混合精度 + 分布式并行 + 显存优化三件套。

本文即围绕这"三件套"展开，它们分别对应数值精度 、显存占用 、算力扩展三个维度。

---

二、混合精度训练（Mixed Precision Training）

2.1 为什么需要混合精度

现代 GPU（NVIDIA Tensor Core）在 FP16/BF16 上的吞吐量通常是 FP32 的 2--16 倍，且显存占用减半。但纯低精度训练会引入两类问题：

问题	成因	后果
上溢（Overflow）	FP16 最大仅约 65504	梯度/激活变成 Inf
下溢（Underflow）	FP16 最小正规数约 6×10⁻⁵	小梯度变成 0，参数不更新
舍入误差累积	FP16 尾数仅 10 位（相对误差 ~1e-3）	训练后期震荡、不收敛

2.2 FP16 与 BF16 的本质差异

FP32: [1位符号 | 8位指数 | 23位尾数]   范围 ≈ ±3.4e38   精度 ≈ 7位十进制
FP16: [1位符号 | 5位指数 | 10位尾数]   范围 ≈ ±65504     精度 ≈ 3位十进制
BF16: [1位符号 | 8位指数 | 7位尾数]    范围 ≈ ±3.4e38   精度 ≈ 2-3位十进制
                                  ▲
                          BF16与FP32"指数位相同"
                          → 动态范围一致，几乎不会溢出
                          → 精度略低，但训练稳定性显著优于FP16

关键结论：BF16 因动态范围与 FP32 一致，是训练大模型的首选低精度格式，基本无需损失缩放；FP16 则需配合损失缩放（Loss Scaling）防止下溢。

2.3 混合精度的标准范式

现代框架（PyTorch AMP、DeepSpeed、Megatron）采用如下分工：

混合精度训练数据流：

  1. Master Weights (FP32)  ← 唯一权威权重，用于更新
        │  cast
        ▼
  2. 副本 Weights (BF16/FP16) ← 参与前向/反向计算
        │
        ▼
  3. 前向 (BF16):  activation = f(X_bf16, W_bf16)
  4. 反向 (BF16):  grad_act, grad_W_bf16
        │
        ▼  (若FP16需 Loss Scaling)
  5. 梯度反缩放 → cast 回 FP32
        │
        ▼
  6. 优化器用 FP32 梯度更新 Master Weights (FP32)

损失缩放（FP16 专用）：在反向前将损失乘以系数 S（如 1024），使小梯度整体放大到 FP16 可表示区间；反向后再除以 S 还原。BF16 因动态范围足够，通常省略此步。

2.4 实战代码（PyTorch AMP）

import torch
from torch.cuda.amp import autocast, GradScaler

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()  # FP16 需要；BF16 可省略

for x, y in dataloader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad()

    # 前向：自动以 BF16/FP16 计算
    with autocast(dtype=torch.bfloat16):
        logits = model(x)
        loss = criterion(logits, y)

    # 反向：scaler 处理缩放（BF16 时 scaler 等效为恒等）
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

实测收益（以 7B 模型为例）：BF16 相比 FP32，显存约降 40--50%，吞吐提升约 1.6--2.3 倍，收敛终点的验证损失差异通常 < 0.1%。

---

三、激活重计算（Gradient Checkpointing）

3.1 显存瓶颈其实在"激活值"

多数初学者会误以为显存主要被参数占据。实际上在前向传播中，每一层的中间激活（activation）都需要保留，供反向传播计算梯度：

单卡训练单样本显存构成（粗略）：
  参数 + 梯度 + 优化器 : 相对固定（与模型规模相关）
  激活值               : 正比于 (batch_size × seq_len × hidden × num_layers)
                         ↑ 长上下文 + 大 batch 时，激活可占显存的 60% 以上

3.2 原理：以时间换空间

标准做法在前向时保存所有 层的激活；激活重计算只保存少数检查点层 的激活，反向时从最近检查点重新前向一遍得到中间激活，再继续反向。

标准反向：
  前向: L1→L2→L3→L4→L5  (全部激活入显存)
  反向: 直接从显存读激活，逐层算梯度
  显存: ████████████ (高)
  计算: 1× 前向 + 1× 反向

激活重计算：
  前向: L1→L2→L3→L4→L5  (仅保存 L1, L3, L5 检查点)
  反向: L5→L4 时, 从L3重算 L3→L4 的激活
        L4→L3 时, 从L3重算 ...
  显存: ████ (低, 约降 1/(checkpoint间隔))
  计算: 1× 前向 + 1× 反向 + 重算 0.3~1× 前向

权衡：显存降幅可达 60--80%，代价是约 20--40% 的额外计算。在显存受限场景下，这是必选项。

3.3 实战代码

# 方式一：PyTorch 内置
from torch.utils.checkpoint import checkpoint

def forward(self, x):
    # 仅对计算密集层启用重计算
    h = checkpoint(self.block1, x)
    h = checkpoint(self.block2, h)
    return self.head(h)

# 方式二：DeepSpeed / Hugging Face Trainer
# training_args = TrainingArguments(
#     gradient_checkpointing=True,   # 一键开启
#     gradient_checkpointing_kwargs={"use_reentrant": False},
# )

---

四、数据并行（Data Parallelism, DP）

4.1 原理

最直观的扩展方式：每张 GPU 持有完整模型副本 ，输入数据切分到各卡，各自前向/反向得到梯度后 AllReduce 求平均，再各自用相同梯度更新。

        GPU0      GPU1      GPU2      GPU3
        ┌────┐    ┌────┐    ┌────┐    ┌────┐
数据 →  │模型│    │模型│    │模型│    │模型│  (副本相同)
        └─┬──┘    └─┬──┘    └─┬──┘    └─┬──┘
          │ 梯度     │ 梯度     │ 梯度     │ 梯度
          └────────►AllReduce◄──────────┘
                     (求平均)
          └──────── 相同梯度更新各副本 ──────┘

4.2 实现与瓶颈

• DDP（DistributedDataParallel） ：梯度在反向时分桶异步 AllReduce，通信与计算重叠，几乎无额外等待。
• 通信瓶颈：AllReduce 总量 = 模型参数量 × 2 字节（FP16）。175B 模型单步需同步约 350GB 梯度，对网络带宽极其敏感。

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group("nccl")
model = DDP(model.cuda(), device_ids=[local_rank])
# 其余训练循环与单卡一致

适用边界：DP 在 8--64 卡内扩展性极佳，但参数量超过单卡显存时无法使用（模型副本放不下）。

---

五、模型并行：切分模型本身

当模型单副本无法放入单卡，必须切分模型参数。主要有两条路线。

5.1 张量并行（Tensor Parallelism, TP）

由 Megatron-LM 提出，将单层内的权重矩阵切分到多卡，按列的线性层（如 FFN 第一层）按列切、按行的线性层（如 FFN 第二层）按行切，使矩阵乘可在多卡上并行且结果可拼接。

以 FFN 为例 (Y = GeLU(X·A)·B)：

  GPU0              GPU1
  ┌──────┐          ┌──────┐
  │ X·A₀ │  ─────►  │ X·A₁ │    A 按列切 [A₀|A₁]
  │GeLU  │          │GeLU  │    (逐元素，各自独立)
  │ ·B₀  │  ◄─────  │ ·B₁  │    B 按行切
  └──┬───┘          └──┬───┘
     └─── 拼接 Y ───────┘

通信：每层仅 1 次 AllReduce（激活/结果的边界）
特点：通信频繁但量小，适合单机多卡（NVLink 高带宽）

5.2 流水线并行（Pipeline Parallelism, PP）

将不同层分配到不同设备，形成流水线。微批次（micro-batch）切分使各设备尽量满负荷。

按层切分 (GPU0: L1-4 | GPU1: L5-8 | GPU2: L9-12 | GPU3: L13-16)

  时间 ──────────────────────────────────────▶
  GPU0: [mb1][mb2][mb3][mb4]
  GPU1:   [mb1][mb2][mb3][mb4]   (等GPU0)
  GPU2:     [mb1][mb2][mb3][mb4]
  GPU3:       [mb1][mb2][mb3][mb4]
                ▲ 气泡(bubble)：流水线启动/结束的空闲

GPipe：在所有微批前向完再统一反向（显存高）
PipeDream：1F1B 交错调度（降低气泡，但需存多版本激活）

气泡问题：流水线启停存在空闲（bubble），micro-batch 越多气泡占比越小，但显存占用上升。

5.3 ZeRO：优化器状态分片

DeepSpeed 的 ZeRO 介于 DP 与 MP 之间，将优化器状态 / 梯度 / 参数分片到各卡，而非每卡全存：

ZeRO 三阶段：
  Stage 1: 优化器状态分片   → 显存降 ~4×
  Stage 2: + 梯度分片        → 显存降 ~8×
  Stage 3: + 参数分片        → 显存降 ~1/N (N=卡数)，等同模型并行效果
特点：通信量与 DP 相当，却获得接近模型并行的显存收益。

---

六、专家并行（Expert Parallelism, EP）

在第11篇《混合专家模型（MoE）》中我们介绍过，MoE 层包含多个专家 FFN。专家并行将不同专家分配到不同设备，路由（router）把 token 分发到对应专家所在卡。

MoE 层 (8 专家) 切到 4 张卡：
  GPU0: Expert 0,1   GPU1: Expert 2,3
  GPU2: Expert 4,5   GPU3: Expert 6,7

  token ──► Router ──► 分发到所属专家所在卡 ──► 计算 ──► 收集
  通信：token 在卡间 all-to-all 分发/收集（容量因子内）

EP 与 TP/PP 正交，是 MoE 模型（如 Mixtral、DeepSeek-V3）训练的必选项。

---

七、三维并行：工程落地范式

工业界训练超大模型普遍采用 DP × TP × PP × EP 的组合，典型如 Megatron-LM + DeepSpeed：

总 GPU 数 = DP 维度 × TP 维度 × PP 维度

示例：训练 70B 模型用 512 张 A100
  配置：TP=8 (单机内NVLink) × PP=4 × DP=16 = 512
  显存：模型被 TP×PP 切薄，每份仅 ~2-3GB 参数级碎片
  通信：
    TP 层内 AllReduce (高频小量, 走 NVLink)
    PP 层间 P2P (中频)
    DP 跨节点 AllReduce (低频大量, 走 IB/以太网)
  优化：ZeRO 可叠加在 DP 维度进一步省显存

设计原则：

• TP 尽量限制在单机内（依赖高带宽 NVLink）；
• PP 跨机但控制在浅流水线（减少气泡）；
• DP 跨节点扩展，配合 ZeRO 减显存。

---

八、实战：用 DeepSpeed / FSDP 启动训练

8.1 DeepSpeed ZeRO-3 配置（节选）

{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 8,
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "none" },
    "overlap_comm": true
  },
  "gradient_checkpointing": true,
  "tensor_parallel": 1,
  "pipe_parallel": 1
}

# 启动 8 卡训练
deepspeed --num_gpus=8 train.py --deepspeed ds_config.json

8.2 PyTorch 原生 FSDP（更轻量替代）

from torch.distributed.fsdp import FSDP, MixedPrecision
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

model = FSDP(
    model,
    auto_wrap_policy=size_based_auto_wrap_policy,
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
    ),
    use_orig_params=True,
)

---

九、并行策略选型决策树

模型能放进单卡吗？
├── 能（< 单卡显存）
│     └── 多卡只为加速 → 纯 DP（DDP）即可，最简单
└── 不能（> 单卡显存）
      ├── 参数量略超 → ZeRO-2/3（DP+分片），无需切模型
      └── 远超（十亿~万亿级）
            ├── 单机多卡、带宽高 → TP 为主（NVLink）
            ├── 跨节点、层数深 → PP 为辅（控气泡）
            ├── MoE 模型       → 叠加 EP
            └── 综合最优       → DP × TP × PP（×EP）+ ZeRO

显存 vs 通信权衡总表：

策略	显存收益	通信开销	适用场景
DP (DDP)	无（每卡全副本）	中（梯度 AllReduce）	模型可单卡放、多卡加速
ZeRO-3	高（~1/N）	中-高	超大模型、显存紧张
TP	高（按层切）	高（层内频繁）	单机多卡、NVLink
PP	高（按层切）	低（层间P2P）	跨节点、深层模型
EP	中（专家切）	中（all-to-all）	MoE 模型专属

---

十、总结与展望

本文核心要点

1. 混合精度：BF16 是训练首选（动态范围=FP32，几乎不溢出）；
   FP16 需 Loss Scaling；Master Weight 始终 FP32。
2. 激活重计算：以 20~40% 额外计算换 60~80% 显存下降，长序列必开。
3. 数据并行：DPP 通信与计算重叠，扩展性好，但受单卡显存上限约束。
4. 模型并行：TP 切层内矩阵（靠NVLink）、PP 切层间（控气泡）、ZeRO 分片优化器。
5. 专家并行：MoE 专属，与 TP/PP 正交叠加。
6. 三维并行 DP×TP×PP（+EP+ZeRO) 是工业界训练超大模型的标配范式。

技术演进方向

2020: ZeRO (DeepSpeed) 重新定义显存分片
2021: Megatron-LM 张量并行 + GPipe 流水线成熟
2022: 3D 并行成为 530B~万亿模型标准
2023: FlashAttention 降低激活显存 + 序列并行（SP）补齐 TP 短板
2024: 通信计算重叠（overlap）算法进一步压缩气泡
2025: 异构训练（CPU/NPU 卸载）、专家并行与 MoE 深度耦合
2026: 国产 NPU（Ascend）适配 + 算法-硬件协同成为新焦点

下期预告（7月16日，周四）：《归一化与激活函数：LLM 中的关键设计选择》------ 为什么 LLM 几乎不用 BatchNorm？LayerNorm、RMSNorm 的差异何在？ReLU、GELU、SwiGLU 各自适合什么场景？我们将从梯度流与训练稳定性角度给出工程答案。

---

参考资料

1. Micikevicius et al. (2018) --- Mixed Precision Training（混合精度训练奠基论文）
2. Rajbhandari et al. (2020) --- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models（DeepSpeed）
3. Shoeybi et al. (2020) --- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
4. Huang et al. (2019) --- GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
5. Narayanan et al. (2021) --- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM（3D 并行）
6. Chen et al. (2016) --- Training Deep Nets with Sublinear Memory Cost（Gradient Checkpointing）
7. Korthikanti et al. (2023) --- Reducing Activation Recomputation in Large Transformer Models（序列并行）

延伸讨论

思考题：

1. 若你有 8 张 A100（单机），要训练一个单副本放不下的 30B 模型，你会优先选 ZeRO-3 还是 TP×PP？为什么？
2. BF16 精度比 FP16 低，为何训练稳定性反而更好？这与"损失曲面"有何关系？

实践作业：

• 用 PyTorch torch.cuda.amp 或 DeepSpeed ZeRO-2 复现一次 7B 模型的混合精度训练，对比 FP32 的显存与吞吐。
• 尝试开启/关闭 gradient_checkpointing，测量单步显存与耗时的变化，验证文中"20~40% 时间换 60~80% 显存"的经验值。