如何理解 GPU

GPU 是什么？

现代机器学习 GPU（例如 H100、B200）基本上由一组专门执行矩阵乘法的计算核心（称为流式多处理器（Streaming Multiprocessor，SM））与一条高速内存（称为高带宽内存（HBM））相连而成。下图展示了它的结构：

与 TPU 的 TensorCore 类似，每个 SM 都有一个专用的矩阵乘法核心（不巧，它也叫 Tensor Core¹）、一个向量算术单元（称为 Warp Scheduler²），以及一块高速片上缓存（称为 SMEM）。TPU 最多只有 2 个独立的“TensorCore”，而现代 GPU 有 100 多个 SM（H100 上有 132 个）。每个 SM 的能力都远弱于一个 TPU TensorCore，但整个系统更加灵活。各个 SM 几乎完全独立，因此 GPU 可以同时执行数百项不同的任务。³

下面更细致地看一下 H100 的一个 SM：

来源：NVIDIA Hopper GH100 概览。

每个 SM 会被划分为 4 个完全相同的象限，NVIDIA 称之为 SM 子分区（SM subpartition）。每个子分区都包含一个 Tensor Core、16k 个 32 位寄存器，以及一个称为 Warp Scheduler 的 SIMD/SIMT 向量算术单元；NVIDIA 将其各条通道（ALU）称为 CUDA 核心（CUDA Core）。每个分区的核心组件可以说是 Tensor Core，它负责执行矩阵乘法，并贡献绝大部分 FLOPs/s；但它并不是唯一值得关注的组件。

• CUDA 核心：每个子分区都包含一组称为 CUDA 核心的算术逻辑单元（ALU），用于执行 SIMD/SIMT 向量运算。每个 ALU 通常能在每个周期内执行 1 次算术操作，例如 f32.add。⁴ 每个子分区包含 32 个 fp32 核心（另有数量较少的 int32 和 fp64 核心），它们在每个周期内都执行同一条指令。与 TPU 的 VPU 类似，CUDA 核心负责 ReLU、逐点向量运算和归约（求和）。⁵

• Tensor Core（TC）：每个子分区都有自己的 Tensor Core，它是类似于 TPU MXU 的专用矩阵乘法单元。Tensor Core 贡献了 GPU 绝大部分 FLOPs/s（例如在 H100 上，bf16 TC 的算力为 990 TFLOP/s，而 CUDA 核心只有 66 TFLOPs/s）。

  • 990 bf16 TFLOPs/s 配合 132 个以 1.76GHz 运行的 SM，意味着每个 H100 TC 每周期可执行 7.5e12 / 1.76e9 / 4 ~ 1024 次 bf16 FLOPs，大致相当于一次 8x8x8 矩阵乘法。⁶
  • 与 TPU 一样，GPU 能以更高吞吐量执行较低精度的矩阵乘法（例如 H100 的 fp8 FLOPs/s 是 fp16 的 2 倍）。低精度训练或推理服务可以快得多。
  • 从 Volta 开始，每一代 GPU 都比上一代增大了 TC 的规模（可参阅这篇不错的文章）。到了 B200，TC 已经大到其输入无法再装进 SMEM，因此 B200 引入了一种名为 TMEM 的新内存空间。⁷

CUDA 核心比 TPU 的 VPU 更灵活：GPU 的 CUDA 核心（从 V100 开始）采用所谓 SIMT（Single Instruction Multiple Threads，单指令多线程）编程模型，而 TPU 采用 SIMD（Single Instruction Multiple Data，单指令多数据）模型。与 TPU VPU 中的 ALU 一样，同一子分区内的 CUDA 核心必须在每个周期执行相同操作（例如，如果一个核心正在把两个浮点数相加，那么该子分区中的其他所有 CUDA 核心也必须执行加法）。不过，与 VPU 不同的是，每个 CUDA 核心（在 CUDA 编程模型中也称“线程”）都有自己的指令指针，可以被独立地编程。当同一个 warp 中的两个线程被要求执行不同操作时，实际上会把两个操作都执行一遍，同时屏蔽掉不需要执行当前分支操作的核心。

来源：NVIDIA Volta 架构白皮书。

这使得线程级编程非常灵活，但代价是：如果 warp 过于频繁地发生分支发散，性能会悄无声息地下降。线程能够访问的内存也更加灵活；VPU 只能操作连续的内存块，而 CUDA 核心可以访问共享寄存器中的单个浮点数，并维护逐线程状态。

CUDA 核心的调度也更加灵活：SM 的运行方式有点像多线程 CPU，因为它们可以并发“调度”许多程序（即 warp，每个 SM 最多 64 个），但每个 Warp Scheduler 在每个时钟周期只会执行一个程序。⁸ Warp Scheduler 会在活跃的 warp 之间自动切换，以隐藏内存加载等 I/O 操作。相比之下，TPU 通常是单线程的。

内存

除了计算单元之外，GPU 还有一套内存层次结构：最大的是 HBM（GPU 的主内存），然后依次是更小的缓存（L2、L1/SMEM、TMEM 和寄存器内存）。

• 寄存器：在 H100/B200 上，每个子分区都有自己的寄存器文件，其中包含 16,384 个 32 位字（每个 SM 为 4 * 16384 * 4 = 256kiB），CUDA 核心可以访问这些寄存器。

  • 每个 CUDA 核心一次最多只能访问 256 个寄存器，因此，虽然每个 SM 最多可以调度 64 个“驻留 warp”，但如果每个线程都使用 256 个寄存器，就只能同时容纳 8 个（256 * 1024 / (4 * 32 * 256)）。

• SMEM（L1 缓存）：每个 SM 都有自己的 256kB 片上缓存，称为 SMEM；它既可以作为“共享内存”由程序员控制，也可以作为片上缓存由硬件使用。SMEM 用来存储激活值和 TC 矩阵乘法的输入。

• L2 缓存：所有 SM 共享⁹一块相对较大的约 50MB L2 缓存，用来减少对主内存的访问。

  • 它的容量与 TPU 的 VMEM 相近，但速度慢得多，而且不能由程序员控制。这就产生了一点“幽灵般的超距作用”：程序员必须调整内存访问模式，才能确保 L2 缓存得到充分利用。¹⁰
  • NVIDIA 并未公布其芯片的 L2 带宽，但实测约为 5.5TB/s。这大约是 HBM 带宽的 1.6 倍，不过它是全双工的，因此有效双向带宽更接近 3 倍。相比之下，TPU 的 VMEM 容量是其 2 倍，带宽也高得多（约 40TB/s）。

