CABINET / The Physical Stack of AI / AI Infra
AI Infra
Physical Stack of AI 系列 · 第四章(完结)
前两章从电路和晶体管讲到 CPU 与 GPU 的架构,再到半导体从沙子到芯片的整条产业链。我们已经理解了”一颗芯片如何工作”和”一颗芯片如何被制造出来”。这是 Physical Stack of AI 系列的最后一章,回答的是系统层面的问题:一颗 GPU(甚至一整个 NVL72 机架)的算力仍然不够训练一个万亿参数的大模型。
AI Infra(AI基础设施)要回答的核心问题是:如何把成千上万颗GPU组织成一个高效协作的计算系统?
这涉及三个层次的工程:硬件如何互联(让GPU之间能快速通信)、软件如何编程(让算法跑在GPU上)、系统如何调度(让训练和推理高效运行)。对应 CPU 一章开头八层技术栈里的 Layer 6(系统互联)、Layer 7(软件栈)、Layer 8(训练与推理系统)。
本篇的路线图:
§18 系统互联与网络层:GPU之间如何通信
§19 软件栈与编程模型:CUDA生态与替代方案
§20 精度控制与量化:从FP32到FP4的工程实践
§21 分布式训练:数千GPU如何协同训练一个模型
§22 推理优化:如何高效地用训练好的模型提供服务
§23 AI数据中心基础设施:电力、散热与存储
§18. 系统互联与网络层
18.1 互联层级总览
一个AI训练集群的互联结构可以分为三个层级,每个层级的带宽差异巨大:
| 层级 | 连接对象 | 典型技术 | 带宽(双向) | 延迟 |
|---|---|---|---|---|
| 芯片内部 | SM ↔ L2 ↔ HBM | 片上互联 + HBM PHY | 数 TB/s | ~ns |
| 节点内部 | GPU ↔ GPU(同一台服务器内) | NVLink + NVSwitch | 900 GB/s ~ 1.8 TB/s(每GPU) | ~μs |
| 节点之间 | 服务器 ↔ 服务器 | InfiniBand / RoCE | 200~800 Gb/s(每端口) | ~1-5 μs |
从芯片内部到节点之间,带宽差了约 三个数量级。这意味着分布式训练的通信策略必须尽量减少跨节点传输。把这三个层级的带宽差异记在心里,后面讲分布式训练的并行策略时会反复用到。
18.2 NVLink 与 NVSwitch
在 GPU 一章 §10 中我们提到,PCIe 是 CPU 与 GPU 之间的标准总线,但 PCIe Gen 5 x16 的 128 GB/s 双向带宽远远跟不上 GPU 间通信的需求。NVIDIA 的解决方案是自研高速互联:NVLink。
NVLink 的本质:一条绕过 CPU 和 PCIe 总线的 GPU 到 GPU 直连通道。
| 代际 | 引入架构 | 每链路带宽 | 每GPU总带宽(双向) | 链路数/GPU |
|---|---|---|---|---|
| NVLink 1.0 | Pascal (2016) | 40 GB/s | 160 GB/s | 4 |
| NVLink 2.0 | Volta (2017) | 50 GB/s | 300 GB/s | 6 |
| NVLink 3.0 | Ampere (2020) | 50 GB/s | 600 GB/s | 12 |
| NVLink 4.0 | Hopper (2022) | 50 GB/s | 900 GB/s | 18 |
| NVLink 5.0 | Blackwell (2024) | 100 GB/s | 1,800 GB/s | 18 |
NVLink 5.0 的每GPU带宽是 PCIe Gen 5 x16 的 14 倍。
但 NVLink 是点对点连接。如果一台服务器内有 8 颗 GPU,要让它们两两直连,需要 $C(8,2) = 28$ 条 NVLink 链路。每颗 GPU 只有 18 条链路,每颗要连其他 7 颗,刚好不够做全互联。
解决方案是 NVSwitch:一颗专用的交换芯片,所有 GPU 的 NVLink 链路都连到 NVSwitch 上,由 NVSwitch 负责任意两颗 GPU 之间的数据路由。
DGX H100 节点的内部互联拓扑:
GPU 0 ──┐ ┌── GPU 4
GPU 1 ──┤ ├── GPU 5
GPU 2 ──┼── NVSwitch ──┼── GPU 6
GPU 3 ──┘ ×4 └── GPU 7
4 颗 NVSwitch 芯片
8 颗 H100 GPU 通过 NVLink 4.0 全互联
任意两颗 GPU 之间: 900 GB/s 双向
在 DGX H100 节点中,4 颗第三代 NVSwitch 芯片提供了总计 7.2 TB/s 的 all-to-all 带宽。对节点内的 8 颗 GPU 来说,数据传输几乎没有拓扑限制,任意一对 GPU 之间都能以满速通信。
Blackwell 的 NVLink Switch:在 GB200 NVL72 机架中,NVIDIA 将 NVSwitch 的概念进一步扩展。72 颗 B200 GPU 通过第五代 NVLink Switch 实现了机架级全互联,而不只是节点级:
NVL72 机架互联结构:
┌─────────────────────────────────────────────┐
│ NVLink Switch Tray │
│ (机架顶部的独立 NVSwitch 交换层) │
│ 总带宽: 130 TB/s │
└─────────┬──────┬──────┬──────┬──────────────┘
│ │ │ │
┌────┴──┐ ┌─┴───┐ ┌┴────┐ ... ×36 个 GB200/GB300
│GB200 │ │GB200│ │GB200│
│#1 │ │#2 │ │#3 │
│(2 GPU │ │ │ │ │
│+1 CPU)│ │ │ │ │
└───────┘ └─────┘ └─────┘
72 颗 GPU 全互联
每 GPU 1.8 TB/s NVLink 带宽
整个机架对软件层呈现为"一颗超大GPU"
这是 NVIDIA 系统级工程能力的集中体现:不仅设计 GPU Die 本身,还设计了 NVLink PHY、NVSwitch 芯片、交换拓扑、机架级系统。这种从芯片到机架的垂直整合,是其他 GPU 厂商短期内难以复制的。
18.3 InfiniBand
NVLink 解决的是节点内部(或 NVL72 机架内部)的 GPU 互联。但训练大模型需要数百甚至数千台服务器协同工作,节点之间的通信需要数据中心级别的网络。
InfiniBand 是 AI 训练集群的首选网络技术。NVIDIA 在 2020 年以 69 亿美元收购了 InfiniBand 的核心厂商 Mellanox,从而掌握了从 GPU 到网络的完整通信栈。
InfiniBand 的核心特性:
| 特性 | 说明 |
|---|---|
| RDMA(Remote Direct Memory Access) | 远程直接内存访问:一台机器可以直接读写另一台机器的内存,不经过 CPU,不经过操作系统内核 |
| 极低延迟 | 端到端延迟约 ~1 μs(以太网通常 5-10 μs) |
| Lossless | 基于信用(Credit-based)的流控机制,不丢包 |
| 硬件卸载 | 网络协议的处理由网卡(HCA)硬件完成,不消耗 CPU 资源 |
InfiniBand 的速率演进:
| 代际 | 每端口速率 | 年份 | 典型应用 |
|---|---|---|---|
| FDR | 56 Gb/s | 2011 | 早期 HPC |
| EDR | 100 Gb/s | 2014 | HPC + 早期 AI |
| HDR | 200 Gb/s | 2019 | A100 训练集群 |
| NDR | 400 Gb/s | 2022 | H100 训练集群 |
| XDR | 800 Gb/s | 2024 | Blackwell 训练集群 |
用具体的数字来感受 InfiniBand 的带宽:NDR 400 Gb/s = 50 GB/s。一颗 H100 的 HBM 带宽是 3.35 TB/s,InfiniBand 带宽只有 HBM 的 1.5%。也就是说,跨节点传输数据比 GPU 内部读取数据慢了约 67 倍。这个差距是分布式训练系统设计的核心约束。
RDMA 与 GPUDirect
RDMA 的全称是 Remote Direct Memory Access(远程直接内存访问)。传统的网络通信需要经过多层软件栈:
传统网络路径(TCP/IP):
GPU 显存 → CPU 内存 → 操作系统内核 → 网络协议栈
→ 网卡 DMA → 网络 → 对端网卡 → 对端内核
→ 对端 CPU 内存 → 对端 GPU 显存
涉及: 多次内存拷贝 + CPU 中断 + 内核态/用户态切换
延迟: ~10-50 μs
RDMA 跳过了所有中间层:
RDMA 路径:
GPU 显存 → 网卡(HCA)→ 网络 → 对端网卡 → 对端 GPU 显存
零拷贝 + 零 CPU 介入 + 旁路内核
延迟: ~1 μs
GPUDirect RDMA(GDR)更进一步:网卡直接访问 GPU 的显存(HBM),连 CPU 内存都不经过。具体来说:
| 技术 | 功能 |
|---|---|
| GPUDirect P2P | 同一节点内,GPU 之间直接访问彼此的显存(通过 NVLink) |
| GPUDirect RDMA | 跨节点时,网卡直接读写 GPU 显存(跳过 CPU 内存) |
| GPUDirect Storage | GPU 直接从 NVMe SSD 读取数据(跳过 CPU 内存) |
在大规模训练集群中,GPUDirect RDMA 是标配。没有它,每次跨节点通信都需要在 GPU 显存和 CPU 内存之间做额外的数据拷贝,延迟和带宽都会大幅恶化。
18.4 RoCE 与以太网方案
InfiniBand 性能最优,但成本也最高(专用网卡、专用交换机、专用线缆)。对于不追求极致延迟的场景,或者已有大规模以太网基础设施的云厂商,RoCE(RDMA over Converged Ethernet) 是另一个选择。
| 维度 | InfiniBand | RoCE(v2) |
|---|---|---|
| 物理层 | InfiniBand 专有 | 标准以太网 |
| RDMA 支持 | 原生 | 在以太网上实现 RDMA |
| 延迟 | ~1 μs | ~2-5 μs |
| 无损保证 | 原生无损(Credit-based) | 需要配置 PFC/ECN(有一定丢包风险) |
| 成本 | 高(专用设备) | 较低(复用以太网设备) |
| 代表用户 | OpenAI、xAI、Meta(AI集群) | 字节跳动、阿里云、Google |
RoCE 的核心挑战是无损保证:以太网本质上允许丢包(靠TCP重传恢复),但 RDMA 要求零丢包。RoCE 通过 PFC(Priority Flow Control) 和 ECN(Explicit Congestion Notification) 来模拟无损行为,但在大规模部署时,PFC 的 Head-of-Line Blocking 和 PFC Storm 是已知的运维难题。
字节跳动和 Google 的经验表明,经过精心的网络调优,RoCE 可以在万卡集群上达到接近 InfiniBand 的实际性能。但这需要深厚的网络工程能力和持续的调优投入。
Ultra Ethernet Consortium(UEC) 是行业对 InfiniBand 垄断的回应。由 AMD、Broadcom、Cisco、Intel、Meta 等联合推动,目标是定义一套面向 AI/HPC 的以太网标准,原生支持集合通信等 AI 训练常用操作。UEC 规范仍在制定中(2025年发布1.0),实际产品和生态的成熟还需要时间。
[!info] 从投资角度看互联
NVIDIA 在 2020 年收购 Mellanox(InfiniBand)的战略意义,在几年后变得清晰:NVIDIA 不只卖 GPU,而是卖整个训练集群的解决方案(GPU + NVLink + NVSwitch + InfiniBand + 软件栈)。客户买了 H100,最自然的网络选择就是 NVIDIA 的 InfiniBand,因为端到端的性能和兼容性最好。
这种垂直整合策略提高了客户的切换成本:即使另一家的 GPU 硬件参数接近,缺少对应的高速互联和软件生态,整体系统性能可能差距巨大。
18.5 集合通信与 NCCL
GPU 之间的通信不是随意的”你给我发一个包”,而是有高度结构化的模式。分布式训练中最常用的通信模式叫集合通信(Collective Communication),核心操作包括:
| 操作 | 功能 | 典型场景 |
|---|---|---|
| All-Reduce | 所有 GPU 把各自的数据聚合(如求和),结果广播给所有 GPU | 数据并行中的梯度同步 |
| All-Gather | 每个 GPU 持有数据的一部分,收集后每个 GPU 都有完整数据 | 张量并行的参数重组 |
| Reduce-Scatter | 每个 GPU 把自己的数据分块发给不同 GPU,每个 GPU 负责聚合一块 | FSDP/ZeRO 中的梯度聚合 |
| Broadcast | 一个 GPU 把数据发给所有其他 GPU | 参数初始化 |
| All-to-All | 每个 GPU 向所有其他 GPU 发送不同的数据 | MoE 模型的 Expert 路由 |
以 All-Reduce 为例,这是数据并行训练中使用频率最高的操作。假设 4 颗 GPU 各自计算了一份梯度(Gradient),需要对这 4 份梯度求平均,然后每颗 GPU 都持有平均后的结果:
Before All-Reduce:
GPU 0: grad_0 = [1, 2, 3, 4]
GPU 1: grad_1 = [5, 6, 7, 8]
GPU 2: grad_2 = [2, 3, 4, 5]
GPU 3: grad_3 = [4, 5, 6, 7]
After All-Reduce (sum):
GPU 0: result = [12, 16, 20, 24]
GPU 1: result = [12, 16, 20, 24]
GPU 2: result = [12, 16, 20, 24]
GPU 3: result = [12, 16, 20, 24]
实现 All-Reduce 有多种算法(Ring All-Reduce、Tree All-Reduce、Recursive Halving-Doubling 等),不同算法在不同拓扑和数据量下的效率差异很大。
NCCL(NVIDIA Collective Communications Library) 是 NVIDIA 为自家 GPU 深度优化的集合通信库。它能自动检测系统拓扑(哪些 GPU 通过 NVLink 连接,哪些通过 PCIe,哪些跨节点),然后选择最优的通信算法和路径。
NCCL 的关键能力:
- 拓扑感知:自动发现 NVLink、NVSwitch、PCIe、InfiniBand 的拓扑结构,避免走低效路径
- 算法选择:根据数据量和拓扑,自动选择 Ring、Tree 或 Direct 算法
- Kernel Fusion:将多个小的通信操作合并成一个大操作,减少启动开销
- 计算通信重叠(Overlap):支持在 GPU 计算的同时进行数据传输,不让 GPU 空等
[!info] 带宽利用率:衡量互联效率的核心指标
理论上,8 颗 H100 通过 NVLink 4.0 全互联,每颗 GPU 的 NVLink 带宽是 900 GB/s。但实际训练中,通信带宽利用率(即实际传输量 / 理论峰值带宽)取决于多个因素:NCCL 算法的效率、数据量大小、通信与计算的重叠程度等。
好的训练系统可以达到 80%+ 的 NVLink 带宽利用率和 70%+ 的 InfiniBand 带宽利用率。剩下的损耗来自协议开销、拓扑不对称和软件栈延迟。
§19. 软件栈与编程模型
19.1 CUDA:NVIDIA 的核心护城河
NVIDIA 在 AI 领域的统治地位,70% 来自软件,30% 来自硬件。硬件上,AMD MI300X 的 HBM 容量(192 GB)甚至超过了 H100(80 GB)。但在实际部署中,绝大多数 AI 工作负载仍然跑在 NVIDIA GPU 上,原因就是 CUDA 生态。
CUDA(Compute Unified Device Architecture)于 2007 年随 GeForce 8800 GTX 一同发布。当时 GPU 编程需要把通用计算伪装成图形渲染操作(通过 OpenGL/DirectX 的 Shader 接口),极其繁琐。CUDA 提供了一套 C/C++ 扩展语言和运行时环境,让程序员可以直接在 GPU 上写通用计算程序。这在当时是革命性的。
到 2026 年,CUDA 生态积累了 19 年的优化库、工具链、教程、开发者社区和用户习惯。这个生态的深度构成了极高的迁移壁垒。
19.2 CUDA 编程模型
CUDA 编程的核心概念可以对应到 GPU 一章讲过的硬件结构。回忆 GPU 一章 §7.1(CUDA 结构和硬件的一一对应):
| CUDA 软件概念 | 对应的硬件 | 说明 |
|---|---|---|
| Thread(线程) | CUDA Core 的一次执行 | 最基本的执行单位 |
| Warp | 32 个 Thread 的硬件调度单位 | SM 的 Warp Scheduler 以 Warp 为粒度调度 |
| Block(线程块) | 一个 SM 上执行的线程组 | Block 内线程可共享 Shared Memory |
| Grid(网格) | 整个 GPU 上执行的所有 Block | 由 GigaThread Engine 分配到各 SM |
| Kernel(核函数) | 一个在 GPU 上运行的函数 | 由 CPU 端发起调用(Launch Kernel) |
一个最简单的 CUDA 程序看起来是这样的(向量加法):
// 定义 Kernel:每个线程处理一个元素
__global__ void vectorAdd(float *A, float *B, float *C, int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
C[i] = A[i] + B[i];
}
}
// CPU 端调用
int main() {
// ... 分配 GPU 内存,拷贝数据 ...
int threadsPerBlock = 256;
int blocksPerGrid = (n + threadsPerBlock - 1) / threadsPerBlock;
vectorAdd<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, n);
// ... 拷贝结果回 CPU ...
}
<<<blocksPerGrid, threadsPerBlock>>> 是 CUDA 特有的语法,指定了 Grid 的大小和每个 Block 包含的线程数。GPU 会自动将这些 Block 分配到各个 SM 上并行执行。
CUDA 内存模型是编写高性能 GPU 代码的关键。对应 GPU 一章 §6.3 中讲过的存储层次:
速度 容量 作用域
Register │ 最快 │ 64KB/SM │ 每线程私有
Shared Memory │ 快 │ ~256KB/SM│ Block 内共享(程序员显式管理)
L1 Cache │ 快 │ 与 Shared Memory 共用 │ SM 内自动缓存
L2 Cache │ 中等 │ ~50MB │ 全 GPU 共享
Global Memory │ 慢 │ ~80GB │ 全 GPU 可见(即 HBM/GDDR)
(HBM/GDDR) │ │ │
关键区别:CPU 的缓存(L1/L2/L3)完全由硬件自动管理,程序员不需要干预。GPU 的 Shared Memory 则由程序员手动控制。你需要显式地把数据从 Global Memory(HBM)搬到 Shared Memory,计算完再写回。这增加了编程复杂度,但给了优化极大的空间。
FlashAttention 的核心优化正是在这里:通过精心设计的 tiling 策略(将大矩阵分成小块,逐块加载到 Shared Memory 中计算),避免了对 HBM 的大量重复访问。我们在 §22 的推理优化中会详细展开。
CUDA Stream(流):CUDA 的异步执行机制。默认情况下,所有操作都在一个 Stream 中顺序执行。但你可以创建多个 Stream,让不同 Stream 中的操作并行执行:
Stream 0: | Kernel A | Kernel C |
Stream 1: | Kernel B | memcpy D2H |
Kernel A 和 Kernel B 可以在不同 SM 上同时执行
Kernel B 和 memcpy 可以同时进行(计算与传输重叠)
计算与通信的重叠(Overlap)是分布式训练性能优化的基础技术。没有 Stream 机制,GPU 在等待数据传输的时候就只能空等。
19.3 CUDA 核心库
日常的 AI 开发中,很少有人直接写 CUDA Kernel。绝大多数计算通过 NVIDIA 提供的优化库完成,这些库是 CUDA 生态的核心资产:
| 库 | 功能 | 重要性 |
|---|---|---|
| cuBLAS | 基础线性代数运算(GEMM 矩阵乘法是核心) | AI 训练/推理中 80%+ 的计算量是 GEMM |
| cuDNN | 深度学习原语(卷积、归一化、激活函数、Attention) | PyTorch/TensorFlow 的底层调用目标 |
| NCCL | 集合通信(§18.5 已详细介绍) | 分布式训练必需 |
| cuSPARSE | 稀疏矩阵运算 | 结构化稀疏(2:4)的硬件加速 |
| cuFFT | 快速傅里叶变换 | 信号处理、部分科学计算 |
| TensorRT | 推理优化引擎(图优化、量化、Kernel Fusion) | 生产环境推理部署的事实标准 |
| Triton(OpenAI) | Python 编写 GPU Kernel 的编程语言 | 降低 Kernel 开发门槛,LLM 推理优化的常用工具 |
以 cuBLAS 为例,理解它为什么重要。矩阵乘法(GEMM:General Matrix Multiply)是 AI 计算中最频繁的操作。一次 Transformer 的前向传播,本质上就是一连串的 GEMM。
一个 naive 的矩阵乘法实现和 cuBLAS 的优化实现之间,性能差距可以达到 10-50 倍。cuBLAS 的优化包括:
- Tiling:将大矩阵分成适合 Shared Memory 的小块
- 数据预取(Prefetching):在计算当前块的同时,提前加载下一块数据
- 寄存器分块(Register Tiling):进一步将 Shared Memory 中的数据分成适合寄存器的微块
- Tensor Core 调用:使用 Tensor Core 的硬件矩阵乘加指令(WMMA/MMA),而不是逐元素用 CUDA Core 计算
- 自动调优(Auto-tuning):根据矩阵大小和精度,自动选择最优的 Kernel 实现
这些优化积累了十几年的工程经验和 GPU 微架构知识,不是一朝一夕能复制的。
19.4 替代软件栈
AMD ROCm 与 HIP
ROCm(Radeon Open Compute) 是 AMD 对标 CUDA 的开源软件栈。HIP(Heterogeneous Interface for Portability) 是 ROCm 中的编程接口,语法与 CUDA 高度相似。
AMD 提供了一个工具 hipify,可以自动将 CUDA 代码转换为 HIP 代码。大部分情况下,转换是直接的 API 名称替换(cudaMalloc → hipMalloc,cudaMemcpy → hipMemcpy)。
PyTorch 从 2.0 起正式支持 ROCm,这是 AMD 生态建设的重要里程碑。开发者可以用同一套 PyTorch 代码在 AMD GPU 上训练和推理,不需要修改上层逻辑。
但 ROCm 的挑战在于深层优化的差距:
| 维度 | CUDA | ROCm |
|---|---|---|
| 库的成熟度 | cuBLAS/cuDNN 经过 15+ 年调优 | rocBLAS/MIOpen 追赶中 |
| Kernel 优化 | 针对每代 GPU 微架构深度优化 | 优化深度不足,部分 Kernel 性能落后 20-40% |
| 第三方库支持 | 几乎所有 AI 库都优先支持 CUDA | 支持范围在扩大但仍有缺口 |
| 调试工具 | Nsight Systems/Compute,成熟完善 | ROCm Profiler 功能在追赶 |
| 文档和社区 | 极其丰富 | 薄弱 |
实际情况:在 MI300X 上跑 PyTorch 训练,大部分标准模型的性能已经接近 H100(得益于 MI300X 更大的 HBM 容量优势)。但涉及到定制 Kernel、复杂推理优化或特殊模型结构时,CUDA 的生态深度仍然明显领先。
国内软件栈
| 公司 | 软件栈 | 策略 |
|---|---|---|
| 寒武纪 | Neuware(CNToolkit + CNNL + CNRT) | 自研编程模型 BANG,学习成本高 |
| 沐曦 | MXMACA | 兼容 CUDA/ROCm 接口,降低迁移成本 |
| 华为 | CANN + MindSpore | 自研算子库 + 框架,与昇腾硬件绑定 |
沐曦的策略值得关注:创始团队来自 AMD,选择兼容 ROCm 生态路线。如果 ROCm 的上层生态持续完善,沐曦可以”搭便车”降低自己的软件生态建设成本。这是一个务实的后发策略。
AI 编译器
AI 编译器解决的问题是:PyTorch 定义的计算图(高层抽象)如何被翻译成 GPU 上高效执行的机器指令(底层具体)。
AI 框架层(PyTorch, JAX)
│ 定义计算图(如: y = Linear(x) → ReLU → Softmax)
↓
AI 编译器
│ 图优化(算子融合、常量折叠、内存规划)
│ 代码生成(针对目标硬件生成 Kernel)
↓
GPU Kernel
│ 在特定硬件上执行
↓
结果
主要的 AI 编译器:
| 编译器 | 开发方 | 定位 |
|---|---|---|
| XLA | TensorFlow/JAX 的后端,TPU 的核心软件栈 | |
| TorchInductor | Meta/PyTorch | PyTorch 2.0 的默认编译后端(torch.compile) |
| Triton | OpenAI | 用 Python 写 GPU Kernel,在手写 CUDA 和自动编译之间找平衡 |
| TVM | Apache 开源社区 | 支持多种硬件后端,自动调优 |
| MLIR | LLVM 基金会 | 编译器基础设施,提供多层中间表示(IR) |
torch.compile 是 PyTorch 2.0 最重要的新特性。只需要在代码中加一行 model = torch.compile(model),PyTorch 会自动将 Python 代码捕获为计算图,然后通过 TorchInductor 生成优化的 GPU Kernel。实际测试中,torch.compile 通常能带来 30-50% 的训练速度提升(取决于模型结构)。
Triton 的价值在于降低了自定义 Kernel 的门槛。传统上,写一个高性能的 CUDA Kernel 需要深入了解 GPU 微架构(Shared Memory 大小、Bank Conflict、Warp Divergence 等)。Triton 用 Python 语法表达 Kernel 逻辑,自动处理内存管理和调度细节。FlashAttention 的核心实现就使用了 Triton。
§20. 精度控制与量化
在 GPU 一章 §6.2.2 中,我们看到 Blackwell GPU 支持从 FP64 到 FP4 的多种数据格式,并且提到”在 AI Infra 中会再次回到量化”。现在兑现这个承诺。
20.1 为什么精度很重要
回顾 GPU 一章的算力表:
| 精度 | B200 Dense 算力 |
|---|---|
| FP32 | 75 TFLOPS |
| FP16/BF16 | 2,250 TFLOPS |
| FP8 | 4,500 TFLOPS |
| FP4 | 9,000 TFLOPS |
从 FP32 到 FP4,算力提升 120 倍。同时,低精度的模型权重占用更少的 HBM 空间(FP4 的权重只有 FP16 的 1/4),这意味着同样的 GPU 可以装下更大的模型。
精度降低的代价是数值精度的损失。精度控制的核心问题是:在算力提升和精度损失之间找到最优的权衡点。
20.2 混合精度训练(Mixed Precision Training)
直接用 FP16 替换 FP32 训练会遇到两个问题:
- 梯度下溢(Underflow):FP16 的最小正数约为 $5.96 \times 10^{-8}$,而训练后期的梯度值常常小于这个数,被直接截断为零,导致训练停滞
- 权重更新精度不足:如果权重是 1.0 而梯度是 0.0001,FP16 无法表示 1.0001(精度不够),权重不会被更新
混合精度训练的标准做法(NVIDIA 在 2018 年提出):
训练循环:
1. 维护一份 FP32 精度的主权重(Master Weights)
2. 前向传播:
FP32 权重 → 转换为 FP16/BF16 → 执行前向计算(FP16/BF16)
→ 输出 FP16/BF16 的 Loss
3. Loss Scaling:
Loss × Scale Factor(如 1024)
→ 放大 Loss,使得后续反向传播的梯度也被放大
→ 避免 FP16 梯度下溢
4. 反向传播:
用放大后的 Loss 计算 FP16/BF16 梯度
5. 梯度反缩放:
FP16 梯度 ÷ Scale Factor → 恢复真实梯度值
→ 转换为 FP32
6. 权重更新:
FP32 主权重 += FP32 梯度 × 学习率
→ 在 FP32 精度下完成更新,避免精度不足
7. 回到步骤 2
整个过程的关键:前向和反向传播用 FP16/BF16(快),权重更新用 FP32(准)。Loss Scaling 确保小梯度不会在 FP16 中被截断为零。
在 H100 和 Blackwell 上,Transformer Engine 自动管理混合精度训练。它会逐层动态选择 FP8 或 FP16 进行计算:对数值分布”友好”的层使用 FP8(算力翻倍),对敏感的层保留 FP16。程序员不需要手动管理每层的精度选择。
20.3 推理量化
训练完成后,模型的权重是固定的。推理阶段不需要计算梯度,也不需要更新权重,因此可以使用更激进的低精度表示。
训练后量化(PTQ:Post-Training Quantization):
原始 FP16 权重: [0.123, -0.456, 0.789, -0.012, ...]
量化为 INT8:
1. 找到权重的最大绝对值: max_abs = 0.789
2. 计算缩放因子: scale = max_abs / 127 = 0.00621
3. 量化: INT8_weight = round(FP16_weight / scale)
→ [20, -73, 127, -2, ...]
4. 推理时: FP16_approx = INT8_weight × scale
→ [0.1243, -0.4536, 0.7888, -0.0124, ...]
原始值 0.123 被近似为 0.1243, 误差 ~1%
原始值 -0.012 被近似为 -0.0124, 误差 ~3%
PTQ 简单快速(不需要重新训练),但精度损失可能较大,尤其在权重分布不均匀(存在 Outlier)时。
量化感知训练(QAT:Quantization-Aware Training):
在训练过程中模拟量化效应。前向传播时插入”伪量化”节点,让模型在训练时就学会适应量化带来的精度损失。QAT 的精度通常优于 PTQ,但需要额外的训练时间。
大模型专用量化方法:
| 方法 | 核心思想 | 典型精度 |
|---|---|---|
| GPTQ | 逐层量化,用 Hessian 信息最小化量化误差 | INT4 / INT3 |
| AWQ(Activation-Aware Weight Quantization) | 保护对激活值影响大的”关键权重”,其余激进量化 | INT4 |
| SqueezeLLM | 利用权重的非均匀分布,对 Outlier 单独处理 | INT4 / INT3 |
| GGUF(llama.cpp) | 多种量化方案的文件格式,支持 Q2 到 Q8 | INT2~INT8 |
| Micro-Tensor Scaling(Blackwell) | 每组数据共享一个缩放因子,FP4 精度下维持模型质量 | FP4 / FP6 |
[!info] 量化的投资含义
量化技术的进步直接降低了 AI 推理的硬件成本。一个 70B 参数的模型:
- FP16 权重占用 ~140 GB HBM(需要 2 张 H100)
- INT4 量化后占用 ~35 GB HBM(1 张 H100 即可)
推理成本直接减半以上。这是推理 Token 价格持续下降的技术基础之一。
从硬件的角度看,Blackwell 支持 FP4 是一个信号:NVIDIA 判断 FP4 精度在适当的 Scaling 技术下足够满足推理需求。如果 FP4 推理成为主流,同样算力的 GPU 可以服务的用户量将进一步翻倍。
§21. 分布式训练
21.1 为什么需要分布式训练
一个直接的计算:训练一个 175B 参数的 GPT-3 级别模型,需要多少算力?
GPT-3 训练估算:
模型参数: 175B (175 × 10⁹)
训练数据: 300B Tokens
每 Token 所需 FLOPs ≈ 6 × 参数量 (近似公式)
总 FLOPs = 6 × 175 × 10⁹ × 300 × 10⁹ = 3.15 × 10²³
用 H100 (FP16 Tensor Core: 989 TFLOPS):
理论最快时间 = 3.15 × 10²³ / (989 × 10¹²) ≈ 3.18 × 10⁸ 秒 ≈ 10 年
实际 MFU 约 50%:
单卡实际时间 ≈ 20 年
用 1000 张 H100,假设 90% 的线性加速比:
实际时间 ≈ 20 年 / (1000 × 0.9) ≈ 8 天
这就是为什么大模型训练需要成千上万张 GPU 并行工作。但”并行”不是简单地把任务切成 1000 份分给 1000 张卡。模型太大,一张卡装不下;数据有依赖关系,梯度需要同步;通信开销可能抵消并行带来的加速。
分布式训练的核心挑战是:如何在成千上万张 GPU 之间高效分配计算任务,同时最小化通信开销。
21.2 数据并行(Data Parallelism)
数据并行是最直观的并行策略。
数据并行的基本流程:
假设 4 张 GPU,训练数据 Batch Size = 1024
1. 每张 GPU 持有模型的完整副本
2. 每张 GPU 分到 1024/4 = 256 条数据
GPU 0: 数据 [0:256] → 前向 → 反向 → 梯度 grad_0
GPU 1: 数据 [256:512] → 前向 → 反向 → 梯度 grad_1
GPU 2: 数据 [512:768] → 前向 → 反向 → 梯度 grad_2
GPU 3: 数据 [768:1024] → 前向 → 反向 → 梯度 grad_3
3. All-Reduce: 4 份梯度求平均 → 每张 GPU 得到相同的平均梯度
4. 每张 GPU 用平均梯度更新自己的模型副本
5. 由于梯度相同,更新后 4 张 GPU 的模型参数仍然一致
6. 回到步骤 2
优点:实现简单,加速比接近线性(只要通信不成为瓶颈)。
缺点:每张 GPU 必须持有完整的模型。对于 175B 参数的模型,FP16 下权重本身占 350 GB,加上优化器状态(Adam 需要额外 2 倍参数量的内存),总内存需求约 1050 GB。而 H100 只有 80 GB HBM,单卡根本装不下。
PyTorch DDP(Distributed Data Parallel) 是数据并行的标准实现。它在反向传播过程中,边计算梯度边发起 All-Reduce 通信(而不是等所有梯度都算完再通信),从而实现计算与通信的重叠。
ZeRO:内存优化的数据并行
ZeRO(Zero Redundancy Optimizer) 是 Microsoft DeepSpeed 提出的优化,解决了”每张卡都存完整模型”的内存浪费问题。
在标准数据并行中,每张 GPU 存储的内容包括:
| 内容 | FP16 模型大小 $\Phi$ 时的内存 |
|---|---|
| 模型参数(FP16) | $2\Phi$ bytes |
| 梯度(FP16) | $2\Phi$ bytes |
| 优化器状态(Adam: FP32 参数 + FP32 动量 + FP32 方差) | $12\Phi$ bytes |
| 总计 | $16\Phi$ bytes |
对于 7.5B 参数模型:$\Phi = 7.5 \times 10^9$,总内存 = 120 GB,单张 H100(80 GB)装不下。
ZeRO 的三个阶段,逐步将这些冗余数据分片到各 GPU 上:
| 阶段 | 分片内容 | 每 GPU 内存(N 张卡) |
|---|---|---|
| ZeRO-1 | 优化器状态 | $2\Phi + 2\Phi + 12\Phi/N$ |
| ZeRO-2 | 优化器状态 + 梯度 | $2\Phi + (2\Phi + 12\Phi)/N$ |
| ZeRO-3 | 优化器状态 + 梯度 + 参数 | $(2\Phi + 2\Phi + 12\Phi)/N = 16\Phi/N$ |
ZeRO-3 在 64 张 GPU 上训练 7.5B 模型时,每 GPU 只需要 $16 \times 7.5 \text{GB} / 64 = 1.875 \text{GB}$(仅模型相关部分),极大释放了 HBM 容量给激活值(Activation)和 KV Cache。
代价是通信量增加:ZeRO-3 在每次前向和反向传播时,都需要通过 All-Gather 临时收集完整参数,计算后丢弃。但在 NVLink 高带宽的节点内,这个开销可以被有效隐藏。
PyTorch FSDP(Fully Sharded Data Parallel) 是 ZeRO-3 思想在 PyTorch 官方的实现。
21.3 张量并行(Tensor Parallelism, TP)
当模型的单层参数就已经很大时(比如一个 12288 × 49152 的权重矩阵),数据并行无法解决问题,因为它要求每张 GPU 都存完整的参数。张量并行的做法是:将一层的权重矩阵切分到多张 GPU 上,每张 GPU 只计算一部分。
以一个简单的线性层 $Y = XW$ 为例,$W$ 是 $[K \times N]$ 的权重矩阵,$X$ 是 $[B \times K]$ 的输入:
列并行(Column Parallelism):
将 W 按列切分到 2 张 GPU
W = [W₁ | W₂] (K×N 切成两个 K×N/2)
GPU 0: Y₁ = X × W₁ (B×K) × (K×N/2) = (B×N/2)
GPU 1: Y₂ = X × W₂ (B×K) × (K×N/2) = (B×N/2)
Y = [Y₁ | Y₂] 拼接得到完整的 B×N 输出
行并行(Row Parallelism):
将 W 按行切分,同时将 X 按列切分
W = [W₁] X = [X₁ | X₂]
[W₂]
GPU 0: Y₁ = X₁ × W₁ (B×K/2) × (K/2×N) = (B×N)
GPU 1: Y₂ = X₂ × W₂ (B×K/2) × (K/2×N) = (B×N)
Y = Y₁ + Y₂ All-Reduce 求和得到完整输出
在 Transformer 中,Megatron-LM 提出了一种精巧的张量并行方案:对 MLP 层的第一个线性层用列并行,第二个线性层用行并行。这样在两层之间不需要额外的通信操作(列并行的输出直接作为行并行的输入),只在第二层结束后做一次 All-Reduce。
Megatron-LM 的 MLP 张量并行(2 GPU):
输入 X
│
│ Broadcast(X 复制到两张 GPU)
↓
GPU 0: X × W₁_col0 → GeLU → × W₂_row0 ──┐
GPU 1: X × W₁_col1 → GeLU → × W₂_row1 ──┤
↓
All-Reduce(求和)
↓
输出 Y
张量并行的带宽需求极高,因为每一层的前向和反向传播都需要通信。因此,张量并行通常只在 NVLink 连接的节点内(同一台机器的 8 张 GPU 之间)使用。跨节点的 InfiniBand 带宽不足以支撑张量并行的通信频率。
21.4 流水线并行(Pipeline Parallelism, PP)
流水线并行的思路不同:将模型的不同层放在不同的 GPU 上。
假设模型有 32 层 Transformer,4 张 GPU:
GPU 0: Layer 0-7 (前 8 层)
GPU 1: Layer 8-15 (中间 8 层)
GPU 2: Layer 16-23 (中间 8 层)
GPU 3: Layer 24-31 (后 8 层)
前向传播: GPU 0 → GPU 1 → GPU 2 → GPU 3
反向传播: GPU 3 → GPU 2 → GPU 1 → GPU 0
问题:流水线气泡(Pipeline Bubble)
在朴素的实现中,当 GPU 0 在做前向传播时,GPU 1/2/3 只能空等。这种空闲时间叫做”气泡”。如果有 4 个流水线阶段,朴素实现的 GPU 利用率只有约 25%(3/4 的时间在等待)。
解决方案是 Micro-Batching:将一个大的 Mini-Batch 切分成多个 Micro-Batch,让多个 Micro-Batch 交替进入流水线:
GPipe 调度(假设 4 个 Micro-Batch,4 个 Pipeline Stage):
时间 →
GPU 0: |F₁|F₂|F₃|F₄| |B₄|B₃|B₂|B₁|
GPU 1: |F₁|F₂|F₃|F₄|B₄|B₃|B₂|B₁|
GPU 2: |F₁|F₂|F₃|F₄|B₃|B₂|B₁|
GPU 3: |F₁|F₂|F₃|F₄|B₂|B₁|
F = 前向传播,B = 反向传播
气泡仍然存在,但比例从 75% 降低到了 ~37.5%
更高级的调度策略(如 1F1B,即交替执行一次前向和一次反向)可以进一步减少气泡:
1F1B 调度:
GPU 0: |F₁|F₂|F₃|F₄|B₁|B₂|B₃|B₄|
GPU 1: |F₁|F₂|F₃|B₁|B₂|F₄|B₃|B₄|
GPU 2: |F₁|F₂|B₁|F₃|B₂|F₄|B₃|B₄|
GPU 3: |F₁|B₁|F₂|B₂|F₃|B₃|F₄|B₄|
气泡比例进一步降低
且内存占用更低(不需要同时保存所有 Micro-Batch 的激活值)
流水线并行的优势:跨层的通信量很小(只需要传递层间的激活值,不需要同步权重),因此适合跨节点部署。
流水线并行的劣势:气泡无法完全消除,且负载均衡(每个阶段的计算量应尽量相等)需要精心设计。
21.5 序列并行、上下文并行与专家并行
除了上述三大并行策略外,还有几种面向特定场景的并行方式:
序列并行(Sequence Parallelism, SP)
在张量并行中,LayerNorm 和 Dropout 等操作不涉及权重矩阵切分,通常在每张 GPU 上独立执行完整的序列。这造成了冗余计算。序列并行将这些操作也按序列维度切分,减少冗余。
上下文并行(Context Parallelism, CP)
专门解决长上下文训练的内存问题。当序列长度达到 128K 甚至 1M 时,Attention 的计算量和 KV Cache 的内存都按序列长度的平方增长。上下文并行将长序列切分到多张 GPU 上,每张 GPU 处理序列的一部分,通过特殊的 Ring Attention 或 Striped Attention 机制在 GPU 之间传递 KV 数据。
专家并行(Expert Parallelism, EP)
专为 MoE(Mixture of Experts)模型设计。MoE 模型有多个”专家”子网络,每个 Token 只激活其中一部分(如 8 选 2)。不同的专家放在不同的 GPU 上,Token 通过 All-to-All 通信路由到对应的专家处理。
DeepSeek-V2/V3 等模型大量使用 MoE 架构,专家并行是其训练效率的关键。
21.6 3D 混合并行与 Megatron-LM
在实际的大模型训练中,以上几种并行策略是同时使用的,这叫做 3D 并行(Data Parallelism + Tensor Parallelism + Pipeline Parallelism)。
典型的 3D 并行配置(以 512 张 H100 训练 175B 模型为例):
总 GPU 数: 512
= 8 (TP) × 8 (PP) × 8 (DP)
TP = 8: 节点内 8 张 GPU 做张量并行(NVLink 高带宽)
PP = 8: 8 个节点串成流水线(InfiniBand 跨节点)
DP = 8: 8 组流水线做数据并行(InfiniBand All-Reduce)
为什么 TP 在节点内,PP 在节点间?
回看 §18.1 的互联带宽表:NVLink 的带宽是 InfiniBand 的 20 倍以上。张量并行每层都要通信,对带宽要求极高,只能放在 NVLink 连接的节点内。流水线并行只在层间传递激活值,通信量小,可以走 InfiniBand。
Megatron-LM(NVIDIA 开发)是 3D 并行训练的事实标准框架。它提供了:
- 高效的张量并行实现(MLP + Attention 的切分策略)
- 流水线并行调度(Interleaved 1F1B)
- 与 NCCL 的深度集成
- 序列并行和上下文并行支持
- 与 PyTorch 和 DeepSpeed 的兼容
DeepSpeed(Microsoft 开发)专注于内存优化,提供 ZeRO 系列优化。实际中 Megatron-LM 和 DeepSpeed 常常配合使用(Megatron-DeepSpeed)。
21.7 训练效率指标
**MFU(Model FLOP Utilization)**是衡量训练效率的核心指标:
$$MFU = \frac{\text{模型实际有效 FLOPs}}{\text{GPU 理论峰值 FLOPs} \times \text{训练时间}}$$
MFU 反映了 GPU 的算力到底有多少被”有效利用”了。其余的算力浪费在:
| 浪费来源 | 说明 |
|---|---|
| 通信等待 | GPU 在等 All-Reduce / All-Gather 完成 |
| 流水线气泡 | Pipeline Parallelism 中的空闲时间 |
| 内存访问 | GPU 计算单元等 HBM 数据 |
| Kernel 启动开销 | CPU 向 GPU 发起 Kernel Launch 的延迟 |
| 重计算(Recomputation) | Gradient Checkpointing 带来的额外计算 |
业界的 MFU 水平:
| 系统 | MFU | 备注 |
|---|---|---|
| Google PaLM(2022,TPU v4) | 46.2% | 当时的标杆 |
| Meta LLaMA-2(2023,H100) | ~42% | 公开报告 |
| 优秀的 Megatron-LM 集群(2025) | 50-55% | 高度优化 |
| 理论上限 | ~60-65% | 受通信和内存墙限制 |
FlashAttention 是近年来对训练效率影响最大的单一优化。
标准 Attention 的计算过程:
标准 Attention(以序列长度 N、头维度 d 为例):
1. Q × K^T → S (N×N 矩阵) ← 这个矩阵存在 HBM 中,占用 O(N²) 空间
2. Softmax(S) → P (N×N 矩阵) ← 再写一次 HBM
3. P × V → O (N×d 矩阵) ← 再读一次 HBM
总 HBM 读写量: O(N² × d)
当 N = 128K 时,S 矩阵大小 = 128K × 128K × 2 bytes = 32 GB
FlashAttention 的核心思想:不把 N×N 的中间矩阵写入 HBM,而是在 Shared Memory 中分块计算。
FlashAttention(Tiling 策略):
将 Q, K, V 各自分成大小为 B 的块
for each Q_block:
for each K_block, V_block:
1. 加载 Q_block, K_block, V_block 到 Shared Memory(快)
2. 在 Shared Memory 中计算局部 Attention
3. 通过在线 Softmax 技巧,增量更新输出
end
end
总 HBM 读写量: O(N × d) (比标准 Attention 的 O(N² × d) 降低了 N/d 倍)
FlashAttention 将 Attention 从 memory-bound(瓶颈在 HBM 读写)变为 compute-bound(瓶颈在 Tensor Core 算力),实际速度提升 2-4 倍,同时内存占用从 O(N²) 降到 O(N)。这使得长上下文训练(128K+)成为可能。
Gradient Checkpointing(梯度检查点) 是另一个重要的内存优化。正常训练中,反向传播需要用到前向传播的中间激活值,因此所有激活值都要保存在内存中。激活值的内存占用往往超过模型参数本身。
Gradient Checkpointing 的做法是只保存部分激活值(检查点),需要用时重新计算。这用 额外约 33% 的计算量换取了 大幅降低的内存占用(通常节省 60-80% 的激活值内存),使得更大的 Batch Size 或更大的模型成为可能。
§22. 推理优化
22.1 推理与训练的本质差异
| 维度 | 训练 | 推理 |
|---|---|---|
| 批量大小 | 大 Batch(数千甚至数万个样本同时处理) | 小 Batch(甚至单条请求) |
| 计算特性 | 计算密集(大矩阵乘法,高算术强度) | 内存带宽密集(小矩阵,低算术强度) |
| 梯度 | 需要(反向传播) | 不需要 |
| 精度 | FP16/BF16(训练需要较高精度) | INT8/INT4/FP8(推理可以更低精度) |
| 延迟要求 | 不敏感(训练一次可以跑几天) | 敏感(用户等待响应) |
| 成本结构 | 一次性(训练完就结束) | 持续性(每次用户请求都消耗算力) |
推理的瓶颈主要在内存带宽,而非算力。
LLM 推理分为两个阶段:
- Prefill(预填充):处理用户输入的所有 Token,计算对应的 KV Cache。这一步可以并行处理所有输入 Token,是 compute-bound 的。
- Decode(解码):逐个生成输出 Token。每生成一个 Token,需要读取整个模型权重和所有历史 Token 的 KV Cache,但只做一次矩阵向量乘法。这一步是 memory-bandwidth-bound 的。
推理的两个阶段:
Prefill:
输入 "What is the capital of France?"
→ 一次性处理所有 8 个 Token
→ 计算 8 个 Token 的 KV Cache
→ compute-bound(大矩阵乘法)
Decode:
生成 "The" → 读取全部权重 + KV Cache → 输出 1 个 Token
生成 "capital" → 读取全部权重 + 更新后的 KV Cache → 输出 1 个 Token
生成 "of" → ...
...
→ 每步只生成 1 个 Token
→ memory-bandwidth-bound(每步读取 GB 级数据,只做极少计算)
以 LLaMA-2 70B 模型在 H100 上推理为例:
模型权重(FP16): 140 GB
H100 HBM 带宽: 3.35 TB/s
每生成一个 Token,至少需要读取全部权重一次:
最快生成速度 = 3.35 TB/s ÷ 140 GB ≈ 24 Tokens/s
H100 的 FP16 Tensor Core 算力: 989 TFLOPS
实际推理时的算力利用率: 通常 < 5%
→ 99%+ 的算力在空转,瓶颈完全在 HBM 带宽
这就解释了为什么推理优化的核心是减少 HBM 访问量(量化、KV Cache 压缩)和提高每次 HBM 访问的有效利用率(Batching)。
22.2 KV Cache 与显存管理
在 Transformer 的自回归推理中,每生成一个新 Token,Attention 机制需要用到所有历史 Token 的 Key 和 Value 向量。如果每次都重新计算,计算量会随序列长度二次增长。KV Cache 的做法是把已经计算过的 Key 和 Value 存在 HBM 中,每次只计算新 Token 的 K、V,然后和缓存的历史 KV 一起做 Attention。
KV Cache 的内存占用:
KV Cache 大小 = 2 × num_layers × num_kv_heads × head_dim × seq_len × precision_bytes
以 LLaMA-2 70B 为例:
num_layers = 80
num_kv_heads = 8 (GQA)
head_dim = 128
seq_len = 4096
precision = FP16 (2 bytes)
KV Cache = 2 × 80 × 8 × 128 × 4096 × 2 = 1.34 GB (per request)
如果同时服务 64 个请求:
KV Cache 总量 = 1.34 × 64 = 85.7 GB
模型权重 = 140 GB
总 HBM 需求 = 225.7 GB → 需要 3 张 H100(80 GB 每张)
KV Cache 的内存管理是推理系统的核心难题。传统做法是为每个请求预分配最大序列长度的 KV Cache 空间,但大多数请求不会用到最大长度,导致大量内存浪费。
PagedAttention(vLLM 的核心创新)借鉴了操作系统的虚拟内存管理思想,将 KV Cache 按固定大小的”页”分配,不再需要连续的内存空间。请求结束后,页被回收。这将 KV Cache 的内存利用率从约 50% 提升到接近 100%,等效于在同一组 GPU 上可以服务接近两倍的并发请求。
22.3 推理框架
| 框架 | 开发方 | 核心技术 | 定位 |
|---|---|---|---|
| vLLM | UC Berkeley | PagedAttention + Continuous Batching | 开源推理框架,LLM 推理最流行的选择 |
| TensorRT-LLM | NVIDIA | TensorRT 图优化 + 定制 CUDA Kernel | NVIDIA 硬件上的极致性能优化 |
| SGLang | Stanford / UC Berkeley | RadixAttention + 结构化生成优化 | 支持复杂推理工作流(如 Tool Use、多轮对话) |
| llama.cpp | Georgi Gerganov | CPU 推理 + 激进量化(GGUF 格式) | 边缘设备和个人电脑推理 |
| MLC-LLM | CMU / OctoAI | 编译优化,多硬件后端 | 跨平台推理(手机、浏览器) |
vLLM 目前是开源 LLM 推理的事实标准。除了 PagedAttention,vLLM 的关键特性还包括:
- Continuous Batching:动态地将新到达的请求插入正在处理的 Batch 中。传统的 Static Batching 需要等一个 Batch 中所有请求都生成完毕才能开始下一个 Batch,短请求被长请求”拖累”。Continuous Batching 让短请求提前完成并释放资源,新请求立即填入。
- Prefix Caching:如果多个请求有相同的前缀(如系统提示词),可以共享这部分的 KV Cache,避免重复计算。
- Tensor Parallelism 支持:多卡推理,将模型切分到多张 GPU 上。
TensorRT-LLM 是 NVIDIA 官方的推理引擎,在 NVIDIA 硬件上通常有 10-30% 的性能优势(相比 vLLM)。但它的使用门槛更高,且与 NVIDIA 硬件深度绑定。
22.4 Speculative Decoding(投机解码)
标准的自回归解码每步只生成一个 Token,严格串行。Speculative Decoding 的思路是:用一个小而快的”草稿模型”(Draft Model)一次性预测多个 Token,然后用大模型并行验证这些预测。
标准解码(Serial):
大模型: → Token 1 → Token 2 → Token 3 → Token 4 → ...
每步: 读取全部权重,生成 1 个 Token
Speculative Decoding:
草稿模型(如 1B 参数):
→ 快速生成 4 个候选 Token: [T1, T2, T3, T4]
大模型(如 70B 参数):
→ 一次性验证 4 个 Token(并行前向传播)
→ 假设前 3 个通过验证,第 4 个被拒绝
→ 接受 [T1, T2, T3],从 T3 处重新开始
效果: 一次大模型前向传播生成了 3 个 Token(而不是 1 个)
加速比取决于草稿模型的"猜中率",通常 2-3 倍
Speculative Decoding 的数学保证:最终生成的 Token 序列与直接用大模型生成的分布完全一致(通过接受/拒绝采样保证)。也就是说,输出质量不受影响,只是速度变快了。
22.5 推理成本分析
以 2025-2026 年的市场价格为参考:
H100 SXM(80GB)的数据中心成本:
硬件租赁: ~$2-3/小时/GPU(云服务商价格)
电费: ~700W × $0.10/kWh ≈ $0.07/小时
冷却等间接成本: 约电费的 30-50%
每 GPU 小时总成本: ~$2.5-3.5
LLaMA-2 70B 推理(FP16, 2 张 H100):
吞吐量: ~2000-4000 output tokens/s(取决于 Batch Size)
成本: ~$5-7/百万 output tokens
INT4 量化后(1 张 H100):
吞吐量: ~3000-6000 output tokens/s
成本: ~$1-2/百万 output tokens
→ 量化使推理成本降低了约 3-4 倍
这就是为什么 API 厂商的 Token 价格在过去两年间持续快速下降:硬件代际升级(H100 → B200)、量化技术进步(FP16 → INT4/FP4)、推理框架优化(vLLM 的 PagedAttention + Continuous Batching)共同推动了推理效率的指数级提升。
§23. AI 数据中心基础设施
23.1 典型训练集群配置
以一个用于训练 ~100B 参数模型的 H100 集群为参考:
节点配置:
DGX H100
├── 8 × H100 SXM5 GPU(80GB HBM3 每张)
├── 2 × Intel Xeon / AMD EPYC CPU
├── 2 TB 系统内存(DDR5)
├── 8 × NVMe SSD(30 TB 本地存储)
├── 4 × NVSwitch(节点内全互联)
├── 8 × ConnectX-7 InfiniBand HCA(每 GPU 一个)
└── 功耗: ~10 kW
集群配置(1024 GPU):
128 个 DGX H100 节点
InfiniBand NDR 400 Gb/s Fat-Tree 网络
并行文件系统: 数 PB 容量,数百 GB/s 吞吐
总功耗: ~1.5-2 MW
估计成本: ~$300-500M(含基础设施)
GB200/GB300 NVL72 时代的新范式:
NVL72 机架将 36 个 GB200 Superchip(72 GPU + 36 Grace CPU)装入一个标准 42U 机架。整个机架对外表现为一个统一的计算节点,72 颗 GPU 通过 NVLink Switch 全互联。
| 指标 | 1024 × H100 集群 | ~15 × NVL72 机架(≈1080 GPU) |
|---|---|---|
| GPU 总数 | 1,024 | 1,080 |
| FP8 总算力 | ~1,000 PFLOPS | ~4,860 PFLOPS |
| HBM 总容量 | ~80 TB | ~200+ TB |
| GPU 间互联 | 节点内 NVLink + 节点间 InfiniBand | 机架内 NVLink 全互联 + 机架间 InfiniBand |
| 功耗 | ~1.5 MW | ~1.8 MW |
NVL72 的核心优势在于:机架内 72 颗 GPU 全部通过 NVLink 互联,消除了”节点内 8 卡 NVLink + 节点间 InfiniBand”的带宽断层。这对张量并行尤其有利:TP 的规模可以从 8 扩展到 72,大幅减少对 InfiniBand 的依赖。
23.2 电力与散热
电力正在成为 AI 扩张的物理约束。
算力需求的增长:
2020: A100 集群, TDP 400W/GPU
2022: H100 集群, TDP 700W/GPU (+75%)
2024: B200 集群, TDP 1000W/GPU (+43%)
一个万卡 B200 集群:
GPU 功耗: 10,000 × 1000W = 10 MW
加上 CPU、网络、存储、散热: 总功耗约 15-20 MW
15 MW ≈ 一个小型城镇的用电量
年电费($0.08/kWh): ~$10M
散热方案的演进:
| 阶段 | 方案 | 适用功耗 | 数据中心改造成本 |
|---|---|---|---|
| 传统 | 风冷(服务器风扇 + 机房空调) | ≤ 300W/GPU | 低 |
| 过渡 | 后门热交换器(风冷 + 液冷辅助) | 300-700W/GPU | 中 |
| 当代 | 直接液冷(DLC)(冷板贴合芯片,循环冷却液) | 700-1200W/GPU | 高(需要冷却管路) |
| 前沿 | 浸没式液冷(整台服务器浸泡在冷却液中) | >1200W/GPU | 极高 |
从 Blackwell 开始,NVIDIA 的数据中心 GPU(B200/B300)标配液冷。这要求数据中心具备冷却液循环管路、泵站、热交换器等设施,对已有数据中心的改造成本很高。
[!info] 从投资角度看电力约束
AI 数据中心的电力需求正在改变能源行业的投资逻辑。北美主要云厂商(Microsoft、Google、Amazon、Meta)在 2024-2025 年签订了大量长期电力采购协议(PPA),包括核电、天然气和可再生能源。
对数据中心运营商(如 Equinix、Digital Realty)而言,能否提供充足且稳定的电力容量,正在取代地理位置成为核心竞争力。一些新建的 AI 数据中心直接选址在电厂旁边。
中国的情况:国内数据中心主要集中在内蒙古、贵州等电力充裕且电价较低的地区。“东数西算”工程的背景之一,就是将高耗能的 AI 训练负载引向西部廉价电力区域。
23.3 存储系统
大模型训练对存储的需求主要在两方面:
- 训练数据加载:数 TB 的训练数据需要以 GB/s 级别的速度持续喂给 GPU,不能让 GPU 等数据
- Checkpoint 保存:训练中定期保存模型快照(Checkpoint),一个 175B 模型的 Checkpoint 约 350 GB(FP16),如果加上优化器状态约 1 TB
典型的存储架构:
┌──────────────────────────────────────┐
│ 并行文件系统 │
│ GPFS (IBM) / Lustre (开源) │
│ / WekaFS / VAST Data │
│ │
│ 容量: 数 PB │
│ 吞吐: 数百 GB/s(聚合) │
│ 通过 InfiniBand 或高速以太网连接 │
└───────────────┬──────────────────────┘
│
┌───────────┴───────────┐
│ NVMe SSD(每节点本地) │
│ 用于缓存热数据 │
│ GPUDirect Storage: │
│ GPU 直接从 SSD 读取 │
└─────────────────────────┘
并行文件系统将数据分散存储在数百个存储节点上,读取时从多个节点同时获取数据(striping),从而实现极高的聚合吞吐量。这和单台 NAS 的工作方式完全不同。
GPUDirect Storage 让 GPU 可以直接从 NVMe SSD 读取训练数据,跳过 CPU 内存中转,减少数据搬运延迟。
全篇回顾:从沙子到大模型
四章覆盖了从物理世界到 AI 模型的完整技术栈:
云 :物理硬件如何被抽象、池化,交付给用户
CPU :逻辑门 → 指令周期 → 存储层次 | 硅砂 → 光刻 → 封装 → 全球供应链
GPU :CUDA Core → Tensor Core → NVL72
AI Infra :互联 → CUDA → 分布式训练 → 推理优化
现在可以用这套框架回答一个入口问题了。
英伟达会被其他芯片厂商颠覆吗?
短期(2-3年)内不会。NVIDIA 的优势不在于单颗 GPU 的硬件参数(AMD MI300X 的 HBM 容量已经超过 H100),而在于:
- CUDA 生态的 19 年积累(§19)
- NVLink + InfiniBand 的垂直整合互联能力(§18)
- NVL72 级别的系统工程能力(§23)
- cuBLAS/cuDNN/TensorRT/NCCL 等核心库的优化深度(§19.3)
这四层壁垒的叠加,使得”硬件参数接近”距离”系统级性能接近”还有很大距离。
中长期的风险点:(1) AMD ROCm 生态持续成熟,加上云厂商(Google TPU、AWS Trainium)自研芯片的规模化部署,逐步侵蚀 NVIDIA 的份额;(2) 推理市场占比持续增大,推理对 CUDA 生态的依赖度低于训练,给替代方案留出空间;(3) 开源编译器(Triton、TVM)降低了硬件切换的软件成本。
AI芯片有GPU之后的下一个范式吗?
目前没有看到能完全替代 GPU 的架构。TPU(Google 的脉动阵列)在 Transformer 场景下性能优异,但通用性弱于 GPU,且仅供 Google 内部使用。存内计算(CIM)和光计算仍处于早期研究阶段,距离商业化有数年的距离。
更现实的趋势是:GPU 本身在特化。Tensor Core 从通用浮点运算单元演变为矩阵乘加专用硬件,FP4/FP6 等极低精度格式的引入使 GPU 越来越”接近 ASIC 的效率”但保留了通用编程能力。GPU 正在吸收 ASIC 的思想,同时维持平台的通用性。
参考资料:
NVIDIA Megatron-LM: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
Microsoft DeepSpeed: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
Tri Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Kwon et al., vLLM: Efficient Memory Management for Large Language Model Serving with PagedAttention
Leviathan et al., Fast Inference from Transformers via Speculative Decoding
NVIDIA H100 Tensor Core GPU Architecture Whitepaper
NVIDIA Blackwell Architecture Technical Overview
NVIDIA NCCL Documentation: NVIDIA Collective Communications Library
InfiniBand Trade Association: InfiniBand Roadmap
SemiAnalysis(半导体与 AI 基础设施深度分析): semianalysis.com