🔗 英文原文： https://jax-ml.github.io/scaling-book/gpus/

✍️ 翻译：北极的树

微信公众号

如何理解 GPU

我们在 Google 热爱 TPU，但 GPU 也很棒。本章深入探讨 NVIDIA GPU 的世界——每个芯片如何工作，它们如何互联，以及这对 LLM 意味着什么，特别是与 TPU 相比。本节建立在第 2 章和第 5 章的基础上，因此建议您先阅读它们。

作者

所属机构

^†Google DeepMind

^*MatX

发布日期

2025 年 8 月 18 日

什么是 GPU？

现代的机器学习 GPU（例如 H100、B200）基本上是一堆专门用于矩阵乘法的计算核心（称为流式多处理器或 SM），连接到一块高速内存（称为 HBM）。下图是一个示意图：

图：展示 H100 或 B200 GPU 抽象布局的示意图。H100 有 132 个 SM，而 B200 有 148 个。我们宽泛地使用术语“Warp 调度器”来描述一组 32 个 CUDA SIMD 核心以及向它们分派工作的调度器。请注意这与 TPU 的相似程度！

每个 SM，就像 TPU 的 TensorCore 一样，都有一个专用的矩阵乘法核心（不幸的是也叫 Tensor CoreGPU 的 Tensor Core 是 SM 的矩阵乘法子单元，而 TPU 的 TensorCore 是包含 MXU、VPU 和其他组件的总称单元。）、一个向量算术单元（称为 Warp 调度器NVIDIA 对此没有一个好的命名，所以我们只是在几个糟糕的选项中选择了最好的一个。Warp 调度器主要是向一组 CUDA 核心分派工作的单元，但我们在这里用它来描述控制单元和它所控制的核心集合。）和一个快速的片上缓存（称为 SMEM）。与 TPU 最多只有 2 个独立的“Tensor Core”不同，现代 GPU 有超过 100 个 SM（H100 上有 132 个）。每个 SM 的功能远不及一个 TPU TensorCore，但整个系统更加灵活。每个 SM 或多或少是完全独立的，因此一个 GPU 可以同时执行数百个独立的任务。虽然 SM 是独立的，但为了达到峰值性能，它们通常被迫进行协调，因为它们都共享一个容量有限的 L2 缓存。

让我们更详细地看一下 H100 的 SM：

图：H100 SM 的示意图（来源），展示了 4 个*子分区*，每个分区包含一个 Tensor Core、Warp 调度器、寄存器文件和不同精度的 CUDA 核心组。底部的“L1 数据缓存”是 256kB 的 SMEM 单元。B200 看起来类似，但增加了大量的 Tensor Memory (TMEM) 来为庞大的 Tensor Core 提供数据。

每个 SM 分为 4 个相同的象限，NVIDIA 称之为 SM 子分区，每个子分区包含一个 Tensor Core、16k 个 32 位寄存器和一个称为 Warp 调度器的 SIMD/SIMT 向量算术单元，其通道（ALU）NVIDIA 称为 CUDA 核心。每个分区的核心组件可以说是 Tensor Core，它执行矩阵乘法并贡献了绝大部分的 FLOPs/s，但它并不是唯一值得注意的组件。

CUDA 核心：每个子分区包含一组称为 CUDA 核心的 ALU，用于执行 SIMD/SIMT 向量算术。每个 ALU 通常每个周期可以执行 1 次算术操作，例如 f32.add。较新的 GPU 支持 FMA（融合乘加）指令，理论上每个周期执行两个 FLOPs，NVIDIA 无情地利用这一事实将其报告的规格翻倍。每个子分区包含 32 个 fp32 核心（以及数量较少的 int32 和 fp64 核心），它们在每个周期内都执行相同的指令。与 TPU 的 VPU 类似，CUDA 核心负责执行 ReLU、逐点向量操作和归约（求和）。历史上，在引入 Tensor Core 之前，CUDA 核心是 GPU 的主要组件，用于渲染，包括光线-三角形相交和着色。在今天的游戏 GPU 上，它们仍然承担大部分渲染工作，而 TensorCore 用于上采样（DLSS），这使得 GPU 可以在较低分辨率下渲染（像素越少 = 工作量越少），然后使用机器学习进行上采样。
Tensor Core (TC)：每个子分区都有自己的 Tensor Core，它是一个专用的矩阵乘法单元，类似于 TPU 的 MXU。Tensor Core 占 GPU FLOPs/s 的绝大部分（例如，在 H100 上，我们有 990 bf16 TC TFLOP/s，而 CUDA 核心只有 66 TFLOPs/s）。
- 990 bf16 TFLOPs/s，在 132 个 SM 以 1.76GHz 运行时，意味着每个 H100 TC 可以执行 7.5e12 / 1.76e9 / 4 ~ 1024 bf16 FLOPs/周期，大约相当于一个 8x8x8 的矩阵乘法。NVIDIA 没有分享很多 TC 硬件细节，所以这更多的是猜测而非确切事实——当然，这并不能说明 TC 是如何实现的。我们知道 V100 每个 TC 每周期可以执行 256 FLOPs。A100 可以做到 512，H100 可以做到 1024，而 B200 的细节尚未公布，但似乎可能是 2048 FLOPs/TC/周期，因为 `2250e12 / (148 * 4 * 1.86e9)` 大约是 2048。更多细节在这里得到确认。
- 与 TPU 类似，GPU 可以以更高的吞吐量执行较低精度的矩阵乘法（例如，H100 的 fp8 FLOPs/s 是 fp16 的 2 倍）。低精度训练或服务可以显著加快速度。
- 自 Volta 以来的每一代 GPU 都增加了 TC 的尺寸（关于这一点的好文章）。到了 B200，TC 已经变得非常大，以至于其输入无法再装入 SMEM，因此 B200 引入了一个名为 TMEM 的新内存空间。在 Ampere 中，Tensor Core 可以由单个 warp 提供数据，而在 Hopper 中则需要一个完整的 SM（warpgroup），在 Blackwell 中则由 2 个 SM 提供数据。在 Blackwell 中，矩阵乘法也变得如此之大，以至于参数（特别是累加器）不再适合寄存器内存/SMEM，因此 Blackwell 增加了 TMEM 来解决这个问题。

CUDA 核心比 TPU 的 VPU 更灵活： GPU 的 CUDA 核心（自 V100 起）使用所谓的 SIMT（单指令多线程）编程模型，而 TPU 使用的是 SIMD（单指令多数据）模型。与 TPU VPU 中的 ALU 类似，一个子分区内的 CUDA 核心必须在每个周期内执行相同的操作（例如，如果一个核心在做两个浮点数相加，那么该子分区中的所有其他 CUDA 核心也必须这样做）。然而，与 VPU 不同的是，每个 CUDA 核心（或在 CUDA 编程模型中称为“线程”）都有自己的指令指针，并且可以被独立编程。当同一 warp 中的两个线程被指令执行不同的操作时，你实际上会执行两种操作，并屏蔽掉那些不需要执行分歧操作的核心。

图：一组线程内 warp 分歧的示例（来源）。白色区域表示至少一部分物理 CUDA 核心的停顿

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

CUDA 核心调度也更加灵活： SM 的运行方式有点像多线程 CPU，因为它们可以并发地“调度”许多程序（warps）（每个 SM 最多 64 个），但每个 Warp 调度器 在每个时钟周期只执行一个程序。在给定 SM 上调度的 Warp 称为“驻留”。 Warp 调度器会自动在活动 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 共享技术上，L2 缓存被分成两半，因此在 H100 上，一半的 SM 可以各自访问 25MB。有一个连接这两半的链接，但带宽较低。一个相对较大的约 50MB 的 L2 缓存，用于减少主内存访问。
- 这在大小上与 TPU 的 VMEM 相似，但它慢得多，并且不受程序员控制。这导致了一些“超距幽灵作用”，程序员需要修改内存访问模式以确保 L2 缓存得到良好利用。L2 缓存被所有 SM 共享这一事实，实际上迫使程序员以一种相当协调的方式运行 SM，尽管原则上它们是独立的单元。
- NVIDIA 没有公布其芯片的 L2 带宽，但据测量约为 5.5TB/s。这大约是 HBM 带宽的 1.6 倍，但它是全双工的，所以有效的双向带宽接近 3 倍。相比之下，TPU 的 VMEM 大 2 倍并且带宽要大得多（约 40TB/s）。
HBM： 主 GPU 内存，用于存储模型权重、梯度、激活值等。
- HBM 的大小从 Volta 的 32GB 大幅增加到 Blackwell (B200) 的 192GB。
- 从 HBM 到 CUDA Tensor Core 的带宽称为 HBM 带宽或内存带宽，在 H100 上约为 3.35TB/s，在 B200 上约为 9TB/s。

GPU 规格摘要

以下是近期 GPU 型号的规格摘要。不同版本的 GPU 的 SM 数量、时钟速度和 FLOPs 会有所不同。以下是内存容量数据：

GPU	代	时钟速度	SM/芯片	SMEM 容量/SM	L2 容量/芯片	HBM 容量/芯片
V100	Volta	1.25GHz/1.38HGz	80	96kB	6MB	32GB
A100	Ampere	1.10GHz/1.41GHz	108	192kB	40MB	80GB
H100	Hopper	1.59GHz/1.98GHz	132	256kB	50MB	80GB
H200	Hopper	1.59GHz/1.98GHz	132	256kB	50MB	141GB
B200	Blackwell	?	148	256kB	126MB	192GB

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