• HBM：GPU 的主内存，用于存储模型权重、梯度、激活值等。

  • 从 Volta 的 32GB 到 Blackwell（B200）的 192GB，HBM 容量增长了很多。
  • 从 HBM 到 CUDA Tensor Core 的带宽称为 HBM 带宽或内存带宽；H100 上约为 3.35TB/s，B200 上约为 9TB/s。

GPU 规格汇总

下面汇总了近期 GPU 型号的规格。同一种 GPU 的不同变体在 SM 数量、时钟速度和 FLOPs 上会略有差异。首先是内存容量数据：

所有代际的每个 SM 都有 256kB 寄存器内存。Blackwell 还为每个 SM 增加了 256kB TMEM。下面是各芯片的 FLOPs 和带宽数据：

这里没有列出 B100，因为它没有大规模量产。¹¹ 有些规格会因 GPU 的具体版本而略有不同，因为 NVIDIA GPU 不像 TPU 那样标准化。

下面这张速查表对比了 GPU 和 TPU 的各个组件：

芯片层面的 GPU 与 TPU 对比

GPU 最初用于渲染电子游戏，但自 2010 年代深度学习兴起以来，它们越来越像专用矩阵乘法机器——换句话说，越来越像 TPU。¹² 在一定程度上，这段历史解释了现代 GPU 为何呈现出如今的样子。它们并非专为 LLM 或机器学习模型而设计，而是通用加速器；其硬件追求某种程度的“通用性”，这既可能是福，也可能是祸。GPU 用于新任务时更常能“直接工作”，对优秀编译器的依赖也远小于 TPU。但这也使 GPU 更难推理、更难达到 Roofline 性能，因为太多编译器特性都可能形成瓶颈。

GPU 更加模块化。TPU 有 1～2 个大型 TensorCore，而 GPU 有数百个小型 SM。类似地，每个 TC 有一个大型 VPU，由 4 个可独立编程的 8x128 单元组成（总计 4096 个 ALU）；相比之下，H100 有 132 * 4 = 528 个独立 SIMD 单元，每个宽度为 32（总计 16k 个 ALU）。下面是一张 GPU 与 TPU 的逐项对比表，突出了这一点：

这种模块化差异一方面让 TPU 的制造成本低得多，也更容易理解；另一方面，它也让编译器承担了更重的责任，必须做出正确选择。TPU 只有单一控制线程，并且只支持 VPU 范围的向量化指令，因此编译器需要手动流水化所有内存加载以及 MXU/VPU 工作，以避免停顿。GPU 程序员则可以直接启动数十个不同的内核，每个内核都在完全独立的 SM 上运行。反过来，这些内核的性能也可能极差，因为它们不断冲刷 L2 缓存，或者未能合并内存访问；由于运行时的许多部分都由硬件控制，理解幕后究竟发生了什么也变得困难。因此，TPU 往往能用更少的工作量，更接近峰值 Roofline 性能。

从历史上看，单个 GPU 比可比的 TPU 更强大（也更昂贵）：单个 H200 的 FLOPs/s 接近 TPU v5p 的 2 倍，HBM 容量是后者的 1.5 倍。与此同时，Google Cloud 上 H200 的标价约为每小时 $10，而 TPU v5p 约为每小时 $4。TPU 通常比 GPU 更依赖于把多个芯片联网协同工作。

TPU 拥有多得多的高速缓存。与 GPU 的 SMEM（加 TMEM）相比，TPU 还拥有多得多的 VMEM。这些内存可以用来存储权重和激活值，使其能够以极高速度被加载和使用。如果能持续把模型权重存入或预取到 VMEM 中，这会让 TPU 在 LLM 推理时更快。

小测 1：GPU 硬件

下面这些题目用于检验前面的内容。答案已经给出，但最好先拿起纸笔，试着自己回答问题，然后再查看答案。

问题 1［CUDA 核心］：H100 有多少个 fp32 CUDA 核心（ALU）？B200 呢？与 TPU v5p 中独立 ALU 的数量相比如何？

问题 2［向量 FLOPs 计算］：单个 H100 有 132 个 SM，时钟速度为 1.59GHz（加速频率最高 1.98GHz）。假设每个 ALU 每周期可以执行一次向量操作。每秒可以执行多少次 fp32 向量 FLOPs？加速频率下呢？与矩阵乘法 FLOPs 相比如何？

问题 3［GPU 矩阵乘法强度］：H100 上 fp16 矩阵乘法的峰值强度是多少？B200 呢？fp8 又是多少？这里的强度是指矩阵乘法 FLOPs/s 与内存带宽之比。

问题 4［矩阵乘法运行时间］：根据问题 3 的答案，你预计单个 B200 上的 fp16[64, 4096] * fp16[4096, 8192] 矩阵乘法要运行多久？fp16[512, 4096] * fp16[4096, 8192] 呢？

问题 5［L1 缓存容量］：H100 的 L1/SMEM 总容量是多少？寄存器内存呢？与 TPU 的 VMEM 容量相比如何？

问题 6［计算 B200 时钟频率］：NVIDIA 在这里报告称，B200 可以执行 80TFLOPs/s 的 fp32 向量计算。已知每个 CUDA 核心在一次 FMA（融合乘加）操作中每周期可执行 2 次 FLOPs，请估算其峰值时钟频率。

问题 7［估算 H100 加法运行时间］：根据上面的数据，计算在单个 H100 上把两个 fp32[N] 向量相加应当需要多长时间。分别计算 $T_\text{math}$ 和 $T_\text{comms}$。该操作的算术强度是多少？如果你能使用相关硬件，也请尝试在 PyTorch 或 JAX 中分别以 N = 1024 和 N=1024 * 1024 * 1024 运行此操作。结果相比如何？

网络互连

网络互连是 GPU 与 TPU 差异最大的领域之一。正如我们已经看到的，TPU 连接成二维或三维环面拓扑，每个 TPU 只与相邻 TPU 直接相连。这意味着，在两个 TPU 之间发送消息时，消息必须经过中间的每一个 TPU，也迫使我们在整个网格上只能使用均匀的通信模式。这在某些方面并不方便，但也意味着每个 TPU 的链路数量保持不变，因此可以扩展到任意大的 TPU“Pod”而不会损失带宽。

GPU 则采用一种更传统的分层树状交换网络。每组 8 个 GPU 称为一个节点（GB200 最多可达 72 个¹⁴），节点内的 GPU 通过名为 NVLink 的高带宽互连在 1 跳之内彼此连接；这些节点再通过连接到每个 GPU 的网络接口卡（NIC），使用带宽较低的 InfiniBand（IB）或以太网连接成更大的单元（称为 SU，即 Scalable Unit）。这些单元还可以借助更高层交换机，进一步连接成任意规模的系统。

