Nenifindo's Home

Appearance

2026-04-02
1078
4m

GPU 与 TPU ​

对应课程 Lecture 5:GPUs, TPUs(Tatsu)

要优化训练，得先理解硬件。这一讲回答：​GPU 为什么擅长深度学习？它的性能极限在哪？

GPU 的组织结构 ​

以 NVIDIA H100 为例，一块 GPU 内部大致是：

text
GPU
├── 数十个 SM(流式多处理器)          ← "CPU 核心们"
│    ├── 张量核心(Tensor Core)       ← 专做矩阵乘法的单位
│    ├── FP32/FP64 计算单元
│    └── 寄存器 + 共享内存(shared memory, ~228KB)  ← SM 内的高速缓存
├── HBM 显存(~80GB, 带宽 ~3TB/s)     ← 主显存
└── NVLink 卡间互连(~900GB/s)         ← 和别的 GPU 说话

三层存储的速度天差地别：寄存器最快、共享内存次之、HBM 最慢（差两个数量级）。​高性能内核的本质就是让数据尽量待在离计算单元近的地方。

关键硬件特性 ​

  • 张量核心（Tensor Core）​：一条指令完成一个小矩阵块的乘加（如 8×8），是 GPU 算力（1000 TFLOP/s bf16）的主要来源。普通 CUDA 核心的 FP32 算力只有它的几十分之一。
  • SIMT 执行模型​：线程按 warp（线程束，32 线程） 为单位调度，一个 warp 里所有线程执行同一条指令。如果同一 warp 内线程走了不同分支（分支发散），性能会掉。
  • 占用率（occupancy）​：SM 上同时驻留的 warp 数。驻留多才能在部分 warp 等显存时，让其他 warp 继续计算——用并行性掩盖访存延迟​。

一个容易犯的错误 ​

把小矩阵乘法交给 GPU，可能反而比 CPU 慢：kernel 启动开销（几微秒）+ 数据搬运时间，远超那点计算量。GPU 的哲学是大批量、少启动、算得密​。

Roofline：性能极限在哪里 ​

前面资源估算的结论在这里落到实处。对每个操作问两个数：

  • 需要多少 FLOPs?（计算时间 = FLOPs ÷ 峰值算力）
  • 需要搬运多少字节？（访存时间 = 字节 ÷ 带宽）

实际时间 ≈ 两者取大。由此得出经验：

操作运算强度瓶颈
大批次矩阵乘法O(d)，高计算受限，利用率可达 60~80%
softmax、激活、残差加≈ 1带宽受限，利用率 1% 量级
小批次解码（权重全读一遍）低带宽受限

Transformer 训练中，矩阵乘法贡献了几乎全部 FLOPs，所以优化目标是让矩阵乘法之外的部分“不拖后腿”——这就是算子融合和手写内核的动机（下一讲）。

通信：多卡系统的新维度 ​

单卡性能之外，集群里还有通信​：

  • PCIe​：通用总线，约 64 GB/s，最慢；
  • NVLink​：NVIDIA 私有互连，H100 上约 900 GB/s，同一节点内首选；
  • InfiniBand / 以太网​：跨节点，200~400 Gbps 量级。

经验法则：​算一算每次同步要搬多少数据，除以链路带宽​。如果通信时间占训练步时间的比例高，算法再优也快不起来。并行策略的设计核心就是最小化、掩盖通信。

TPU：另一条技术路线 ​

Google 的 TPU(Tensor Processing Unit) 是专为矩阵计算设计的 ASIC：

  • 核心是 MXU 脉动阵列（systolic array）​，如 128×128 的乘加单元网格，数据像流水一样在阵列中“脉动”传递，一个数据被读入后参与多次计算，​天然高运算强度&#8203；
  • 结构规整，编译器（XLA）驱动，不需要手写 CUDA 那样的底层代码；
  • 大模型时代 TPU 训练了 Google 系的多数旗舰模型（PaLM、Gemini 等）,latest TPU 的 bf16 集群算力与顶级 GPU 集群在同一量级。

GPU vs TPU 的取舍​：GPU 生态成熟、灵活、可手写内核；TPU 软件栈统一、集群扩展性好，但仅通过云租用。理解它们共享的 roofline 逻辑，比记住产品参数更重要。

小结 ​

  • GPU = 大量 SM + 张量核心 + 多层存储；性能来自并行性掩盖延迟、数据靠近计算。
  • 用 roofline 判断每个操作的瓶颈：矩阵乘法计算受限，逐元素操作带宽受限。
  • 多卡系统里，通信带宽是第三种资源，与算力、显存并列。
  • TPU 用脉动阵列走规整化路线，与 GPU 在大模型训练中并驾齐驱。
最近更新