""" ============================================================================= 第 9 讲 · 代码 2：数据并行、ring all-reduce 与 ZeRO ============================================================================= 对应讲义：lectures/lecture09_distributed_training.md 第 3 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 数据并行是最简单的并行：每张卡一份完整的模型，各算一部分数据，最后把梯度加起来。 "加起来"怎么做最快？卡多了以后还划算吗？ZeRO 把模型状态切开，结果还和原来一样吗？ 演示 1：ring all-reduce 一步一步地走：4 张卡，每一步谁把哪一块发给谁。 演示 2：通信量与时间：环形 vs 参数服务器；卡数增加时，通信时间和计算时间各怎样变化。 演示 3：在第 1 讲的 base model 上做"4 张卡的数据并行"：和单卡用同样的大 batch 训练比较—— 数学上是同一件事。 演示 4：ZeRO-1 / 2 / 3：reduce-scatter 梯度 → 每张卡只更新自己那一块 → all-gather 参数。 参数和普通数据并行逐位相同；每张卡存的东西少了，通信量是多少。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 02_data_parallel.py # 约 10 秒；需要 numpy、PyTorch 依赖本目录的 dist_sim.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py （只用第 1 讲的数据和模型结构，从随机初始化开始训练；不读也不写第 1 讲的缓存）。 【PyTorch 小知识】 - torch.nn.utils.parameters_to_vector(model.parameters())：把所有参数拼成一个一维向量； vector_to_parameters(向量, model.parameters())：反过来，把向量写回各个参数。 我们把"整个模型"当成一个长向量来切块、通信、更新——真实的 ZeRO / FSDP 也是这样做的（"扁平化"）。 """ import os import sys import numpy as np import torch import torch.nn.functional as F from torch.nn.utils import parameters_to_vector, vector_to_parameters 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 (CommLog, all_reduce_time, parameter_server_all_reduce, ring_all_gather, # noqa: E402 ring_all_reduce, ring_reduce_scatter) from fm_world import ALPHABET, build_world # noqa: E402 from tiny_gpt import BASE_CONFIG, CharTokenizer, get_batch, make_model # noqa: E402 # ============================================================================= # 演示 1：ring all-reduce 一步一步 # ============================================================================= def demo1_ring_trace(): print_header("演示 1：ring all-reduce，4 张卡，每张卡一个长度 8 的梯度（切成 4 块，每块 2 个数）") p = 4 arrays = [np.arange(8, dtype=float) * (r + 1) for r in range(p)] # 第 r 张卡的"梯度"是 (r+1)·[0..7] for r, a in enumerate(arrays): print(f" 卡{r} 的梯度：{a.astype(int).tolist()}") log = CommLog(p) out = ring_all_reduce(arrays, log) print() for line in log.trace: print(" " + line) print(f"\n结束后每张卡上的结果：{out[0].astype(int).tolist()}（= 1+2+3+4 = 10 倍的 [0..7]）" f"，4 张卡完全相同：{all(np.array_equal(o, out[0]) for o in out)}") print(f"每张卡发出的数：{log.sent}，= 2(p − 1)/p × 8 = {2 * (p - 1) / p * 8:.0f}；一共 {log.steps} 步 = 2(p − 1)") print("\n观察 Observations：") print(" 1. 前 3 步（reduce-scatter）：每张卡把一块发给环上的下一张卡，下一张卡把它加到自己的同一块上。") print(" 3 步之后，每张卡恰好有一块是 4 张卡之和。后 3 步（all-gather）：把这些完整的块传一圈。") print(" 2. 每一步，每条链路上都只传 1/p 的数据，而且所有链路同时在传——没有哪张卡是瓶颈。") # ============================================================================= # 演示 2：通信量与时间 # ============================================================================= def demo2_scaling(): print_header("演示 2：卡数增加时，每张卡要发多少数据、要花多久") n = 1200 print(f"(a) 实际模拟：每张卡一个长度 {n} 的数组") print(f"{pad('卡数 p', 8, '>')}{pad('环形：每张卡发', 16, '>')}{pad('步数', 6, '>')}" f"{pad('参数服务器：卡 0 收 + 发', 26, '>')}{pad('其他卡发', 10, '>')}") rng = np.random.default_rng(0) for p in [2, 4, 8, 16]: arrays = [rng.normal(size=n) for _ in range(p)] lr, lp = CommLog(p), CommLog(p) ring_all_reduce(arrays, lr) parameter_server_all_reduce(arrays, lp) print(f"{pad(p, 8, '>')}{pad(max(lr.sent), 16, '>')}{pad(lr.steps, 6, '>')}" f"{pad(lp.sent[0] + lp.received[0], 26, '>')}{pad(lp.sent[1], 10, '>')}") print("\n(b) α–β 模型：一个 1.3B 参数的模型（16 位梯度 2.6 GB），每张卡 100 Gbit/s（12.5 GB/s）、延迟 5 µs；") print(" 每一步的 batch 固定为 100 万个 token；每张卡按 A100 峰值的 50% 计算（156 TFLOP/s）") grad_bytes, bw, lat = 2.6e9, 12.5e9, 5e-6 flops_per_step = 6 * 1.3e9 * 1e6 print(f"{pad('卡数 p', 8, '>')}{pad('计算', 12, '>')}{pad('环形 all-reduce', 18, '>')}{pad('参数服务器', 14, '>')}" f"{pad('不重叠时的效率', 16, '>')}{pad('完美重叠时', 12, '>')}") for p in [8, 32, 128, 512, 2048]: comp = flops_per_step / (p * 156e12) ring = all_reduce_time(grad_bytes, p, bw, lat) ps = all_reduce_time(grad_bytes, p, bw, lat, "ps") eff_serial = comp / (comp + ring) eff_overlap = comp / max(comp, ring) print(f"{pad(p, 8, '>')}{pad(f'{comp:.3f} 秒', 12, '>')}{pad(f'{ring:.3f} 秒', 18, '>')}{pad(f'{ps:.1f} 秒', 14, '>')}" f"{pad(f'{eff_serial:.0%}', 16, '>')}{pad(f'{eff_overlap:.0%}', 12, '>')}") print("\n观察 Observations：") print(" 1. 环形算法每张卡发的数据 = 2(p − 1)/p × 数组大小：p 从 2 到 16，只从 1 倍涨到 1.9 倍，上限是 2 倍。") print(" 参数服务器的第 0 张卡要收发 2(p − 1) 份，随卡数线性增长——它是瓶颈。") print(" 2. 但是环形 all-reduce 的时间也不会随卡数**减少**：梯度有多大，就要传多少（约 0.4 秒）。") print(" 而 batch 固定时，每张卡的计算时间随卡数减少。8 张卡时通信只占 6%；512 张卡时计算只剩 0.1 秒，") print(" 不重叠的话效率只有 19%。把通信和反向传播重叠起来（一层的梯度算完就开始传）能撑得更久，但终究会被通信追上。") print(" 3. 所以数据并行的扩展受两个东西限制：通信（梯度的大小不随卡数变小）和 batch 大小") print(" （卡多了要么每张卡分到的数据太少，要么总 batch 太大、超过临界 batch——第 4 讲 §8.2）。") # ============================================================================= # 演示 3：数据并行 = 大 batch # ============================================================================= def flat_grad(model): """把模型所有参数的梯度拼成一个 float64 的 numpy 向量。""" return parameters_to_vector([p.grad for p in model.parameters()]).detach().double().numpy() def worker_grad(model, x, y): """在一份数据 (x, y) 上算交叉熵（在这份数据上取平均）并反向，返回扁平的梯度向量和损失。""" model.zero_grad() loss = F.cross_entropy(model(x).view(-1, 44), y.view(-1)) loss.backward() return flat_grad(model), loss.item() def set_flat_grad(model, g): """把扁平的梯度向量写回各个参数的 .grad。""" vector_to_parameters(torch.tensor(g, dtype=torch.float32), [p.grad for p in model.parameters()]) def demo3_dp_equals_big_batch(data): print_header("演示 3：4 张卡的数据并行（每张卡 8 个序列）vs 单卡 batch 32——同样的 40 步") p, steps = 4, 40 single = make_model(BASE_CONFIG, seed=0) dp = make_model(BASE_CONFIG, seed=0) opt_s = torch.optim.AdamW(single.parameters(), lr=3e-3, weight_decay=0.1, betas=(0.9, 0.95)) opt_d = torch.optim.AdamW(dp.parameters(), lr=3e-3, weight_decay=0.1, betas=(0.9, 0.95)) gen = torch.Generator().manual_seed(1) rows = [] for step in range(steps): x, y = get_batch(data, 32, 96, gen) # (a) 在**同一组参数**（单卡模型当前的参数）上，比较"整个 batch 的梯度"和"4 份梯度 all-reduce 后的平均" g_full, loss_s = worker_grad(single, x, y) g_parts = [worker_grad(single, x[r * 8:(r + 1) * 8], y[r * 8:(r + 1) * 8])[0] for r in range(p)] same_param_diff = np.abs(ring_all_reduce(g_parts)[0] / p - g_full).max() set_flat_grad(single, g_full) # 单卡用整个 batch 的梯度更新 opt_s.step() # (b) 另一个模型真的按数据并行训练：它的参数会和单卡模型慢慢分开 grads, losses = [], [] for r in range(p): g, l = worker_grad(dp, x[r * 8:(r + 1) * 8], y[r * 8:(r + 1) * 8]) grads.append(g) losses.append(l) set_flat_grad(dp, ring_all_reduce(grads)[0] / p) opt_d.step() param_diff = (parameters_to_vector(single.parameters()) - parameters_to_vector(dp.parameters())).abs().max().item() if step % 10 == 0 or step == steps - 1: rows.append((step, same_param_diff, loss_s, float(np.mean(losses)), param_diff)) print(f"{pad('步', 6, '>')}{pad('同一组参数上的梯度差', 22, '>')}{pad('单卡损失', 12, '>')}{pad('4 卡平均损失', 14, '>')}" f"{pad('两条轨迹的参数差', 18, '>')}") for step, gd, a, b, pdiff in rows: print(f"{pad(step, 6, '>')}{pad(f'{gd:.1e}', 22, '>')}{pad(f'{a:.6f}', 12, '>')}{pad(f'{b:.6f}', 14, '>')}{pad(f'{pdiff:.1e}', 18, '>')}") print("\n观察 Observations：") print(" 1. 在同一组参数上，4 张卡梯度的平均和单卡在整个 batch 上的梯度只差 1e-7 量级——float32 求和顺序不同带来的舍入。") print(" 数据并行在数学上就是'用一个更大的 batch'：它不改变算法，只改变谁来算。") print(" 2. 但两条真实的训练轨迹并不相同：第 0 步更新之后参数就差 3.7e-4，40 步后差 3.1e-3。") print(" 原因是 Adam 按 √v 归一化：梯度很小的那些分量，1e-7 的舍入误差相对它们自己并不小，") print(" 归一化之后就变成了 1e-4 量级的更新差别，再被后面的训练继续放大。损失曲线仍然几乎重合（差 < 1e-3）。") print(" 这在大模型训练里是真实存在的问题：换一种并行方式、换一个卡数，训练就不再逐位可复现。") print(" PaLM 特意把整个训练做成'从任何检查点重跑都逐位相同'，第 10 讲会讲为什么这很有用。") print(" 3. 反过来，数据并行不能让一个大 batch 学得比小 batch 更快。卡多了要用更大的 batch，") print(" 超过临界 batch 之后，再加卡也换不来同样多的进展（第 4 讲 §8.2）。") # ============================================================================= # 演示 4：ZeRO # ============================================================================= class FlatAdamW: """ 手写的 AdamW，作用在一个一维 numpy 向量（或它的一段）上。 【思路】和 torch.optim.AdamW 的公式相同：先做权重衰减 p ← p − lr·wd·p，再 m ← β1 m + (1 − β1) g；v ← β2 v + (1 − β2) g²；p ← p − lr · m̂ / (√v̂ + ε)。 写成"只管一段"的形式，ZeRO 就能让每张卡只为自己那一段存 m、v——这就是"切优化器状态"。 【参数】n：这一段有多少个参数；lr、wd、betas、eps：超参数 """ def __init__(self, n, lr=3e-3, wd=0.1, betas=(0.9, 0.95), eps=1e-8): self.m, self.v = np.zeros(n), np.zeros(n) self.lr, self.wd, self.b1, self.b2, self.eps = lr, wd, betas[0], betas[1], eps self.t = 0 def step(self, p, g): """原地更新 p（一段参数）；g 是这一段的梯度。""" self.t += 1 p -= self.lr * self.wd * p self.m = self.b1 * self.m + (1 - self.b1) * g self.v = self.b2 * self.v + (1 - self.b2) * g * g mh = self.m / (1 - self.b1 ** self.t) vh = self.v / (1 - self.b2 ** self.t) p -= self.lr * mh / (np.sqrt(vh) + self.eps) def run_dp_or_zero(data, stage, p=4, steps=15): """ 用 p 张"卡"训练 steps 步，stage = 0 是普通数据并行，1/2/3 是 ZeRO 的三个阶段。 【思路】所有卡共用一个 PyTorch 模型对象来算梯度（因为它们的参数本来就相同）； "每张卡存了什么"用 numpy 数组显式地表示： stage 0：每张卡一份完整参数、完整梯度、完整的 m、v；梯度用 ring all-reduce 求和。 stage ≥ 1：把扁平的参数向量切成 p 段（np.array_split），第 r 张卡只为第 r' 段存 m、v； 梯度用 ring reduce-scatter，第 r 张卡只拿到自己那一段的梯度和， 更新自己那一段参数，再 ring all-gather 把新参数传给所有卡。 stage 3 额外：参数平时也只存一段，前向、反向之前各 all-gather 一次（只影响通信量的记账）。 stage 1 和 2 的计算完全相同，区别只在"梯度算完后是否只保留自己那一段"（显存记账不同）。 【返回】dict(params=最终的扁平参数, comm=每张卡发出的元素个数, per_card=每张卡常驻的元素个数) """ model = make_model(BASE_CONFIG, seed=0) params = parameters_to_vector(model.parameters()).detach().double().numpy().copy() n = params.size bounds = np.array_split(np.arange(n), p) # 每段的下标 gen = torch.Generator().manual_seed(1) log = CommLog(p) if stage == 0: opts = [FlatAdamW(n)] else: # ring_reduce_scatter 结束时第 r 张卡拥有第 (r + 1) % p 段，所以第 r 张卡负责那一段的优化器状态 owner = [(r + 1) % p for r in range(p)] opts = {r: FlatAdamW(bounds[owner[r]].size) for r in range(p)} for _ in range(steps): vector_to_parameters(torch.tensor(params, dtype=torch.float32), model.parameters()) if stage == 3: # 前向、反向前各 all-gather 一次参数（只记账） for _ in range(2): ring_all_gather([np.array_split(params, p) for _ in range(p)], owner, log) x, y = get_batch(data, 32, 96, gen) grads = [worker_grad(model, x[r * 8:(r + 1) * 8], y[r * 8:(r + 1) * 8])[0] for r in range(p)] if stage == 0: g = ring_all_reduce(grads, log)[0] / p opts[0].step(params, g) else: chunks, own = ring_reduce_scatter(grads, log) new_shard = {} # 段号 → 更新后的那一段参数 for r in range(p): seg = own[r] shard = params[bounds[seg]].copy() opts[r].step(shard, chunks[r][seg] / p) # 第 r 张卡只更新自己那一段 new_shard[seg] = shard if stage in (1, 2): # 把新参数传给所有卡：第 r 张卡一开始只有自己那一段 start = [[new_shard[seg] if seg == own[r] else np.zeros_like(new_shard[seg]) for seg in range(p)] for r in range(p)] params = ring_all_gather(start, own, log)[0] else: # stage 3：参数保持切开；这里只是把各段拼回来，方便下一步算梯度 params = np.concatenate([new_shard[seg] for seg in range(p)]) per_card = {0: 4 * n, 1: 2 * n + 2 * n / p, 2: n + 3 * n / p, 3: 4 * n / p}[stage] # 参数 + 梯度 + m + v（fp32 记账） return dict(params=params, comm=max(log.sent), per_card=per_card, n=n) def demo4_zero(data): print_header("演示 4：ZeRO-1 / 2 / 3 与普通数据并行（4 张卡，15 步，同样的数据）") base = run_dp_or_zero(data, 0) n = base["n"] names = {0: "普通数据并行", 1: "ZeRO-1（切优化器状态）", 2: "ZeRO-2（再切梯度）", 3: "ZeRO-3（再切参数）"} print(f"参数个数 Ψ = {n:,}；每张卡常驻的数 = 参数 + 梯度 + Adam 的 m、v（fp32 记账，普通数据并行 4Ψ）") print(f"{pad('做法', 26)}{pad('参数与普通 DP 的最大差', 24, '>')}{pad('每张卡常驻的数', 18, '>')}{pad('每张卡每步发出', 16, '>')}{pad('÷ Ψ', 8, '>')}") for stage in range(4): r = base if stage == 0 else run_dp_or_zero(data, stage) diff = np.abs(r["params"] - base["params"]).max() comm_per = r["comm"] / 15 per_card = f"{r['per_card']:,.0f}" # 先算成字符串，避免 f-string 里嵌套同种引号 print(f"{pad(names[stage], 26)}{pad(f'{diff:.1e}', 24, '>')}{pad(per_card, 18, '>')}" f"{pad(f'{comm_per:,.0f}', 16, '>')}{pad(f'{comm_per / n:.2f}', 8, '>')}") print("\n观察 Observations：") print(" 1. ZeRO 的三个阶段和普通数据并行的参数**逐位相同**（差 0.0）：all-reduce 本来就是 reduce-scatter + all-gather，") print(" ZeRO 只是在两半之间插进了'每张卡只更新自己那一段'——AdamW 是逐元素的，分段更新和整体更新完全一样。") print(" 2. 每张卡常驻的东西：普通 DP 4Ψ → ZeRO-1 2.5Ψ → ZeRO-2 1.75Ψ → ZeRO-3 Ψ（4 张卡；卡越多省得越多）。") print(" 3. 通信量：普通 DP 和 ZeRO-1/2 每张卡每步都发 2(p − 1)/p · Ψ = 1.5Ψ；ZeRO-3 要在前向、反向前各 all-gather") print(" 一次参数，再 reduce-scatter 梯度，共 3(p − 1)/p · Ψ = 2.25Ψ——正好是 1.5 倍（ZeRO 论文 §7）。") if __name__ == "__main__": world, tok = build_world(), CharTokenizer(ALPHABET) data = torch.tensor(tok.encode(world.train_text()), dtype=torch.long) demo1_ring_trace() demo2_scaling() demo3_dp_equals_big_batch(data) demo4_zero(data)