节点层级

GPU 节点是一个小型单元，通常包含 8 个 GPU（GB200 最多可达 72 个），这些 GPU 通过全互连、全带宽、低延迟的 NVLink 互连连接。¹⁵ 每个节点都包含多个高带宽 NVSwitch，负责在所有本地 GPU 之间交换数据包。节点层级的实际拓扑随时间发生了很大变化，包括每个节点的交换机数量；不过在 H100 中，每个节点有 4 个 NVSwitch，GPU 按 5 + 4 + 4 + 5 的链路模式与之相连，如下图所示：

在 Hopper 代（NVLink 4.0）中，每条 NVLink 链路都有 25GB/s 全双工¹⁶带宽（B200 为 50GB/s），因此每个 GPU 进入网络的全双工带宽为 18 * 25=450GB/s。巨大的 NVSwitch 最多拥有 64 个 NVLink 端口，这意味着一个带 4 个交换机的 8xH100 节点最多可以处理 64 * 25e9 * 4=6.4TB/s 带宽。下面概览了这些数据随 GPU 代际的变化：

Blackwell（B200）的节点包含 8 个 GPU。GB200NVL72 支持由 72 个 GPU 组成的更大 NVLink 域。这里同时列出了 8 GPU 系统和 72 GPU 系统的详细信息。

小测 2：GPU 节点

下面还有几道关于网络互连的问答题。我觉得这些题特别值得亲手算一遍，因为它们会迫使你实际推演通信模式。

问题 1［H100 节点的总带宽］：在一个带 4 个交换机的 8xH100 节点中，每个节点共有多少带宽？提示：需要同时考虑 NVLink 和 NVSwitch 带宽。

问题 2［对分带宽］：对分带宽的定义，是把网络均匀划分时任意一种划分方式所能获得的最小带宽。换句话说，如果把一个网络分成大小相等的两半，有多少带宽能够跨越这两半？你能计算 8x H100 节点的对分带宽吗？提示：对分带宽通常包括两个方向的流量。

问题 3［AllGather 开销］：给定一个包含 B 字节的数组，在一个 8xH100 节点上，一次（受吞吐量限制的）AllGather 要运行多久？请对 bf16[D_X, F] 进行计算，其中 D=4096、F=65,536。回答之前，值得先阅读 TPU 集合通信的章节。请先自行推演，下一节我们会详细讨论集合通信。

跨越节点层级

跨出节点之后，GPU 网络拓扑就没有那么标准化了。NVIDIA 发布了一套参考 DGX SuperPod 架构，它使用 InfiniBand 连接比单个节点更多的 GPU；但客户和数据中心提供商可以根据自身需求自由定制。¹⁷

下图展示了一套参考的 1024 GPU H100 系统。底行中的每个方框都是一个 8xH100 节点，包含 8 个 GPU、8 个 400Gbps CX7 NIC（每个 GPU 一个）和 4 个 NVSwitch。

可扩展单元：每组 32 个节点称为一个“可扩展单元”（Scalable Unit，简称 SU），位于同一组 8 个 InfiniBand 叶交换机之下。一个 SU 有 256 个 GPU，每个节点配 4 个 NVSwitch，整个 SU 配 8 个 Infiniband 叶交换机。图中所有线缆均为 InfiniBand NDR（50GB/s 全双工），交换机则是 64 端口 NDR IB 交换机（每端口同样为 50GB/s）。请注意，IB 交换机的带宽是 NVSwitch 的 2 倍（64 个端口，使用 400 Gbps 链路）。

SuperPod：整个 SuperPod 再通过 16 个顶层“脊”IB 交换机连接 4 个这样的 SU，从而得到 1024 个 GPU、512 个节点层级 NVSwitch、32 个叶 IB 交换机和 16 个脊 IB 交换机，总计 512 + 32 + 16 = 560 个交换机。叶交换机按每组 32 个节点的方式连接节点，因此每组 256 个 GPU 配有 8 个叶交换机。所有叶交换机都连接到所有脊交换机。

我们有多少带宽？InfiniBand 网络（称为“横向扩展网络”，即 scale-out network）的整体拓扑是一棵胖树；线缆和交换机保证在节点层级之上拥有完整的对分带宽（这里为 400GB/s）。这意味着，如果把节点分成两半，每个节点都可以同时以 400GB/s 向另一分区中的一个节点送出数据。更重要的是，这意味着在横向扩展网络中，AllReduce 带宽应当大致保持恒定！实际实现未必如此，但你可以设想在横向扩展网络中的任意多个节点上执行环形归约，因为总能构造出一个包含所有节点的环。

作为对比，TPU v5p 每条链路约有 90GB/s 出方向带宽，或者说沿三维环面拓扑所有轴合计有 540GB/s 出方向带宽。它并非点到点互连，因此只能用于受限且均匀的通信模式；但它仍能提供高得多的 TPU 到 TPU 带宽，而且可以扩展到任意大的拓扑（至少能达到 8960 个 TPU）。

理论上，通过增加额外交换机或增加间接层级，可以把 GPU 交换结构扩展到任意规模；代价则是额外延迟和昂贵的网络交换机。

要点：在 H100 节点内，每个 GPU 都拥有 450GB/s 的完整胖树带宽；跨越节点之后，这一带宽下降为节点间 400GB/s。事实会证明，这一点对通信原语至关重要。

GB200 NVL72：NVIDIA 最近开始生产新的 GB200 NVL72 GPU 集群，在单个 NVLink 域中组合 72 个 GPU，并提供完整的 900GB/s GPU 到 GPU 带宽。随后可以把这些域连接成更大的 SuperPod，并配备按比例提高（9 倍）的 IB 胖树带宽。其拓扑如下图所示：

计算单个节点的出方向带宽（上图中的橙色线），可以得到通往叶层级的带宽为 4 * 18 * 400 / 8 = 3.6TB/s，是 H100 的 9 倍（节点中的 GPU 数量也恰好是 9 倍）。这意味着，关键的节点出方向带宽要高得多得多，而跨节点集合通信带宽实际上可能低于节点内带宽。 更多讨论请参阅附录 A。

要点：GB200 NVL72 SuperPod 大幅提高了节点规模和单个节点的出方向带宽，从而显著改变了我们的 Roofline。

小测 3：跨越节点层级

问题 1［胖树拓扑］：使用上面的 DGX H100 示意图，计算整个 1024 GPU Pod 在节点层级的对分带宽。证明每条链路的带宽都经过选择，可以确保完整的对分带宽。提示：请务必同时计算链路带宽和交换机带宽。

