并行训练

一个 7B 参数的模型，权重用 bf16 存只要 14 GB，一张 80 GB 的 H100 绰绰有余。可一旦开始用 AdamW 训练，每个参数要背上约 16 字节的状态，7B × 16 = 112 GB，激活还没算，一张卡已经装不下了。就算装得下，按资源核算与 Scaling Laws 的 6ND 估算，用一张卡训完几万亿 token 要几十年。并行训练要回答的就是：模型、状态和计算怎么切到很多张卡上，切完以后卡和卡之间要传什么、传多少、能不能藏在计算后面。

前置：GPU 与 Kernel（显存层级与带宽）、Transformer 架构（一层由哪些矩阵组成）

一张卡上装了什么：混合精度与 AdamW

先把单卡的账算清楚，因为后面所有切法都是在切这本账。现代训练用混合精度：前向和反向的矩阵乘用 bf16（GPU 与 Kernel 里讲过，Tensor Core 的低精度路径更快、字节更少），但参数的”正本”用 fp32 保存。为什么不全用 bf16？看一个最小的例子：权重是 1.0，每步减去 1e-3，走 1000 步。

import torch
w = torch.tensor(1.0, dtype=torch.bfloat16)
for _ in range(1000): w = w + torch.tensor(-1e-3, dtype=torch.bfloat16)
master = torch.tensor(1.0, dtype=torch.float32)
for _ in range(1000): master -= 1e-3
print("bf16 weight after 1000 steps of -1e-3:", w.item())
print("fp32 master after 1000 steps:        ", round(master.item(), 4))
print("bf16 spacing near 1.0:", torch.finfo(torch.bfloat16).eps)

bf16 weight after 1000 steps of -1e-3: 1.0
fp32 master after 1000 steps:         0.0
bf16 spacing near 1.0: 0.0078125

bf16 只有 7 位尾数，1.0 附近相邻两个数相差 0.0078，减 0.001 每次都被舍入回 1.0，一千步白走。学习率乘梯度通常就是这么小的数（位布局见信息的表示）。所以要有一份 fp32 的 master weights 接收更新，每步再转一份 bf16 给前向用。同样的道理，归约（sum、softmax、LayerNorm）和 loss 在 fp32 里算；用 fp16 而非 bf16 时，因为指数范围小，还要做 loss scaling 防止梯度下溢，bf16 的指数位和 fp32 一样多，通常不需要。

再看优化器。AdamW 给每个参数维护两个状态：一阶矩 m（梯度的指数滑动平均）和二阶矩 v（梯度平方的滑动平均），更新时用 m / (sqrt(v) + eps) 当步长方向，相当于每个参数自己调学习率（Adam 的来历见训练深度网络与 CNN）。AdamW 的 W 指解耦的权重衰减：不是把 λθ 加进梯度（那样它会被 sqrt(v) 除掉，梯度大的参数反而衰减少），而是直接把参数乘以 1 − lr·λ。CS336 作业里的实现核心只有几行：

m.mul_(beta1).add_(grad, alpha=1 - beta1)
v.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
# 两个偏差修正合并进步长：代数上等价，少一个临时张量
step_size = lr * math.sqrt(1 - beta2**t) / (1 - beta1**t)
p.addcdiv_(m, v.sqrt().add_(eps), value=-step_size)
p.mul_(1 - lr * weight_decay)          # 解耦衰减：直接缩参数，不经过 v

偏差修正是因为 m、v 从 0 开始，前几步的滑动平均偏小，除以 1 − β^t 把它拉回来。拿它和 torch.optim.AdamW 对比，在同一个二次函数上跑 100 步：

max |ours - torch| after 100 steps: 0.0008425712585449219
state tensors per param: ['m', 'v']

两者接近但不完全相同，差异来自衰减的顺序：PyTorch 先衰减再做 Adam 更新，这份实现先更新再衰减，每步差一个 lr²·λ 量级的项。两种写法都常见，读代码时要留意。

