第一章：Roofline 模型详解

🎯 跑算法就像开车送货

速度取决于三件事：

bg right:40%

📋 今日议程

  1. 🧮 为什么你的模型跑得慢？
  2. ⏱️ 三个时间杀手
  3. 📊 算术强度——一个数字看穿本质
  4. 📉 点积的悲剧
  5. 🏠 Roofline 图怎么看
  6. ✖️ 矩阵乘法的秘密
  7. 🔗 从单机到分布式
  8. 🎯 总结与实战

你的模型，卡在哪了？

核心问题

为什么同一个模型，跑 50ms 而不是 50s 或者 5ms？

三个可能的瓶颈：

瓶颈 描述 类比
🧮 计算 CPU/GPU 算不过来 工人干活慢
📦 内存带宽 数据搬运太慢 送货太慢
🔗 多卡通信 卡与卡之间传输慢 车间之间协调差

历史小故事

1995年，Patterson 和 Williams 发现：

超级计算机越来越快，但程序速度提升却不如预期。

为什么？

问题不在于 CPU 有多快，而在于 数据能不能及时送到 CPU 手里。

打字员的比喻

🧑‍💻 超级打字员：1000 字/分钟 📝 你递稿子：10 字/分钟

他再快也没用，大部分时间都在干等！

时间杀手 ① 计算时间

公式

\[T_{\text{计算}} = \frac{\text{FLOPs 总量}}{\text{芯片算力 (FLOPs/s)}}\]

主流芯片算力对比

芯片 bf16 算力
NVIDIA H100 ~9.89×10¹⁴ FLOPs/s
Google TPU v6e ~9.1×10¹⁴ FLOPs/s

例子

做 1万亿次（10¹²）运算： \(\frac{10^{12}}{9.89 \times 10^{14}} \approx 1\text{ms}\)

时间杀手 ② 芯片内通信

公式

\[T_{\text{内存}} = \frac{\text{数据量 (Bytes)}}{\text{HBM 带宽 (Bytes/s)}}\]

HBM 带宽对比

芯片 HBM 带宽
NVIDIA H100 3.35 TB/s
Google TPU v6e 1.6 TB/s

工厂流水线比喻

🏭 机器（计算核心）再快 📦 原材料（数据）送不过来 ⏸️ 机器就得停工等待！

时间杀手 ③ 芯片间通信

多卡连接方式

连接类型 说明
NVLink NVIDIA 卡间连接（快）
ICI TPU 卡间连接（快）
PCIe 通用接口（相对慢）

⚠️ 重要发现

卡间带宽比 HBM 低 1-2 个数量级

→ 多卡通信往往是更大的瓶颈！

好消息：计算与通信可重叠！

时间估算

最好情况（完美重叠）： \(T = \max(T_{\text{计算}}, T_{\text{通信}})\)

最坏情况（完全不重叠）： \(T = T_{\text{计算}} + T_{\text{通信}}\)

实际情况

算术强度 —— 一个数字看穿本质

定义

\[\text{算术强度} = \frac{\text{FLOPs}}{\text{Bytes}}\]

直觉：每搬运 1 字节数据，能做多少次运算？

判断法则

条件 结果
算法强度 > 临界强度 ✅ 计算受限（好事！）
算法强度 < 临界强度 ❌ 通信受限（浪费算力）

临界算术强度

公式

\[\text{临界强度} = \frac{\text{芯片算力}}{\text{带宽}}\]

TPU v5e MXU 的临界值

\[\frac{1.97 \times 10^{14}}{8.2 \times 10^{11}} \approx \boxed{240 \text{ FLOPs/Byte}}\]

金句

🎯 算术强度就像”性价比” —— 每搬运 1 字节数据能做多少活。 性价比越高，硬件利用率越高。

点积的悲剧 😢

场景：两个 bf16 向量点积（长度 N）

\[\mathbf{x} \cdot \mathbf{y} = x_1 y_1 + x_2 y_2 + \cdots + x_N y_N\]
项目 数量
加载数据 4N + 2 字节
运算次数 2N - 1 FLOPs

算术强度

\[\text{强度} = \frac{2N - 1}{4N + 2} \approx \boxed{0.5}\]

对比临界强度 240…… 差了 480 倍！💀

点积为什么这么惨？

超级数学家的比喻