问题 2［扩展到更大的 DGX Pod］：假设我们想使用 2048 个而非 1024 个 GPU 进行训练。修改上述 DGX 拓扑以支持这种规模时，最简单或最好的办法是什么？4096 个 GPU 时又该怎么做？提示：不存在唯一正确答案，但请设法控制成本，并注意链路容量。这份文档可能会有帮助。

GPU 上的集合通信如何工作？

GPU 可以执行与 TPU 完全相同的集合通信操作：ReduceScatter、AllGather、AllReduce 和 AllToAll。与 TPU 不同，这些操作的工作方式会因其是在节点层级（通过 NVLink）执行，还是在更高层级（通过 InfiniBand）执行而发生变化。NVIDIA 在 NVSHMEM 和 NCCL（读作“nickel”）库中实现了这些集合通信操作。NCCL 的开源代码在这里。NCCL 会根据延迟要求和拓扑采用多种实现（详情）；从这里开始，我们将讨论交换式树状结构上的理论最优模型。

节点内集合通信

AllGather 或 ReduceScatter：在节点层级执行 AllGather 或 ReduceScatter 时，可以像 TPU 一样沿环形拓扑完成，并在每一跳使用完整的 GPU 到 GPU 带宽。以任意顺序排列 GPU，然后利用完整的 GPU 到 GPU 带宽，沿环发送数组的一部分。¹⁸ 每一跳的开销为 $T_\text{hop} = \text{bytes} / (N * \text{GPU egress bandwidth})$，所以总开销为

你会注意到，这与 TPU 上的情况完全相同。对于 AllReduce，可以像往常一样组合一次 RS 和一次 AG，开销翻倍。

如果担心延迟（例如数组非常小），可以执行树状归约：先在每 2 个一组的 GPU 内执行 AllReduce，再扩展到每 4 个、每 8 个一组；这样总共只需 $\log(N)$ 跳，而不是 $N - 1$ 跳，不过总开销仍然相同。

要点：在单个节点内对 B 字节数组执行 AllGather 或 ReduceScatter 的开销约为 $T_\text{comms} = B * (8 - 1) / (8 * W_\text{GPU egress}) \approx B / W_\text{GPU egress}$。理论上，在 H100 上约为 $B / \text{450e9}$，在 B200 上约为 $B / \text{900e9}$。除非启用网内归约，否则 AllReduce 的开销是其 2 倍。

随堂小测 1［AllGather 时间］：使用一个全双工带宽为 450 GB/s 的 8xH100 节点时，AllGather(bf16[B_X, F]) 需要多长时间？令 $B=1024$、$F=16,384$。

AllToAll：节点内 GPU 彼此全互连，这使 AllToAll 变得相当简单——嗯，确实如此。每个 GPU 只需直接向目标节点发送数据。对于节点内的 B 字节数据，每个 GPU 持有 $B / N$ 字节，并向目标节点各发送 $(B / N^2)$ 字节，目标节点共有 $N - 1$ 个，因此总开销为

相比之下，TPU 上的开销是 $B / (4W)$。因此，在单个节点内，理论运行时间可获得 2 倍加速（$B / 4W$ 对比 $B / 8W$）。

对于混合专家模型（Mixture of Experts，MoE），我们经常希望执行一种稀疏或变长（ragged）AllToAll；我们保证输出维度上最多有 $k$ 个非零分片（共 $N$ 个）。也就是说，$T_\text{AllToAll} \rightarrow K[B, N]$，并且每条轴上最多有 $k$ 个非零条目（共 $N$ 个）。其开销会按 $k/N$ 的比例下降，总计约为 $\min(k/N, 1) \cdot B / (W \cdot N)$。在 MoE 中，我们通常会独立、随机地选择非零值，因此实际非零数可能少于 $k$，从而近似得到 $(N-1)/N \cdot \min(k/N, 1) \cdot B / (W \cdot N)$。¹⁹

随堂小测 2［AllToAll 时间］：使用一个单向带宽为 450 GB/s 的 8xH100 节点时，AllToAll_X->N(bf16[B_X, N]) 需要多长时间？如果已知 8 个条目中只有 4 个非零，结果又如何？

要点：在单个节点内的 GPU 上，对 $B$ 字节数组执行 AllToAll 的开销约为 $T_\text{comms} = (B \cdot (8 - 1)) / (8^2 \cdot W_\text{GPU egress}) \approx B / (8 \cdot W_\text{GPU egress})$。对于变长（top-$k$）AllToAll，开销会进一步下降到 $(B \cdot k) / (64 \cdot W_\text{GPU egress})$。

实测结果：下面是在一个 8xH100 节点上测得的 AllReduce 带宽。Algo BW 是实测带宽（字节数 / 运行时间），Bus BW 则按 $2 \cdot W \cdot (8 - 1) / 8$ 计算，理论上衡量的是实际链路带宽。可以看到，我们确实达到了接近 370GB/s；虽然低于 450GB/s，但已相当接近，不过这要到每设备约 10GB 的消息大小才能实现。这意味着，尽管这些估算在理论上正确，却需要很大的消息才能真正达到相应性能。

这是一个实实在在的问题，因为它使我们能够提出的任何理论断言都明显变复杂。例如，即便对大小合理的数组执行 AllReduce——比如 LLaMA-3 70B 的 MLP（大小为 bf16[8192, 28672]，或者在 8 路模型分片下为 bf16[8192, 3584] = 58MB）——与 450GB/s 峰值相比也只能达到约 150GB/s。相比之下，TPU 在消息小得多时就能达到峰值带宽（参见附录 B）。

要点：虽然 NVIDIA 宣称 H100 NVLink 的带宽约为 450GB/s，但实践中很难超过 370 GB/s，因此需要相应调整前面的估算。

网内归约：从 Hopper 代开始，NVIDIA 交换机支持 “SHARP”（Scalable Hierarchical Aggregation and Reduction Protocol），从而可以执行“网内归约”。这意味着网络交换机本身能够执行归约操作，并把结果多路复用或“MultiCast”到多个目标 GPU：

理论上，这几乎能把 AllReduce 开销减半：每个 GPU 可以把数据发给顶层交换机，由交换机自身执行归约并把结果广播到每个 GPU，无需让数据两次从每个 GPU 送出，同时还能降低网络延迟。

请注意，这是精确结果，并没有差一个 $1/N$ 因子：每个 GPU 首先送出 $B \cdot (N - 1) / N$，然后接收其本地分片的部分归约版本（接收 $B/N$），完成归约后再次送出 $B/N$，最后接收完全归约的结果（接收 $B \cdot (N - 1) / N$），最终接收的总量恰好是 $B$ 字节。