GPU	代	HBM 带宽/芯片	FLOPs/s/芯片 (bf16/fp16)	FLOPs/s/芯片 (fp8/int8)	FLOPs/s/芯片 (fp4)
V100	Volta	9.0e11	—	—	—
A100	Ampere	2.0e12	3.1e14	6.2e14	—
H100	Hopper	3.4e12	9.9e14	2.0e15	—
H200	Hopper	4.8e12	9.9e14	2.0e15	—
B200	Blackwell	8.0e12	2.3e15	4.5e15	9.0e15

我们排除了 B100，因为它没有大规模生产。虽然 NVIDIA 制造了 B100 代，但它们只短暂销售和生产，据称是由于设计缺陷导致它们无法接近其声称的规格运行。它们在不因散热和功耗问题而降频的情况下难以达到峰值 FLOPs。一些规格会因 GPU 的具体版本而略有不同，因为 NVIDIA GPU 不像 TPU 那样标准化。

这是一个有用的 GPU 和 TPU 组件对比备忘单：

GPU	TPU	它是什么？
流式多处理器 (SM)	TensorCore	包含其他单元的核心“单元”
Warp 调度器	VPU	SIMD 向量算术单元
CUDA 核心	VPU ALU	SIMD ALU
SMEM (L1 缓存)	VMEM	快速片上缓存内存
Tensor Core	MXU	矩阵乘法单元
HBM (又名 GMEM)	HBM	高带宽大容量内存

芯片层面上的 GPU 与 TPU 对比

GPU 最初是用于渲染视频游戏的，但自从深度学习在 2010 年代兴起以来，它们越来越像专用的矩阵乘法机器——换句话说，越来越像 TPU。在深度学习热潮之前，GPU（“图形处理单元”）做的是，嗯，图形处理——主要是为了视频游戏。视频游戏用数百万个小三角形来表示物体，游戏将这些三角形渲染（或“光栅化”）成一个二维图像，每秒在屏幕上显示 30-60 次（这个频率称为帧率）。光栅化涉及将这些三角形投影到相机的坐标系中，并计算哪些三角形与哪些像素重叠，每秒数十亿次。可以想象，这是非常昂贵的，而这仅仅是开始。然后你必须通过组合可能与光线相交的几个半透明三角形的颜色来为每个像素着色。GPU 被设计用来极快地执行这些操作，并着眼于通用性；你需要同时运行许多不同的 GPU 工作负载（称为“着色器”），而没有任何单一操作占主导地位。因此，面向消费者的图形 GPU 可以进行矩阵乘法，但这并不是它们的主要功能。在某种程度上，这段历史解释了为什么现代 GPU 是现在这个样子。它们并非纯粹为 LLM 或 ML 模型设计，而是作为通用加速器，硬件追求一定程度的“通用性”，这既是福也是祸。GPU 在应用于新任务时更常“开箱即用”，并且对优秀编译器的依赖远低于 TPU。但这也使得它们更难推理或获得Roofline性能，因为太多的编译器特性可能导致瓶颈。

GPU 更加模块化。 TPU 有 1-2 个大的 TensorCore，而 GPU 有数百个小的 SM。同样，每个 TensorCore 有 4 个大的 VPU，每个 VPU 有 1024 个 ALU，而 GPU 的 H100 有 132 * 4 = 528 个小的独立 SIMD 单元。以下是 GPU 与 TPU 的 1:1 比较，突出了这一点：

GPU	TPU	H100 #	TPU v5p #
SM (流式多处理器)	TensorCore	132	2
Warp 调度器	VPU	528	8
SMEM (L1 缓存)	VMEM	32MB	128MB
寄存器	向量寄存器 (VRegs)	32MB	256kB
Tensor Core	MXU	528	8

这种模块化上的差异一方面使得 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 有更多的快速缓存内存。 TPU 的 VMEM 也比 GPU 的 SMEM (+TMEM) 多得多，这些内存可以用来存储权重和激活值，使得它们可以被极快地加载和使用。如果能够持续地将模型权重存储或预取到 VMEM 中，这可以使它们在 LLM 推理方面更快。

测验 1：GPU 硬件

这里有一些练习题，用于测试上述内容。提供了答案，但在查看答案之前，最好手持纸笔尝试回答问题。

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

点击此处查看答案。

答案： 一个 H100 有 132 个 SM，每个 SM 有 4 个子分区，每个子分区包含 32 个 fp32 CUDA 核心，所以我们有 132 * 4 * 32 = 16896 个 CUDA 核心。一个 B200 有 148 个 SM，所以总共有 18944 个。一个 TPU v5p 有 2 个 TensorCore（通常通过 Megacore 连接），每个都有一个 VPU，具有 (8, 128) 个通道和每个通道 4 个独立的 ALU，所以有 2 * 4 * 8 * 128 = 8192 个 ALU。这大约是 H100 向量通道数量的一半，运行频率大致相同。

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

点击此处查看答案。

答案： 132 * 4 * 32 * 1.59e9 = 26.9TFLOPs/s。使用 boost 时为 33.5 TFLOPs/s。这是规格表中报告的一半，因为技术上我们可以在一个周期内执行一次 FMA（融合乘加），这算作两个 FLOPs，但这在大多数情况下并不实用。我们可以执行 990 bfloat16 矩阵乘法 TFLOPs/s，所以忽略 FMA，Tensor Core 的 FLOPs/s 大约是其 30 倍。

问题 3 [GPU 矩阵乘法强度]： H100 上的峰值 fp16 矩阵乘法强度是多少？B200 呢？fp8 呢？强度我们指的是矩阵乘法 FLOPs/s 与内存带宽的比率。

点击此处查看答案。

答案： 对于 H100，我们有峰值 990e12 fp16 FLOPs 和 3.35e12 字节/秒的带宽。所以临界强度是 990e12 / 3.35e12 = 295，与 TPU 的 240 相当接近。对于 B200，它是 2250e12 / 8e12 = 281，非常相似。这意味着，与 TPU 类似，我们需要大约 280 的批处理大小才能在矩阵乘法中达到计算密集型。

对于 H100 和 B200，我们的 fp8 FLOPs 都是 2 倍，所以峰值强度也翻倍到 590 和 562，尽管在某种意义上它保持不变，如果我们考虑到我们的权重也可能以 fp8 加载的话。

问题 4 [矩阵乘法运行时]： 使用问题 3 的答案，你预计一个 fp16[64, 4096] * fp16[4096, 8192] 的矩阵乘法在单个 B200 上需要多长时间？fp16[512, 4096] * fp16[4096, 8192] 呢？

点击此处查看答案。

从上面我们知道，当批处理大小低于 281 个 token 时，我们会受通信限制。因此第一个完全受带宽限制。我们读取或写入 2BD + 2DF + 2BF 字节 (2*64*4096 + 2*4096*8192 + 2*64*8192=69e6)，带宽为 8e12 字节/秒，所以大约需要 69e6 / 8e12 = 8.6us。实际上我们可能只能获得总带宽的一部分，所以可能需要接近 10-12us。当我们增加批处理大小时，我们完全受计算限制，所以我们预计 T=2*512*4096*8192/2.3e15=15us。我们同样只期望获得总 FLOPs 的一部分，所以我们可能会看到接近 20us。

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

点击此处查看答案。

答案： 每个 SM 有 256kB 的 SMEM 和 256kB 的寄存器内存，所以每种大约 33MB (132 * 256kB)。加在一起，总共大约 66MB。这大约是现代 TPU 120MB VMEM 的一半，尽管一个 TPU 总共只有 256kB 的寄存器内存！TPU VMEM 的延迟低于 SMEM 的延迟，这也是为什么 TPU 上的寄存器内存不那么关键的原因之一（向 VMEM 的溢出和填充成本很低）。

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

点击此处查看答案。

答案： 我们知道我们有 148 * 4 * 32 = 18944 个 CUDA 核心，所以我们可以执行 18944 * 2 = 37888 FLOPs / 周期。因此 80e12 / 37888 = 2.1GHz，这是一个很高但合理的峰值时钟速度。B200 通常是液冷的，所以更高的时钟周期更合理。

问题 7 [估算 H100 加法运行时]： 使用上面的数据，计算在单个 H100 上将两个 fp32[N] 向量相加需要多长时间。计算 T_\text{math} 和 T_\text{comms}。这个操作的算术强度是多少？如果你能访问到，也尝试在 PyTorch 或 JAX 中对 N = 1024 和 N=1024 * 1024 * 1024 运行这个操作。结果如何？

点击此处查看答案。

答案： 首先，将两个 fp32[N] 向量相加执行 N FLOPs，需要加载 4 * N * 2 字节并写回 4 * N 字节，总共 3 * 4 * N = 12N 字节。计算它们的比率，我们得到 总 FLOPs / 总字节数 = N / 12N = 1 / 12，这相当糟糕。

正如我们上面计算的，忽略 FMA，我们可以达到大约 33.5 TFLOPs/s 的 boost。这只有在所有 CUDA 核心都被使用的情况下才能实现。对于 N = 1024，我们最多只能使用 1024 个 CUDA 核心或 8 个 SM，这将花费更长的时间（假设我们受计算限制，大约长 16 倍）。我们还有 3.35e12 字节/秒的内存带宽。因此，我们的峰值硬件强度是 33.5e12 / 3.35e12 = 10。值得注意的是，这个强度在最近几代 GPU 中保持不变。对于 H100s 是 33.5 / 3.5，对于 B200 是 80 / 8。为什么会这样尚不清楚，但这是一个有趣的观察。所以我们将严重受通信限制。因此我们的运行时就是

T = max (T_{comms}, T_{math}) = \frac{12 \cdot N}{3.35e12} = \frac{N}{2.8e11}

