""" ============================================================================= CS349E · 第 4 讲 · 02_megatron_tensor_parallel.py —— 把第 2 讲的 Transformer 切到几块虚拟 GPU 上（张量并行） ============================================================================= 【这个脚本要回答的问题】 原课作业 3 的主题：张量并行（tensor parallelism，Megatron-LM 的切法）。每一层： 注意力：按头切 Wq / Wk / Wv（列）和 Wo（行）→ 每块设备算自己那几个头 → 一次 all-reduce； MLP ：按 F 切 W_gate / W_up（列）和 W_down（行）→ 一次 all-reduce。 这里真的把第 2 讲的模型切开、在虚拟 GPU 上跑，检查输出和单卡一样，并数一数通信。 演示 1：正确性与通信量：1 / 2 / 4 / 8 块设备的 logits 和单卡的差；每块设备存多少权重；每次前向几次 all-reduce。 演示 2：带 KV cache 的张量并行 decode：生成的 token 和单卡一样；每块设备的 KV cache—— 设备数超过 KV 组数时，KV 要被复制。 演示 3：为什么是"先列后行"：一个 MLP 的三种切法各需要多少通信（Llama-3.1-70B 的尺寸）；以及 lm_head 的 logits。 【怎么读输出】全部是确定性的计算，每次运行结果相同。小模型的通信时间几乎全是延迟 α（消息只有几 KB）， 所以这里重点看"次数"和"字节"；真实尺寸的时间在 04 脚本里算。 【运行】python3 02_megatron_tensor_parallel.py （约 1 秒） 【依赖】numpy；同目录的 shard_lab；../lecture02/transformer_lab；../lecture01 的 console、llm_math。 """ 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 llm_math import LLAMA31_70B # noqa: E402 from shard_lab import kv_group_of_device, make_comm, shard_params_tp, tp_caches, tp_forward # noqa: E402 from transformer_lab import TINY, count_params, forward, generate, init_params # noqa: E402 def weights_on_device(dp): """一块设备上的参数个数（复制的词嵌入、norm 也算上）。""" n = dp["embed"].size + dp["final_norm"].size + dp["lm_head"].size for lp in dp["layers"]: n += sum(a.size for a in lp.values()) return n # --------------------------------------------------------------------------- # 演示 1：正确性与通信量 # --------------------------------------------------------------------------- def demo1_correctness(): print_header("演示 1：张量并行的前向 = 单卡前向；每块设备存多少、传多少") cfg = TINY p = init_params(cfg, seed=0) toks = np.random.default_rng(1).integers(0, cfg.vocab, size=(2, 16)) ref = forward(p, cfg, toks) total = count_params(p) rows = [] for n in (1, 2, 4, 8): dev = shard_params_tp(p, cfg, n) comm = make_comm(n) out = tp_forward(dev, cfg, toks, comm) ar_steps = [r for r in comm.records if "all-gather" not in r.label] ag_steps = [r for r in comm.records if "all-gather" in r.label] rows.append([n, f"{np.max(np.abs(out - ref)):.1e}", f"{weights_on_device(dev[0]):,}", f"{weights_on_device(dev[0]) / total * 100:.1f}%", 0 if n == 1 else 2 * cfg.n_layers, fmt_bytes(sum(r.bytes_total for r in ar_steps) / n), fmt_bytes(sum(r.bytes_total for r in ag_steps) / n), fmt_time(comm.time)]) print(f"小模型（{cfg.n_layers} 层，H = {cfg.n_heads}，G = {cfg.n_kv_heads}，F = {cfg.d_ff}，V = {cfg.vocab}），" f"2 个序列 × 16 个 token 的前向（prefill）：") print_table(["设备数", "logits 与单卡的最大差", "每块设备的参数", "占全部", "all-reduce 次数", "每设备发送（all-reduce）", "每设备发送（logits all-gather）", "通信时间"], rows) print_observations([ "切到 2、4、8 块设备后，logits 和单卡只差 1e-15 量级：差别只来自 all-reduce 时求和的顺序不同（浮点加法不满足结合律，前置讲 P3 §4.2）。", "每块设备的参数没有降到 1/n：词嵌入和 RMSNorm 是复制的，n > G = 2 时 K、V 的投影也被复制。" "这个小模型词表占比大，复制的部分显得多；大模型里绝大部分参数都被切开了。", "每层固定两次 all-reduce（注意力一次、MLP 一次），和设备数无关。最后还有一次 logits 的 all-gather（lm_head 按词表切）。", "小模型的消息只有几 KB，通信时间几乎全是 α：设备越多，ring all-reduce 的步数越多（2(n−1) 步），时间反而越长。", ]) # --------------------------------------------------------------------------- # 演示 2：带 KV cache 的张量并行 decode # --------------------------------------------------------------------------- def tp_generate(dev, cfg, prompt, n_new): """张量并行的贪心生成：prefill 一次写入各设备的 KV cache，之后每步 1 个 token。返回 (tokens, comm)。""" n = len(dev) comm = make_comm(n) caches = tp_caches(cfg, dev, prompt.shape[0]) logits = tp_forward(dev, cfg, prompt, comm, caches=caches)[:, -1] toks = prompt.copy() for i in range(n_new): nxt = np.argmax(logits, axis=-1)[:, None] toks = np.concatenate([toks, nxt], axis=1) if i == n_new - 1: break logits = tp_forward(dev, cfg, nxt, comm, caches=caches)[:, -1] return toks, comm, caches def demo2_decode_kv(): print_header("演示 2：带 KV cache 的张量并行 decode——每块设备只存自己那几组 KV（直到 KV 组不够分）") cfg = TINY p = init_params(cfg, seed=0) prompt = np.random.default_rng(2).integers(0, cfg.vocab, size=(3, 12)) ref, _ = generate(p, cfg, prompt, 16) single_kv = 2 * cfg.n_layers * cfg.n_kv_heads * cfg.d_head * 2 # BF16 字节 / token rows = [] for n in (1, 2, 4, 8): dev = shard_params_tp(p, cfg, n) toks, comm, caches = tp_generate(dev, cfg, prompt, 16) groups = [kv_group_of_device(cfg, n, d) for d in range(n)] per_dev = caches[0].nbytes() / (prompt.shape[0] * caches[0].length) rows.append([n, "是" if np.array_equal(toks, ref) else "否", str(groups), f"{per_dev:.0f} B", f"{per_dev / single_kv * 100:.0f}%", f"{per_dev * n / single_kv:.1f}×"]) print(f"3 个序列，提示 12 个 token，贪心生成 16 个；H = {cfg.n_heads} 个 query 头、G = {cfg.n_kv_heads} 组 KV；" f"单卡每 token 的 KV（BF16）= {single_kv} B") print_table(["设备数", "token 与单卡相同", "每块设备的 KV 组", "每设备 KV / token", "相对单卡", "全部设备合计 / 单卡"], rows) m = LLAMA31_70B print(f"\nLlama-3.1-70B 有 {m.n_kv_heads} 组 KV：张量并行度 ≤ 8 时 KV 完全切开；16 时每组 KV 要存 2 份。") print_observations([ "带 KV cache 的张量并行生成和单卡生成的 token 完全相同：KV cache 也跟着头一起被切开了，每块设备只存、只读自己那几组。", "设备数 ≤ KV 组数（这里 2）时，每块设备的 KV 是单卡的 1/n——张量并行同时摊薄了**权重**和 **KV 的带宽、容量**。", "设备数 > KV 组数时，一组 KV 要服务分到不同设备上的几个 query 头，只能每块设备各存一份：n = 8 时全部设备合计存了 4 份。" "这就是 GQA / MQA 和张量并行的冲突（MQA 只有 1 组，切几份就复制几份）。", "Pope et al. 的解法：注意力部分**按批切**（每块设备负责一部分序列的完整 KV），用 all-to-all 在'按头'和'按批'之间换布局（03 脚本演示 3）。", ]) # --------------------------------------------------------------------------- # 演示 3：为什么是"先列后行" # --------------------------------------------------------------------------- def demo3_why_column_then_row(): print_header("演示 3：一个 MLP 的三种切法——为什么 Megatron 选'先列后行'（Llama-3.1-70B，8 块设备，BF16）") m = LLAMA31_70B d, f, n = m.d_model, m.d_ff, 8 rows = [] for B, label in ((1, "decode，B = 1"), (64, "decode，B = 64"), (4096, "prefill，4096 个 token")): x, h = B * d * 2, B * f * 2 # 激活 [B, d] 和中间值 [B, f] 的字节 ar = lambda v: 2 * (n - 1) / n * v # ring all-reduce 每块设备发送的字节 ag = lambda v: (n - 1) / n * v # all-gather 每块设备发送的字节 plans = [ ("先列后行（Megatron）", "up 按列 → down 按行 → all-reduce [B, d]", ar(x), 1), ("先行后列", "up 按行 → all-reduce [B, f] → down 按列 → all-gather [B, d]", ar(h) + ag(x), 2), ("两次都按列", "up 按列 → all-gather [B, f] → down 按列 → all-gather [B, d]", ag(h) + ag(x), 2), ] for name, how, by, k in plans: rows.append([label, name, how, k, fmt_bytes(by)]) print_table(["场景", "切法", "过程", "集合通信次数", "每块设备发送"], rows) lg = 64 * m.vocab * 2 print(f"\nlm_head 按词表切时，B = 64 的 logits all-gather 要传 {fmt_bytes(lg)}，而一次层内 all-reduce 的消息只有 " f"{fmt_bytes(64 * d * 2)}——只需要下一个 token 时，可以在每块设备上先取本地的最大值 / top-k，只交换候选。") print_observations([ "'先列后行'：第一个矩阵按输出（列）切，每块设备得到中间值的 F/n 列，正好是第二个矩阵按输入（行）切时需要的那部分——" "中间一个字节都不用传，最后只 all-reduce 一个 [B, d] 的结果。SwiGLU 的逐元素运算也在本地完成（它只作用在自己那几列上）。", "反过来'先行后列'要在中间 all-reduce [B, f]（F = 3.5·d，大 3.5 倍），最后还得 all-gather；'两次都按列'要 all-gather 整个中间值。" "两种都要两次集合通信，字节是 Megatron 的 2.25–4 倍。", "注意力是同一个道理：Wq / Wk / Wv 按头（列）切，每块设备完整地算自己那几个头的注意力，Wo 按行切，最后 all-reduce。" "所以每层两次 all-reduce 就是 Megatron 张量并行的全部通信。", ]) if __name__ == "__main__": demo1_correctness() demo2_decode_kv() demo3_why_column_then_row()