""" ============================================================================= CS349E · 前置讲 P5 · 02_collectives_semantics.py —— 集合通信：每一种在做什么、彼此什么关系、要传多少字节 ============================================================================= 【这个脚本要回答的问题】 "集合通信"（collective communication）是一组设备一起完成的数据交换：广播、归约、all-reduce、 reduce-scatter、all-gather、all-to-all……名字很多，但每一种的"输入是什么、输出是什么"非常具体。本脚本： 演示 1：在 4 个设备上用一眼就能看懂的小数组，打印每种集合通信之前和之后每个设备手里的数据。 演示 2：几个恒等式（用真实数据在 8 个设备上验证）：all-reduce = reduce-scatter + all-gather； all-to-all 做两次回到原样；广播 = scatter + all-gather。 演示 3：每个设备要发、要收多少字节？和"无论什么算法都至少要传多少"的下界比较——ring 算法为什么是带宽最优的。 演示 4：LLM 推理里，每种并行方式用到哪种集合通信。 【怎么读输出】全部是确定性的计算，每次运行结果相同。 【运行】python3 02_collectives_semantics.py （不到 1 秒） 【依赖】numpy；从 ../lecture01 导入 console；同目录的 comm_sim。 """ import os import sys from collections import defaultdict import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import print_header, print_observations, print_table # noqa: E402 from comm_sim import (NVLINK4, Comm, allgather_ring, allreduce_ring, alltoall_direct, broadcast_tree, # noqa: E402 reduce_scatter_ring, reduce_tree, switch) def show(arr): """把一个小整数数组写成 [a b c d]。""" return "[" + " ".join(f"{int(x):3d}" for x in arr) + "]" # --------------------------------------------------------------------------- # 演示 1：每种集合通信在做什么 # --------------------------------------------------------------------------- def demo1_semantics(p=4): """ 设备 d 一开始持有 [10d+0, 10d+1, 10d+2, 10d+3]（4 个数，方便切成 4 块，每块 1 个数）。 依次做每种集合通信，打印"之后"每个设备手里的数据。所有算法用 comm_sim 里真的搬运数据的实现。 """ print_header(f"演示 1：每种集合通信在做什么（{p} 个设备，设备 d 一开始持有 [10d, 10d+1, 10d+2, 10d+3]）") bufs = [np.arange(4, dtype=np.int64) + 10 * d for d in range(p)] sw = switch(p, NVLINK4) results = {} results["之前"] = bufs tmp = [bufs[0]] + [np.zeros(4, dtype=np.int64)] * (p - 1) results["broadcast（根 = 0）"] = broadcast_tree(Comm(sw), tmp, 0) red = reduce_tree(Comm(sw), bufs, 0) results["reduce（求和到 0）"] = [red] + [None] * (p - 1) results["all-reduce"] = allreduce_ring(Comm(sw), bufs) rs, owner = reduce_scatter_ring(Comm(sw), bufs) results["reduce-scatter"] = rs pieces = [np.array([10 * d], dtype=np.int64) for d in range(p)] own = {d: d for d in range(p)} results["all-gather（设备 d 先只有 [10d]）"] = allgather_ring(Comm(sw), pieces, own) results["all-to-all"] = alltoall_direct(Comm(sw), bufs) rows = [] for name, out in results.items(): cells = [name] for d in range(p): if out[d] is None: cells.append("（不关心）") elif name == "reduce-scatter": cells.append(f"第 {owner[d]} 块 = {show(out[d])}") else: cells.append(show(out[d])) rows.append(cells) print_table(["操作"] + [f"设备 {d}" for d in range(p)], rows) print_observations([ "broadcast：根的数据复制给所有人。reduce：所有人的数据逐元素相加，结果只在根上。", "all-reduce：所有人都得到逐元素的和（这里每个位置是 0+10+20+30 = 60 再加上位置号）。张量并行每层都要做它。", "reduce-scatter：和也求了，但每个设备只拿到**一块**（设备 d 拿到第 owner[d] 块）——all-reduce 的'前一半'。", "all-gather：每个设备贡献一块，最后每个人都有完整的拼接结果——all-reduce 的'后一半'。", "all-to-all：设备 i 的第 j 块发给设备 j。把 4 个设备的数据排成 4×4 的矩阵（行 = 设备，列 = 块），" "all-to-all 就是**转置**。MoE 的专家并行用它把 token 送到各自的专家那里。", ]) # --------------------------------------------------------------------------- # 演示 2：恒等式 # --------------------------------------------------------------------------- def demo2_identities(p=8, n=1 << 12, seed=0): """用真实的随机数据在 p 个设备上验证几个恒等式。""" print_header(f"演示 2：几个恒等式（{p} 个设备，每个 {n:,} 个数，真实数据）") rng = np.random.default_rng(seed) bufs = [rng.standard_normal(n) for _ in range(p)] total = sum(bufs) sw = switch(p, NVLINK4) rows = [] c1 = Comm(sw) ar = allreduce_ring(c1, bufs) c2 = Comm(sw) rs, owner = reduce_scatter_ring(c2, bufs) t_rs = c2.time ag = allgather_ring(c2, rs, owner) rows.append(["all-reduce = reduce-scatter + all-gather", "是" if all(np.allclose(a, total) for a in ag) and all(np.array_equal(a, b) for a, b in zip(ar, ag)) else "否", f"RS {t_rs * 1e6:.2f} µs + AG {(c2.time - t_rs) * 1e6:.2f} µs = AR {c1.time * 1e6:.2f} µs"]) once = alltoall_direct(Comm(sw), bufs) twice = alltoall_direct(Comm(sw), once) rows.append(["all-to-all 做两次 = 原样（转置的转置）", "是" if all(np.array_equal(a, b) for a, b in zip(twice, bufs)) else "否", ""]) bc_times = {} for label, root_data in (("小消息", bufs[0]), ("大消息", rng.standard_normal(1 << 21).astype(np.float32))): c3 = Comm(sw) scattered = np.array_split(root_data, p) c3.step([(0, d, scattered[d].nbytes) for d in range(1, p)], "scatter") bc = allgather_ring(c3, scattered, {d: d for d in range(p)}) c4 = Comm(sw) broadcast_tree(c4, [root_data] + [np.zeros_like(root_data)] * (p - 1), 0) bc_times[label] = (c3.time, c4.time) rows.append([f"broadcast = scatter + all-gather（{label}，{root_data.nbytes / 1e6:.2f} MB）", "是" if all(np.array_equal(b, root_data) for b in bc) else "否", f"scatter + AG {c3.time * 1e6:.2f} µs；二项树广播 {c4.time * 1e6:.2f} µs"]) print_table(["恒等式", "成立？", "（在 NVLink 交换机上的模拟时间）"], rows) print_observations([ "all-reduce 可以拆成 reduce-scatter + all-gather，而且 ring all-reduce 正是这么实现的。拆开以后，" "两半之间可以插入别的计算——'序列并行'（sequence parallelism）就把 LayerNorm 这类逐 token 的运算放在中间，" "只在 1/p 的数据上算（第 4 讲）。", "all-to-all 是转置，做两次回到原样：MoE 里先 all-to-all 把 token 发给专家（dispatch），专家算完再 all-to-all 发回来（combine）。", "广播也可以拆成 scatter + all-gather（Van de Geijn 的算法）：每条链路上只传约 2 份数据，而二项树要传 log2 p 份完整数据。" f"小消息时它更慢（{bc_times['小消息'][0] * 1e6:.1f} vs {bc_times['小消息'][1] * 1e6:.1f} µs，步数多、α 多），" f"大消息时更快（{bc_times['大消息'][0] * 1e6:.0f} vs {bc_times['大消息'][1] * 1e6:.0f} µs）。" "同一个语义，不同的实现，适合不同的消息大小——MPI、NCCL 都会按消息大小切换算法（03 脚本）。", ]) # --------------------------------------------------------------------------- # 演示 3：每个设备要传多少字节 # --------------------------------------------------------------------------- def demo3_traffic(p=8, n_elem=1 << 16, seed=0): """ 统计每种算法里**每个设备发出的字节数**（从模拟器记录的消息里数），和下界比较。 【下界】设每个设备有 n 字节： all-gather（结果共 n·p 字节）：每个设备至少要收到别人的 (p−1)·n 字节 → 收 ≥ (p−1)/p × 结果大小； reduce-scatter：每个设备最后的那一块需要另外 p−1 个设备的贡献 → 每个设备至少收 (p−1)/p · n； all-reduce：可以证明每个设备至少要发 2(p−1)/p · n（Patarasuk & Yuan 2009）； all-to-all：每个设备要把 (p−1)/p · n 发给别人。 """ print_header(f"演示 3：每个设备要发多少字节？（{p} 个设备，每个 n = {n_elem * 8 // 1024} KiB）") rng = np.random.default_rng(seed) bufs = [rng.standard_normal(n_elem) for _ in range(p)] n = n_elem * 8 sw = switch(p, NVLINK4) class Counting(Comm): """在 Comm 的基础上，把每个设备发出的字节数累加起来。""" def __init__(self, topo): super().__init__(topo) self.sent = defaultdict(float) def step(self, msgs, label=""): for s, _, nb in msgs: if isinstance(s, int): self.sent[s] += nb return super().step(msgs, label) from comm_sim import allreduce_naive, allreduce_rabenseifner, allreduce_recursive_doubling, allreduce_tree cases = [ ("all-reduce：ring", lambda c: allreduce_ring(c, bufs), 2 * (p - 1) / p), ("all-reduce：Rabenseifner", lambda c: allreduce_rabenseifner(c, bufs), 2 * (p - 1) / p), ("all-reduce：递归倍增", lambda c: allreduce_recursive_doubling(c, bufs), 2 * (p - 1) / p), ("all-reduce：二叉树", lambda c: allreduce_tree(c, bufs), 2 * (p - 1) / p), ("all-reduce：朴素", lambda c: allreduce_naive(c, bufs), 2 * (p - 1) / p), ("reduce-scatter：ring", lambda c: reduce_scatter_ring(c, bufs), (p - 1) / p), ("all-to-all：直接", lambda c: alltoall_direct(c, bufs), (p - 1) / p), ] rows = [] for name, fn, lb in cases: c = Counting(sw) fn(c) sent = np.array([c.sent[d] for d in range(p)]) / n rows.append([name, f"{sent.mean():.3f} n", f"{sent.max():.3f} n", f"{lb:.3f} n", len(c.records)]) print_table(["算法", "平均每设备发出", "最多的那个设备发出", "下界（每设备）", "步数"], rows) print_observations([ f"ring 和 Rabenseifner 每个设备都正好发 2(p−1)/p · n = {2 * (p - 1) / p:.3f} n，等于下界：它们是**带宽最优**的。" "区别在步数：ring 要 2(p−1) 步，Rabenseifner 只要 2·log2 p 步（α 少得多）。", "递归倍增每步都发完整的 n，一共 log2 p 步：发出 log2 p · n，比下界多；但步数最少，适合小消息。", "二叉树和朴素算法的问题是**不均匀**：一部分设备（根、树的上层）发得多、其他设备闲着。" "朴素算法里根要发 (p−1) 份——它是所有算法里最慢的（03 脚本）。", "all-reduce 的字节数约是 reduce-scatter / all-gather 的 2 倍——所以 all-reduce 的时间约是 all-gather 的 2 倍" "（How to Scale Your Model 第 3 章的经验法则）。", ]) # --------------------------------------------------------------------------- # 演示 4：LLM 推理里的集合通信 # --------------------------------------------------------------------------- def demo4_where_in_llm(): """一张对照表：并行方式 → 切什么 → 每层用什么通信 → 传多少字节的量级（B = 这一步的 token 数，d = 隐藏维度）。""" print_header("演示 4：LLM 推理里，每种并行方式用到哪种集合通信") rows = [ ["张量并行 TP（Megatron）", "每个矩阵切成 p 份，每块芯片算一部分", "每层 2 次 all-reduce（注意力后、MLP 后）", "B × d × 2 字节 / 次", "第 4 讲"], ["TP + 序列并行 SP", "在 all-reduce 中间插入逐 token 的运算", "每层 2 次 reduce-scatter + 2 次 all-gather", "同上（总字节相同）", "第 4 讲"], ["专家并行 EP（MoE）", "不同的专家放在不同芯片上", "每个 MoE 层 2 次 all-to-all（分发、合并）", "B × k × d × 字节 / 次", "第 3、4 讲"], ["流水线并行 PP", "不同的层放在不同芯片上", "相邻两段之间点对点发送", "B × d × 2 字节 / 次", "第 4 讲"], ["数据并行 DP（多副本）", "每个副本一份完整的模型", "推理时副本之间不通信", "—", "第 4、5 讲"], ["解耦式推理（P/D 分离）", "prefill 和 decode 在不同的机器上", "prefill 完成后把 KV cache 发给 decode 机器", "提示长度 × 每 token 的 KV 字节", "第 5、7 讲"], ] print_table(["并行方式", "怎么切", "通信", "每次传多少", "详见"], rows) print_observations([ "张量并行的 all-reduce 在每一层的关键路径上：decode 时 B 很小，一次只有几 KB 到几百 KB，是延迟主导的（05 脚本）。", "all-to-all 的通信量和 top-k 成正比，而且要跨越所有持有专家的芯片——它常常是 MoE 推理的瓶颈，" "所以 DeepSeek-V3 等系统专门为它写通信内核、用 FP8 发送。", "流水线和解耦的点对点发送不需要所有设备协调，但 KV cache 很大（一条 4096 token 的提示，Llama-3.1-70B 的 KV 超过 1 GB），" "跨节点搬运要算账（05 脚本）。", ]) if __name__ == "__main__": demo1_semantics() demo2_identities() demo3_traffic() demo4_where_in_llm()