对于 N = 65,536，这大约是 0.23us。实际上我们在 JAX 中看到的运行时大约是 1.5us，这没问题，因为我们预计在这里会受到严重的延迟限制。对于 N = 1024 * 1024 * 1024，我们的Roofline大约是 3.84ms，我们看到的是 4.1ms，这很好！

网络

网络是 GPU 和 TPU 差异最大的领域之一。正如我们所见，TPU 以 2D 或 3D 环面连接，每个 TPU 只连接到其邻居。这意味着在两个 TPU 之间发送消息必须经过所有中间的 TPU，并迫使我们只能在网格上使用统一的通信模式。虽然在某些方面不方便，但这也意味着每个 TPU 的链接数量是恒定的，我们可以扩展到任意大的 TPU “pod” 而不损失带宽。

另一方面，GPU 使用更传统的基于交换机的分层树状网络。一组 8 个 GPU 称为节点（对于 GB200 最多 72 个术语“节点”是重载的，可以指两件事：NVLink 域，即通过 NVLink 互连完全连接的 GPU 集合，或者连接到单个 CPU 主机的 GPU 集合。在 B200 之前，这两者通常是相同的，但在 GB200 NVL72 中，我们有一个包含 72 个 GPU 的 NVLink 域，但每个主机仍然只连接 8 个 GPU。我们在这里使用术语“节点”来指代 NVLink 域，但这有争议。），它们通过称为 NVLink 的高带宽互连在 1 跳内连接，这些节点通过连接到每个 GPU 的 NIC 使用较低带宽的 InfiniBand (IB) 或以太网网络连接成更大的单元（称为 SU 或可扩展单元）。这些单元又可以通过更高级别的交换机连接成任意大的单元。

图：一个典型 H100 网络的示意图。一组 8 个 GPU 通过 NVSwitches（也称为 NVLink 交换机）连接成一个节点或 NVLink 域，这些节点通过交换式 InfiniBand 结构相互连接。在 NVLink 域中，每个 H100 约有 450GB/s 的出口带宽，每个节点有 400GB/s 的出口带宽进入 IB 网络。

节点层面

一个 GPU 节点是一个小单元，通常由 8 个 GPU（对于 GB200 最多 72 个）组成，通过全对全、全带宽、低延迟的 NVLink 互连连接。有人向我描述 NVLink 就像一个增强版的 PCIe 连接，具有低延迟和协议开销，但不是为可扩展性/容错性设计的，而 InfiniBand 更像以太网，专为更大的有损网络设计。每个节点包含几个高带宽的 NVSwitches，用于在所有本地 GPU 之间交换数据包。实际的节点级拓扑随时间变化很大，包括每个节点的交换机数量，但对于 H100，我们每个节点有 4 个 NVSwitches，GPU 以 5 + 4 + 4 + 5 的链接模式连接到它们，如图所示：

图：从 Pascall (P100) 开始的节点（即 NVLink 域）示意图。自 Volta (V100) 以来，我们通过一组交换机在节点内实现了全对全连接。H100 节点有 4 个 NVSwitches，通过 25GB/s 的链接连接到所有 8 个 GPU。

对于 Hopper 代（NVLink 4.0），每个 NVLink 链接具有 25GB/s 的全双工这里的全双工意味着每个方向 25GB/s，两个方向相互独立。你可以在链路上总共发送 50GB/s，但每个方向最多 25GB/s。带宽（B200 为 50GB/s），这使得每个 GPU 进入网络的带宽为 18 * 25=450GB/s 全双工。巨大的 NVSwitches 最多有 64 个 NVLink 端口，这意味着一个带有 4 个交换机的 8xH100 节点最多可以处理 64 * 25e9 * 4=6.4TB/s 的带宽。以下是这些数字随 GPU 代次变化的概述：

NVLink 代	NVSwitch 代	GPU 代	NVLink 带宽 (GB/s, 全双工)	NVLink 端口 / GPU	节点 GPU 间带宽 (GB/s 全双工)	节点大小 (NVLink 域)	NVSwitches / 节点
3.0	2.0	Ampere	25	12	300	8	6
4.0	3.0	Hopper	25	18	450	8	4
5.0	4.0	Blackwell	50	18	900	8/72	2/18

Blackwell (B200) 有 8 个 GPU 的节点。GB200NVL72 支持更大的 72 个 GPU 的 NVLink 域。我们展示了 8 GPU 和 72 GPU 系统的详细信息。

测验 2：GPU 节点

这里有更多关于网络的问题/解答。我发现亲自做这些特别有用，因为它们让你真正理解通信模式。

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

点击此处查看答案。

答案： 我们有 Gen4 4xNVSwitches，每个具有 64 * 25e9=1.6TB/s 的单向带宽。这将在交换机层面给我们 4 * 1.6e12=6.4e12 的带宽。然而，请注意每个 GPU 只能处理 450GB/s 的单向带宽，这意味着我们最多有 450e9 * 8 = 3.6TB/s 的带宽。由于这个数字更小，峰值带宽是 3.6TB/s。

问题 2 [对分带宽]：对分带宽定义为网络任意均分后可用的最小带宽。换句话说，如果将网络分成相等的两半，两半之间有多少带宽？你能计算出 8x H100 节点的对分带宽吗？提示：对分带宽通常包括双向流量。

点击此处查看答案。

答案： 任何均分都会在每半边有 4 个 GPU，每个 GPU 可以向另一半输出 4 * 450GB/s 的流量。考虑到双向流量，这使得跨越分区的字节数为 8 * 450GB/s，即 3.6TB/s 的对分带宽。这与 NVIDIA 报告的相符，例如这里。

问题 3 [AllGather 成本]：给定一个 B 字节的数组，一个（吞吐量受限的）AllGather 在 8xH100 节点上需要多长时间？对 bf16[D_X, F] 进行计算，其中 D=4096，F=65,536。在回答这个问题之前，值得阅读 TPU 集合通信部分。在这里思考一下，但我们接下来会更详细地讨论集合通信。

点击此处查看答案。

答案： 每个 GPU 可以输出 450GB/s，每个 GPU 有 B / N 字节（其中 N=8，节点大小）。我们可以想象每个节点将其字节一个接一个地发送给其他 N - 1 个节点，总共导致 (N - 1) 轮，每轮 T_\text{comms} = (B / (N * W_\text{unidirectional}))，或者 T_\text{comms} = (N - 1) * B / (N * W_\text{unidirectional})。这大约是 B / (N * W_\text{uni}) 或 B / \text{3.6e12}，即对分带宽。

对于给定的数组，我们有 B=4096 * 65536 * 2=512MB，所以总时间是 536e6 * (8 - 1) / 3.6e12 = 1.04ms。这可能是延迟受限的，所以实际上可能需要更长的时间（实际上大约需要 1.5ms）。

超越节点层面

在节点层面之上，GPU 网络的拓扑结构不那么标准化。NVIDIA 发布了一个参考 DGX SuperPod 架构，该架构使用 InfiniBand 连接比单个节点更多的 GPU，但客户和数据中心提供商可以根据自己的需求进行定制。例如，Meta 在一个与此描述显著不同的数据中心网络上训练了 LLaMA-3，该网络使用以太网、一个三层交换结构，并且在顶层有一个超额订阅的交换机。

这是一个参考的 1024 GPU H100 系统的图表，其中底行的每个方框都是一个包含 8 个 GPU、8 个 400Gbps CX7 NIC（每个 GPU 一个）和 4 个 NVSwitches 的 8xH100 节点。

图：参考的 1024 H100 DGX SuperPod 图，包含 128 个节点（有时是 127 个），每个节点有 8 个 H100 GPU，连接到一个 InfiniBand 扩展网络。每 32 个节点（256 个 GPU）的集合称为“可扩展单元”或 SU。叶交换机和脊交换机 IB 交换机为节点间提供了足够的对分带宽。

可扩展单元： 每组 32 个节点称为一个“可扩展单元”（或 SU），隶属于一组 8 个叶 InfiniBand 交换机。这个 SU 有 256 个 GPU，每个节点有 4 个 NVSwitches，还有 8 个 Infiniband 叶交换机。图中所示的所有布线都是 InfiniBand NDR（50GB/s 全双工），配备 64 端口 NDR IB 交换机（每个端口也是 50GB/s）。请注意，IB 交换机的带宽是 NVSwitches 的 2 倍（64 个端口，每个端口 400 Gbps 链接）。

SuperPod： 整个 SuperPod 随后连接 4 个这样的 SU，配备 16 个顶层“脊”IB 交换机，总共给我们 1024 个 GPU，512 个节点级 NVSwitches，32 个叶 IB 交换机和 16 个脊 IB 交换机，总共有 512 + 32 + 16 = 560 个交换机。叶交换机以 32 个节点为一组连接到节点，所以每组 256 个 GPU 有 8 个叶交换机。所有叶交换机都连接到所有脊交换机。

我们有多少带宽？ InfiniBand 网络（称为“扩展网络”）的整体拓扑结构是一个胖树，其电缆和交换机保证了节点级以上的完全对分带宽（这里是 400GB/s）。这意味着如果我们将节点分成两半，每个节点可以同时向另一分区中的节点输出 400GB/s。更重要的是，这意味着我们在扩展网络中应该有大致恒定的 AllReduce 带宽！虽然可能不是这样实现的，但你可以想象在任意数量的扩展网络节点上进行环形规约，因为你可以构建一个包含所有节点的环。

