Skip to content

ZeRO 与 FP8 训练账:显存拆分的实战手册 ​

70B 模型训练,光权重就是 140GB,加上梯度和优化器状态,单卡 80G 根本放不下。行业答案是两个"省"字诀:ZeRO 把状态拆到各卡,FP8 把精度降下来省显存。本文把这两笔账算清楚,衔接 训练资源账 与 分布式并行。

1. 训练的四份状态(先回顾) ​

状态大小(70B,FP16 训练)说明
权重 W140 GB模型参数
梯度 G140 GB反向计算的梯度
优化器状态 O280-420 GBAdam 的 m/v(2-3 份副本)
激活 A视 batch 而定前向保存,反向重算

结论:权重只占训练显存的 20-30%,优化器状态才是大头——这是 ZeRO 的出发点。

2. ZeRO 三阶段:状态从"每人全量"到"每人一份" ​

阶段切分内容每卡显存(相对全量)通信开销
DDP 基线无100%All-Reduce 梯度
ZeRO-1优化器状态~50%不变
ZeRO-2优化器+梯度~25%不变
ZeRO-3优化器+梯度+权重~1/N(8 卡≈12%)权重 All-Gather

关键判断:

  • ZeRO-1/2 不增加通信量(还是那一次梯度 All-Reduce),白捡显存——默认就该开;
  • ZeRO-3 用额外 All-Gather 换显存,通信增加约 1.5×——显存极限时才用;
  • 每步通信时间可用 通信账 复算,评估 ZeRO-3 的扩展代价。

3. ZeRO 与并行策略的组合 ​

DP(数据并行)── 梯度 All-Reduce,卡越多越好
TP(张量并行)── 超节点内,通信频繁
PP(流水并行)── 切层,通信稀疏
ZeRO(状态分片)── 与 DP 叠加,省显存不增通信(1/2 阶段)

组合经验:小集群(8-32 卡)DP+ZeRO-2 最省心;大模型(70B+)"DP×TP×PP×ZeRO-3"四维组合,先固定 TP/PP(通信拓扑约束),再用 ZeRO-3 榨显存。

4. FP8 混合精度:算力翻倍的另一半账 ​

FP8 训练(Hopper+)的收益与代价:

项目FP16 基线FP8收益
Tensor Core 算力1×~2×吞吐翻倍
权重/梯度显存2B/参数1B/参数显存减半
动态范围大小(指数位少)需 loss scaling

工程要点:

  • FP8 不是"全程 FP8"——通常 FP8 做前向 GEMM,反向与优化器保持 FP16/FP32(混合精度);
  • 需要 loss scaling + 溢出检测(见 混合精度检查点);
  • 收益上限受"算力账"约束:只有算力受限的训练才吃得下 FP8 的翻倍(见 算力账)。

5. 可复算 Python:70B 训练显存预算 ​

python
def zero_mem(n_params_b, world, stage, fp8=False):
    """ZeRO 各阶段每卡显存(GB,忽略激活)"""
    w = n_params_b * 1e9 * (1 if fp8 else 2) / 1e9      # 权重 GB
    g = n_params_b * 1e9 * (1 if fp8 else 2) / 1e9      # 梯度
    o = n_params_b * 1e9 * 4 / 1e9 * 2                   # Adam m+v, FP32
    parts = {"0": w+g+o, "1": w+g+o/world, "2": (w+g)/world+o/world,
             "3": (w+g+o)/world}
    return round(parts[str(stage)], 1)

zero_mem(n_params_b=70, world=64, stage=3, fp8=True)
# ≈ 3.3 GB/卡(纯权重/梯度/优化器;激活另计)→ 64 卡可训 70B

6. 交付清单 ​

  1. 训练默认开 ZeRO-1/2——不增通信白省显存;
  2. 显存仍不够再上 ZeRO-3,并实测通信占比(>30% 考虑改拓扑);
  3. Hopper+ 卡上 FP8 训练:前向 FP8 + 反向/优化器 FP16/32 + loss scaling;
  4. 显存预算用上面的函数先算,再对照 验收基线 实测;
  5. 每个并行维度改动跑一次小规模基准,确认扩展效率不崩。

本文为本站整理思考后撰写,用于知识库学习交流;具体版本号与实测数据以现场为准。

AI-HPC 智算超算专家 · Contact: openaihpc@gmail.com