现在可以把每个参数的字节数列出来：bf16 参数 2 + bf16 梯度 2 + fp32 master 4 + fp32 m 4 + fp32 v 4 = 16 字节。其中后三项合称优化器状态，占了 12 字节，是最大的一块。有的配置把 m、v 存成 bf16，总数变成 12 字节，读别人的数字前要先确认用的是哪种设定。激活另算，它随 batch 和序列长度增长，怎么估算见资源核算与 Scaling Laws。

停一下，检查理解单选

混合精度训练里，参数已经有一份 bf16 副本参与前向和反向，为什么还要再存一份 fp32 的 master weights？

bf16 的矩阵乘结果不准，需要用 fp32 重新算一遍前向每步的更新量 lr × 梯度 往往小于 bf16 在参数附近的间距，直接加到 bf16 上会被舍入掉fp32 副本是为了 checkpoint 时文件更小Adam 的二阶矩只能和 fp32 参数相乘

训练循环与 checkpoint

有了模型和优化器，一个训练步是固定的几行：

for it in range(start_iter, max_iters):
    lr = get_lr_cosine_schedule(it, max_lr, min_lr, warmup_iters, cosine_iters)
    for g in opt.param_groups: g["lr"] = lr
    x, y = get_batch(data, batch_size, context_len, device)
    loss = cross_entropy(model(x).reshape(-1, vocab), y.reshape(-1))
    opt.zero_grad(set_to_none=True)
    loss.backward()
    gradient_clipping(model.parameters(), max_norm=1.0)   # 防止偶发的大梯度把参数打飞
    opt.step()
    if it % ckpt_every == 0:
        save_checkpoint(model, opt, it, path)

（示意，函数名沿用 CS336 Assignment 1。）学习率先线性 warmup：训练刚开始时 Adam 的 v 估计还很不准，步子太大容易发散；之后按余弦从 max_lr 降到 min_lr。梯度裁剪把所有参数梯度拼成一个向量，L2 范数超过阈值就整体等比缩小，方向不变。batch 太大放不下时用梯度累积：连续做 k 次前向反向、梯度自动累加，再做一次 step，效果等价于 k 倍的 batch（loss 要除以 k）。

checkpoint 必须保存的不只是模型权重：还有优化器状态（m、v、步数 t）和当前迭代号，最好还有随机数状态和数据读取位置。只存权重的话，恢复后 m、v 从 0 开始，偏差修正从 t = 1 重来，相当于在训练中途重新做了一次 warmup，loss 曲线会出现一个台阶。大规模训练里这不是边角问题：Llama 3 405B 的 54 天预训练中断了几百次，其中约 78% 来自确认或疑似的硬件故障，checkpoint 和自动恢复是基础设施的一部分。

激活太大时还有一个单卡就能用的手段：activation checkpointing，前向只保存每层的输入，反向时把这一层重新算一遍得到中间激活。多花约一次前向的计算，换来激活显存从”每层所有中间量”降到”每层一个输入”。它和下面所有的切分方法都可以叠加。

卡和卡怎么说话：集合通信

多卡之间的数据流动不用一条条点对点消息描述，而用几个集合通信原语，一次说清所有设备之间的通信模式。设备叫 rank，总数叫 world size。用 PyTorch 的 gloo 后端在 CPU 上起 4 个进程，看最常用的三个：

import os, torch, torch.distributed as dist, torch.multiprocessing as mp
def run(rank, W):
    os.environ["MASTER_ADDR"], os.environ["MASTER_PORT"] = "127.0.0.1", "29512"
    dist.init_process_group("gloo", rank=rank, world_size=W)   # GPU 上用 "nccl"
    x = torch.arange(W, dtype=torch.float32) + rank
    a = x.clone(); dist.all_reduce(a)                          # 默认求和，原地修改
    part = torch.empty(1); dist.reduce_scatter_tensor(part, x)
    full = torch.empty(W); dist.all_gather_into_tensor(full, part)
    print(f"rank{rank} in={x.tolist()} all_reduce={a.tolist()} "
          f"reduce_scatter={part.tolist()} all_gather={full.tolist()}", flush=True)
    dist.destroy_process_group()
if __name__ == "__main__":
    mp.spawn(run, args=(4,), nprocs=4)

