""" ============================================================================= CS349E · 第 4 讲 · 01_sharded_matmuls.py —— 切分记号与"切分矩阵乘法的四种情形" ============================================================================= 【这个脚本要回答的问题】 How to Scale Your Model 第 3 章把所有"切开的矩阵乘法"归成四种情形，每种情形需要什么通信一目了然。 这里在 8 块虚拟 H100（NVLink 交换机）上用真实的 numpy 数组把四种情形各做一遍，结果和单卡比较，并给通信记账。 演示 1：切分记号：A[I_X, J] 是什么意思；同一个矩阵的几种切法，每块设备各存多少。 演示 2：四种情形：谁需要 all-gather、谁需要 all-reduce / reduce-scatter、谁根本不需要通信； 每块设备的运算量和通信时间（用 byte_scale 把小数组"放大"成 Llama-3.1-70B 的 MLP 尺寸）。 演示 3：网络 roofline：把被求和的维度切开（情形 3）什么时候值得——通信时间 vs 计算时间。 演示 4：三种集合通信的代价：all-gather ≈ reduce-scatter，all-reduce ≈ 两者之和；小消息被延迟主导。 【怎么读输出】全部是确定性的计算（固定种子），每次运行结果相同。通信时间来自 comm_sim 的 α–β 模拟 （NVLink 4：每卡单向 450 GB/s，每条消息 α = 2 µs，α 是量级上的假设值，见前置讲 P5）。 【运行】python3 01_sharded_matmuls.py （< 1 秒） 【依赖】numpy；同目录的 shard_lab（它再导入 ../prereq5/comm_sim 等）；../lecture01 的 console、hw_specs。 """ 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_num, fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 from shard_lab import (NVLINK4, all_gather, all_reduce, make_comm, matmul_case1, matmul_case2, # noqa: E402 matmul_case3, matmul_case4, reduce_scatter, shard, t_formula, unshard) N_DEV = 8 # --------------------------------------------------------------------------- # 演示 1：切分记号 # --------------------------------------------------------------------------- def demo1_notation(): print_header("演示 1：切分记号——同一个矩阵的几种切法（8 块设备，BF16）") D, F, B = 8192, 28672, 64 rows = [] for spec, shape_fn, note in ( ("W[D, F]", lambda: (D, F), "完全复制：每块设备一整份"), ("W[D, F_X]", lambda: (D, F // N_DEV), "按列切：每块设备 F/8 列"), ("W[D_X, F]", lambda: (D // N_DEV, F), "按行切：每块设备 D/8 行"), ("A[B, D]", lambda: (B, D), "激活复制"), ("A[B_X, D]", lambda: (B // N_DEV, D), "按批切（数据并行的样子）"), ): s = shape_fn() rows.append([spec, f"{s[0]:,} × {s[1]:,}", fmt_bytes(s[0] * s[1] * 2), note]) print(f"Llama-3.1-70B 的 MLP 上投影 W[D = {D}, F = {F}]，以及一批 B = {B} 个 token 的激活 A[B, D]：") print_table(["记号", "每块设备上的形状", "每块设备的字节", "含义"], rows) print_observations([ "下标 X 表示'这一维被切到网格轴 X 上'：每块设备只有 1/|X| 那一段。没有下标的维度是完整复制的。", "这个记号同时说明了两件事：每块设备存多少（显存），以及做矩阵乘法时需要什么通信（演示 2）。", "一个 70B 模型的一个 MLP 矩阵就有 470 MB；完全复制到 8 块卡上是 8 份，切开就只有 1 份。" "推理时切权重的首要原因往往是**放得下**，其次才是**更快**（更多卡的 HBM 带宽一起读权重）。", ]) # --------------------------------------------------------------------------- # 演示 2：四种情形 # --------------------------------------------------------------------------- def demo2_four_cases(): print_header("演示 2：切分矩阵乘法的四种情形（8 块虚拟 H100，NVLink 交换机）") s = 32 # 实际数组的每一维是"模拟尺寸"的 1/s I, J, K = 512 // s, 8192 // s, 28672 // s byte_scale = s * s * 2 / 8 # 数组每维小 s 倍 → 字节小 s² 倍；float64（8 字节）→ BF16（2 字节） rng = np.random.default_rng(0) A = rng.standard_normal((I, J)) Bm = rng.standard_normal((J, K)) C = A @ Bm flops_full = 2 * (I * s) * (J * s) * (K * s) peak = H100.peak_flops("bf16") rows = [] def run(label, fn, per_dev_frac, rule): comm = make_comm(N_DEV, NVLINK4, byte_scale=byte_scale) out, ok = fn(comm) fl = flops_full * per_dev_frac rows.append([label, rule, "是" if ok else "否", fmt_bytes(sum(r.bytes_total for r in comm.records) / N_DEV), fmt_time(comm.time), fmt_num(fl), fmt_time(fl / peak)]) run("情形 1：A[I_X, J]·B[J, K] → C[I_X, K]", lambda c: (None, np.allclose(unshard(matmul_case1(shard(A, 0, N_DEV), Bm), 0), C)), 1 / N_DEV, "不需要通信") run("情形 2：A[I, J_X]·B[J, K] → C[I, K]", lambda c: (None, all(np.allclose(x, C) for x in matmul_case2(c, shard(A, 1, N_DEV), Bm))), 1.0, "all-gather A，每块设备算全部") run("情形 3：A[I, J_X]·B[J_X, K] → C[I, K]", lambda c: (None, all(np.allclose(x, C) for x in matmul_case3(c, shard(A, 1, N_DEV), shard(Bm, 0, N_DEV)))), 1 / N_DEV, "本地部分和 + all-reduce") run("情形 3'：同上 → C[I, K_X]", lambda c: (None, np.allclose(unshard(matmul_case3(c, shard(A, 1, N_DEV), shard(Bm, 0, N_DEV), 1), 1), C)), 1 / N_DEV, "本地部分和 + reduce-scatter") run("情形 4：A[I_X, J]·B[J, K_X] → C[I_X, K]", lambda c: (None, np.allclose(unshard(matmul_case4(c, shard(A, 0, N_DEV), shard(Bm, 1, N_DEV)), 0), C)), 1 / N_DEV, "先 all-gather B，变成情形 1") print(f"模拟尺寸：I = {I * s}（token），J = {J * s}，K = {K * s}（Llama-3.1-70B 的 MLP 上投影）；BF16；" f"实际运行的数组每维缩小 {s} 倍，通信字节按比例放大") print_table(["情形", "做法", "结果 = 单卡", "每块设备发送", "通信时间", "每块设备 FLOPs", "计算时间"], rows) print_observations([ "情形 1：被求和的维度 J 没被切——每块设备用自己的 I 段算自己的那几行，**一个字节都不用传**。", "情形 2：只有 A 的 J 被切，B 是完整的——必须先把 A all-gather 回来，然后每块设备做**全部**运算（算力没被分摊）。", "情形 3：两边的 J 切法一致——每块设备算一个'部分和'，再 all-reduce（或 reduce-scatter）把它们加起来。" "运算被 8 块设备分摊了，代价是传输输出 C[I, K]：这里 C 比 A 大 3.5 倍，all-reduce 又是 all-gather 的两倍，" "所以通信是情形 2 的 7 倍。", "reduce-scatter 的通信只有 all-reduce 的一半（all-reduce = reduce-scatter + all-gather），结果是切开的 C[I, K_X]——" "如果下一步本来就要按 K 切着用，就没必要 all-reduce。", "情形 4：两个'不被求和'的维度切在同一条轴上，每块设备只能算出 C 的对角块，C[I_X, K_X] 不合法；" "先 all-gather 其中一个，变回情形 1。这里 all-gather 的是权重 B（411 MB、0.93 ms）；如果改为 all-gather 激活 A，" "只要传 7 MB——**能搬激活就别搬权重**，这是推理切分的基本原则（How to Scale Your Model 第 7 章）。", ]) # --------------------------------------------------------------------------- # 演示 3：网络 roofline # --------------------------------------------------------------------------- def demo3_network_roofline(): print_header("演示 3：把被求和的维度切开值不值——通信 vs 计算（情形 3，ring all-reduce）") peak = H100.peak_flops("bf16") rows = [] for B in (16, 512): for D, F in ((4096, 14336), (8192, 28672), (16384, 53248)): for n in (2, 8): flops = 2 * B * D * F / n t_math = flops / peak t_comm = t_formula("ring", n, B * F * 2, NVLINK4.alpha, NVLINK4.bw) rows.append([B, f"{D} × {F}", n, fmt_time(t_math), fmt_time(t_comm), "计算" if t_math > t_comm else "通信", f"{D / n:,.0f}"]) print(f"A[B, D_X]·W[D_X, F] → all-reduce C[B, F]；H100 BF16 {H100.peak_flops('bf16') / 1e12:.0f} TFLOP/s，" f"NVLink {NVLINK4.bw / 1e9:.0f} GB/s、α = {NVLINK4.alpha * 1e6:.0f} µs") print_table(["B", "D × F", "设备数 n", "计算时间", "通信时间", "谁更慢", "每块设备分到的 D"], rows) crit = peak / NVLINK4.bw print(f"\nH100 的'网络临界强度' = 峰值算力 / 链路带宽 = {crit:,.0f} FLOP/byte（第 1 讲 §3.10）；" f"本情形的带宽项临界点 D / (n − 1) ≈ 2 × {crit:,.0f} = {2 * crit:,.0f}") print_observations([ "情形 3 每块设备的运算是 2·B·D·F / n，ring all-reduce 的带宽项约 2·(n−1)/n·(B·F·2 字节) / 带宽（输出的大小，和 D 无关）。" "两者之比 = D / (n − 1) ÷ (2 × 网络临界强度)：切得越细（n 越大）、模型越窄（D 越小），越容易被通信拖住。", f"所以要 D / (n − 1) ≳ {2 * crit:,.0f} 才是计算更慢：Llama-3.1-70B（D = 8192）切 2 份可以（8192），" "切 8 份就不行了（8192 / 7 ≈ 1170）；连 D = 16384 的模型切 8 份也是通信更慢。B = 512 的六行都符合这个判据；B = 16 时还要加上延迟（下一条）：" "D = 8192 切 2 份按带宽算应该是计算更慢，实际却是通信更慢。", "B = 16（decode 那种小批）时，通信时间几乎全是 α：ring all-reduce 有 2(n−1) 步，每步 2 µs，n = 8 时 28 µs，" "而计算只有几微秒。**decode 的张量并行几乎总是被通信的延迟主导**（02、04 脚本还会看到）。", ]) # --------------------------------------------------------------------------- # 演示 4：三种集合通信的代价 # --------------------------------------------------------------------------- def demo4_collectives(): print_header("演示 4：all-gather、reduce-scatter、all-reduce 的代价（8 块 H100，NVLink 交换机）") rng = np.random.default_rng(1) rows = [] for total in (16 * 1024, 64 * 2 ** 20): elems = 4096 scale = total / (elems * 8) # 实际数组 4096 个 float64，按比例放大到 total 字节 arrs = [rng.standard_normal(elems) for _ in range(N_DEV)] res = {} for name, fn in (("all-gather", lambda c: all_gather(c, [a[:elems // N_DEV] for a in arrs], 0)), ("reduce-scatter", lambda c: reduce_scatter(c, arrs, 0)), ("all-reduce（ring）", lambda c: all_reduce(c, arrs, "ring")), ("all-reduce（递归倍增）", lambda c: all_reduce(c, arrs, "rd"))): comm = make_comm(N_DEV, NVLINK4, byte_scale=scale) fn(comm) res[name] = comm.time f_ag = t_formula("ag", N_DEV, total, NVLINK4.alpha, NVLINK4.bw) for name, t in res.items(): ref = {"all-gather": f_ag, "reduce-scatter": f_ag, "all-reduce（ring）": t_formula("ring", N_DEV, total, NVLINK4.alpha, NVLINK4.bw), "all-reduce（递归倍增）": t_formula("rd", N_DEV, total, NVLINK4.alpha, NVLINK4.bw)}[name] rows.append([fmt_bytes(total), name, fmt_time(t), fmt_time(ref), f"{t / res['all-gather']:.2f}"]) print("数据量 = 完整数组的大小（all-gather 的结果、reduce-scatter / all-reduce 的输入）：") print_table(["数据量", "集合通信", "模拟时间", "α–β 公式", "相对 all-gather"], rows) sum_check = [] rng = np.random.default_rng(2) arrs = [rng.standard_normal((8, 16)) for _ in range(N_DEV)] c1 = make_comm(N_DEV) rs = reduce_scatter(c1, arrs, 0) ag = all_gather(c1, rs, 0) c2 = make_comm(N_DEV) ar = all_reduce(c2, arrs, "ring") sum_check.append(np.allclose(ag[3], sum(arrs)) and np.allclose(ar[3], sum(arrs))) print(f"\n检查：reduce-scatter 之后再 all-gather 的结果 = all-reduce 的结果 = 直接求和：{sum_check[0]}；" f"时间 {fmt_time(c1.time)} vs {fmt_time(c2.time)}") print_observations([ "大消息（64 MB）时，all-gather 和 reduce-scatter 的时间相同（都约等于 (n−1)/n × 数据量 / 带宽），" "ring all-reduce 恰好是它们的两倍——因为 ring all-reduce 就是 reduce-scatter + all-gather。" "递归倍增每一步都传完整的数据，大消息时反而慢（3.14 倍）。", "小消息（16 KB）时三者都被延迟 α 主导：ring 要 2(n−1) = 14 步，递归倍增只要 log2 n = 3 步，快得多" "（前置讲 P5 §5.3）。张量并行 decode 每层的 all-reduce 恰好是这种小消息。", "这些关系就是切分矩阵乘法的'价目表'：情形 2 / 4 付一次 all-gather，情形 3 付一次 all-reduce（或半价的 reduce-scatter）。", ]) if __name__ == "__main__": demo1_notation() demo2_four_cases() demo3_network_roofline() demo4_collectives()