然而在实践中，启用 SHARP 后观察到的带宽增幅约为 30%，而非预测的 75%。这只能让有效集合通信带宽达到约 480GB/s，远未接近 2 倍。

要点：理论上，NVIDIA SHARP（多数 NVIDIA 交换机均可用）应能把对 $B$ 字节执行 AllReduce 的开销从约 $2 * B / W$ 降低到 $B / W$。然而实践中，带宽只提升约 30%。由于纯 AllReduce 在 LLM 中相当少见，因此这并不是特别有用。

跨节点集合通信

跨越节点层级之后，开销会微妙一些。在树状结构上执行归约时，可以把它看成从底向上进行：先在节点内归约，再在叶层级归约，最后在脊层级归约，每个层级都采用常规算法。对于 AllReduce 尤其如此：这种方式可以减少总通信数据量，因为在节点层级完成 AllReduce 后，向叶层级送出的只需 $B$ 字节，而不是 $B * N$。

开销有多大？作为一阶近似，由于我们拥有完整的对分带宽，AllGather 或 ReduceScatter 的开销大致等于缓冲区字节数除以节点出方向带宽（H100 上为 400GB/s），而且与树状归约的任何具体细节都无关。

对于上述 H100 网络，$W_\text{node}$ 出方向带宽通常为 400GB/s（每个节点有 8 条向外连接的 400Gbps IB 链路）。最直观的理解方式，是设想在集群中的每个节点上执行环形归约。由于采用胖树拓扑，我们总能在任意两个节点之间构造一个拥有 $W_\text{node}$ 出方向带宽的环，并执行常规归约。节点层级归约（几乎）永远不会成为瓶颈，因为它拥有更高的总带宽和更好的延迟；不过一般而言，其开销为

其他集合通信：除非启用 SHARP，否则 AllReduce 的开销仍然是上述数值的 2 倍。NVIDIA 也销售支持 SHARP 的 IB 交换机，不过并非所有提供商都会部署。跨节点时，AllToAll 的变化要大得多，因为它不像 AllReduce 那样具有“分层”特性。如果希望把数据从每个 GPU 发送到其他每个 GPU，就无法充分利用节点层级的完整对分带宽。这意味着，如果一次 N 路 AllToAll 跨越 $M = N / 8$ 个节点，其开销为

其有效带宽是 50GB/s，而不是 400GB/s。在单个 H100 节点内，开销为 $B / (8 * \text{450e9})$；跨越 2 个节点时则变为 $B / (2 \cdot \text{400e9})$，性能下降超过 4 倍。

下面汇总了 1024-GPU DGX H100 SuperPod 的架构：

这里使用“集合通信带宽”一词，表示 GPU 或节点能够向外发送数据的有效带宽。它也等于 $\text{bisection bandwidth} * 2 / N$。

要点：在节点层级之上，对 B 字节执行 AllGather 或 ReduceScatter 的开销约为 $B / W_\text{node egress}$；在 H100 DGX SuperPod 上即为 $B / \text{400e9}$。除非启用 SHARP，否则 AllReduce 的开销是其两倍。整体拓扑是一棵胖树，旨在让任意两对节点之间都能获得恒定带宽。

数组沿另一条轴分片时的归约：考虑如下归约的开销：

这里，我们正沿 X 对一个本身又沿另一条轴 $Y$ 分片的数组执行 AllReduce。在 TPU 上，由于每条轴发送的数据量都是原来的 $1 / Y$，所以与未分片版本相比，这项操作的总开销会缩小 $1 / Y$ 倍。在 GPU 上，开销取决于哪条轴是“内轴”（节点内还是节点间），以及每个分片是否跨越不止一个节点。假设 $Y$ 是内轴，并且数组总共有 $\text{bytes}$ 字节，则总开销实际上会缩小 $Y$ 倍，但这只在 $Y$ 跨越多个节点时成立：

其中 N 是 GPU 数量，$D_\text{node}$ 仍是一个节点中的 GPU 数量（即节点的度）。可以看到，当 $Y < D_\text{node}$ 时，我们在节点层级获得了收益，但总运行时间通常不会减少；当 $Y > D_\text{node}$ 时，则会获得与跨越节点数成正比的加速。

如果要精确考虑环形归约，对于树状 AllGather_X(A_Y { U_X })（假设 Y 是内轴），一般规律为

其中，$S_i$ 是 M * N * …，也就是树中第 i 层以下子节点的规模。粗略来说，这表示我们跨越的 GPU 或节点越多，可用带宽就越高，但这种提升只发生在相应节点内。

随堂小测 3［沿两条轴分片］：假设要在单个 SU（256 个芯片）上执行 $\text{AllGather}_X(\text{bf16}[D_X, F_Y])$，其中 $Y$ 是内轴。用 $D$、$F$ 和 $Y$ 表示时，这项操作需要多长时间？

要点：当沿多条轴进行分片时，外层归约的开销会除以内轴所跨越的节点数。

小测 4：集合通信

问题 1［SU AllGather］：只考虑一个包含 M 个节点、每节点有 N 个 GPU 的 SU。在一次 AllGather 期间，节点层级交换机准确地接收和送出了多少字节？顶层交换机呢？

问题 2［单节点 SHARP AR］：考虑一个每节点有 N 个 GPU 的单节点。使用 SHARP（网内归约）执行 AllReduce 时，交换机准确地接收和送出了多少字节？

问题 3［跨节点 SHARP AR］：考虑一个在单节点 N 个 GPU 上分片的数组 bf16[D_X, F_Y]。AllReduce(bf16[D, F_Y] { U_X }) 需要多长时间？可以假设执行网内归约。请解释跨越多个节点后会有何不同。

问题 4［脊层级 AR 开销］：考虑与上面相同的设置，但令 $Y = 256$（因此 AR 发生在脊层级）。AllReduce 需要多长时间？同样，可以假设执行网内归约。

问题 5［2 路 AllGather 开销］：计算恰好跨越 2 个节点、对 $B$ 字节执行 AllGather 的精确开销。务必计算精确开销而不是近似值，并同时考虑节点内和跨节点开销。

GPU 上 LLM 扩展的 Roofline

现在，让我们来看看前面的所有内容最终是为了什么：理解 GPU 上 LLM 扩展的 Roofline。这部分是对这里的 TPU 训练章节的补充。和在那里一样，我们的目标是考察不同并行策略下的总 $T_\text{math}$ 和 $T_\text{comms}$，并理解 $T_\text{comms} > T_\text{math}$ 会在什么条件下发生。与之前一样，我们只考虑具有以下运算的 MLP 块：

其中 $B$ 是以 token 计的全局批大小（即 $B = \text{batch size} \cdot \text{sequence length}$）。