级别	GPU 数量	每单元交换机数	交换机类型	每单元带宽 (TB/s, 全双工)	GPU 间带宽 (GB/s, 全双工)	胖树带宽 (GB/s, 全双工)
节点	8	4	NVL	3.6	450	450
叶	256	8	IB	12.8	50	400
脊	1024	16	IB	51.2	50	400

相比之下，一个 TPU v5p 每个链接约有 90GB/s 的出口带宽，或者在 3D 环面的所有轴向上有 540GB/s 的出口带宽。这不是点对点的，所以只能用于受限的、统一的通信模式，但这仍然给我们提供了更高的 TPU 间带宽，可以扩展到任意大的拓扑结构（至少高达 8960 个 TPU）。

理论上，GPU 交换结构可以通过增加额外的交换机或间接层来扩展到任意大小，但代价是增加延迟和昂贵的网络交换机。

要点：在一个 H100 节点内，我们每个 GPU 有 450GB/s 的全胖树带宽，而在节点之外，这个带宽下降到 400GB/s 的节点间带宽。这对于通信原语来说至关重要。

GB200 NVL72s： NVIDIA 最近开始生产新的 GB200 NVL72 GPU 集群，将 72 个 GPU 组合在一个 NVLink 域中，具有完整的 900GB/s GPU 间带宽。这些域随后可以连接成更大的 SuperPod，具有成比例增加的（9 倍）IB 胖树带宽。以下是该拓扑的图表：

图：一个包含 576 个 GPU 的 GB200 DGX SuperPod 的示意图。底层每个机架包含 72 个 GB200 GPU。

计算单个节点（上图中的橙色线）的出口带宽，我们有 4 * 18 * 400 / 8 = 3.6TB/s 到叶级别的带宽，这比 H100 多 9 倍（正如节点包含 9 倍多的 GPU）。这意味着关键的节点出口带宽要高得多，我们的跨节点集合通信带宽实际上可能比节点内的要低。更多讨论请参见附录 A。

节点类型	每节点 GPU 数	GPU 出口带宽	节点出口带宽
H100	8	450e9	400e9
B200	8	900e9	400e9
GB200 NVL72	72	900e9	3600e9

要点：GB200 NVL72 SuperPods 极大地增加了节点大小和给定节点的出口带宽，这显著改变了我们的Roofline模型。

测验 3：超越节点层面

问题 1 [胖树拓扑]： 使用上面的 DGX H100 图，计算整个 1024 GPU pod 在节点级别的对分带宽。证明每个链接的带宽选择是为了确保完全的对分带宽。提示：确保计算链接带宽和交换机带宽。

点击此处查看答案。

答案： 让我们逐个组件来分析：

首先，每个节点有 8x400Gbps NDR IB 电缆连接到叶交换机，给每个节点 8 * 400 / 8 = 400 GB/s 的到叶的带宽。我们有 8 个叶交换机，每个 3.2TB/s（64 个 400 GBps 链接），但我们只能使用 64 个端口中的 32 个从 SU 输入，所以那是 32 * 400 / 8 = 12.8TB/s 用于 32 个节点，同样是每个节点 400GB/s。
然后在脊级别，我们有 8 * 16 * 2 400Gbps NDR IB 电缆连接每个 SU 到脊，给每个 SU 8 * 16 * 2 * 400 / 8 = 12.8 TB/s 的到叶的带宽。同样，这是每个节点 400GB/s。我们有 16 个脊交换机，每个 3.2TB/s，给我们 16 * 3.2 = 51.2 TB/s，对于 128 个节点，这又是 400GB/s。

因此，如果我们以任何方式对分我们的节点，它们之间将有每个 GPU 400GB/s 的带宽。每个组件都恰好有确保胖树所需的带宽。

问题 2 [扩展到更大的 DGX pod]： 假设我们想在 2048 个 GPU 而不是 1024 个上进行训练。修改上述 DGX 拓扑以处理这个问题的最简单/最好的方法是什么？4096 个呢？提示：没有唯一的正确答案，但尽量降低成本。记住链接容量。这篇文档可能会有帮助。

点击此处查看答案。

答案： 一个选项是保持 SU 结构不变（32 个节点在 8 个交换机下），然后用更多的顶层交换机增加更多的 SU。我们需要 2 倍多的脊交换机，所以我们将有 8 个 SU 和 32 个脊交换机，给我们足够的带宽。

这样做的一个问题是，每个叶交换机只有 64 个端口，我们在上面的图表中已经全部使用了。但实际上，每个脊使用 1x 400 Gbps NDR 电缆而不是 2x 是很容易的，这提供了相同的总带宽但为我们节省了一些端口。

对于 4096 个 GPU，我们实际上用完了端口，所以我们需要增加另一个间接层，也就是说，在层次结构中增加一个级别。NVIDIA 称这些为“核心交换机”，并用 128 个脊交换机和 64 个核心交换机构建一个 4096 GPU 集群。你可以计算一下，这提供了足够的带宽。

集合通信操作在 GPU 上如何工作？

GPU 可以执行与 TPU 相同的所有集合通信操作：ReduceScatters、AllGathers、AllReduces 和 AllToAlls。与 TPU 不同，这些操作的工作方式取决于它们是在节点级别（通过 NVLink）还是在更高级别（通过 InfiniBand）执行。这些集合通信由 NVIDIA 在 NVSHMEM 和 NCCL（发音为“nickel”）库中实现。NCCL 是开源的，在这里。虽然 NCCL 根据延迟要求/拓扑使用多种实现（详情），但从现在开始，我们将讨论一个在交换树结构上的理论最优模型。

节点内集合通信

AllGather 或 ReduceScatter： 对于节点级别的 AllGather 或 ReduceScatter，你可以像 TPU 一样围绕一个环执行它们，在每一跳都使用完整的 GPU 间带宽。任意排序 GPU，并使用完整的 GPU 间带宽将数组的一部分沿环发送。你也可以认为每个 GPU 将其大小为 \text{bytes} / N 的块发送给其他 N - 1 个 GPU，总共通信了 (N - 1) * N * bytes / N 字节，这给了我们每一跳的成本是 T_\text{hop} = \text{bytes} / (N * \text{GPU 出口带宽})，所以总成本是

T_{AG or RS comms} = \frac{bytes \cdot (N - 1)}{N \cdot GPU egress bandwidth} \to \frac{bytes}{GPU egress bandwidth}

你会注意到这和 TPU 上完全一样。对于 AllReduce，你可以像往常一样组合一个 RS + AG，成本是两倍。

图：带宽最优的 1D 环形 AllGather 算法。对于 B 字节，这将在顶层交换机上传输 V / X 字节 X - 1 次。

如果你担心延迟（例如，如果你的数组很小），你可以做一个树形规约，即在 2、4、8 对内进行 AllReduce，总共有 \log(N) 跳而不是 N - 1 跳，尽管总成本仍然相同。

要点：在一个节点内对 B 字节的数组进行 AllGather 或 ReduceScatter 的成本大约是 T_\text{comms} = B * (8 - 1) / (8 * W_\text{GPU egress}) \approxeq B / W_\text{GPU egress}。理论上，在 H100 上这大约是 B / \text{450e9}，在 B200 上是 B / \text{900e9}。除非启用了网络内规约，否则 AllReduce 的成本是这个的两倍。

即时测验 1 [AllGather 时间]：使用一个具有 450 GB/s 全双工带宽的 8xH100 节点，AllGather(bf16[B_X, F]) 需要多长时间？设 B=1024，F=16,384。

点击此处查看答案。

答案： 我们总共有 2 \cdot B \cdot F 字节，单向带宽为 450e9。这大约需要 T_\text{comms} = (2 \cdot B \cdot F) / \text{450e9}，或者更精确地说是 (2 \cdot B \cdot F \cdot (8 - 1)) / (8 \cdot \text{450e9})。使用给定的值，这大约是 (2 \cdot 1024 \cdot 16384) / \text{450e9} = \text{75us}，或者更精确地说是 \text{65us}。

AllToAlls： 节点内的 GPU 具有全对全连接性，这使得 AllToAlls 非常容易。每个 GPU 只需直接发送到目标节点。在一个节点内，对于 B 字节，每个 GPU 有 B / N 字节，并向 N - 1 个目标节点发送 (B / N^2) 字节，总共是

T_{AllToAll comms} = \frac{B \cdot (N - 1)}{W \cdot N^{2}} \approx \frac{B}{W \cdot N}

与 TPU 相比，TPU 的成本是 B / (4W)。因此，在单个节点内，我们在运行时间上获得了 2 倍的理论加速（B / 4W vs. B / 8W）。

对于专家混合 (MoE) 模型，我们经常需要进行稀疏或不规则的 AllToAll，我们保证输出维度上的 N 个分片中最多有 k 个是非零的，也就是说 T_\text{AllToAll}_X \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)。真实成本实际上是 $(1 - {(\frac{Z - 1}{Z})}^{K}) \cdot \frac{Z - 1}{Z}$ 在 K 次掷骰子中不同结果的期望数量，但它非常接近给出的近似值。更多细节请参见附录。

即时测验 2 [AllToAll 时间]：使用一个具有 450 GB/s 单向带宽的 8xH100 节点，AllToAll_X->N(bf16[B_X, N]) 需要多长时间？如果我们知道 8 个条目中只有 4 个是非零的呢？

点击此处查看答案。

答案： 从上面我们知道，在密集情况下，成本是 B \cdot N / (W \cdot N)，或者 B / W。如果我们知道只有 \frac{1}{2} 的条目是非填充的，我们可以发送 B \cdot N \cdot k / (W \cdot N^2) = B \cdot k/N / W = B / (2 \cdot W)，大约是总成本的一半。

