这是 DeepMind Scaling Book 系列的第 12 部分（额外章节）。
关于 GPU 的一切 (How to Think About GPUs)
How To Scale Your Model Part 12 (Part 11: Conclusion | The End)
在 Google 我们热爱 TPU，但 GPU 也很棒。本章深入探讨 GPU 的世界——每个芯片如何工作，它们如何联网，以及这对 LLM 意味着什么，特别是与 TPU 相比。虽然有来自 NVIDIA, AMD, Intel 等的众多 GPU 架构，但在这里我们将重点关注 NVIDIA GPU。本节建立在第 2 章和第 5 章的基础上，因此鼓励你先阅读它们。
GPU 是什么？ (What Is a GPU?)
现代 ML GPU（例如 H100, B200）基本上是一堆专门用于矩阵乘法的计算核心（称为 Streaming Multiprocessors 或 SMs）连接到一根快速内存（称为 HBM）。
Figure: H100 或 B200 GPU 的抽象布局示意图。H100 有 132 个 SM，而 B200 有 148 个。我们使用术语“Warp Scheduler”来宽泛地描述一组 32 个 CUDA SIMD 核心以及向它们分派工作的调度器。注意这看起来多像 TPU！
每个 SM，像 TPU 的 Tensor Core 一样，都有一个专用的矩阵乘法核心（不幸的是也叫 Tensor Core [1]），一个向量算术单元（叫 Warp Scheduler [2]），和一个快速的片上缓存（叫 SMEM）。与 TPU 不同（TPU 最多有 2 个独立的“Tensor Cores”），现代 GPU 有超过 100 个 SM（H100 上有 132 个）。这些 SM 中的每一个都远不如 TPU Tensor Core 强大，但整个系统更灵活。每个 SM 或多或少是完全独立的，因此 GPU 可以同时执行数百个单独的任务 [3]。
让我们更详细地看一下 H100 SM：
Figure: H100 SM 示意图，显示 4 个 subpartitions，每个包含一个 Tensor Core, Warp Scheduler, Register File, 和一组 CUDA Cores。底部的 ‘L1 Data Cache’ 是 256kB SMEM 单元。
- CUDA Cores: 每个子分区包含一组用于 SIMD/SIMT 向量算术的 ALU，称为 CUDA Cores。每个 ALU 通常每个周期可以执行 1 个算术运算，例如 f32.add。每个子分区包含 32 个 fp32 核心。
- Tensor Core (TC): 每个子分区都有自己的 Tensor Core，这是一个专用的矩阵乘法单元，就像 TPU MXU 一样。Tensor Core 代表了 GPU 的绝大多数 FLOPs/s（例如，在 H100 上，我们有 990 bf16 TC TFLOP/s，而 CUDA 核心只有 66 TFLOPs/s）。
CUDA 核心比 TPU 的 VPU 更灵活：GPU CUDA 核心使用称为 SIMT (Single Instruction Multiple Threads) 的编程模型，相比之下 TPU 使用 SIMD (Single Instruction Multiple Data) 模型。
CUDA 核心调度也更灵活：SM 运行有点像多线程 CPU，因为它们可以并发地“调度”许多程序（warps）（每个 SM 最多 64 个），但每个 Warp Scheduler 在每个时钟周期只执行单个程序。
内存 (Memory)
除了计算单元，GPU 还有内存层次结构：
- Registers: 每个子分区有自己的寄存器文件，包含 16,384 个 32-bit 字。
- SMEM (L1 Cache): 每个 SM 有自己的 256kB 片上缓存，称为 SMEM。
- L2 Cache: 所有 SM 共享一个相对较大的 ~50MB L2 缓存。
- HBM: 主 GPU 内存，用于存储模型权重、梯度、激活等。
GPU 规格摘要 (Summary of GPU specs)
| GPU | Generation | Clock Speed | SMs/chip | SMEM capacity/SM | L2 capacity/chip | HBM capacity/chip |
|---|---|---|---|---|---|---|
| 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 |
| GPU | Generation | HBM BW/chip | FLOPs/s/chip (bf16/fp16) | FLOPs/s/chip (fp8/int8) | FLOPs/s/chip (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 |
芯片层面的 GPU vs. TPU (GPUs vs. TPUs at the chip level)
- GPU 更模块化。TPU 有 1-2 个大 Tensor Cores，而 GPU 有数百个小 SM。
- 历史上，单个 GPU 更强大（也更贵）。
- TPU 有更多快速缓存内存。TPU 有比 GPU SMEM 多得多的 VMEM。
网络 (Networking)
网络是 GPU 和 TPU 差异最大的领域之一。TPU 连接在 2D 或 3D 环面中，而 GPU 使用更传统的分层树状交换网络。
节点层面 (At the node level)
GPU 节点是一个小单元，通常有 8 个 GPU（GB200 最多 72 个），通过全对全、全带宽、低延迟的 NVLink 互连连接。
Figure: 从 Pascall (P100) 开始的节点即 NVLink 域图。
对于 Hopper 一代 (NVLink 4.0)，每个 NVLink 链路具有 25GB/s 的全双工带宽，给我们每个 GPU 18 * 25=450GB/s 的全双工带宽进入网络。
| NVLink Gen | NVSwitch Gen | GPU Generation | NVLink Bandwidth (GB/s, full-duplex) | NVLink Ports / GPU | Node GPU to GPU bandwidth (GB/s full-duplex) | Node size (NVLink domain) | NVSwitches per node |
|---|---|---|---|---|---|---|---|
| 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 |
节点之外 (Beyond the node level)
在节点级别之外，GPU 网络的拓扑结构不那么标准化。NVIDIA 发布了一个参考 DGX SuperPod 架构，使用 InfiniBand (IB) 将更大的 GPU 组连接起来。
Figure: 参考 1024 H100 DGX SuperPod 示意图，有 128 个节点。
我们有多少带宽？ InfiniBand 网络（称为“scale out network”）的总体拓扑是 Fat Tree (胖树)，电缆和交换机保证节点级别以上的全对分带宽。这意味着如果我们将节点分成两半，每个节点可以同时向另一个分区中的节点输出 400GB/s。
| Level | GPUs | Switches per Unit | Switch Type | Bandwidth per Unit (TB/s, full-duplex) | GPU-to-GPU Bandwidth (GB/s, full-duplex) | Fat Tree Bandwidth (GB/s, full-duplex) |
|---|---|---|---|---|---|---|
| Node | 8 | 4 | NVL | 3.6 | 450 | 450 |
| Leaf | 256 | 8 | IB | 12.8 | 50 | 400 |
| Spine | 1024 | 16 | IB | 51.2 | 50 | 400 |
集体通信在 GPU 上是如何工作的？ (How Do Collectives Work on GPUs?)
我们使用 NCCL (NVIDIA Collective Communication Library) 库。
节点内集体通信 (Intra-node collectives)
AllGather 或 ReduceScatter: 在节点级别，你可以像 TPU 一样在环周围执行它们，在每一跳使用完整的 GPU 到 GPU 带宽。
\[T_\text{AG or RS comms} = \frac{\text{bytes}}{\text{GPU egress bandwidth}}\]AllToAlls: 节点内的 GPU 具有全对全连接，这使得 AllToAlls 非常容易。
\[T_\text{AllToAll comms} = \frac{B \cdot (N - 1)}{W \cdot N^2} \approx \frac{B}{W \cdot N}\]网络内归约 (In network reductions): 自 Hopper 一代以来，NVIDIA 交换机支持 SHARP，允许网络交换机本身执行归约操作。理论上，这应该将 AllReduce 的成本减半。
跨节点集体通信 (Cross-node collectives)
当我们超越节点级别时，成本有点微妙。
\[T_\text{total} = \max(T_\text{comms at node}, T_\text{comms in scale-out network}) = \max\left[\frac{\text{bytes}}{W_\text{GPU egress}}, \frac{\text{bytes}}{W_\text{node egress}}\right]\]一般来说，节点外 AllGather 或 ReduceScatter 的成本大约是 $B / W_\text{node egress}$，在 H100 DGX SuperPod 上是 $B / \text{400e9}$。
GPU 上 LLM Scaling 的 Rooflines (Rooflines for LLM Scaling on GPUs)
数据并行 (Data Parallelism)
对于纯数据并行或 FSDP 没有网络内归约，我们需要：
\[\frac{B}{X} > \frac{C}{W_\text{collective}}\]- 节点内，我们需要每个 GPU token batch size > $\text{990e12} / \text{450e9} = 2200$。
- SU 内或 spine 级别，BS > $\text{990e12} / \text{400e9} = 2475$。
张量并行 (Tensor Parallelism)
张量并行需要在激活上进行 AllGather 和 ReduceScatter。
\[Y < \frac{F \cdot W_\text{collective}}{C}\]在节点内，这给我们大约 $F / 2200$ 或节点外 $F / 2475$。对于 LLaMA-3 ($F=28000$)，这大约是 11 路 TP（或者向下取整，大约 8 路，这是节点的大小）。
专家并行 (Expert Parallelism)
对于混合专家 (MoE) 模型，如果 $F < 8 * C / W_\text{node}$，专家并行可以跨越 1-2 个节点，成本与 TP 相似（略低）；或者如果 $F > 8 * C / W_\text{node}$，我们可以进行大量的专家并行（直到 $E$ 个节点），成本相对较低。
流水线并行 (Pipeline Parallelism)
流水线并行的总体通信成本很小。
\[T_\text{per-layer comms} \approx 1.5 \cdot \frac{2BD}{W \cdot N_\text{layers}}\]从通信的角度来看，流水线基本上是免费的。
总结：GPU 上 LLM 扩展的要点 (TLDR of LLM scaling on GPUs)
- 数据并行或 FSDP (ZeRO-1/3) 需要每个 GPU 约 2500 tokens 的本地批量大小。
- 张量并行在计算受限的情况下最多可达约 8 路，这主要限制我们在单个 NVLink 域内（即单节点）。
- 任何跨越多节点的模型并行形式都可以进一步降低 FSDP 的成本，因此我们经常希望混合 PP + EP + TP 来跨越许多节点并降低 FSDP 成本。
- 如果能处理零气泡流水线的代码复杂性并保持较大的批量大小，流水线并行效果很好。
脚注
[1] 令人困惑的命名。 [2] 实际上包含 Dispatch Unit, Register File, CUDA Cores 等。 [3] 通过 CUDA Streams。