这里我们会再次给出上面的表格，展示 GPU 和节点层级的有效带宽：

注意： GPU 和节点的出方向带宽都会决定 LLM 的 Roofline。我们将使用 $W_\text{collective}$ 这个术语来表示 GPU 或节点带宽，具体取决于我们是在节点内部还是在节点层级以上运行。

让我们像分析 TPU 时那样，考察数据并行、张量并行、流水线并行、专家并行及其组合的计算—通信 Roofline。在本节其余部分的具体计算中，我们将重点关注 H100 的 Roofline。GB200-NVL72 的总体 Roofline 相同，但由于其节点出方向带宽更大，我们有时可能反而会在节点层级遇到瓶颈。

数据并行

如前所述，DP 和 ZeRO 分片在反向传播中需要进行权重 AllReduce，或者一次 ReduceScatter 加一次 AllGather。由于二者开销相同，对于不使用网内归约的纯数据并行或 FSDP，要做到计算受限，在大小为 X 的轴上，反向传播期间每层有：

因此，要满足 $T_\text{math} > T_\text{comms}$，我们需要 $B / (XC) > 1 / W_\text{collective}$，也就是

其中 $W_\text{collective}$ 是 GPU 层级还是节点层级的出方向带宽，取决于我们是在节点内还是跨节点进行分片。因此：

• 在节点内，我们只需要每 GPU 的 token 批大小 > $\text{990e12} / \text{450e9} = 2200$。
• 在 SU 内或脊层级，BS > $\text{990e12} / \text{400e9} = 2475$。

这比 TPU 上高得多；在 TPU 上，使用全部三个轴时，这个数值是 850。例如，在 16000 个 H100 上训练的 LLaMA-3 至少需要 40M token 的批大小（作为参考，他们实际使用的是 16M）。DeepSeek v3 在 2048 个 H800 GPU 上训练，其带宽较低，只有 300GB/s（而不是 H100 上的 450GB/s），因此每个 GPU 需要 $\text{990e12} / \text{300e9} = 3300$ 个 token，也就是总共大约 6.7M（实践中，他们使用的是 4M）。

启用网内归约并使用纯数据并行时，理论上 AllReduce 带宽会提高到 2 倍，从而把这两个数值都减半。但在实践中，收益更接近 30%，这其实只够弥补我们通常难以达到标称数值的差距。此外，由于纯数据并行很少有用，这一点在实践中基本无关紧要。

MoE 模型： 对于一个混合专家（Mixture of Experts，MoE）模型，设有 E 个专家、每个 token 激活 k 个专家，上式变为

这会让每 GPU 的 token 批大小增大 $E/k$ 倍，即

例如，对于新的 OpenAI OSS 模型，其 $k=4$、$E=128$，跨节点时这个数值会增至 32 * 2475 = 79,200，高得有些荒谬。

X 很小时会怎样？ 当我们只进行例如 2 节点数据并行时，会受益于 $(X - 1) / X$ 的缩放关系，从而得到

其中 X 是节点数，且 $N = 8 \cdot X$。那么对于稠密模型，有 $B / N > \alpha \cdot (X - 1) / X$；例如 $B / N > \text{1237}$，只有上述数值的一半。正因如此，你会相当频繁地看到 2 路数据并行。

要点： 假设通信与计算完美重叠、FLOPs 利用率也达到理想水平，数据并行和 ZeRO 分片要在 H100 或 B200 上达到计算受限，每 GPU 需要约 2500 个 token 的批大小。对于 MoE 模型，这一数值会增大 $E / k$ 倍，也就是总参数量与激活参数量之比。只进行少量数据并行时，临界批大小会下降。

张量并行

张量并行需要对激活值执行一次 AllGather 和一次 ReduceScatter，并且需要让这些操作与 MLP 的 FLOPs 重叠。换句话说，在前向传播中，有

要做到计算受限，由此可以得到规则

在节点内，这给出大约 $F / 2200$；超过一个节点后则为 $F / 2475$。对于像 LLaMA-3 那样的 $F=\text{28000}$，这大约是 11 路 TP（或者向下取整，约为 8 路，也就是一个节点所含 GPU 的数量）。与上面一样，当恰好跨越 2 个节点时，我们会额外获得 2 倍带宽，所以一般可以进行 16 路张量并行（$F > 2475 \cdot (Y - 8)$）；理论上，这允许我们达到最高 19 路模型并行。

要点： 对于大小为 Y、前馈维度为 F 的轴，张量并行会在 $Y > F / 2475$ 时变为通信受限，这通常将我们限制在节点内 TP，或者最多 2 节点 TP。

专家并行

正如上面已经指出的，混合专家（MoE）模型的模型权重多了 E 倍，而 FLOPs 只多了 k 倍，这让数据并行变得困难得多。我们可以沿专家维度对权重进行分片，即 W_in[E_Z, D, F]，从而在一定程度上缓解这个问题。为了执行 MLP 块，我们需要引入 2 次 AllToAll，把激活值发送给相应的专家。

如上所述，如果 AllToAll_Z->k([B, D, k]) 跨越多个节点，其开销大约为 $T_\text{AllToAll} = 2 \cdot B \cdot D \cdot (Z-8)/Z \min(8 * k / Z, 1)$，因此对于纯专家并行，我们需要

我们要么需要在 $K > Z/8$ 时满足 $F > \alpha \cdot (Z - 8)/k$，要么需要在 $Z \gg K$ 时满足 $F > 8 \cdot \alpha$，其中 $\alpha = C/W$。这给出了两个可以采用专家并行的区域：一是使用少量专家并行（大约 2 个节点）且 $F$ 较小；二是 $F$ 较大，而 $Z$ 可以任意大（最高可达到 E 路专家并行）。

实践中这两种情况都能见到：要么采用少量专家并行（例如 DeepSeek v3，它的 F 非常小，跨节点专家并行也相对较小且受到限制），要么模型具有较大的 F，此时可以在 TP 之外再进行大量跨节点 EP。

要点： 如果 $F < 8 * C / W_\text{node}$，专家并行可以跨越 1–2 个节点，开销与 TP 相近（但略低）；或者，如果 $F > 8 * C / W_\text{node}$，我们可以用相对较低的开销进行大量专家并行（最高可达 $E$ 个节点）。

流水线并行

流水线并行把不同层拆分到不同节点上，其通信成本极低，因为我们只需每隔几层发送一小批激活值。过去，流水线一直受到“流水线气泡”的困扰，但借助新的零气泡流水线方法，通常可以消除这一问题。

