""" ============================================================================= CS349E · 前置讲 P5 · 03_allreduce_algorithms.py —— all-reduce 的算法：延迟主导 vs 带宽主导 ============================================================================= 【这个脚本要回答的问题】 all-reduce 是推理里最常见的集合通信（张量并行每层两次）。实现它的算法有好几种，各自的"步数"和"每步字节数"不同： ring：2(p−1) 步，每步 n/p 字节——字节最少（带宽最优），步数最多； 递归倍增：log2 p 步，每步 n 字节——步数最少，字节多； Rabenseifner：2·log2 p 步，字节和 ring 一样——两头兼顾； 二叉树、朴素：负载不均匀；交换机内归约（SHARP 类）：交换机帮你加。 哪个最快，取决于消息大小和设备数。本脚本： 演示 1：8 个 GPU（NVLink 交换机），消息从 1 KB 到 1 GB，六种算法的模拟时间（每个都和 α–β 公式核对），以及每种大小的最优算法。 演示 2：设备数从 8 到 512（IB 网络），小消息和大消息下各算法的时间——ring 的 2(p−1) 步在大规模时变成灾难。 演示 3："算法带宽"和"总线带宽"——NCCL 报告的指标；和公开的 8×H100 实测数字对照，看 α–β 模型乐观了多少。 演示 4：广播的三种算法（朴素、二项树、流水线）——切块流水把大消息的时间压到接近 n/β。 【怎么读输出】全部是模拟器 + 公式，每次运行结果相同。α 是假设值，结论看"相对大小"和"交叉点的数量级"。 【运行】python3 03_allreduce_algorithms.py （约 3 秒） 【依赖】numpy；从 ../lecture01 导入 console；同目录的 comm_sim。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 from comm_sim import (ALLREDUCE_ALGOS, IB_NDR, ICI_V5E, NVLINK4, Comm, broadcast_naive, # noqa: E402 broadcast_pipeline, broadcast_tree, switch, t_formula) SIM_MAX_BYTES = 4e6 # 模拟时真的分配数据的上限；更大的消息只用公式（演示 1 已经验证两者完全一致） def simulate(key, p, nbytes, link): """ 在 p 个设备的交换机上，用真实数据跑一次 all-reduce，返回模拟的时间。 【省内存的技巧】模拟器的时间只取决于消息的字节数，不取决于数值，所以用 float32、元素个数 = 字节数 / 4。 消息超过 SIM_MAX_BYTES 时直接用公式（返回值标记为公式）。 """ if nbytes > SIM_MAX_BYTES: return t_formula(key, p, nbytes, link.alpha, link.bw), "公式" n = max(int(nbytes // 4), p) rng = np.random.default_rng(0) bufs = [rng.standard_normal(n).astype(np.float32) for _ in range(p)] comm = Comm(switch(p, link)) out = ALLREDUCE_ALGOS[key][1](comm, bufs) ref = sum(b.astype(np.float64) for b in bufs) if not np.allclose(out[p // 2], ref, atol=1e-3): raise AssertionError(f"{key} 结果不对") return comm.time, "模拟" # --------------------------------------------------------------------------- # 演示 1：消息大小 # --------------------------------------------------------------------------- def demo1_message_size(p=8): """8 个设备、NVLink 交换机：六种算法 × 消息大小。模拟的时间和公式逐个核对（误差 < 1e-9）。""" print_header(f"演示 1：{p} 个 GPU（NVLink 交换机，α = {NVLINK4.alpha * 1e6:.0f} µs），all-reduce 的时间 vs 消息大小") sizes = [1e3, 16e3, 256e3, 1e6, 4e6, 64e6, 1e9] rows, best, max_dev = [], {}, 0.0 for n in sizes: cells = [fmt_bytes(n)] times = {} for key in ALLREDUCE_ALGOS: t, how = simulate(key, p, n, NVLINK4) if how == "模拟": max_dev = max(max_dev, abs(t - t_formula(key, p, n, NVLINK4.alpha, NVLINK4.bw)) / t) times[key] = t cells.append(fmt_time(t)) fastest = min((k for k in times if k != "innet"), key=times.get) best[n] = fastest cells.append(ALLREDUCE_ALGOS[fastest][0]) rows.append(cells) print_table(["消息大小"] + [ALLREDUCE_ALGOS[k][0] for k in ALLREDUCE_ALGOS] + ["最快（不算交换机内归约）"], rows) print(f"\n≤ {fmt_bytes(SIM_MAX_BYTES)} 的消息用真实数据模拟，和 α–β 公式的最大相对差 {max_dev:.1e}；更大的消息直接用公式。") # 交叉点：递归倍增 vs ring a, b = NVLINK4.alpha, NVLINK4.bw L = np.log2(p) n_cross = (2 * (p - 1) - L) * a / (L - 2 * (p - 1) / p) * b print(f"递归倍增和 ring 的交叉点：log2 p·(α + n/β) = 2(p−1)α + 2(p−1)/p·n/β → n ≈ {fmt_bytes(n_cross)}") small = [n for n in sizes if best[n] == "naive"] print_observations([ f"小消息：步数少的赢。最小的几个消息（≤ {fmt_bytes(max(small)) if small else '—'}）连朴素算法都赢了——它只有 2 步；" f"再大一点（几百 KB 到 1 MB），递归倍增（log2 {p} = {int(L)} 步）最快；ring 要 {2 * (p - 1)} 步，" "每步一个 α，光延迟就是递归倍增的好几倍。", f"大消息（≥ 64 MB）：字节少的赢。ring 和 Rabenseifner 每个设备只发 2(p−1)/p·n ≈ 1.75 n，递归倍增发 {int(L)} n，" "二叉树的根附近要发 2·⌈log2 p⌉ 份完整数据。", "Rabenseifner 在两头都接近最优：步数是 2·log2 p（和树一样少），字节和 ring 一样少。" "NCCL、MPI 都会按消息大小（和拓扑）在几种算法之间切换，没有一个算法在所有大小上都最好。", "交换机内归约（SHARP 一类）在所有大小上都最快：上行时交换机就把数加好了，每个设备只发一次 n、收一次 n。" "理论上把大消息的时间减半；How to Scale Your Model 第 12 章报告实际（NCCL 2.27.5）只快了约 30%。", "朴素算法在大消息时最慢：根的端口要收 p−1 份再发 p−1 份——瓶颈集中在一个设备上。", ]) return best # --------------------------------------------------------------------------- # 演示 2：设备数 # --------------------------------------------------------------------------- def demo2_scale(): """ 设备数 8 … 512，挂在一个无阻塞的 InfiniBand 网络上（每个设备 50 GB/s，α = 5 µs）。只用公式（演示 1 已核对）。 看两种消息：64 KB（像 decode 时的张量并行激活）和 256 MB（像训练时的梯度）。 """ print_header(f"演示 2：设备数变多（{IB_NDR.name}，α = {IB_NDR.alpha * 1e6:.0f} µs）") algos = ["ring", "rd", "rab", "tree", "innet"] rows = [] for n in (64e3, 256e6): for p in (8, 64, 512): cells = [fmt_bytes(n), p] times = {k: t_formula(k, p, n, IB_NDR.alpha, IB_NDR.bw) for k in algos} for k in algos: cells.append(fmt_time(times[k])) ring_lat = 2 * (p - 1) * IB_NDR.alpha cells.append(f"{ring_lat / times['ring'] * 100:.0f}%") rows.append(cells) print_table(["消息", "设备数 p"] + [ALLREDUCE_ALGOS[k][0] for k in algos] + ["ring 时间里 α 的占比"], rows) print_observations([ "64 KB、512 个设备：ring 要 1022 步，光是延迟就超过 5 ms，几乎全部时间都在等；递归倍增只要 9 步。" "设备越多，'步数'越要命——这就是为什么大规模训练的 NCCL 会用树形或分层的算法。", "256 MB 时，ring 和 Rabenseifner 的时间几乎不随 p 变化（约 2n/β）：每个设备发的字节数趋于 2n，与 p 无关。" "How to Scale Your Model 第 3 章说'带宽主导时，集合通信的时间只取决于数组大小和带宽，不取决于设备数'，就是这个意思。", "递归倍增和二叉树在大消息时随 log2 p 变慢：每一步都要传完整的数据。", ]) # --------------------------------------------------------------------------- # 演示 3：算法带宽与总线带宽 # --------------------------------------------------------------------------- def demo3_bus_bandwidth(p=8): """ NCCL 的两个指标： 算法带宽 algbw = n / 时间（"每个设备的数据多快被 all-reduce 完"）； 总线带宽 busbw = algbw × 2(p−1)/p（换算成"每条链路实际承载的带宽"，和链路规格可比）。 对照：How to Scale Your Model 第 12 章给出的 8×H100（NVLink，关闭 SHARP）实测： 约 10 GB 的消息总线带宽约 370 GB/s；58 MB 的消息（Llama-3-70B 一个 MLP 权重的 1/8）只有约 150 GB/s。 """ print_header(f"演示 3：总线带宽——α–β 模型 vs 公开的 8×H100 实测（ring，α = {NVLINK4.alpha * 1e6:.0f} µs）") rows = [] measured = {58e6: 150e9, 10e9: 370e9} for n in (1e6, 16e6, 58e6, 256e6, 10e9): t = t_formula("ring", p, n, NVLINK4.alpha, NVLINK4.bw) algbw = n / t busbw = algbw * 2 * (p - 1) / p meas = measured.get(n) rows.append([fmt_bytes(n), fmt_time(t), f"{algbw / 1e9:.0f} GB/s", f"{busbw / 1e9:.0f} GB/s", f"约 {meas / 1e9:.0f} GB/s" if meas else ""]) print_table(["消息", "模型时间", "算法带宽", "模型的总线带宽", "公开实测的总线带宽"], rows) n = 58e6 t_meas = n * 2 * (p - 1) / p / measured[n] bw_part = 2 * (p - 1) / p * n / NVLINK4.bw alpha_eff = (t_meas - bw_part) / (2 * (p - 1)) print(f"\n要让模型在 58 MB 时也只有 150 GB/s，每一步的'有效 α'要大约 {alpha_eff * 1e6:.0f} µs（我们假设的是 2 µs）。") print_observations([ "α–β 模型在大消息（10 GB）时的总线带宽接近链路的 450 GB/s，和实测的约 370 GB/s 差不太多（实测达不到规格的约 80%）。", f"但在 58 MB 这样的'中等'消息上，模型预测的总线带宽远高于实测的约 150 GB/s。要用 α–β 模型解释实测，" f"每一步的有效延迟得是 {alpha_eff * 1e6:.0f} µs 左右，比链路本身的延迟大一个数量级：真实实现里还有同步、分块、" "协议、调度占用的 SM 等开销。", "教训：模型告诉你'该往哪个方向优化'（少步数还是少字节），但绝对时间要用实测校准。How to Scale Your Model 也提醒：" "要达到标称带宽，需要非常大的消息。", ]) # --------------------------------------------------------------------------- # 演示 4：广播 # --------------------------------------------------------------------------- def demo4_broadcast(p=8, n_bytes=8e6): """ 把 8 MB 从 0 号广播给其他 7 个设备：朴素（根一次发 7 份）、二项树（3 步，每步完整数据）、 流水线链（切成 k 块，沿一条链传：p − 2 + k 步，每步 n/k）。 【为什么用 TPU ICI 的链路参数】流水线的好处要在消息远大于半带宽点 n½ = αβ 时才明显。NVLink 的 n½ 约 900 KB， 要上百 MB 的消息才看得出来（模拟器要真的分配这么多数据）；ICI 的 n½ 只有 45 KB，8 MB 就够了。 【最优块数】T(k) = (p − 2 + k)(α + n/(kβ))，对 k 求导：k* = √((p − 2)·n/(αβ))， T* = (√((p − 2)α) + √(n/β))²——k 越大流水越满，但每块一个 α。 """ link = ICI_V5E print_header(f"演示 4：广播 {fmt_bytes(n_bytes)} 给 {p - 1} 个设备——朴素、二项树、流水线（{link.name}，α = {link.alpha * 1e6:.0f} µs）") rng = np.random.default_rng(1) data = rng.standard_normal(int(n_bytes // 4)).astype(np.float32) bufs = [data] + [np.zeros_like(data) for _ in range(p - 1)] rows = [] c = Comm(switch(p, link)) out = broadcast_naive(c, bufs) rows.append(["朴素（根一次发给所有人）", len(c.records), fmt_time(c.time), "✓" if all(np.array_equal(o, data) for o in out) else "✗"]) c = Comm(switch(p, link)) out = broadcast_tree(c, bufs) t_tree = c.time rows.append(["二项树", len(c.records), fmt_time(c.time), "✓" if all(np.array_equal(o, data) for o in out) else "✗"]) best = None for k in (1, 4, 16, 32, 64, 256): c = Comm(switch(p, link)) out = broadcast_pipeline(c, bufs, 0, k) rows.append([f"流水线链，切成 {k} 块", len(c.records), fmt_time(c.time), "✓" if all(np.array_equal(o, data) for o in out) else "✗"]) if best is None or c.time < best[1]: best = (k, c.time) print_table(["算法", "步数", "时间", "结果正确"], rows) lower = n_bytes / link.bw k_star = np.sqrt((p - 2) * n_bytes / (link.alpha * link.bw)) t_star = (np.sqrt((p - 2) * link.alpha) + np.sqrt(n_bytes / link.bw)) ** 2 print(f"\n下界：每个设备至少要收 n 字节 → n/β = {fmt_time(lower)}；公式的最优块数 k* ≈ {k_star:.0f}，" f"对应时间 {fmt_time(t_star)}（下界的 {t_star / lower:.2f} 倍）。") nv_tree = np.ceil(np.log2(p)) * NVLINK4.time(4e6) nv_pipe = min((p - 2 + k) * NVLINK4.time(4e6 / k) for k in range(1, 257)) print_observations([ "朴素广播只要 1 步，但根的端口要发 7 份完整数据；二项树 3 步、每步 1 份，根只发 3 份。", f"流水线把数据切成小块沿链传，所有链路同时在传不同的块：最好的 k = {best[0]}（{fmt_time(best[1])}），" f"是二项树的 {best[1] / t_tree:.0%}、下界的 {best[1] / lower:.2f} 倍。块太少，流水线'灌满'之前的等待太长；" "块太多，每块一个 α，延迟又占了上风——和公式的 k* 一致。", f"消息不够大时流水线反而吃亏：同样的算法用在 NVLink 上、4 MB 的消息（只有 n½ 的 {4e6 / NVLINK4.half_bw_size():.1f} 倍），" f"最好的流水线要 {fmt_time(nv_pipe)}，二项树只要 {fmt_time(nv_tree)}。什么算法最好，永远要看 n 和 n½ = αβ 的比值。", "'切块 + 流水'是通信里最常用的技巧之一：ring all-reduce、NCCL 的分块协议、流水线并行的 micro-batch，" "都是同一个思想——让所有链路和所有设备同时忙起来。", ]) if __name__ == "__main__": demo1_message_size() demo2_scale() demo3_bus_bandwidth() demo4_broadcast()