rank1 in=[1.0, 2.0, 3.0, 4.0] all_reduce=[6.0, 10.0, 14.0, 18.0] reduce_scatter=[10.0] all_gather=[6.0, 10.0, 14.0, 18.0]
rank0 in=[0.0, 1.0, 2.0, 3.0] all_reduce=[6.0, 10.0, 14.0, 18.0] reduce_scatter=[6.0] all_gather=[6.0, 10.0, 14.0, 18.0]
rank2 in=[2.0, 3.0, 4.0, 5.0] all_reduce=[6.0, 10.0, 14.0, 18.0] reduce_scatter=[14.0] all_gather=[6.0, 10.0, 14.0, 18.0]
rank3 in=[3.0, 4.0, 5.0, 6.0] all_reduce=[6.0, 10.0, 14.0, 18.0] reduce_scatter=[18.0] all_gather=[6.0, 10.0, 14.0, 18.0]

（还有两行 deprecation 警告，已省略；打印顺序乱是因为进程异步。）all-reduce 让每个 rank 都拿到全体之和；reduce-scatter 求和但每个 rank 只拿自己那一段（rank i 拿第 i 个位置的和）；all-gather 把每个 rank 的一段拼起来发给所有人。最后一列说明了贯穿后文的等式：

all-reduce = reduce-scatter + all-gather

另外还有 broadcast（一个 rank 发给所有人，比如加载初始 checkpoint）和 all-to-all（rank j 收到每个 rank 的第 j 段，相当于转置，MoE 路由用它）。

成本怎么算？NCCL 在带宽受限时常用 ring 算法：W 张卡连成环，数据切成 W 块，先走 W−1 步 reduce-scatter，再走 W−1 步 all-gather，每步每张卡发送 size / W。所以 all-reduce 每张卡发送 2 × size × (W−1)/W，W 大时接近 2 倍数据量，与卡数基本无关；reduce-scatter 或 all-gather 单独做是一半。小消息另当别论，几 KB 到几 MB 的消息主要被每步的启动延迟限制，所以 DDP 会把小梯度打包成 bucket（PyTorch 默认 25 MB）再发。

传多快取决于走哪条线。以 H100 SXM5 节点为例：节点内 8 张卡经 NVLink 4 / NVSwitch 互连，每卡 900 GB/s（双向合计，单向 450 GB/s），这已经比本卡 HBM 的 3.35 TB/s 慢好几倍；节点之间走 InfiniBand，NDR 每卡一块 400 Gb/s 的网卡约 50 GB/s，又慢一个数量级。B200 的 NVLink 5 是每卡 1.8 TB/s。NVLink 和 InfiniBand 都支持 RDMA，一张卡直接读写另一张卡的显存，不经过 CPU。节点内 vs 节点间的带宽落差，决定了后面每种并行该放在哪一层。

停一下，检查理解单选

8 张卡做 ring all-reduce，张量大小 800 MB。每张卡一共要发送多少数据？

800 MB100 MB1400 MB6400 MB

数据并行与 ZeRO：按 batch 切，再把状态切碎

最简单的并行是数据并行（DDP）：每张卡一份完整模型，batch 切成 W 份，每张卡算自己那份的梯度，然后对梯度做一次 all-reduce 取平均，大家用同样的梯度做同样的更新，参数始终一致。为什么平均各分片的梯度就等于全 batch 的梯度？因为 loss 是样本 loss 的平均，梯度是线性的：

import torch
torch.manual_seed(0)
W1 = torch.randn(64, 32, dtype=torch.float64, requires_grad=True)
X, Y = torch.randn(128, 64, dtype=torch.float64), torch.randn(128, 32, dtype=torch.float64)
def loss_fn(x, y): return ((torch.tanh(x @ W1) - y) ** 2).mean()
g_full, = torch.autograd.grad(loss_fn(X, Y), W1)
shards = [torch.autograd.grad(loss_fn(x, y), W1)[0] for x, y in zip(X.chunk(4), Y.chunk(4))]
print("DDP avg grad == full-batch grad:", torch.allclose(torch.stack(shards).mean(0), g_full))

DDP avg grad == full-batch grad: True

