分布式训练扩展

分布式训练扩展

本节摘要:你的 124M 模型在单卡上训完了。现在试 70 亿参数——模型放不进内存,数据在单机上要跑几周。大规模训练不是可选项,是唯一出路。本节带你吃透让大规模训练成为可能的四大策略:数据并行(每卡一份完整模型,梯度取平均)、张量并行(把单层切到多卡)、流水线并行(按层分阶段,微批次填气泡)、完全分片数据并行(FSDP/ZeRO,每卡只存 1/N 的参数/梯度/优化器状态)。你会算清「70B 模型到底要几张 A100」,理解为什么 Adam 优化器状态是显存杀手,以及为什么 Llama 3 405B 要用 16,384 张 H100。

学习目标

阅读完本节,你应当能够:

1. 解释三种并行(数据、张量、流水线)各自何时必需,依据模型与集群规模判断。
2. 用 PyTorch DDP 实现数据并行训练,跨多卡同步梯度。
3. 为给定模型大小算清显存预算(权重 + 优化器状态 + 梯度 + 激活),定出最低硬件。
4. 配置 FSDP 或 DeepSpeed ZeRO 各阶段,把模型状态分片到多卡,装下超单卡的模型。

一、问题与直觉

一个 7B 参数模型,FP16 下光是权重就 14GB。Adam 优化器为每个参数存两份额外拷贝(一阶矩、二阶矩),又是 28GB。反向传播的梯度再加 14GB。一个激活还没存,你已经 56GB 了。

NVIDIA A100 显存 80GB。56/80 已用,剩 24GB 给激活——前向计算、反向要保留的中间值。2048-token 序列、4096 维模型,单层激活约 64MB;32 层,每样本 2GB;批 8 需 16GB。你只有 24GB,批 12 就爆。

现在试 70B。光权重 FP16 就 140GB,单卡放不下。至少要 2 张 A100(160GB)装权重,加上优化器状态和梯度远不止——至少 3 张,现实是 8~16 张视分片策略而定。Llama 3 405B 在 16,384 张 H100 上训练,估算算力成本约 1 亿美元。DeepSeek V3 训了个相当模型只花约 560 万美元——靠架构巧妙(MoE 每 token 只激活部分参数)和训练效率。

为什么必须分布式

真实模型的显存账,每个数字都是算出来的,不是估的。

「Adam 状态」那列是杀手。Adam 为每个参数存一份运行均值(m)和方差(v),都是 FP32。70B 模型即 70B×4 字节×2 = 560GB——光优化器就要 7 张 A100。单张 H100 80GB,Llama 3 405B 光装权重+优化器+梯度至少要 61 张 H100,加激活更多。Meta 用 16,384 张 GPU 不是想用,是不得不用。

数据并行

最简单的分布式策略。把整个模型复制到 N 卡,每批切成 N 等份,每卡在自己数据片上前向+反向,反向后跨卡平均梯度,每卡用同样的平均梯度更新自己的权重副本,保持同步。好处:吞吐线性扩展,N 卡每步处理 N 倍数据,通信仅限梯度平均(可与计算重叠)。坏处:每卡都装一份完整模型+优化器+梯度,70B 模型每卡要 840GB——数据并行毫不降低单卡显存,只降训练时间。有效批大小 = 每卡批 × N。

张量并行

把单层切到多卡。一次矩阵乘法分给多卡,各算部分结果。8192×8192 的权重矩阵,4 路张量并行每卡持 8192×2048 的片,各乘输入得部分结果,再 all-reduce 或 all-gather 合成完整输出。好处:降低单卡权重显存,70B 分到 8 卡每卡约 8.75B 参数。坏处:每层后都要快速卡间通信(NVLink 节点内 900 GB/s 可行,跨节点 InfiniBand 约 50 GB/s 不行),几乎总限于单节点内(8 卡)。Megatron-LM 开创此法,Llama 3 405B 节点内用 8 路张量并行。

流水线并行

按层切。GPU 1 跑 18 层,GPU 2 跑 916 层,依此类推。数据流过流水线:GPU 1 算完它的层把激活发给 GPU 2。好处:卡间通信极少——只有层边界的激活(比梯度或权重小得多),适合跨节点。坏处:流水线气泡——GPU 4 算微批 1 的前向时,GPU 1、2、3 在闲(已算完自己的部分)。朴素流水线 GPU 利用率只有 1/N。GPipe 和 PipeDream 把批切成微批解决:GPU 1 一算完微批 1 就开始微批 2,跨阶段重叠计算。M 个微批、N 阶段,气泡占比降到 (N-1)/M——16 微批 4 阶段,气泡 3/16 = 18.75%。