🧮 数学家能力：每秒算 240 道题 📝 你送题速度：每秒 1 道题

他 99.5% 的时间都在摸鱼！

启示

这就是为什么 GPU/TPU 上很多”逐元素运算”效率不高：

它们的算术强度都很低！

Roofline 图 —— 一张图看懂性能

   FLOPs/s
      ↑
  峰值 ├────────────────── 天花板（算力极限）
      │              ╱
      │            ╱
      │          ╱  斜坡（带宽限制）
      │        ╱
      │      ╱
      │    ╱
      └───────────────────→ 算术强度
            ↑ 临界点

Roofline 图解读

两个区域

区域 位置 含义 性能表现
🔺 斜坡 临界点左边 通信受限 性能随强度线性增长
⬜ 平台 临界点右边 计算受限 性能达到峰值

跑步比赛比喻

矩阵乘法 —— 深度学习的心脏

场景：$X_{[B,D]} \times Y_{[D,F]} \rightarrow Z_{[B,F]}$ (bf16)

项目 计算
加载数据 2(BD + DF + BF) 字节
运算次数 2BDF FLOPs

算术强度

\[\text{强度} = \frac{2BDF}{2(BD + DF + BF)} = \frac{BDF}{BD + DF + BF}\]

矩阵乘法强度的简化

当 D, F » B 时

\[\text{强度} \approx \frac{BDF}{DF} = \boxed{B}\]

🎯 黄金法则

bf16 矩阵乘法要跑满算力，batch size > 240！

这解释了很多现象

  1. ✅ 训练时要用大 batch —— 保证每卡 batch 够大
  2. ⚠️ 推理时小 batch 效率低 —— 强度不够
  3. 📈 模型越大，batch 可以相对越小 —— D 和 F 变大了

从单机到分布式

场景：两张 TPU 联合做矩阵乘法

分布式策略

  1. X 和 Y 沿 D 维度各存一半
  2. TPU 0 算：X[:, :D/2] @ Y[:D/2, :]
  3. TPU 1 算：X[:, D/2:] @ Y[D/2:, :]
  4. 交换结果，求和得到 Z

多卡时间对比

  单卡 双卡
计算时间 $\frac{2BDF}{\text{算力}}$ $\frac{BDF}{\text{算力}}$ ⬇️
通信来源 HBM带宽 卡间带宽 🔄
通信量 2(BD+DF+BF) 2BF

临界条件变化

🎯 关键洞察

多卡时，瓶颈从 batch size 变成了模型宽度！

这意味着

  1. 模型越宽（D 越大），多卡扩展效果越好
  2. 窄模型即使加再多卡，也可能被卡间通信拖累
  3. 这就是为什么大模型训练要用”够宽“的架构！

思考题 🤔

如果有 4 张卡呢？8 张卡呢？

总结：Roofline 三大公式

① 时间估算

\[T = \max(T_{\text{计算}}, T_{\text{通信}}) \quad \text{到} \quad T_{\text{计算}} + T_{\text{通信}}\]

② 算术强度

\[\text{强度} = \frac{\text{FLOPs}}{\text{Bytes}}\]

③ 临界判断

\[\text{强度} \gtrless \frac{\text{算力}}{\text{带宽}} \Rightarrow \text{计算/通信 受限}\]

黄金法则速查表

场景 要满足的条件
单卡 bf16 matmul (TPU) Batch size > 240
单卡 bf16 matmul (H100) Batch size > 300
多卡并行 模型宽度 D > 临界值
int8 量化 临界 batch size 更低（约 120）

下一步

第 2 章：TPU 内部结构详解

第 3 章：分片与多卡矩阵乘法

课后练习

  1. 如果用 fp32 代替 bf16，算术强度会怎么变？

  2. 注意力机制的算术强度是多少？它是计算受限还是通信受限？

  3. 如果你的模型只有 D=1024，在多卡上扩展效果会好吗？

  4. H100 的临界 batch size 是多少？（提示：查规格表）

参考资源

推荐阅读

📄 Roofline 论文原文 Williams, Waterman, Patterson (2009) “Roofline: An Insightful Visual Performance Model”

🔗 硬件规格

谢谢！

Q&A 时间 🙋

联系方式


下节课见！我们一起打开 TPU 的”黑盒子”！ 🔓