流水线的总体通信成本很小：设有 $N_\text{MB}$ 个微批次和 $N_\text{stages}$ 个阶段，则有 $T_\text{comms per hop} = 2 \cdot B \cdot D / (W \cdot N_\text{MB})$，并且需要经过 $N_\text{MB} + N_\text{stages} - 2$ 跳，因此大致为

由于这里除以了 $N_\text{layers}$，这一成本远小于其他任何成本。换句话说，从通信角度看，流水线基本上是免费的。那为什么不干脆只用流水线呢？原因有几个：

(1) 代码复杂度： 与其他方法相比，流水线不太容易融入自动并行框架（例如 XLA 的 GSPMD）。由于它会引入微批次来隐藏流水线气泡，因此会改变程序结构；而自定义零气泡流水线调度还要求以复杂方式交错执行前向传播与反向传播，进一步加剧了这个问题。

(2) 流水线让数据并行和 FSDP 变得困难： 不采用流水线的最大理由，很可能是它与 FSDP 和数据并行配合得不好。尤其是 ZeRO-3 分片表现很差，因为它要求我们对每个微批次都执行权重 AllGather，而此时只有 $B / N_\text{microbatches}$ 个 token 可用于摊销 AllGather 的成本。此外，在反向传播期间，只有当最后一个微批次通过某个阶段后，我们才能对梯度执行 AllReduce 或 ReduceScatter，这意味着会有大量无法与计算重叠的通信时间。

(3) 流水线气泡与步骤不均衡： 从上面这个（糟糕的）流水线调度可以看出，朴素流水线调度很容易产生大量气泡（也就是浪费的计算）。在上图中，第二个阶段在第 0 步处于空闲状态，第一个阶段从第 2 步到第 3 步处于空闲状态，而第二个阶段在最后一步又再次处于空闲状态。虽然通过谨慎调度可以在一定程度上避免这些气泡，但通常还是会留下一些。我们还必须在关键路径上把激活值从一个阶段传递到下一个阶段，这会增加额外开销：

这些问题各自都有变通办法，但通常实现复杂且难以维护；相较于其他方法，流水线仍然是一种通信成本较低的技术。

关于延迟的提醒： 如前所述，即使消息相当大，GPU 也很难达到完整的 AllReduce 带宽。这意味着，即便从理论上看，我们可以把专家并行 AllToAll 等操作扩展到多个节点，实际也可能连总带宽的 50% 都难以达到。因此，我们确实会尝试把 TP 或 EP 控制在较少的节点内，以尽量降低延迟开销。

示例

DeepSeek 是怎么做的？ 作为参考，DeepSeek V3 使用 2048 个 H800 GPU 训练，并采用：

• 跨越 8 个节点的 64 路专家并行（EP）
• 16 路流水线并行（PP）
• 2 路 ZeRO-1 数据并行（DP）

其稳态批大小为 4096 * 15360 = 62,914,560 个 token，也就是每 GPU 约 30k 个 token。可以看出，这已经相当大了；但其模型也非常稀疏（k=8，E=256），所以确实需要相当大的批大小。采用 64 路 EP 和 16 路 PP 后，总共会得到 1024 路模型并行，这意味着 AllReduce 在脊层级执行；又因为它只有 2 路，所以在实践中会得到 $2 / (2 - 1) = 2$ 倍的带宽。这也有助于降低最终数据并行 AllReduce 与最后几个流水线阶段重叠时的成本。

LLaMA-3 是怎么做的？ LLaMA-3 在 16k 个 GPU 上以 16M token 的 BS 进行训练，即每个 GPU 约 1k 个 token。其配置为：

• 节点内 8 路张量并行（TP）
• 16 路流水线并行（PP）
• 128 路 ZeRO-1 数据并行

这也是一个稠密模型，所以总体而言这些问题都相当简单。16 路 PP 把数据并行 AllReduce 的成本降低了 16 倍，有助于降低临界批大小。

GPU 上 LLM 扩展总结

让我们退一步，对目前所学内容做一个总体总结：

• 数据并行或 FSDP（ZeRO-1/3）要求每 GPU 的局部批大小约为 2500 个 token，不过理论上，网内归约加纯 DP 可以在一定程度上降低这一数值。
• 张量并行在最高约 8 路时是计算受限的，但如果继续扩展，我们会在变成通信受限之前缺少足够的带宽。这基本上把我们限制在单个 NVLink 域内（也就是单节点，或者需要使用最多包含 72 个 GPU 的 GB200NVL72）。
• 任何跨越多个节点的模型并行形式都可以进一步降低 FSDP 的成本，因此我们常常希望混合 PP、EP 和 TP，跨越许多节点并降低 FSDP 成本。
• 如果你能处理零气泡流水线的代码复杂度，并保持相当大的批大小以避免数据并行瓶颈，那么流水线并行会很好用。 流水线通常会让 ZeRO-3 变得不可行（因为你需要在每个流水线阶段执行 AllGather），但可以改用 ZeRO-1。

从高层次看，这为我们给出了在 GPU 上分片大型模型的一套方法：

• 对于相对较小的稠密模型，只要批大小足够，激进的 FSDP 就非常有效；如有需要，可以再加一定程度的流水线并行或张量并行。
• 对于更大的稠密模型，组合使用 1–2 节点 TP、多节点 PP 和纯 DP 会很有效。
• 对于 MoE，上述规则同样适用，不过我们还可以使用专家并行，而且通常比起 TP，我们更偏好专家并行。如果 $F > 8 * C / W_\text{node}$，就可以进行大量多节点专家并行；否则大致只能采用 2 节点 EP。

小测 5：LLM Roofline

问题 1［B200 Roofline］： B200 DGX SuperPod（不是 GB200 NVL72）的节点内带宽是原来的 2 倍（900GB/s 出方向），但横向扩展网络的带宽不变（400GB/s）（来源）。其总 FLOPs 已在上文给出。这会如何改变模型并行和数据并行的 Roofline？

问题 2 [如何对 LLaMA-3 70B 进行分片]： 考虑 LLaMA-3 70B，使用 bfloat16 训练，并为 Adam 使用 fp32 优化器状态。

1. 仅仅为了存储权重和优化器，我们至少需要多少个 H100？
2. 假设我们希望在 4096 个 H100 GPU 上训练 15T 个 token。假设达到了 45% 的 MFU（模型 FLOPs 利用率）。训练需要多长时间？
3. LLaMA-3 70B 的 F = 28,672，训练时使用的批大小约为 4M token。在不受通信限制的前提下，我们最多可以进行多少路模型并行？加上纯 DP 后，能否在 4k 个芯片上训练 LLaMA-3，同时保持计算受限？ZeRO-3 呢？再加上 8 路流水线又如何？注意：需要同时考虑通信成本和 GPU 内存用量。