FSDP:完全分片数据并行

FSDP 结合数据并行的可扩展性与分片的省显存。每卡不持完整模型,而只持 1/N 的参数、梯度、优化器状态。某层前向前,FSDP 跑 all-gather 收集全卡完整参数到每卡内存,前向后丢弃非本地参数;反向时再 all-gather 重建参数算梯度;反向后 reduce-scatter 把梯度片分给各卡,每卡只存 1/N 梯度。

70B 模型在 8 卡上的账:

无 FSDP 单张 80GB 卡装不下 70B;有 FSDP 8 卡每卡 105GB 仍超——至少要 16 卡才降到 80GB 以下,或配合激活检查点(反向时重算而非存储)。

DeepSpeed ZeRO

DeepSpeed 的 ZeRO 在概念上与 FSDP 相同(Microsoft 独立开发),定义三阶段:ZeRO-1(只分片优化器状态,约 4 倍省显存)、ZeRO-2(+梯度,约 8 倍)、ZeRO-3(+参数,约 N 倍)。ZeRO-3 等价于 FSDP。DeepSpeed 还有 ZeRO-Offload(把优化器状态卸到更便宜更大的 CPU 内存)和 ZeRO-Infinity(卸到 NVMe SSD),用速度换容量。

混合精度训练

现代训练同时用多种浮点格式:前向用 FP16/BF16(16 位,内存减半,tensor core 上矩阵乘快 2 倍);主权重用 FP32(优化器为数值精度维护);损失缩放(反向前把损失乘大常数防 FP16 梯度下溢,优化器步前再除回)。BF16 与 FP32 同指数范围(8 位指数)但精度低(7 位尾数 vs FP32 的 23),几乎不需损失缩放;FP16 指数 5 位尾数 10 位,能表示精细值但极端幅度会溢出/下溢。业界基本转向 BF16。

Megatron-LM 与 3D 并行

真实大规模训练三者并用:跨节点组数据并行(扩批)、节点内张量并行(8 卡切单层)、跨节点流水线并行(切层组)。Llama 3 405B 在 16,384 张 H100 上:节点内 8 路张量并行、跨节点 16 路流水线并行、剩余维度 128 路数据并行(16384/8/16=128)。这种 3D 分解(8×16×128=16384)就是扩展到数千卡的方法。DeepSeek V3 走另一条路:MoE 每 token 只激活 37B(总 671B),2048 张 H800——不到 Meta 的 1/8——花了 560 万美元 vs Meta 估的 1 亿。

二、从零实现

Step 1:模拟数据并行

import numpy as np
def simulate_data_parallelism(data, num_gpus, model_fn):
    shard = len(data) // num_gpus
    gpu_losses, gpu_grads = [], []
    for i in range(num_gpus):
        s = data[i*shard:(i+1)*shard]
        loss, grad = model_fn(s)
        gpu_losses.append(loss); gpu_grads.append(grad)
    return np.mean(gpu_losses), np.mean(gpu_grads, axis=0)   # all-reduce 平均梯度
​

all-reduce(平均梯度)是数据并行唯一的通信,实践用 NVIDIA 的 NCCL 库做环形 all-reduce。

Step 2:模拟张量并行

def simulate_tensor_parallelism(x, W, num_gpus):
    d_in, d_out = W.shape
    shard = d_out // num_gpus
    partials = [x @ W[:, i*shard:(i+1)*shard] for i in range(num_gpus)]
    full = np.concatenate(partials, axis=-1)                  # 列并行:拼接
    return full, np.abs(full - x@W).max()                     # 误差应为 0
​

张量并行数学上精确——产出与单卡算完整矩阵乘相同。列并行(切输出维)用拼接,行并行(切输入维)用求和。Transformer FFN 第一线性(扩张)用列并行,第二(收缩)用行并行,避免两层间 all-reduce。

Step 3:模拟流水线并行

模拟层跨虚拟 GPU,展示气泡问题:4 阶段 1 微批,气泡 75%(4 卡中 3 卡闲);16 微批降到约 19%。消除气泡的代价是内存——要同时存所有在途微批的激活。

Step 4:显存计算器

