“训练 AI 要用 GPU” 几乎成了常识，但常识背后的工程逻辑常被一句"GPU 并行快"糊弄过去。真正的答案要拆三层：硬件架构为什么适合，神经网络计算为什么吃这套，以及为什么很多时候瓶颈根本不是算力而是显存带宽。理解后者，你才知道为什么 kernel fusion、FlashAttention 这类优化能带来数倍加速。

直觉：吞吐机器 vs 延迟机器

CPU 是为低延迟设计的：少量强核心、巨大的乱序执行窗口、多级缓存、复杂分支预测，目标是让单条指令流尽快跑完。GPU 反过来，是高吞吐机器：成千上万个简单核心，靠海量并行掩盖单个操作的延迟。

一个常被引用的对比：CPU 可能十几到几十个核，GPU 有数千个 CUDA core。但核数只是表象，关键是执行模型——GPU 用 SIMT（Single Instruction, Multiple Threads）：一组线程（NVIDIA 的 warp，通常 32 个）锁步执行同一条指令，只是数据不同。这恰好契合神经网络的计算特征。

机制：神经网络 = 一堆大矩阵乘法

把前向传播剥开，核心就是矩阵乘法。一层全连接是 Y=XWY = XW，卷积可通过 im2col 展开成矩阵乘，注意力里的 QK⊤QK^\top、softmax⋅V\text{softmax}\cdot V 也都是矩阵乘。而矩阵乘是极度数据并行的：

Cij=∑kAikBkjC_{ij} = \sum_k A_{ik} B_{kj}

每个输出元素 CijC_{ij} 的计算彼此独立，可以分给不同线程同时算。一个 1024×10241024\times1024 矩阵乘有上百万个独立输出——正好喂饱 GPU 的几千个核心。这就是"神经网络天生适合 GPU"的本质：不是因为算法神秘，而是因为它的主要算子是规则、独立、可平铺的密集线性代数。

NVIDIA 的 Tensor Core 把这一步推到极致：它不是标量乘加单元，而是直接做小矩阵块的乘加（如 4×44\times4 块），单条指令完成一整块 MMA（Matrix Multiply-Accumulate），并原生支持 fp16/bf16/int8 等低精度。这是现代 GPU 在深度学习上相对纯 FP32 算力暴涨的主因。

被忽视的真相：瓶颈常在显存带宽

很多人以为 GPU 慢是因为"算力不够"，其实大量算子是显存带宽受限（memory-bound）而非算力受限（compute-bound）。判断标准是算术强度（arithmetic intensity）：

\text{AI} = \frac{\text{FLOPs}}{\text{Bytes moved}}\ \ (\text{每搬一字节做多少次浮点运算})

这是 Roofline 模型的核心。每块 GPU 有两个上限：峰值算力（FLOP/s）和峰值显存带宽（Byte/s）。它们的比值给出一个临界算术强度 I∗=peak FLOP/speak Byte/sI^* = \frac{\text{peak FLOP/s}}{\text{peak Byte/s}}：

  • 算子的 AI <I∗< I^* → 带宽受限，算力空转，瓶颈在搬数据。
  • 算子的 AI >I∗> I^* → 算力受限，才真正吃满计算单元。
1
2
性能 = min( 峰值算力 ,  显存带宽 × 算术强度 )
                          └────── 斜线区：带宽决定一切 ──────┘

逐元素算子（ReLU、加法、LayerNorm、GELU）算术强度极低——读一个数、做一两次运算、写回去，几乎纯搬运。大矩阵乘的算术强度随维度增长，才有机会进入算力受限区。这解释了为什么"算子融合（kernel fusion）"如此关键：把 matmul → bias → GELU → dropout 这一串本来各自读写显存的算子合并成一个 kernel，中间结果留在片上寄存器/共享内存里，省掉多次往返 HBM 的搬运，对带宽受限段是数倍加速。FlashAttention 的核心思想也是如此——不把巨大的 N×NN\times N 注意力矩阵物化到显存，而是分块在片上算完。

显存层级：为什么"放不下"比"算不动"更常见

GPU 显存（HBM，几十到上百 GB）容量大但相对慢；真正快的是片上的共享内存/L1 和寄存器，但极小（每个 SM 仅几十到上百 KB）。优化的艺术就是最大化片上数据复用、最小化 HBM 访问。训练时的显存占用大致是：

1
显存 ≈ 模型参数 + 梯度 + 优化器状态 + 激活值(activations)

其中激活值常是最大头，且随 batch size 和序列长度线性甚至平方增长。这是 OOM 最常见的来源，也催生了 gradient checkpointing（用重算换显存）、激活值量化、序列并行等技术。换句话说，工程上很多时候不是"训不快"，而是"装不下"。

工程权衡与常见误区

  • 小 batch 喂不饱 GPU。 并行度不足时，几千个核心大量闲置，吞吐远低于峰值。这也是推理（batch 常为 1）难以打满 GPU、要靠 continuous batching 把多请求拼起来的原因。
  • CPU↔GPU 拷贝是隐形杀手。 PCIe 带宽远低于 HBM，频繁 .cpu()/.cuda()、在训练循环里同步打印 loss 标量，都会逼停 GPU。尽量让数据驻留显存、异步传输。
  • GPU 是异步的。 kernel launch 后 CPU 立即返回，计时若不 synchronize 会测到假的"极快"。这是基准测试最经典的坑。
  • 低精度不是免费午餐。 fp16 动态范围窄易溢出，需 loss scaling；bf16 范围大但尾数少。能用就用（Tensor Core 提速明显），但要理解数值影响。
  • 不是所有任务都该上 GPU。 强分支、强串行依赖、数据量小的逻辑，GPU 的 SIMT 会因 warp divergence（同一 warp 内走不同分支被迫串行）而退化，未必比 CPU 快。

小结

GPU 之于 AI，不是"更快的 CPU"，而是一类吞吐导向、靠海量并行掩盖延迟的架构，恰好咬合神经网络"密集、独立、可平铺的矩阵乘"这一计算特征，再由 Tensor Core 把低精度矩阵乘推向极致。但真正决定你能不能榨干它的，往往是 Roofline 上那条带宽斜线——算术强度、算子融合、显存层级与激活值占用。把"算力"和"带宽"分开看，你才能解释清楚为什么有些优化立竿见影，有些则纯属安慰。