问题 3 [Megatron-LM 超参数]： 考虑 Megatron-LM 仓库中的这张图，它突出了其很高的 MFU 数值。

请注意，所有配置的序列长度均为 4096。对于 16B、70B 和 314B 模型，每 GPU 的 token 批大小是多少？假设数据并行是最外层轴，并假设使用 bfloat16 归约，请判断各配置理论上是计算受限还是通信受限，以及是否存在更优的配置？

致谢与延伸阅读

本章在很大程度上得益于许多知识渊博的 GPU 专家的帮助，包括：

• Adam Paszke，他帮助解释了 GPU 内核编程的现实情况。
• Swapnil Patil，他最先向我解释了 GPU 网络的工作方式。
• Stas Bekman，他指出 GPU 的实测情况往往与宣称的规格不同。
• Reiner Pope，他帮助澄清了 GPU 和 TPU 在硬件层面的比较方式。
• Frédéric Bastien，他对芯片层面的叙述提供了详细反馈。
• Nouamane Tazi，他在 GPU 上训练 LLM 的经验帮助改进了 Roofline 一节。
• Sanford Miller，他帮助我理解 GPU 的联网方式，以及 NVIDIA 的规格与实际部署情况之间的差异。

关于 GPU，有大量优秀的阅读材料；以下是我最喜欢的一些：

• SemiAnalysis’ History of the NVIDIA Tensor Core：一篇精彩的文章，讲述 GPU 如何从电子游戏引擎转变为 ML 加速器。
• SemiAnalysis’ Analysis of Blackwell Performance：值得阅读，有助于理解 NVIDIA GPU 的下一代产品。
• H100 DGX SuperPod Reference：内容枯燥，但对于理解大型 GPU 集群如何联网很有帮助。这里还有一份关于 GB200 系统的类似文档。
• Hot Chips Talk about the NVLink Switch：一份介绍 NVLink 和 NCCL 集合通信的有趣材料，尤其涵盖了网内归约。
• DeepSeek-V3 Technical Report：一份大型半开放 LLM 训练报告的优秀示例，描述了其分片配置的选择方式。
• How to Optimize a CUDA Matmul：一篇很棒的博客，介绍如何使用 CUDA Cores 实现高效矩阵乘法，并着眼于 GPU 上的缓存一致性。
• HuggingFace Ultra-Scale Playbook：一份 GPU 上 LLM 并行的指南，本章的部分灵感来自它。
• Making Deep Learning Go Brrrr From First Principles：一篇更聚焦于 GPU 和 PyTorch 的 LLM Roofline 与性能工程教程。
• Cornell Understanding GPU Architecture site：一份与本书类似的指南，更具体地比较了 GPU 和 CPU 的内部结构。

附录 A：使用 GB200 时有何变化？

Blackwell 引入了许多重大的网络变化，包括总体 NVLink 带宽翻倍（900GB/s）的 NVLink 5。B200 和 H100 一样，仍采用 8-GPU 节点；但 GB200 系统（将 B200 GPU 与 Grace CPU 结合）引入了大得多的 NVLink 域（NVL72 中为 72 个 GPU，理论上最多可达 576 个）。更大的 NVLink 域实际上也提高了节点出方向带宽，从而降低节点层级以上的集合通信成本。

节点内带宽从 450GB/s 增至 900GB/s，但这并没有带来太大差异，因为每个 GPU 的总 FLOPs/s 也翻了一倍。我们的 Roofline 基本保持不变，不过 NVLink 的带宽高得多，使得专家并行变得更容易。

在节点之外，变化更大。下面是一张来自这里的 SuperPod 示意图。

如图所示，每节点出方向带宽增至 4 * 18 * 400 / 8 = 3.6TB/s，高于 H100 的 400GB/s。由于每芯片 FLOPs 也翻了一倍，这会让有效的跨节点 Roofline 改善约 4 倍。现在，我们可能要开始担心瓶颈会不会出现在节点层级，而不是横向扩展层级。

Grace Hopper： NVIDIA 还销售把若干 GPU 与 Grace CPU 配对的 GH200 和 GB200 系统。例如，一个 GH200 包含 1 个 H200 和 1 个 Grace CPU，而一个 GB200 系统包含 2 个 B200 和 1 个 Grace CPU。这种系统的一个优势是，CPU 通过全带宽 NVLink 连接（称为 NVLink C2C）与 GPU 相连，因此 CPU 到 GPU 的带宽非常高，适合把参数卸载到主机 RAM。换句话说，对于任意一个给定 GPU，访问主机内存的带宽与访问另一个 GPU 的 HBM 相同。

附录 B：更多网络互连细节

下面是一张 NVLink 4 交换机的示意图。它总共有 64 个 NVLink4 端口（每个端口使用 2 条物理通道），还有一个负责通道间交换的大型交叉开关。相比之下，TPU 使用带反射镜的光交换机，可以动态重新配置。

在每个层级上，可用链路带宽或交换机总带宽都可能成为瓶颈。

• 节点层级： 在节点层级，我们有 4 * 1.6TB/s = 6.4TB/s 的 NVSwitch 带宽，但 8 个 GPU 中的每一个都只能以 450GB/s 的速度向交换机传出数据，这意味着节点内的峰值带宽实际上是 450e9 * 8 = 3.6TB/s（全双工）。
• SU/叶层级： 在 SU 层级，有 8 个交换机以全互连方式连接 32 个节点，使用 1x400 Gbps InfiniBand。由此得到节点的出方向带宽为 8 * 32 * 400 / 8 = 12.8TB/s，而交换机层级有 8 * 1.6TB/s = 12.8TB/s，因此两者恰好一致。
• 脊层级： 在脊层级，有 16 个交换机通过 2x400 Gbps 链路连接 32 个叶交换机，因此出方向带宽为 32 * 16 * 400 * 2 / 8 = 51.2TB/s。与叶交换机不同，脊交换机的全部 64 个端口都朝向下层，因此每个交换机可以传输 64 * 400 / 8 = 3.2TB/s 的流量，在交换机层级总计得到 16 * 3.2TB/s = 51.2TB/s，再次恰好一致。

平均到每个 GPU，这意味着节点层级的 GPU 到 GPU 带宽为 450GB/s，而 SU 和脊层级均为 50GB/s。

GPU 实测 AR 带宽：

TPU v5p 带宽（1 个轴）：

下面还有 AllGather 带宽：

关于 AllToAll 成本的更多信息：

这里可以比较近似值 $\min(K / Z) * (Z - 1) / Z$ 与真实值 $(1 - ((Z - 1) / Z) ** K) * (Z - 1) / Z$。除 $Z$ 很小时外，两者很相似。