（前提是每个分片一样大；LayerNorm 这类逐样本的归一化不跨样本，所以不影响这个等式。）通信量是每步一次梯度 all-reduce，约 2 × 梯度字节，与 batch 无关。7B 模型、bf16 梯度 14 GB、8 卡时每卡发送约 24.5 GB，NVLink 单向 450 GB/s 下约 54 ms，跨节点 50 GB/s 下约 0.5 s。真正的 DDP 不等反向结束才通信：最后一层的梯度最先算完，一个 bucket 满了就异步发起 all-reduce，和前面层的反向计算重叠。

DDP 的问题在内存：每张卡都存全部 16 字节/参数，卡再多也不能训更大的模型。ZeRO 的想法是：既然 all-reduce = reduce-scatter + all-gather，那就别让每张卡都持有全部状态。

ZeRO-1 切优化器状态。反向后对梯度做 reduce-scatter，每张卡只拿到自己负责那 1/W 参数的梯度之和，只更新这一片，再 all-gather 更新后的参数。通信总量还是 2 倍参数量，和 DDP 一样，所以应该总是开。
ZeRO-2 再切梯度。反向时一层的梯度一算完就 reduce 给负责的卡，然后释放，通信量不变。
ZeRO-3（PyTorch 叫 FSDP）连参数也切。前向每到一层先 all-gather 这层参数，算完释放；反向再 all-gather 一次，算完对梯度 reduce-scatter。通信变成 3 倍参数量（DDP 的 1.5 倍），靠”算第 k 层时预取第 k+1 层参数”藏在计算后面。

按 ZeRO 论文的例子（7.5B 参数、64 卡，优化器状态每参数 K = 12 字节）算每卡显存：

baseline  120.0 GB/GPU
ZeRO-1     31.4 GB/GPU
ZeRO-2     16.6 GB/GPU
ZeRO-3      1.9 GB/GPU

只有 ZeRO-3 随卡数线性下降。但它切不掉激活：每张卡仍要对自己那份数据跑完整的前向，激活一点没少。另外，数据并行的 batch 也不能无限加：每卡至少要分到一些样本，而全局 batch 超过 critical batch size 后，再加样本对优化的帮助不如多走一步。

停一下，检查理解单选

FSDP（ZeRO-3）和流水线并行都把参数分到了多张卡上。两者最本质的区别是什么？

FSDP 每张卡只跑一部分层，流水线每张卡跑完整模型FSDP 每张卡都跑完整模型、卡间传的是参数；流水线每张卡只跑自己那几层、卡间传的是激活两者都传激活，只是 FSDP 用 all-gather，流水线用 send / recvFSDP 没有通信，流水线有 bubble

张量并行与序列并行：按宽度切

当一层的矩阵本身太大，或者想让激活也随卡数下降，就要切进层的内部。张量并行（TP）利用的是矩阵乘可以拆成小矩阵乘。Megatron-LM 对 MLP 的切法：第一个矩阵 A 按列切，第二个矩阵 B 按行切。

D, F_ = 16, 64
x = torch.randn(8, D, dtype=torch.float64)
A, B = torch.randn(D, F_, dtype=torch.float64), torch.randn(F_, D, dtype=torch.float64)
ref = torch.nn.functional.gelu(x @ A) @ B
parts = [torch.nn.functional.gelu(x @ a) @ b               # 每张卡：自己那几列 A、几行 B
         for a, b in zip(A.chunk(4, dim=1), B.chunk(4, dim=0))]
print("sum of 4 partials == full MLP:", torch.allclose(sum(parts), ref))
print("partial shape", tuple(parts[0].shape))

sum of 4 partials == full MLP: True
partial shape (8, 16)

按列切 A 时，每张卡得到隐藏层的一部分列，GeLU 是逐元素的，可以直接在分片上做，中间不需要任何通信；按行切 B 后，每张卡得到一个完整形状的部分和，最后一次 all-reduce 加起来。attention 同理：按头切（每卡算 H/W 个头），输出投影按行切。所以每个 Transformer 层前向 2 次 all-reduce、反向 2 次，每次的大小是激活 b × s × h。这个通信在关键路径上，下一层要等它，很难重叠，量又大，所以 TP 只在节点内的 NVLink 域里用，通常 ≤ 8。注意 LayerNorm、softmax 这种要看完整隐藏维的运算不能在列分片上做，这正是 Megatron 把它们放在切分边界之外的原因。