要点：在单个节点内的 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) / (864 \cdot W_\text{GPU egress})。

实证测量： 这是一个在 8xH100 节点上 AllReduce 带宽的实证测量。Algo BW 是测量的带宽（字节/运行时），而 Bus BW 计算为 2 \cdot W \cdot (8 - 1) / 8，理论上是实际链接带宽的度量。你会注意到我们确实接近 370GB/s，低于 450GB/s 但相当接近，尽管每个设备只有大约 10GB。这意味着尽管这些估计在理论上是正确的，但需要一个大的消息才能实现它。

图：一个 8xH100 节点（禁用 SHARP）的 AllReduce 吞吐量。蓝色曲线是经验链接带宽，根据经验测量计算为 2 * \text{bytes} * (N - 1) / (N * \text{runtime})。请注意，即使使用巨大的 10GB 数组，我们也没有特别接近声称的 450GB/s 带宽。

这是一个真正的问题，因为它有意义地复杂化了我们可以提出的任何理论主张，因为例如，即使是在一个合理大小的数组上进行 AllReduce，比如 LLaMA-3 70B 的 MLP（大小为 bf16[8192, 28672]，或者 8 路模型分片后为 bf16[8192, 3584] = 58MB），也只能达到大约 150GB/s，而峰值是 450GB/s。相比之下，TPU 在更小的消息大小下就能达到峰值带宽（见附录 B）。

要点：尽管 NVIDIA 声称 H100 NVLink 的带宽约为 450GB/s，但在实践中很难超过 370 GB/s，因此请相应调整上述估计。

网络内规约： 自 Hopper 代以来，NVIDIA 交换机支持“SHARP”（可扩展分层聚合和规约协议），它允许“网络内规约”。这意味着网络交换机本身可以执行规约操作，并将结果多路复用或“多播”到多个目标 GPU：

图：没有 SHARP 的 AllReduce 理论成本是原来的 2 倍，因为它必须两次通过每个 GPU。实际上，速度提升只有大约 30%（来自 NCCL 2.27.5）。

理论上，这几乎将 AllReduce 的成本减半，因为它意味着每个 GPU 可以将其数据发送到一个顶层交换机，该交换机本身执行规约并将结果广播到每个 GPU，而无需两次从每个 GPU 出口，同时也减少了网络延迟。

T_{SHARP AR comms} = \frac{bytes}{GPU egress bandwidth}

请注意，这是精确的，而不是差一个 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 的 AllReduce 算法带宽的实证测量。在峰值时，增益约为 30% 的吞吐量提升，尽管从算法上讲，它应该能够实现接近 75% 的增益。

要点：理论上，NVIDIA SHARP（在大多数 NVIDIA 交换机上可用）应将对 B 字节的 AllReduce 成本从大约 2 * B / W 降低到 B / W。然而，在实践中，我们只看到大约 30% 的带宽提升。由于纯 AllReduce 在 LLM 中相当罕见，这并不是特别有用。

跨节点集合通信

当我们超越节点级别时，成本变得更加微妙。当在树上进行规约时，你可以想象从下往上规约，首先在节点内，然后在叶级别，然后在脊级别，在每个级别都使用常规算法。特别是对于 AllReduce，你可以看到这使我们能够总体上传输更少的数据，因为在节点级别进行 AllReduce 后，我们只需要向叶级别出口 B 字节，而不是 B * N。

这有多昂贵？ 作为一阶近似，因为我们有完全的对分带宽，AllGather 或 ReduceScatter 的成本大约是缓冲区大小（字节）除以节点出口带宽（H100 上为 400GB/s），无论树形规约的任何细节如何。

T_{AG or RS comms} = \frac{bytes}{W_{node egress}} \underset{H 100}{=} \frac{bytes}{400e9}

其中 W_\text{node} 出口带宽对于上述 H100 网络通常是 400GB/s（每个节点出口 8x400Gbps IB 链接）。想象这一点的最清晰方式是想象在集群中的每个节点上进行环形规约。由于胖树拓扑结构，我们总能构建一个在任意两个节点之间具有 W_\text{node} 出口带宽的环，并进行正常的规约。节点级规约（几乎）永远不会是瓶颈，因为它具有更高的总带宽和更好的延迟，尽管一般而言成本是

T_{total} = max (T_{comms at node}, T_{comms in scale-out network}) = max [\frac{bytes}{W_{GPU egress}}, \frac{bytes}{W_{node egress}}]

你可以在这里看到更精确的推导。

我们可以更精确地指出，我们实际上是在网络的每一层进行环形规约，这些规约大部分可以重叠，所以我们有：

T_{AG or RS comms} = bytes \cdot m a x_{depth i} [\frac{D_{i} - 1}{D_{i} \cdot W_{link i}}]

其中 D_i 是深度 i 的度（深度 i 的子节点数量），W_\text{link i} 是连接每个子节点到节点 i 的链接带宽。

使用这个，我们可以计算可用的 AllGather/AllReduce 带宽为 min_\text{depth i}(D_i * W_\text{link i} / (D_i - 1)) 对于给定的拓扑。在上面的情况下，我们有：

节点： D_\text{node} = 8，因为我们节点中有 8 个 GPU，Wlink i = 450GB/s。因此我们的 AG 带宽是 450e9 * 8 / (8 - 1) = 514GB/s。
叶： D_\text{leaf} = 32，因为我们 SU 中有 32 个节点，Wlink i = 400GB/s（8x400Gbps IB 链接）。因此我们的带宽是 400e9 * 32 / (32 - 1) = 413GB/s。
脊： D_\text{spine} = 4，因为我们有 4 个 SU，W_\text{link i} = 12.8TB/s（来自上面的 8 * 16 * 2 * 400Gbps 链接）。我们的带宽是 12.8e12 * 4 / (4 - 1) = 17.1TB/s。

因此我们的整体 AG 或 RS 带宽是 min(514GB/s, 413GB/s, 17.1TB/s) = 413GB/s 在叶级别，所以实际上 T_\text{AG or RS comms} = B / \text{413GB/s}，即即使在最高级别，我们也有大约 413GB/s 的 AllReduce 带宽。对于带 SHARP 的 AllReduce，它会略低于这个值（大约 400GB/s），因为我们没有 (N - 1) / N 因子。尽管如此，450GB/s 和 400GB/s 作为近似值足够接近。

其他集合通信： 除非启用 SHARP，否则 AllReduces 的成本仍然是上述的两倍。NVIDIA 也销售支持 SHARP 的 IB 交换机，尽管并非所有提供商都有。AllToAlls 在跨节点时变化很大，因为它们不像 AllReduces 那样是“分层”的。如果我们想将数据从每个 GPU 发送到每个其他 GPU，我们无法利用节点级别的完全对分带宽。这意味着如果我们有一个跨越 M = N / 8 个节点的 N 路 AllToAll，成本是

T_{AllToAll comms} = \frac{B \cdot (M - 1)}{M^{2} \cdot W_{node egress}} ≊ \frac{B}{M \cdot W_{node egress}}

这实际上只有 50GB/s 而不是 400GB/s 的带宽。我们从单个 H100 节点内的 B / (8 * \text{450e9}) 变成了跨越 2 个节点时的 B / (2 \cdot \text{400e9})，性能下降超过 4 倍。

以下是 1024-GPU DGX H100 SuperPod 架构的摘要：

