""" ============================================================================= CS349E · 前置讲 P5 · 04_topology_and_hierarchy.py —— 拓扑与映射、all-to-all 与对分带宽、层次化 all-reduce ============================================================================= 【这个脚本要回答的问题】 03 脚本假设所有设备挂在一个理想的交换机上。真实的网络有形状：TPU 是 torus，GPU 集群是"节点内 NVLink + 节点间 InfiniBand"的两级网络。同一个算法放到不同的网络上，表现可以差很多。本脚本： 演示 1：ring all-reduce 在 16 个设备上——交换机、物理环、4×4 torus（三种"逻辑环 → 物理位置"的映射）， 以及按维度分两步做的二维 all-reduce。映射不好会怎样？ 演示 2：all-to-all——交换机、环、torus 上的时间，和对分带宽给出的下界比较。 演示 3：两级网络（2–8 个节点 × 8 块 H100）：扁平 ring vs 层次化 all-reduce。 【怎么读输出】全部是模拟（真实数据 + 逐消息记账），每次运行结果相同。为了模拟大消息又不占太多内存， 用小数组 + Comm 的 byte_scale 放大字节数（正确性在小数组上检查，时间按放大后的字节数算）。 【运行】python3 04_topology_and_hierarchy.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 (IB_NDR, ICI_V5E, Comm, allgather_ring, allreduce_hierarchical, # noqa: E402 allreduce_ring, alltoall_direct, bisection_bw, reduce_scatter_ring, ring, switch, torus2d, two_level) N_ELEM = 1 << 12 # 每个设备真实持有的元素数（小数组） def make_bufs(p, seed=0): rng = np.random.default_rng(seed) return [rng.standard_normal(N_ELEM) for _ in range(p)] def scale_for(nbytes): """要让 N_ELEM 个 float64 代表 nbytes 字节，byte_scale 应取多少。""" return nbytes / (N_ELEM * 8) def check(out, ref): return all(np.allclose(o, ref) for o in out) def allreduce_2d(comm, bufs, a, b): """ 二维（按维度）all-reduce，a × b 个设备排成网格（编号 x·b + y）： 1. 每一行做 ring reduce-scatter（a 个环同时，各 b 个设备）； 2. 每一列在自己那一块上做 ring reduce-scatter（b 个环同时）； 3. 每一列 all-gather；4. 每一行 all-gather。 步数 2(a − 1) + 2(b − 1)，比一个 p = ab 的大环的 2(ab − 1) 少得多；每一步都只和同一行 / 同一列的邻居说话， 在 torus 上正好走物理链路。 """ rows = [[x * b + y for y in range(b)] for x in range(a)] cols = [[x * b + y for x in range(a)] for y in range(b)] p1, own1 = reduce_scatter_ring(comm, bufs, rows, "行 RS") p2, own2 = reduce_scatter_ring(comm, p1, cols, "列 RS") g2 = allgather_ring(comm, p2, own2, cols, "列 AG") return allgather_ring(comm, g2, own1, rows, "行 AG") # --------------------------------------------------------------------------- # 演示 1：拓扑与映射 # --------------------------------------------------------------------------- def demo1_mapping(n_bytes=64e6, seed=0): """ 16 个设备、每条链路都是 45 GB/s / 1 µs（ICI 的参数），ring all-reduce 64 MB。 逻辑环的顺序决定了"环上的邻居"在物理上隔几跳： 蛇形：0 1 2 3 7 6 5 4 8 9 10 11 15 14 13 12（每一对相邻都是物理邻居，最后 12 → 0 走环绕链路）； 按行：0 1 2 … 15（每行末尾跳到下一行开头要 2 跳）； 随机：打乱的顺序（很多对相邻要走好几跳，而且路径互相重叠——争用）。 """ p, a = 16, 4 print_header(f"演示 1：ring all-reduce {fmt_bytes(n_bytes)}，16 个设备，每条链路 45 GB/s——拓扑和映射") bufs = make_bufs(p, seed) ref = sum(bufs) sc = scale_for(n_bytes) ideal = 2 * (p - 1) / p * n_bytes / ICI_V5E.bw snake = [0, 1, 2, 3, 7, 6, 5, 4, 8, 9, 10, 11, 15, 14, 13, 12] rng = np.random.default_rng(seed) rand = [int(x) for x in rng.permutation(p)] cases = [("交换机（每个设备 1 个端口）", switch(p, ICI_V5E), list(range(p))), ("物理双向环，按环的顺序", ring(p, ICI_V5E), list(range(p))), ("4×4 torus，蛇形映射", torus2d(a, a, ICI_V5E), snake), ("4×4 torus，按行映射", torus2d(a, a, ICI_V5E), list(range(p))), ("4×4 torus，随机映射", torus2d(a, a, ICI_V5E), rand)] rows = [] res = {} for name, topo, order in cases: comm = Comm(topo, sc) out = allreduce_ring(comm, bufs, [order]) hops = [topo.hops(order[i], order[(i + 1) % p]) for i in range(p)] res[name] = comm.time rows.append([name, f"{np.mean(hops):.2f} / {max(hops)}", fmt_time(comm.time), f"{comm.time / ideal:.2f}", "✓" if check(out, ref) else "✗"]) comm = Comm(torus2d(a, a, ICI_V5E), sc) out = allreduce_2d(comm, bufs, a, a) res["2d"] = comm.time rows.append(["4×4 torus，按维度的二维 all-reduce", "1.00 / 1", fmt_time(comm.time), f"{comm.time / ideal:.2f}", "✓" if check(out, ref) else "✗"]) print_table(["拓扑与映射", "环上相邻两个设备的跳数（平均 / 最多）", "时间", "÷ 单环理想时间", "结果"], rows) print(f"\n单环理想时间 = 2(p−1)/p · n/β = {fmt_time(ideal)}（每个设备一条 45 GB/s 的出口在用）。") small = 64e3 c_ring = Comm(torus2d(a, a, ICI_V5E), scale_for(small)) allreduce_ring(c_ring, bufs, [snake]) c_2d = Comm(torus2d(a, a, ICI_V5E), scale_for(small)) allreduce_2d(c_2d, bufs, a, a) res["small"] = (c_ring.time, c_2d.time) print(f"换成 {fmt_bytes(small)} 的小消息：蛇形单环 {fmt_time(c_ring.time)}（{len(c_ring.records)} 步），" f"二维 all-reduce {fmt_time(c_2d.time)}（{len(c_2d.records)} 步）。") print_observations([ "交换机、物理环、蛇形映射的 torus 都达到了单环的理想时间：环上的每一对邻居都是物理上直接相连的，没有争用。" "ring 算法只需要'一个环'——这也是它在 TPU 的 torus 上很好用的原因。", f"按行映射时，每行最后一个设备要隔 2 跳才到下一行的开头，但多走的那一跳用的链路恰好没有别人用，只慢了 " f"{(res['4×4 torus，按行映射'] / ideal - 1) * 100:.0f}%（多了一跳的延迟）。随机映射就严重了" f"（{res['4×4 torus，随机映射'] / ideal:.2f} 倍）：环上的邻居平均隔 2 跳以上，很多消息的路径重叠在同一条链路上——" "'哪个逻辑编号放在哪块物理芯片上'是一个实实在在的性能问题（MPI、JAX 的 mesh 都要考虑设备的物理排布）。", f"按维度做的二维 all-reduce 在 64 MB 时和单环一样快（{res['2d'] / ideal:.2f} 倍）：传的字节数相同，" f"大消息只看字节。它的好处在小消息：步数从 2(p−1) = 30 降到 2(a−1) + 2(b−1) = 12，" f"64 KB 时从 {fmt_time(res['small'][0])} 降到 {fmt_time(res['small'][1])}。" "它也还没用上 torus 的全部 4 个端口——行、列两个阶段先后进行；真实实现会让两个方向、两个维度的环同时跑" "（NCCL 的多 channel、TPU 的多维集合通信）。", ]) # --------------------------------------------------------------------------- # 演示 2：all-to-all 与对分带宽 # --------------------------------------------------------------------------- def demo2_alltoall(n_bytes=16e6, seed=0): """ 每个设备有 n 字节，切成 p 块发给所有人（直接算法：p − 1 步，第 s 步发给编号 +s 的设备）。 【两个下界，取较大的】 对分带宽：一半设备要发给另一半的数据，每个设备有 n/2 要过中线，一共 p·n/4 字节（单方向）→ 时间 ≥ (p·n/4) / 对分带宽； 出口带宽：每个设备自己要发出 (p−1)/p · n 字节，只能走自己的端口 → 时间 ≥ (p−1)/p · n / (端口数 × β)。 """ print_header(f"演示 2：all-to-all（每个设备 {fmt_bytes(n_bytes)}），和对分带宽给出的下界比较") rows = [] ports = {"switch": 1, "ring": 2, "torus": 4} res = {} for p, cases in ((16, [("交换机", switch(16, ICI_V5E), "switch"), ("双向环", ring(16, ICI_V5E), "ring"), ("4×4 torus", torus2d(4, 4, ICI_V5E), "torus")]), (64, [("交换机", switch(64, ICI_V5E), "switch"), ("双向环", ring(64, ICI_V5E), "ring"), ("8×8 torus", torus2d(8, 8, ICI_V5E), "torus")])): bufs = make_bufs(p, seed) for name, topo, kind in cases: comm = Comm(topo, scale_for(n_bytes)) out = alltoall_direct(comm, bufs) ok = all(np.array_equal(out[j], np.concatenate([np.array_split(bufs[i], p)[j] for i in range(p)])) for j in range(p)) bis = bisection_bw(kind, p, ICI_V5E.bw) lb_bis = p * n_bytes / 4 / bis lb_inj = (p - 1) / p * n_bytes / (ports[kind] * ICI_V5E.bw) lb = max(lb_bis, lb_inj) res[(p, kind)] = (comm.time, lb) rows.append([p, name, f"{bis / 1e9:,.0f} GB/s", fmt_time(lb_bis), fmt_time(lb_inj), fmt_time(comm.time), f"{comm.time / lb:.2f}", "✓" if ok else "✗"]) print_table(["设备数", "拓扑", "对分带宽", "下界①：对分", "下界②：出口", "模拟时间", "÷ 较大的下界", "结果"], rows) print_observations([ f"交换机上，all-to-all 的时间几乎等于'出口'下界（{res[(64, 'switch')][0] / res[(64, 'switch')][1]:.2f} 倍）：" "每个设备的端口一直满速在发，对分带宽绰绰有余。", f"环上，all-to-all 被对分带宽卡死：64 个设备时 {fmt_time(res[(64, 'ring')][0])}，是交换机的 " f"{res[(64, 'ring')][0] / res[(64, 'switch')][0]:.0f} 倍——只有两条链路跨过中线，'对面'的数据都要从这里挤过去。", "torus 介于两者之间：对分带宽随 √p 增长，设备越多，相对交换机越吃亏。How to Scale Your Model 第 12 章提到，" "GPU 节点内（交换机）做 all-to-all 比 TPU 的 torus 快——这正是 MoE 的专家并行偏爱 NVSwitch 域" "（以及 GB200 NVL72 这种 72 块 GPU 的大交换域）的原因（05 脚本）。", "环和 torus 上，模拟时间比下界高出约 2–4 倍：直接算法分 p − 1 步、每步之间要同步，而且每一步的流量在拓扑上不均匀" "（有的链路忙、有的闲，最短路由遇到'两边一样远'时只走一边）。更好的 all-to-all 算法会按拓扑重新编排每一步的发送对象。", ]) # --------------------------------------------------------------------------- # 演示 3：层次化 all-reduce # --------------------------------------------------------------------------- def demo3_hierarchical(seed=0): """ 两级网络：每个节点 8 块 H100（NVLink 450 GB/s），节点间每块 GPU 一块 50 GB/s 的网卡。 扁平 ring：所有 GPU 按编号排成一个大环，跨节点的那几步走网卡——最慢的链路决定每一步的时间。 层次化：节点内 RS（NVLink）→ 节点间 all-reduce（8 个环同时，每个只处理 1/8 的数据，走网卡）→ 节点内 AG。 """ print_header("演示 3：两级网络（N 个节点 × 8 块 H100）：扁平 ring vs 层次化 all-reduce") rows = [] res = {} for n_bytes in (1e6, 256e6): for n_nodes in (2, 4, 8): p = 8 * n_nodes bufs = make_bufs(p, seed) ref = sum(bufs) topo = two_level(n_nodes, 8) c1 = Comm(topo, scale_for(n_bytes)) o1 = allreduce_ring(c1, bufs) c2 = Comm(topo, scale_for(n_bytes)) o2 = allreduce_hierarchical(c2, bufs, 8) res[(n_bytes, n_nodes)] = (c1.time, c2.time) ideal_ib = 2 * (n_nodes - 1) / n_nodes * (n_bytes / 8) / IB_NDR.bw rows.append([fmt_bytes(n_bytes), f"{n_nodes} × 8 = {p}", fmt_time(c1.time), fmt_time(c2.time), f"{c1.time / c2.time:.1f}×", fmt_time(ideal_ib), "✓" if check(o1, ref) and check(o2, ref) else "✗"]) print_table(["消息", "GPU 数", "扁平 ring", "层次化", "加速", "节点间部分的理想时间", "结果"], rows) print_observations([ f"256 MB、8 个节点：扁平 ring 用了 {fmt_time(res[(256e6, 8)][0])}，层次化只要 {fmt_time(res[(256e6, 8)][1])}" f"（{res[(256e6, 8)][0] / res[(256e6, 8)][1]:.1f} 倍）。扁平 ring 的每一步都要等跨节点的那条 50 GB/s 链路，" "相当于整个环都按网卡的速度走；层次化让每块 GPU 只把 1/8 的数据送出节点，8 块网卡同时工作。", "节点间部分的理想时间（最后一列）说明：层次化之后，跨节点的流量只剩 2(N−1)/N · n/8 字节，" "加上两段很快的 NVLink 阶段——它几乎就是两级网络上 all-reduce 能做到的最好。", f"1 MB 时加速反而随节点数变大（{res[(1e6, 2)][0] / res[(1e6, 2)][1]:.1f} → {res[(1e6, 8)][0] / res[(1e6, 8)][1]:.1f} 倍）：" "小消息的时间由步数主导。扁平 ring 有 2(8N − 1) 步，而且几乎每一步都有消息跨节点、要付网卡更大的 α；" "层次化只有 2·7 步走 NVLink、2(N − 1) 步走网卡。消息大时（256 MB），加速随节点数略微下降" f"（{res[(256e6, 2)][0] / res[(256e6, 2)][1]:.1f} → {res[(256e6, 8)][0] / res[(256e6, 8)][1]:.1f} 倍）：" "节点间部分 2(N−1)/N · n/8 随 N 变大，而扁平 ring 已经接近 2n/β_网卡 的上限。", "这就是 NCCL 在多节点上用的思路（以及'rail-optimized'网络的设计：每个节点里第 k 块 GPU 的网卡接到同一组交换机上，" "让节点间那 8 个环各走各的路）。推理里的张量并行一般放在节点内，就是为了完全避开跨节点的慢链路（05 脚本）。", ]) if __name__ == "__main__": demo1_mapping() demo2_alltoall() demo3_hierarchical()