TP 切掉了权重和大部分激活，但留下了一块：Korthikanti 等（2022）算出每层激活约为 sbh(34 + 5as/h) 字节（s 序列长、b micro-batch、h 隐藏维、a 头数），开 t 路 TP 后变成 sbh(10 + 24/t + 5as/(ht))。那个 10 来自 LayerNorm、dropout 和残差输入，它们在每张卡上都是完整的一份。这些都是沿序列逐位置的运算，于是序列并行（SP）沿序列轴把它们切开，进入 TP 区域前 all-gather，离开时 reduce-scatter，代替原来的 all-reduce，通信量不变。再对 attention 的平方项做选择性重算（就是 FlashAttention 做的事），激活降到 34sbh / t，完全随 t 线性下降。

停一下，检查理解多选

在按列切的张量并行里，每张卡手上是隐藏维的一个分片。哪些运算可以直接在分片上做、不需要先通信？

流水线并行：按层切

流水线并行（PP）把层切成 p 段（stage），每张卡一段，激活从前一段发给后一段，反向时梯度反着走。通信是点对点的 send / recv，每个边界只传 b × s × h 的激活，比 TP 每层 4 次 all-reduce 少得多，所以 PP 能放在慢的节点间链路上。

代价是 bubble。朴素地一次喂一整个 batch，任意时刻只有一张卡在算，p 张卡每张只有 1/p 的时间在工作。解法是把 batch 切成 m 个 micro-batch，第 1 段算完第一个就发出去、接着算第二个。最朴素的调度（GPipe：全部前向、再全部反向）下，bubble 占总时间的比例是 (p − 1) / (m + p − 1)：

bubble (p-1)/(m+p-1): {1: 0.75, 4: 0.429, 8: 0.273, 32: 0.086}

p = 4 时，m = 32 才能把空闲压到 9% 以下。（也有文献用”bubble 时间 / 有用计算时间” (p − 1)/m，分母不同，m 远大于 p 时两者接近。）所以 PP 吃 batch：它和数据并行争的是同一份全局 batch。

GPipe 还要同时存下全部 m 个 micro-batch 的激活。1F1B 调度在热身阶段做 p 个前向后，每做一个前向就紧跟一个反向，反向完就释放那份激活，bubble 不变，但每段最多同时存约 p 份激活，m 可以放心加大。interleaved 1F1B 让每张卡负责 v 个不相邻的小段，bubble 再除以 v，代价是通信多 v 倍。zero bubble 调度则利用反向的一个不对称：反向里”把梯度传给上一段”（B）有人在等，“算本层权重梯度”（W）没人等，先做 B、把 W 推迟到空闲时段，几乎能填满流水线。

停一下，检查理解单选

p = 4 个 stage、m = 32 个 micro-batch，从 GPipe 换成 1F1B，主要改变了什么？

bubble 从约 9% 降到接近 0bubble 不变，但每个 stage 同时要存的激活从 32 份降到约 4 份通信量减半，因为前向和反向交替进行不再需要切 micro-batch

长序列与 MoE：上下文并行和专家并行

剩下两个维度对应两种特殊的”宽”。上下文并行（CP，又叫 ring attention）把超长序列按 token 切到多张卡上，每张卡持有自己那段的 Q，K、V 的块沿环在卡间传递，每张卡用传来的 KV 块和自己的 Q 做分块 attention，用的正是 FlashAttention 里那个可合并的 online softmax 状态。它主要用在长上下文扩展阶段。

专家并行（EP）用于 MoE（路由机制见注意力变体与 MoE）：不切矩阵，而是把整个 expert 放到不同卡上，router 决定每个 token 去哪个 expert，用 all-to-all 把 token 发过去、算完再 all-to-all 发回来。和 TP 相比，EP 保留了每个 expert 完整的大矩阵，GPU 利用率更好；难点是 all-to-all 的目的地事先不知道，各卡负载可能不均，而且计算在等 token 到达，必须低延迟。DeepSeek 为此专门写了 DeepEP 通信库。attention 层没有 expert，用不了 EP，所以 Megatron 让 attention 和 MLP 用不同的并行配置：attention 走 TP/CP/DP，MoE 层走 EP 和专家内的 TP/DP。