精确计算任意模型大小的显存需求——回答每个 ML 工程师都问的问题:「我要几张卡?」

def memory_calculator(params_b, precision_bytes=2, optimizer="adam",
                      num_gpus=1, sharding="none"):
    p = params_b * 1e9
    weights = p * precision_bytes
    optim = p * 4 * 2 if optimizer == "adam" else 0           # Adam 的 m 与 v,FP32
    grads = p * precision_bytes
    if sharding in ("fsdp","zero3"):
        weights /= num_gpus; optim /= num_gpus; grads /= num_gpus
    elif sharding == "zero2":
        optim /= num_gpus; grads /= num_gpus
    elif sharding == "zero1":
        optim /= num_gpus
    per_gpu = weights + optim + grads
    return {"per_gpu_gb": per_gpu/1e9, "fits_80gb": per_gpu/1e9 <= 80}
​

喂入模型大小,调分片策略直到单卡总量降到 80GB 以下。

Step 5:混合精度对比

最大意外:混合精度不减半显存。7B 模型 FP32 训练用 112GB,混合精度(BF16+FP32 主权重)用 84GB——降 25% 而非 50%,因为优化器状态(Adam 的 m 和 v)无论精度都留 FP32,而优化器占大头。

三、框架对比

PyTorch FSDP 是原生实现,DeepSpeed 是 Microsoft 的等价框架。配置上一行 --fsdp_strategy full_shard 即可。本节模拟器用纯 Python 让你理解机制——生产中你写的是配置文件,不是分布式原语。

四、可复用产物

本节产出 outputs/prompt-distributed-training-planner.md——一个提示,接收模型大小与可用硬件,产出完整分布式训练计划:并行策略、显存预算、通信开销、预期吞吐。

五、练习

1. (Easy) 给显存计算器加激活检查点:每 K 层只存一次激活(K=1 即全重算),展示省多少显存、慢多少(约 33% 额外计算)。

2. (Medium) 扩展流水线模拟,实现 PipeDream 的 1F1B 调度,4 阶段 8 微批对比朴素调度的气泡占比——1F1B 峰值内存更小因为更早开始反向。

3. (Medium) 实现梯度累积模拟器:不在每个微批后 all-reduce,而是本地累积 K 步再 all-reduce,展示通信减 K 倍但最终梯度相同。

4. (Hard) 构建成本估算器:给定模型大小、目标 token 数、GPU 型号(A100 $2/时、H100 $3.5/时)、并行策略,估美元总成本,对照已知成本(Llama 3 405B 约 1 亿、DeepSeek V3 约 560 万)验证。

5. (Hard) 给显存计算器加 ZeRO-Offload:设 CPU 内存每节点 512GB、NVMe 2TB,展示把优化器状态卸到 CPU 让 70B 在 4 卡而非 16 卡上训练,代价是优化器步慢 30~50%。

本节要点回顾

1. 数据并行复制模型到每卡:每卡处理不同数据片,反向后 all-reduce 平均梯度;线性提速但不降单卡显存。
2. 张量并行切单层:把权重矩阵按列/行分到多卡,各算部分矩阵乘;需要 NVLink 快速卡间通信,几乎限于节点内。
3. 流水线并行按层分:每卡跑一组层,数据流过流水线;气泡问题是主要代价,微批次填气泡。
4. FSDP 每卡只存 1/N:all-gather 重建全参数、reduce-scatter 分发梯度片;70B 从单卡 840GB 降到 8 卡每卡 105GB。
5. ZeRO 三阶段递进分片:ZeRO-1 优化器、ZeRO-2 +梯度、ZeRO-3 +参数(=FSDP)。
6. Adam 状态是显存杀手:每参数两份额外 FP32 拷贝,70B 即 560GB,光优化器就要 7 张 A100。
7. 混合精度只省 25% 非 50%:优化器状态无论精度都留 FP32 且占大头。
8. BF16 优于 FP16:同指数范围无需损失缩放,业界已转向 BF16。
9. 3D 并行缩到数千卡:节点内张量并行 × 跨节点流水线并行 × 跨组数据并行,Llama 3 405B 即 8×16×128=16384。
10. DeepSeek 用 MoE 省一个数量级:每 token 只激活 37B/671B,2048 卡花 560 万 vs Meta 16384 卡花 1 亿。

下一节,我们把这个基座模型变成有用的助手——指令微调(SFT)是「token 预测器」与「能回答问题的助手」之间的桥梁。