级别	GPU 数量	度 (# 子节点)	交换机带宽 (全双工, TB/s)	电缆带宽 (全双工, TB/s)	集合通信带宽 (GB/s)
节点	8	8	6.4	3.6	450
叶 (SU)	256	32	25.6	12.8	400
脊	1024	4	51.2	51.2	400

我们使用术语“集合通信带宽”来描述我们可以从 GPU 或节点出口的有效带宽。它也是 \text{对分带宽} * 2 / N。

要点：在节点级别之上，对 B 字节进行 AllGather 或 AllReduce 的成本大约是 B / W_\text{node egress}，在 H100 DGX SuperPod 上是 B / \text{400e9}。整体拓扑是一个胖树，旨在为任意两对节点之间提供恒定的带宽。

当数组在单独的轴上分片时的规约： 考虑像这样的规约成本

{AllReduce}_{X} (A [I_{Y}, J] {U_{X}})

其中我们在一个本身沿另一个轴 Y 分片的数组上进行 AllReduce。在 TPU 上，此操作的总成本比未分片版本降低了 1 / Y 的因子，因为我们每个轴发送的数据量是 1 / Y。在 GPU 上，成本取决于哪个轴是“内部”轴（节点内 vs. 节点间）以及每个分片是否跨越多个节点。假设 Y 是这里的内部轴，总成本有效地降低了 Y，但前提是 Y 跨越多个节点：

T_{comms at node} = \frac{bytes \cdot D_{node}}{min (Y, D_{node}) \cdot W_{GPU egress}}

T_{comms in scale-out network} = \frac{bytes \cdot N}{Y \cdot W_{node egress}}

T_{total} = max (T_{comms at node}, T_{comms in scale-out network})

其中 N 是 GPU 的数量，D 再次是节点中的 GPU 数量（节点的度）。如你所见，如果 Y < D_\text{node}，我们在节点级别获得了胜利，但通常不会看到总运行时间的减少，而如果 Y > D_\text{node}，我们会得到一个与跨越的节点数量成比例的加速。

如果我们想精确地进行环形规约，对于树形 AllGather_X(A_Y { U_X })（假设 Y 是内部轴）的一般规则是

T_{AR or RS comms} = bytes \cdot max_{depth i} [\frac{D_{i} - 1}{D_{i} \cdot max (Y, S_{i - 1}) \cdot W_{link i}}]

其中 S_i 是 M * N * …，即树中 i 级以下子节点的大小。这大致是说，我们跨越的 GPU 或节点越多，我们可用的带宽就越大，但仅限于该节点内。

即时测验 3 [沿 2 个轴分片]： 假设我们想在单个 SU（256 个芯片）上执行 \text{AllGather}_X(\text{bf16}[D_X, F_Y])，其中 Y 是内部轴。这将花费多长时间，作为 D、F 和 Y 的函数？

点击此处查看答案。

答案： 我们可以将其分为两种情况，Y <= 8 和 Y > 8。当 Y <= 8 时，我们仍然受叶交换机的限制，所以答案像往常一样是 T_\text{comms} = 2 * D * F * (32 - 1) / (32 * 400e9)。当 Y > 8 时，我们从上面得到，大约是

T_{comms} = \frac{2 \cdot D \cdot F \cdot 256}{Y \cdot 12.8e12} = \frac{2 D F}{Y \cdot 50GB/s}

对于 D = 8192，F = 32,768，我们有：

请注意，如果我们恰好进行 8 路模型并行，我们确实将节点级规约的成本降低了 8 倍，但总成本保持不变，所以这是免费的，但对提高总带宽没有帮助。

要点：当我们有多个分片轴时，外部规约的成本会因内部轴跨越的节点数量而降低。

测验 4：集合通信

问题 1 [SU AllGather]： 仅考虑一个具有 M 个节点和每个节点 N 个 GPU 的 SU。在 AllGather 期间，节点级交换机精确地入口和出口多少字节？顶层交换机呢？

点击此处查看答案。

答案： 让我们一步一步来，分析规约的组成部分：

每个 GPU 向交换机发送 B / MN 字节，总入口量为 NB / MN = B / M 字节。
我们将全部 B / M 字节出口到脊交换机。
我们从脊交换机入口 B * (M - 1) / M 字节。
我们出口 B - B / MN 字节 N 次，总共是 N * (B - B / MN) = NB - B / M。

总共是 B 入口和 BN 出口，所以我们应该受出口的瓶颈限制，总时间将是 T_\text{AllGather} = BN / W_\text{node} = B / \text{450e9}。

对于脊交换机，数学实际上更简单。我们必须有 B / M 字节入口 M 次（总共 B 字节），然后 B (M - 1) / M 出口 M 次，总共是 B * (M - 1) 出口。由于这个数字明显更大，成本是 T_\text{AllGather} = B \cdot (M - 1) / (M \cdot W_\text{node}) = B \cdot (M - 1) / (M \cdot \text{400e9})。

问题 2 [单节点 SHARP AR]： 考虑一个具有 N 个 GPU 的单节点。在使用 SHARP（网络内规约）进行 AllReduce 期间，交换机精确地入口和出口多少字节？

点击此处查看答案。

答案： 和以前一样，让我们一步一步来。

每个 GPU 发送 B * (N - 1) / N 字节，所以我们有 N * B * (N - 1) / N = B * (N - 1) 字节入口。
我们累加部分和，然后向每个 GPU 发回 B / N 字节，所以 N * B / N = B 字节出口。
我们在本地对残差进行部分求和，然后将其发送回交换机。这总共是 N * B / N = B 字节入口。
我们捕获所有分片并进行多播，向 N 个目的地发送 B * (N - 1) / N，总共是 B * (N - 1) / N * N = B * (N - 1) 字节出口。

因此，总共有 B * (N - 1) + B = BN 字节入口和出口。这支持总吞吐量恰好是 B / W_\text{egress}。

问题 3 [跨节点 SHARP AR]： 考虑一个在单个 N GPU 节点上分片的数组 bf16[D_X, F_Y]。AllReduce(bf16[D, F_Y] { U_X }) 需要多长时间？你可以假设我们进行网络内规约。解释一下如果我们有多个节点，这会有什么不同？

点击此处查看答案。

答案： 我们可以尝试修改上面问题的答案。基本上，我们首先从每个 GPU 出口 B * (X - 1) / XY 字节，然后向每个 GPU 发回 B / XY，然后将相同的量发回交换机，然后向每个 GPU 发回 B * (X - 1) / XY。总共是 NB / Y 入口和出口，所以总时间是 T_\text{comms} = NB / (Y * N * W_\text{link}) = N * 2DF / (Y * N * W_\text{link}) = 2 * D * F / (Y * W_\text{link})，所以总时间确实随着 Y 的增加而减少。

如果我们超越单个节点，我们可以进行与上面大致相同的规约，但是当我们从节点级交换机出口时，我们需要发送所有 B 字节，而不仅仅是 B / Y。这是因为我们需要保持每个分片的分离。

问题 4 [脊级别 AR 成本]： 考虑与上面相同的设置，但是 Y = 256（所以 AR 发生在脊级别）。AllReduce 需要多长时间？同样，可以假设网络内规约。

点击此处查看答案。

答案： 这让我们能够利用脊级别相当惊人的带宽。我们在 4 个节点上有 25.6TB/s 的带宽，所以 AllReduce 带宽是 6.4TB/s。使用 SHARP，这可能只需要 2 * D * F / 6.4e12 秒。

问题 5 [2 路 AllGather 成本]： 考虑在恰好 2 个节点上进行 AllGather 的成本。精确地说，它是什么？确保计算精确成本而不是近似值。

点击此处查看答案。

答案： 在节点级别，我们有 T_\text{comms} = B * 7 / (8 * \text{450e9}) = B / \text{514e9}，而在节点之外，我们实际上有 T_\text{comms} = B * (2 - 1) / (2 * \text{400e9}) = B / \text{800e9}。因此，我们实际上受节点级规约的限制，而不是叶级别！这激励了例如 DeepSeek v3，它进行 2 路数据并行。

GPU 上 LLM 扩展的Roofline模型

现在让我们看看这一切都是为了什么：理解 GPU 上 LLM 扩展的Roofline模型。这是对 TPU 训练章节这里的补充。和那里一样，这里的目标是查看不同并行策略的总 T_\text{math} 和 T_\text{comms}，并理解在什么点上 T_\text{comms} > T_\text{math}。和以前一样，我们只考虑 MLP 块，操作为

MLP (x) \equiv x [B, D] *_{D} W_{in} [D, F] \cdot_{F} W_{out} [F, D]

其中 B 是全局批处理大小（以 token 为单位），即 B = \text{batch size} \cdot \text{sequence length}。

这里我们将重现上面的表格，显示 GPU 和节点级别的有效带宽：

节点类型	每节点 GPU 数	GPU 出口带宽	节点出口带宽
H100	8	450e9	400e9
B200	8	900e9	400e9
GB200 NVL72	72	900e9	3600e9

注意： GPU 和节点的出口带宽都决定了我们 LLM 的Roofline模型。我们将使用术语 W_\text{collective} 来描述 GPU 或节点的带宽，具体取决于我们是在节点内还是节点之上操作。

让我们像为 TPU 做的那样，为数据并行、张量并行、流水线并行、专家并行以及它们的组合，看看计算通信的Roofline模型。本节的其余部分，我们将专注于 H100 的Roofline模型进行具体计算。GB200-NVL72 具有相同的通用Roofline模型，但因为我们有更大的节点出口带宽，我们有时可能会在节点级别遇到瓶颈。

数据并行性

如前所述，DP 和 ZeRO 分片在反向传播中涉及权重 AllReduce 或 ReduceScatter + AllGather。由于这两者的成本相同，要使纯数据并行或 FSDP 没有网络内规约的情况下达到计算密集型，我们每层在反向传播中，对于大小为 X 的轴有：

T_{math} = \frac{2 \cdot 2 \cdot 2 \cdot B D F}{X \cdot C}

T_{comms} = \frac{2 \cdot 2 \cdot 2 \cdot D F}{W_{collective}}

因此，要使 T_\text{math} > T_\text{comms}，我们需要 B / (XC) > 1 / W_\text{collective} 或

\frac{B}{X} > \frac{C}{W_{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）。在 2048 个 H800 GPU 上训练的 DeepSeek v3，带宽较低，为 300GB/s（而不是 H100 上的 450GB/s），将需要每个 GPU \text{990e12} / \text{300e9} = 3300 token，或约 6.7M（实际上，他们使用了 4M）。

启用网络内规约并使用纯数据并行，理论上我们的 AllReduce 带宽是 2 倍，这将使这两个数字减半。然而，实际上收益接近 30%，这仅仅弥补了我们通常难以达到报告数字的事实。此外，由于纯数据并行很少有用，这在实践中基本上无关紧要。

MoE 模型： 对于专家混合 (MoE) 模型，我们有 E 个专家，每个 token k 个专家，这增加到

T_{math} = \frac{2 \cdot 2 \cdot 2 \cdot k \cdot B D F}{X \cdot C}

T_{comms} = \frac{2 \cdot 2 \cdot 2 \cdot E D F}{W_{collective}}

这将每个 GPU 的 token 批处理大小增加了 E/k 的因子，即

\frac{B}{X} > \frac{E}{k} \frac{C}{W_{collective}}

例如，对于新的 OpenAI OSS 模型，k=4 和 E=128，这增加到跨节点 32 * 2475 = 79,200，这是一个相当高的数字。

当 X 很小时会发生什么？ 当我们只做例如 2 节点数据并行时，我们受益于 (X - 1) / X 的缩放，这给了我们

T_{math} = \frac{2 \cdot 2 \cdot 2 \cdot B D F}{N * C}

T_{comms} = \frac{2 \cdot 2 \cdot 2 \cdot D F \cdot (X - 1)}{X \cdot W_{collective}}

其中 X 是节点数，N = 8 \cdot X。那么对于一个密集模型，我们有 B / N > \alpha \cdot (X - 1) / X，或者例如 B / N > \text{1237}，是上面值的一半。你会因此经常看到 2 路数据并行。

要点：数据并行和 ZeRO 分片需要每个 GPU 大约 2500 个 token 的批处理大小才能在 H100 或 B200 上达到计算密集型，假设完美的重叠和 FLOPs 利用率。对于 MoE 模型，这个数字增加了 E / k 的因子，即总参数与激活参数的比率。当进行少量数据并行时，临界批处理大小会减小。

张量并行性

张量并行需要在激活值上进行 AllGather 和 ReduceScatter，我们需要将其与 MLP FLOPs 重叠。换句话说，在前向传播中，我们有

T_{math} = \frac{2 \cdot 2 \cdot B D F}{Y \cdot C}

T_{comms} = \frac{2 \cdot 2 \cdot B D}{W_{collective}}

要达到计算密集型，我们得到规则

Y < \frac{F \cdot W_{collective}}{C}

在节点内，这大约是 F / 2200 或在节点之外是 F / 2475。对于像 LLaMA-3 那样的 F=\text{28000}，这大约是 11 路 TP（或者向下取整，大约 8 路，这是一个节点的大小）。和上面一样，当我们恰好跨越 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 块，我们需要引入 2x AllToAll 将我们的激活值发送到相应的专家。

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

T_{math} = \frac{4 \cdot B \cdot k \cdot D \cdot F}{Z \cdot C}

T_{comms} = \frac{4 \cdot B \cdot D \cdot (Z - 8)}{W \cdot Z} \cdot min (\frac{8 \cdot 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 很大的模型，在这种情况下我们可以进行显著的跨节点 EP 和 TP。

要点：如果 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 跳，所以大约是

T_{total PP comms} = \frac{2 B D}{W \cdot N_{microbatches}} \cdot (N_{microbatches} + N_{stages} - 2)

T_{per-layer comms} \approx 1.5 \cdot \frac{2 B D}{W \cdot N_{layers}}

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

(1) 代码复杂性： 流水线不像其他方法那样能很好地融入自动并行框架（如 XLA 的 GSPMD）。因为它引入了微批次来隐藏流水线气泡，它改变了程序的结构，而定制的零气泡流水线调度通过要求前向和后向传播的复杂交错而加剧了这个问题。

(2) 流水线使得数据并行和 FSDP 变得困难： 可能不做流水线的最大原因是它与 FSDP 和数据并行的兼容性不好。特别是 ZeRO-3 分片效果很差，因为它要求我们在每个微批次上 AllGather 权重，而当我们只有 B / N_\text{microbatches} 个 token 来摊销 AllGather 成本时，这是行不通的。此外，在反向传播期间，我们无法 AllReduce 或 ReduceScatter 梯度，直到最后一个微批次通过给定阶段，这意味着我们有显著的非重叠通信时间。

图：一个 2 阶段、2 微批次流水线的示例。F 表示阶段前向传播，B 是阶段后向传播（成本是 2 倍）。G 表示数据并行 AllReduces，它可能比单个微批次的时间长得多。

(3) 流水线气泡和步骤不平衡： 正如你在上面（糟糕的）流水线调度中看到的，在一个简单的流水线调度中很容易出现显著的气泡（意味着计算浪费）。在上面，第二阶段在步骤 0 是空闲的，第一阶段在步骤 2 到 3 是空闲的，第二阶段在最后一步再次是空闲的。虽然我们可以通过仔细的调度来在一定程度上避免这些，但我们仍然经常有一些气泡。我们还必须在关键路径上将激活值从一个阶段传递到下一个阶段，这会增加开销：

图：一个示例流水线，显示了红色的传输成本。这会使阶段相互移动，并增加流水线气泡开销。

每个问题都有解决方法，但它们往往实现复杂且难以维护，但流水线仍然是一种相对于其他方法通信成本较低的技术。

关于延迟的警告： 如前所述，即使消息相当大，GPU 也很难达到完整的 AllReduce 带宽。这意味着即使我们理论上可以在多个节点上扩展例如专家并行的 AllToAlls，我们也可能难以达到总带宽的 50%。这意味着我们确实试图将 TP 或 EP 保持在较少数量的节点内，以最小化延迟开销。

示例

DeepSeek 是如何做的？ 作为参考，DeepSeek V3 是用 2048 个 H800 GPU 训练的，配置如下：

64 路专家并行 (EP)，跨越 8 个节点
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 的批处理大小进行训练，或每个 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 域内（即单节点或需要使用 GB200NVL72，最多 72 个 GPU）。
任何跨越多个节点的模型并行形式都可以进一步降低 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模型有什么影响？

点击此处查看答案。

答案： 我们的 bfloat16 FLOPs/s 从 990 增加到 2250 TFLOPs，增加了 2.25 倍。在节点内，带宽是 2 倍，我们的Roofline模型大致保持不变。例如，对于 TP，临界强度上升到 2250e12 / 900e9 = 2500，所以我们的限制是 Y < F / 2500，只高了一点点（而且除非节点大小增加，否则这对我们没有帮助）。

然而，在节点之外，缺乏额外的带宽实际上使得我们更难达到计算受限！例如，对于数据并行，我们的临界批处理大小增加到 2250e12 / 400e9 = 5625，因为我们的 GPU 可以在相同带宽下执行显著更多的 FLOPs。

具有 72-GPU 节点的 GB200 SuperPods 通过增加更多的出口带宽来改变这一点（来源）。

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

至少需要多少个 H100 才能存储权重和优化器？
假设我们想在 4096 个 H100 GPU 上训练 15T token。假设我们达到了 45% 的 MFU（模型 FLOPs 利用率）。训练需要多长时间？
LLaMA-3 70B 的 F = 28,672，训练时的批处理大小约为 4M token。在不成为通信受限的情况下，我们最多可以做多少模型并行？用这个加上纯 DP，我们能在保持计算受限的情况下训练 LLaMA-3 吗？ZeRO-3 呢？8 路流水线呢？

点击此处查看答案。

我们需要 2 字节用于权重，8 字节用于优化器状态，所以至少 700GB。每个 H100 有 80GB 的 DRAM，所以我们至少需要 9 个 GPU，或者（向上取整）至少 2 个 8xH100 节点。这需要很长时间来训练，并且无法保存梯度检查点，但这是一个下限。
这总共需要 6 * 70e9 * 15e12 = 6.3e24 bf16 FLOPs。每个 GPU 可以执行 990e12 FLOPs，所以在 40% MFU 下我们可以执行 1.6e18 FLOPs/s。因此整个过程需要 3.9e6 秒，或 45 天。
在节点内，我们有 450GB/s 的带宽，所以限制大约是 F / 1995 = 28672 / 1995 = 14.372。由于这不会跨越 2 个节点，实际意味着我们会做到 8 路模型并行。
1. 这需要我们做 512 路 DP。首先，我们需要看是否有足够的内存。由于我们的模型只分片 8 路，这意味着 700GB / 8 = 87.5GB / GPU，这放不下，所以不行！
2. 使用 ZeRO-3 和 8 路 TP，我们将进行 512 路 ZeRO-3。这不会有任何内存问题，因为我们积极地分片了一切。我们将有每个 GPU 4e6 / 4096 = 976 的批处理大小。这相当低，甚至低于我们的纯 DP 限制，而且这是该限制的两倍，因为我们必须移动我们的权重。所以不行。
3. 使用 8 路流水线，每个模型并行分片现在跨越 8 个节点。正如我们所见，这将我们叶级别的 AllGathers 成本降低了 8 倍，所以那里的总 AllReduce/AllGather 带宽从 400GB/s 增加到 8 * 400GB/s = 3200GB/s。那么Roofline模型就是 989e12 / 3200e9 = 309，所以我们应该没问题！我们只需要高效地实现流水线。

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

请注意，他们的序列长度到处都是 4096。对于 16B、70B 和 314B 模型，每个 GPU 的 token 批处理大小是多少？假设数据并行是最外层的轴，并假设 bfloat16 规约，确定这些模型中的每一个在理论上是计算受限还是通信受限，以及是否有更优的配置可用？

点击此处查看答案。

答案： 让我们从每个 GPU 的批处理大小开始。

16B: 192 * 4096 / 192 = 4096 token / GPU
70B: 384 * 4096 / 768 = 2048 token / GPU
314B: 1536 * 4096 / 3072 = 2048 token / GPU

这意味着除了第一个，这些都徘徊在每个批次 2k token 左右，这恰好是我们为 FSDP 计算的临界阈值附近。我们计算出该界限为 2,472 token / GPU，基于脊级别的规约，这应该在这里大致适用。然而，对于 70B 和 314B，因为我们分别有 16 路和 64 路模型分片，我们在脊级别获得了 2 倍和 8 倍的吞吐量提升，这意味着我们应该分别在大约 1k 和 300 token / 步时达到计算受限。

致谢与延伸阅读

本章严重依赖于许多知识渊博的 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 的 NVIDIA Tensor Core 历史：一篇精彩的文章，描述了 GPU 如何从视频游戏引擎转变为 ML 加速器。
SemiAnalysis 对 Blackwell 性能的分析：值得一读，以了解下一代 NVIDIA GPU。
H100 DGX SuperPod 参考：关于大型 GPU 集群如何联网的枯燥但有用的读物。这里是关于 GB200 系统的类似文档。
关于 NVLink Switch 的 Hot Chips 演讲：关于 NVLink 和 NCCL 集合通信的有趣读物，特别包括网络内规约。
DeepSeek-V3 技术报告：一个大型半开放 LLM 训练报告的好例子，描述了他们如何选择分片设置。
如何优化 CUDA 矩阵乘法：一篇很棒的博客，描述了如何使用 CUDA 核心实现高效的矩阵乘法，并着眼于 GPU 上的缓存一致性。
HuggingFace 超大规模手册：一个关于 GPU 上 LLM 并行性的指南，本章部分灵感来源于此。
从第一性原理让深度学习飞速发展：一个更侧重于 GPU 和 PyTorch 的关于 LLM Roofline模型和性能工程的教程。

附录 A：GB200 会带来哪些变化？

Blackwell 引入了一系列重大的网络变化，包括 NVLink 5，其总 NVLink 带宽是原来的两倍（900GB/s）。B200 仍然有 8-GPU 节点，就像 H100s 一样，但 GB200 系统（将 B200 GPU 与 Grace CPU 结合）引入了更大的 NVLink 域（NVL72 中有 72 个 GPU，理论上最多 576 个）。这个更大的 NVLink 域也有效地增加了节点出口带宽，从而降低了节点之上的集合通信成本。

图：一个展示 GB200 NVL72 单元如何构建的图表，包含 18 个交换机和 72 个 GPU。

在节点内，这种增加的带宽（从 450GB/s 到 900GB/s）并没有太大区别，因为我们也使每个 GPU 的总 FLOPs/s 翻倍。我们的Roofline模型基本保持不变，尽管因为 NVLink 有更好的带宽，专家并行变得更容易。

在节点之外，情况变化更大。这是来自这里的 SuperPod 图。

图：一个包含 576 个 GPU 的 GB200 DGX SuperPod 的示意图。

如你所见，每个节点的出口带宽增加到 4 * 18 * 400 / 8 = 3.6TB/s，而 H100 中是 400GB/s。由于我们的 FLOPs/芯片也翻倍，这使得有效的跨节点Roofline模型提高了约 4 倍。现在我们可能开始担心我们是否在节点级别而不是扩展级别遇到瓶颈。

Grace Hopper： NVIDIA 还销售 GH200 和 GB200 系统，它们将一定数量的 GPU 与一个 Grace CPU 配对。例如，一个 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 个交换机连接 32 个叶交换机，使用 2x400 Gbps 链接，所以我们有 32 * 16 * 400 * 2 / 8 = 51.2TB/s 的出口带宽。16 个交换机给了我们 16 * 1.6TB/s = 25.6TB/s 的带宽，所以这是这个级别的瓶颈。

每个 GPU，这给了我们节点级别 450GB/s 的 GPU 间带宽，SU 级别 50GB/s，脊级别 25 GB/s。

GPU 经验 AR 带宽：

图：一个 8xH100 集群上的 AllReduce 带宽（节点内，禁用 SHARP）。

TPU v5p 带宽（1 轴）：

图：一个 TPU v5p 4x4x4 集群上的 AllReduce 带宽（沿一个轴）。

这里还有 AllGather 带宽：

图：一个 TPU v5e 8x16 集群上的 AllGather 带宽（沿一个轴）。

更多关于 AllToAll 成本的信息：

这里我们可以比较近似值 \min(K / Z) * (Z - 1) / Z 和真实值 (1 - ((Z - 1) / Z) ** K) * (Z - 1) / Z。除了 Z 值较小时，它们是相似的。

脚注

GPU 的 Tensor Core 是 SM 的矩阵乘法子单元，而 TPU 的 TensorCore 是包含 MXU、VPU 和其他组件的总称单元。[↩]
NVIDIA 对此没有一个好的命名，所以我们只是在几个糟糕的选项中选择了最好的一个。Warp 调度器主要是向一组 CUDA 核心分派工作的单元，但我们在这里用它来描述控制单元和它所控制的核心集合。[↩]
虽然 SM 是独立的，但为了达到峰值性能，它们通常被迫进行协调，因为它们都共享一个容量有限的 L2 缓存。[↩]
较新的 GPU 支持 FMA（融合乘加）指令，理论上每个周期执行两个 FLOPs，NVIDIA 无情地利用这一事实将其报告的规格翻倍。[↩]
历史上，在引入 Tensor Core 之前，CUDA 核心是 GPU 的主要组件，用于渲染，包括光线-三角形相交和着色。在今天的游戏 GPU 上，它们仍然承担大部分渲染工作，而 TensorCore 用于上采样（DLSS），这使得 GPU 可以在较低分辨率下渲染（像素越少 = 工作量越少），然后使用机器学习进行上采样。[↩]
NVIDIA 没有分享很多 TC 硬件细节，所以这更多的是猜测而非确切事实——当然，这并不能说明 TC 是如何实现的。我们知道 V100 每个 TC 每周期可以执行 256 FLOPs。A100 可以做到 512，H100 可以做到 1024，而 B200 的细节尚未公布，但似乎可能是 2048 FLOPs/TC/周期，因为 `2250e12 / (148 * 4 * 1.86e9)` 大约是 2048。更多细节在这里得到确认。[↩]
在 Ampere 中，Tensor Core 可以由单个 warp 提供数据，而在 Hopper 中则需要一个完整的 SM（warpgroup），在 Blackwell 中则由 2 个 SM 提供数据。在 Blackwell 中，矩阵乘法也变得如此之大，以至于参数（特别是累加器）不再适合寄存器内存/SMEM，因此 Blackwell 增加了 TMEM 来解决这个问题。[↩]
在给定 SM 上调度的 Warp 称为“驻留”。[↩]
技术上，L2 缓存被分成两半，因此在 H100 上，一半的 SM 可以各自访问 25MB。有一个连接这两半的链接，但带宽较低。[↩]
L2 缓存被所有 SM 共享这一事实，实际上迫使程序员以一种相当协调的方式运行 SM，尽管原则上它们是独立的单元。[↩]
虽然 NVIDIA 制造了 B100 代，但它们只短暂销售和生产，据称是由于设计缺陷导致它们无法接近其声称的规格运行。它们在不因散热和功耗问题而降频的情况下难以达到峰值 FLOPs。[↩]
在深度学习热潮之前，GPU（“图形处理单元”）做的是，嗯，图形处理——主要是为了视频游戏。视频游戏用数百万个小三角形来表示物体，游戏将这些三角形渲染（或“光栅化”）成一个二维图像，每秒在屏幕上显示 30-60 次（这个频率称为帧率）。光栅化涉及将这些三角形投影到相机的坐标系中，并计算哪些三角形与哪些像素重叠，每秒数十亿次。可以想象，这是非常昂贵的，而这仅仅是开始。然后你必须通过组合可能与光线相交的几个半透明三角形的颜色来为每个像素着色。GPU 被设计用来极快地执行这些操作，并着眼于通用性；你需要同时运行许多不同的 GPU 工作负载（称为“着色器”），而没有任何单一操作占主导地位。因此，面向消费者的图形 GPU 可以进行矩阵乘法，但这并不是它们的主要功能。[↩]
值得注意的是，这个强度在最近几代 GPU 中保持不变。对于 H100s 是 33.5 / 3.5，对于 B200 是 80 / 8。为什么会这样尚不清楚，但这是一个有趣的观察。[↩]
术语“节点”是重载的，可以指两件事：NVLink 域，即通过 NVLink 互连完全连接的 GPU 集合，或者连接到单个 CPU 主机的 GPU 集合。在 B200 之前，这两者通常是相同的，但在 GB200 NVL72 中，我们有一个包含 72 个 GPU 的 NVLink 域，但每个主机仍然只连接 8 个 GPU。我们在这里使用术语“节点”来指代 NVLink 域，但这有争议。[↩]
有人向我描述 NVLink 就像一个增强版的 PCIe 连接，具有低延迟和协议开销，但不是为可扩展性/容错性设计的，而 InfiniBand 更像以太网，专为更大的有损网络设计。[↩]
这里的全双工意味着每个方向 25GB/s，两个方向相互独立。你可以在链路上总共发送 50GB/s，但每个方向最多 25GB/s。[↩]
例如，Meta 在一个与此描述显著不同的数据中心网络上训练了 LLaMA-3，该网络使用以太网、一个三层交换结构，并且在顶层有一个超额订阅的交换机。[↩]
你也可以认为每个 GPU 将其大小为 \text{bytes} / N 的块发送给其他 N - 1 个 GPU，总共通信了 (N - 1) * N * bytes / N 字节，这给了我们[↩]
真实成本实际上是 $(1 - {(\frac{Z - 1}{Z})}^{K}) \cdot \frac{Z - 1}{Z}$ 在 K 次掷骰子中不同结果的期望数量，但它非常接近给出的近似值。更多细节请参见附录。[↩]

其他

^*在 Google DeepMind 完成的工作，现就职于 MatX。

引用

在学术背景下引用，请按如下格式引用本作品：

    Austin et al., "How to Scale Your Model", Google DeepMind, online, 2025.

或作为 BibTeX 条目：

    @article{scaling-book,
      title = {How to Scale Your Model},
      author = {Austin, Jacob and Douglas, Sholto and Frostig, Roy and Levskaya, Anselm and Chen, Charlie and Vikram, Sharad
      and Lebron, Federico and Choy, Peter and Ramasesh, Vinay and Webson, Albert and Pope, Reiner},
      publisher = {Google DeepMind},
      howpublished = {Online},
      note = {Retrieved from https://jax-ml.github.io/scaling-book/},
      year = {2025}
    }