""" ============================================================================= 第 9 讲 · 代码 4：流水线并行的"气泡"，以及 3D 并行怎样组合 ============================================================================= 对应讲义：lectures/lecture09_distributed_training.md 第 5、6 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 流水线并行把模型的层分给不同的卡：第 1 张卡算前几层，把激活传给第 2 张卡……问题是后面的卡要等前面的卡， 前面的卡要等后面的卡把梯度传回来——中间有大量的空闲时间（"气泡"）。 演示 1：画出时间线：不切 micro-batch、GPipe（先做完所有前向再做反向）、1F1B（一个前向一个反向交替）。 演示 2：气泡占多少时间：模拟结果 vs 公式 (p − 1)/(m + p − 1)；micro-batch 越多，气泡越小。 演示 3：同时要存几个 micro-batch 的激活：GPipe 存 m 个，1F1B 最多存 p 个。 演示 4：在第 1 讲的 base model 上真的切成 2 个阶段、4 个 micro-batch：梯度和不切时一样； 阶段之间传的数据有多大。 演示 5：BLOOM 的 3D 并行（张量 4 × 流水线 12 × 数据 8 = 384 张卡）：每张卡的模型状态、气泡、 数据并行的通信量；换几种组合比一比。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 04_pipeline_and_3d.py # 约 2 秒；需要 PyTorch 依赖本目录的 dist_sim.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py （第 1 讲的 base model 只读加载，先 copy.deepcopy 再算梯度）。 【Python 小知识】 - itertools.product(A, B, C)：A、B、C 的所有组合（笛卡尔积），等价于三重 for 循环。 """ import copy import itertools import os import sys import torch import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import pad, print_header # noqa: E402 from dist_sim import fmt_bytes, render_timeline, simulate_pipeline # noqa: E402 from fm_world import ALPHABET, build_world # noqa: E402 from tiny_gpt import CharTokenizer, get_batch, load_or_train_base_model # noqa: E402 def demo1_timelines(): print_header("演示 1：4 个阶段的流水线，时间线（每格一个时间单位；F = 前向，b = 反向，·· = 空闲）") print("每个 micro-batch 在一个阶段上前向 1 个单位、反向 2 个单位（反向约是前向的 2 倍）\n") cases = [("(a) 不切：整个 batch 一次过（相当于 1 个 micro-batch，前向 4、反向 8）", 4, 1, "naive", 4, 8), ("(b) GPipe，切成 4 个 micro-batch", 4, 4, "gpipe", 1, 2), ("(c) GPipe，切成 8 个 micro-batch", 4, 8, "gpipe", 1, 2), ("(d) 1F1B，切成 8 个 micro-batch", 4, 8, "1f1b", 1, 2)] for title, p, m, kind, tf, tb in cases: r = simulate_pipeline(p, m, kind, tf, tb) print(f"{title}：总时间 {r['makespan']:.0f}，气泡占 {r['bubble']:.1%}") print(render_timeline(r)) print() print("观察 Observations：") print(" 1. (a) 不切的时候，任何时刻只有一个阶段在干活，其余 3 个都闲着：气泡占 75%——4 张卡只干了 1 张卡的活。") print(" 2. (b)(c) GPipe 把 batch 切成 micro-batch，前一个 micro-batch 交给下一个阶段后，马上开始算下一个。") print(" 气泡只出现在开头（后面的阶段等第一个 micro-batch 到达）和中间（前向做完、等第一个反向传回来）。") print(" 切成 8 份时气泡从 75% 降到 27%。") print(" 3. (d) 1F1B 的总时间和 GPipe 完全一样——它没有减少气泡，减少的是**同时要存的激活**（演示 3）。") def demo2_bubble(): print_header("演示 2：气泡占总时间的比例（GPipe；1F1B 相同）") ms = [1, 2, 4, 8, 16, 32, 64] print(f"{pad('阶段数 p', 10)}" + "".join(pad(f"m = {m}", 10, ">") for m in ms)) for p in (4, 8, 16): sim = [simulate_pipeline(p, m, "gpipe")["bubble"] for m in ms] formula = [(p - 1) / (m + p - 1) for m in ms] assert all(abs(a - b) < 1e-12 for a, b in zip(sim, formula)) print(f"{pad(p, 10)}" + "".join(pad(f"{b:.1%}", 10, ">") for b in sim)) print("（模拟结果与公式 (p − 1)/(m + p − 1) 逐个相等）") print("\n观察 Observations：") print(" 1. 气泡比例 = (p − 1)/(m + p − 1)：开头有 p − 1 个时间片后面的阶段在等，结尾有 p − 1 个前面的阶段在等。") print(" 要让气泡小，micro-batch 的个数 m 要远大于阶段数 p：p = 8 时 m = 64 才把气泡压到 10% 以下。") print(" Narayanan et al. (2021) 写成'气泡时间 ÷ 理想时间 = (p − 1)/m'，是同一件事的另一种写法。") print(" 2. 但 m 不能无限大：一个 batch 的总大小是有限的（太大会超过临界 batch），切得越碎，每个 micro-batch 越小，") print(" 每个阶段的矩阵乘法越'瘦'，硬件利用率越低（第 8 讲 §3.3）。流水线并行总是在气泡和效率之间折中。") def demo3_activation_memory(): print_header("演示 3：每个阶段同时存着几个 micro-batch 的激活（p = 4，m = 16）") g = simulate_pipeline(4, 16, "gpipe") f = simulate_pipeline(4, 16, "1f1b") print(f"{pad('阶段', 8)}{pad('GPipe', 10, '>')}{pad('1F1B', 10, '>')}") for s in range(4): print(f"{pad(s, 8)}{pad(g['peak_act'][s], 10, '>')}{pad(f['peak_act'][s], 10, '>')}") print(f"总时间：GPipe {g['makespan']:.0f}，1F1B {f['makespan']:.0f}") print("\n观察 Observations：") print(" 1. GPipe 先做完全部 16 个前向，每个阶段都要同时存 16 个 micro-batch 的激活。") print(" 2. 1F1B 一旦进入'一前一后'的稳定阶段，每做完一个反向就能扔掉一份激活：第 s 个阶段最多存 p − s 份") print(" （第 0 个阶段 4 份，最后一个阶段 1 份），而且和 m 无关。总时间完全相同。") print(" 这就是为什么 Megatron、DeepSpeed 默认用 1F1B（PipeDream-Flush）而不是 GPipe 的原始调度。") def demo4_real_pipeline(base, data): print_header("演示 4：把 base model 切成 2 个阶段，batch 32 切成 4 个 micro-batch") model = copy.deepcopy(base) model.train() gen = torch.Generator().manual_seed(0) x, y = get_batch(data, 32, 96, gen) def stage0(idx): # 阶段 0：嵌入 + 第 1 个 Block T = idx.shape[1] return model.blocks[0](model.tok_emb(idx) + model.pos_emb(torch.arange(T))) def stage1(h): # 阶段 1：第 2 个 Block + 最后的 LayerNorm + 输出层 return model.head(model.ln_f(model.blocks[1](h))) # 不切：整个 batch 一次前向、一次反向 model.zero_grad() F.cross_entropy(stage1(stage0(x)).view(-1, 44), y.view(-1)).backward() ref = {n: p.grad.clone() for n, p in model.named_parameters()} # 流水线：4 个 micro-batch，梯度累加（每个 micro-batch 的损失除以 4，加起来就是整个 batch 的平均） model.zero_grad() m, sent = 4, 0 for i in range(m): h = stage0(x[i * 8:(i + 1) * 8]) sent += h.numel() * h.element_size() # 阶段 0 → 阶段 1：传这个 micro-batch 的激活 loss = F.cross_entropy(stage1(h).view(-1, 44), y[i * 8:(i + 1) * 8].reshape(-1)) / m loss.backward() # 反向时阶段 1 → 阶段 0：传同样大小的梯度 diff = max((p.grad - ref[n]).abs().max().item() for n, p in model.named_parameters()) n_params = sum(p.numel() for p in model.parameters()) print(f"累加 4 个 micro-batch 的梯度 vs 整个 batch 的梯度：最大差 {diff:.1e}") print(f"阶段之间每一步传的数据：前向 {sent:,} 字节（4 × 8 × 96 × 64 个 fp32 数），反向同样大小") print(f"对照：数据并行每一步要 all-reduce 的梯度 = 参数个数 {n_params:,} × 4 字节 = {n_params * 4:,} 字节") print("\n观察 Observations：") print(" 1. 流水线并行 + 梯度累加在数学上和不切完全一样（差 1e-8 量级，来自求和顺序）——它也不改变算法。") print(" 2. 阶段之间只传'切口'处的激活（和它的梯度），是点对点通信，不是 all-reduce。在这个玩具模型上") print(" 它比参数还大，但对大模型，一个切口的激活（token 数 × 宽度）远小于整个模型的参数——") print(" 这正是流水线并行适合放在**带宽较低的机器之间**的原因。") def demo5_bloom_3d(): print_header("演示 5：BLOOM 176B 的 3D 并行：384 张 A100 80GB = 张量 t × 流水线 p × 数据 d") psi, layers, gbs, mbs = 176e9, 72, 2048, 2 # 72 = 70 层 + 两端的嵌入层也按 Block 计（训练日志的说法） configs = [(t, p, 384 // (t * p)) for t, p in itertools.product((1, 2, 4, 8), (6, 12, 24, 48)) if 384 % (t * p) == 0 and layers % p == 0] print(f"Ψ = 176B，全局 batch {gbs} 个序列，micro-batch {mbs}；数据并行上用 ZeRO-1（只切优化器状态）") print(f"{pad('t × p × d', 14)}{pad('每张卡的参数', 14, '>')}{pad('模型状态/卡', 14, '>')}{pad('m', 6, '>')}" f"{pad('气泡', 8, '>')}{pad('DP 每卡每步发', 16, '>')} 备注") rows = [] for t, p, d in configs: per_gpu_params = psi / (t * p) states = (2 + 2 + 12 / d) * per_gpu_params # bf16 参数、梯度各 2 字节；fp32 主参数 + Adam 12 字节切 d 份 m = gbs // (d * mbs) bubble = (p - 1) / (m + p - 1) dp_bytes = 2 * (d - 1) / d * 2 * per_gpu_params notes = [] if t > 8: notes.append("张量并行跨机器") if states > 80e9: notes.append("放不下") if (t, p, d) == (4, 12, 8): notes.append("← BLOOM 实际用的") rows.append((t, p, d)) print(f"{pad(f'{t} × {p} × {d}', 14)}{pad(f'{per_gpu_params / 1e9:.2f}B', 14, '>')}{pad(fmt_bytes(states), 14, '>')}" f"{pad(m, 6, '>')}{pad(f'{bubble:.1%}', 8, '>')}{pad(fmt_bytes(dp_bytes), 16, '>')} {'，'.join(notes)}") print("\n观察 Observations：") print(" 1. 张量并行 t 和流水线 p 一起决定'每张卡分到多少参数'；数据并行 d 决定'同时训练几份'。") print(" t × p 太小（例如 1 × 6），每张卡要放 29B 参数，模型状态 123 GB，放不下。") print(" 2. BLOOM 的 4 × 12 × 8：每张卡 3.67B 参数、模型状态约 20 GB，给激活留出了大半张卡；") print(" 每条流水线 128 个 micro-batch，气泡约 8%；张量并行 4 份留在一台机器（8 张卡）里。") print(" 3. 换成 1 × 24 × 16（不用张量并行、靠更长的流水线）：每张卡 7.33B 参数、35 GB 模型状态，气泡升到 26%——") print(" 流水线越长、每条流水线分到的 micro-batch 越少，气泡越大。换成 8 × 6 × 8：气泡只有 4%，但张量并行用满了一台机器的 8 张卡，") print(" 每层的通信更多、每张卡的矩阵更小（第 8 讲：矩阵越小，利用率越低）。") print(" 4. 没有'最优'的公式答案：要在显存、气泡、通信、矩阵大小之间折中，最后靠实测吞吐来选。") print(" BLOOM 的训练日志也是这样做的：试了几种 micro-batch 大小，选了实测最快的。") if __name__ == "__main__": world, tok = build_world(), CharTokenizer(ALPHABET) base, _ = load_or_train_base_model(world, tok) data = torch.tensor(tok.encode(world.train_text()), dtype=torch.long) demo1_timelines() demo2_bubble() demo3_activation_memory() demo4_real_pipeline(base, data) demo5_bloom_3d()