""" ============================================================================= CS349E · 第 4 讲 · 03_pope_partitioning.py —— Pope et al.：大模型推理的分区方式与通信代价模型 ============================================================================= 【这个脚本要回答的问题】 Pope et al., *Efficiently Scaling Transformer Inference*（2022，本讲阅读材料）在 TPU v4 上给 PaLM 540B 做推理， 用一个简单的解析模型挑选分区方式。这里把论文的几个公式算出来，复现它的主要结论： 演示 1：前馈层的 1D vs 2D weight-stationary：1D 的通信不随芯片数下降，2D 按 1/√n 下降；芯片数超过 16 时 2D 更省。 演示 2：weight-stationary vs weight-gathered：批里的 token 多到一定程度，搬权重反而比搬激活便宜（prefill）。 演示 3：多查询注意力（MQA）怎么切：按头切会把唯一的那组 KV 复制到每块芯片上；按批切才能真正省下 KV—— 复现论文表 1 的"最大上下文长度"。 演示 4：PaLM 540B 的 decode 时间：roofline 下界 vs 论文实测（29 ms / token）。 【怎么读输出】全部是公式计算，每次运行结果相同。TPU v4 的规格取自论文 §4：每芯片 BF16 275 TFLOP/s、 32 GiB HBM、1200 GB/s HBM 带宽、270 GB/s 互连带宽（3D torus）。 论文的通信时间公式写作"元素个数 / 带宽"，这里乘上每个数的字节（BF16 = 2）换算成秒。 【运行】python3 03_pope_partitioning.py （< 1 秒） 【依赖】同目录的 shard_lab；../lecture01 的 console、hw_specs。 """ import os import sys 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 hw_specs import Chip # noqa: E402 from shard_lab import pope_ffn_comm # noqa: E402 TPU_V4 = Chip("TPU-v4", "TPU", peak={"bf16": 275e12}, mem_bw=1.2e12, mem_cap=32 * 2 ** 30, mem_kind="HBM2", link_bw=270e9, year=2021, notes="Pope et al. §4：275 TFLOP/s、32 GiB、1200 GB/s、270 GB/s 互连") # PaLM 540B（Chowdhery et al. 2022；Pope et al. 把 48 个头补到 64 个，多出约 18B 参数） PALM = {"layers": 118, "E": 18432, "F": 73728, "heads": 64, "d_head_mqa": 256, "d_head_mha": 128, "params": 540e9} # --------------------------------------------------------------------------- # 演示 1：1D vs 2D weight-stationary # --------------------------------------------------------------------------- def demo1_1d_vs_2d(): print_header("演示 1：前馈层的 1D vs 2D weight-stationary（PaLM 540B，decode B = 512）") E, F, L = PALM["E"], PALM["F"], PALM["layers"] B = 512 rows = [] for n in (8, 16, 32, 64, 128, 256): t1 = pope_ffn_comm("1d", B, E, F, n, TPU_V4.link_bw) * L t2 = pope_ffn_comm("2d", B, E, F, n, TPU_V4.link_bw) * L w = 2 * E * F * 2 * L / n / TPU_V4.mem_bw # 前馈层权重（两个矩阵，BF16）每块芯片的读取时间 fl = 2 * B * 2 * E * F * L / n / TPU_V4.peak_flops("bf16") rows.append([n, fmt_time(t1), fmt_time(t2), "2D" if t2 < t1 else ("相同" if abs(t1 - t2) < 1e-12 else "1D"), fmt_time(w), fmt_time(fl)]) print("118 层前馈层的总通信时间（只算带宽项），以及同样 118 层的读权重、计算时间（作对照）：") print_table(["芯片数 n", "1D 通信", "2D 通信", "更省的", "读权重（每芯片）", "计算（每芯片）"], rows) print_observations([ "1D（Megatron 式，只按 F 切）：每块芯片每层都要 all-gather / reduce-scatter 一整份 [B, E] 的激活，" "通信时间 2·BL·E / 带宽，**和芯片数无关**——芯片越多，读权重和计算越快，通信却不变，迟早成为瓶颈。", "2D（E、F 两个方向都切，每块芯片拿一个近似方形的块）：激活也只需要沿一个轴聚合，通信按 1/√n 下降。" "两者相等的点是 √n = 4（F = 4E 时），即 n = 16：论文说'芯片数超过 16 时 2D 更省'，表里正是如此。", "n = 64 时 1D 的通信（16.5 ms）已经接近计算（18.7 ms），是读前馈层权重（8.4 ms）的两倍；2D 只有 8.3 ms，n = 256 时 4.1 ms。" "**通信时间能不能随芯片数下降，决定了加芯片还能不能降低延迟**。", ]) # --------------------------------------------------------------------------- # 演示 2：weight-stationary vs weight-gathered # --------------------------------------------------------------------------- def demo2_gathered(): print_header("演示 2：搬激活还是搬权重——weight-stationary vs weight-gathered（论文图 3 的设置）") E, F, n = 16384, 65536, 64 rows = [] for tokens in (2048, 8192, 16384, 65536, 262144, 1048576): t_ws = pope_ffn_comm("2d", tokens, E, F, n, TPU_V4.link_bw) t_wg, N = pope_ffn_comm("gathered", tokens, E, F, n, TPU_V4.link_bw) comp = 2 * tokens * 2 * E * F / n / TPU_V4.peak_flops("bf16") rows.append([f"{tokens:,}", fmt_time(t_ws), fmt_time(t_wg), f"{N:.1f}", "gathered" if t_wg < t_ws else "stationary", fmt_time(comp), f"{comp / (comp + min(t_ws, t_wg)) * 100:.0f}%"]) print(f"一个前馈层，E = {E}、F = {F}、{n} 块芯片（X = Y = Z = 4）；N = 权重在几块芯片之间 all-gather（最优值 √(BL·n / F)）：") print_table(["这一步的 token 数 BL", "2D stationary 通信", "weight-gathered 通信", "最优 N", "更省的", "计算时间", "计算占比（不重叠时）"], rows) print_observations([ "weight-stationary 的通信 ∝ BL（激活随 token 数增长），weight-gathered 的通信 ∝ √BL（多搬一些权重，少搬一些激活）。" f"两者相等在 BL = F / 4 = {F // 4:,} 个 token：比这少时搬激活便宜，比这多时搬权重便宜。", "decode 的批（几百个 token）远在交叉点左边，所以 decode 用 2D weight-stationary；prefill 动辄几十万个 token，" "换成 weight-gathered。同一套权重布局（E_x F_yz）两种都能用，只是运行时选不同的通信方式。", "token 越多，通信的占比越小（最后一行计算占 94%，而且通信还能和计算重叠）：论文在 prefill 时达到 76% 的 MFU，靠的就是这个。", "【系统联系】这和第 1 讲的 roofline 是同一个思路：比较'搬的字节'和'算的次数'，只不过这里的屋顶是芯片之间的网络，" "而且可以选择搬哪一种数据（激活 vs 权重）。", ]) # --------------------------------------------------------------------------- # 演示 3：MQA 怎么切 # --------------------------------------------------------------------------- def demo3_mqa_partition(): print_header("演示 3：多查询注意力怎么切——按头切会复制 KV，按批切才省（复现论文表 1）") n, L = 64, PALM["layers"] budget = n * TPU_V4.mem_cap * 0.30 # 论文：预留 30% 的显存给 KV cache mha_tok = 2 * L * PALM["heads"] * PALM["d_head_mha"] * 2 # MHA 变体：64 头 × 128 维（参数量和 MQA 版相同） mqa_tok = 2 * L * 1 * PALM["d_head_mqa"] * 2 # MQA：1 组 KV，256 维 variants = [("多头注意力（MHA）", mha_tok, 1, "1,320 / 330"), ("MQA，按头切（KV 复制到每块芯片）", mqa_tok, n, "660 / 165"), ("MQA，按批切（Pope 的做法）", mqa_tok, 1, "43,000 / 10,700")] rows = [] for name, per_tok, copies, paper in variants: lens = [int(budget / (B * per_tok * copies)) for B in (128, 512)] rows.append([name, fmt_bytes(per_tok), copies, f"{lens[0]:,} / {lens[1]:,}", paper]) print(f"PaLM 540B，{n} 块 TPU v4，KV 预算 = 总显存的 30% = {fmt_bytes(budget)}；BF16 KV") print_table(["注意力的切法", "每 token 的 KV（一份）", "存几份", "最大上下文：B = 128 / 512", "论文表 1"], rows) B, ctx = 256, 8192 rows = [] for name, per_tok, copies, _ in variants: kv_chip = B * ctx * per_tok * copies / n rows.append([name, fmt_bytes(kv_chip), fmt_time(kv_chip / TPU_V4.mem_bw), "能" if kv_chip <= TPU_V4.mem_cap * 0.3 else "放不下"]) a2a = 2 * B * PALM["heads"] * PALM["d_head_mqa"] * 2 * L / n / TPU_V4.link_bw print(f"\n批大小 {B}、上下文 {ctx} 的一步 decode，每块芯片读 KV 的时间（118 层）：") print_table(["切法", "每块芯片的 KV", "读 KV 的时间", "放得进 30% 预算吗"], rows) print(f"按批切多出来的 all-to-all（每层前后各一次，搬 Q 和注意力输出）：约 {fmt_time(a2a)}（只算带宽项）") print_observations([ "MQA 只有一组 K/V。如果像 MHA 那样按头切，每块芯片都需要这唯一的一组 K/V，只能各存一份：64 块芯片存 64 份，" "MQA 省下的显存全被复制吃掉了——比 MHA 还差一半（MQA 版的 d_head 是 256）。", "按批切：每块芯片只负责一部分序列，存这些序列的完整 KV，不复制。最大上下文长度提高了 32–64 倍，" "和论文表 1（43,000 vs 660 / 1,320）一致。", "代价是一次 all-to-all：把 Q 从'按头切'换成'按批切'，算完再换回去。decode 时 Q 只有一个 token，" "比 KV（几千个 token）小几个数量级——**花一点通信，省大量 KV 读取**。prefill 时 Q 很长，就不划算了，论文 prefill 时按头切。", "GQA（Llama 3 的 8 组 KV）介于两者之间：张量并行度 ≤ 8 时按头切不复制，超过 8 就要考虑按批切（02 脚本演示 2）。", ]) # --------------------------------------------------------------------------- # 演示 4：PaLM 540B 的 decode 时间 # --------------------------------------------------------------------------- def demo4_palm_latency(): print_header("演示 4：PaLM 540B 的 decode——roofline 下界 vs 论文实测") n, L, E = 64, PALM["layers"], PALM["E"] ctx = 2048 mqa_tok = 2 * L * PALM["d_head_mqa"] * 2 rows = [] for w_bytes, B, measured in ((1, 64, 28.5e-3), (2, 64, 36.9e-3)): t_w = PALM["params"] * w_bytes / (n * TPU_V4.mem_bw) t_kv = B * ctx * mqa_tok / n / TPU_V4.mem_bw t_fl = 2 * PALM["params"] * B / (n * TPU_V4.peak_flops("bf16")) t_comm = pope_ffn_comm("2d", B, E, PALM["F"], n, TPU_V4.link_bw) * L bound = max(t_w + t_kv, t_fl) + t_comm rows.append(["int8" if w_bytes == 1 else "bf16", B, fmt_time(t_w), fmt_time(t_kv), fmt_time(t_fl), fmt_time(t_comm), fmt_time(bound), fmt_time(measured), f"{measured / bound:.1f}×", f"{n * measured / B * 1e3:.1f}"]) print(f"{n} 块 TPU v4，上下文 {ctx}，2D weight-stationary（并行的注意力 / 前馈层，每层一组通信）：") print_table(["权重", "B", "读权重", "读 KV", "计算", "通信（带宽项）", "下界", "论文实测", "实测 / 下界", "成本（芯片·毫秒 / token）"], rows) print_observations([ "下界里最大的一项是读权重：540B 个 int8 参数分到 64 块芯片上，每块读 8.4 GB，要 7 ms；bf16 翻倍。" "int8 权重让论文的实测从 36.9 ms 降到 28.5 ms——小批量 decode 只看权重字节（第 3 讲 §5.1）。", "实测是下界的 2.4–3.4 倍。差距来自下界没算的东西：每次集合通信的启动延迟（118 层 × 每层若干次）、" "通信和计算没能完全重叠、注意力和 layernorm 等小运算的开销、芯片达不到峰值带宽。**roofline 告诉你'不可能更快'，" "不告诉你'能做到多快'**。", "成本按'芯片·毫秒 / token'（= 芯片数 × 每步时间 / 批大小）算：B = 64 时约 28 芯片·毫秒。论文说批大小 512 左右成本最低，" "而最低延迟（小批）比批大小 512 时的延迟低约 3 倍——这就是延迟和成本的取舍（04 脚本）。", ]) if __name__ == "__main__": demo1_1d_vs_2d() demo2_gathered() demo3_mqa_partition() demo4_palm_latency()