把它们组合起来

每种方法都有短板，没有哪一种全面胜出：

方法	切什么	每步主要通信	能和计算重叠吗	放在哪层网络
DDP / ZeRO-1/2	batch	梯度，约 2 倍参数量	能，藏在反向里	任意
FSDP / ZeRO-3	batch，参数也分片	参数 + 梯度，约 3 倍参数量	能，靠预取	节点间也可
TP（+ SP）	每层的宽度	每层 4 次激活大小的集合通信	很难，在关键路径上	节点内 NVLink
PP	层	每边界每 micro-batch 一次激活	部分，靠调度	节点间
CP	序列	KV 块沿环传递	部分	节点内或相邻
EP	expert	每个 MoE 层两次 all-to-all	部分	节点内为主

经验规则可以压成三条。第一，先让模型放进显存：TP（和 EP）最多开到一台机器的卡数，通常 8；还放不下就跨节点加 PP，或者用 FSDP，看节点间带宽。第二，放下之后，剩下的卡全给数据并行（配合 ZeRO-1 切优化器状态）。第三，每卡分到的 batch 太小、通信盖不住时，用梯度累积换更大的 batch。每多开一个维度，就多一类集合通信、多一次同步、多一个出故障的地方，所以维度能少开就少开。

Llama 3 405B（dense 模型，所以没有 EP）的实际配置印证了这些规则：

阶段	GPU 数	TP	CP	PP	DP	序列长度	MFU
初期	8,192	8	1	16	64	8,192	43%
主预训练	16,384	8	1	16	128	8,192	41%
长上下文	16,384	8	16	16	8	131,072	38%

GPU 数 = TP × CP × PP × DP。维度从内到外按带宽需求排：TP 在最里层（同一台服务器），DP（实际是 FSDP）在最外层，因为它能异步预取参数、异步归约梯度，最能容忍延迟。长上下文阶段把 CP 加到 16、DP 降到 8，卡数不变。

一条贯穿全篇的取舍：同一份数据可以重算、可以存在本卡、也可以存在别的卡上用时再传。activation checkpointing 是前两者之间的取舍，ZeRO 和各种模型并行把第三种也加了进来。硬件会变快，模型总会变得更大，这三者之间的平衡点会一直移动。

停一下，检查理解多选

用 1024 张 GPU（每节点 8 张）训练一个 dense 模型，模型要 TP 8 × PP 4 才放得下。哪些做法合理？

TP 设为 8，放在一个节点的 NVLink 域内剩下的维度给数据并行：DP = 1024 / (8 × 4) = 32为了少用流水线，把 TP 加到 32，跨 4 个节点做张量并行如果每卡分到的 batch 太小，用梯度累积换更大的 batch

记住这三件事

混合精度 + AdamW 下每个参数约 16 字节，其中优化器状态占 12 字节；ZeRO 靠 all-reduce = reduce-scatter + all-gather，在几乎不加通信的前提下把这些状态切到各卡。
数据并行传梯度、可以藏在反向后面；张量并行每层都传激活、只能放在 NVLink 节点内；流水线只在相邻段传激活、能跨慢链路，但要用大 batch 填 bubble。
组合的顺序是：先用 TP（≤ 8）和 PP 让模型放进显存，再把剩下的卡全给数据并行；checkpoint 要连优化器状态和迭代号一起存。

CS 与 LLM 学习知识库

探索

并行训练

一张卡上装了什么：混合精度与 AdamW

训练循环与 checkpoint

卡和卡怎么说话：集合通信

数据并行与 ZeRO：按 batch 切，再把状态切碎

张量并行与序列并行：按宽度切

流水线并行：按层切

长序列与 MoE：上下文并行和专家并行

把它们组合起来

记住这三件事

延伸阅读

目录

反向链接