""" ============================================================================= 第 5 讲 · 代码 1：Roofline 模型 —— 这个运算受计算限制，还是受内存限制？ ============================================================================= 对应讲义：lectures/lecture05_gpus_tpus.md 第 2、3 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 一个运算在 GPU 上要花多少时间，取决于两件事中较慢的那一个： - 计算：FLOPs ÷ 峰值算力 - 搬数据：字节数 ÷ 内存带宽 把"每搬一个字节做多少次运算"叫做**算术强度**（arithmetic intensity，单位 FLOP/字节）。 硬件也有一个"平衡点"：峰值算力 ÷ 内存带宽。 算术强度 < 平衡点 → 内存受限（大部分时间在等数据） 算术强度 > 平衡点 → 计算受限（能跑满算力） 画成图就是一个"屋顶"形状：先斜着上升（受带宽限制），再变平（受算力限制）——这就是 Roofline 模型。 这个文件： 1. 列出 A100 / H100 / B200 的关键规格，算出它们的平衡点； 2. 计算常见运算的算术强度，判断它们在 H100 上是哪一类； 3. 解释为什么"推理解码"是内存受限的，而"训练"是计算受限的； 4. 在你自己的电脑上实测带宽和算力，画一条属于你的 roofline，并验证不同运算落在哪里。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 01_roofline.py （几秒到十几秒，主要在计时） """ import time import torch def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：硬件规格 # ============================================================================= # 数据来自 NVIDIA 公开规格（SXM 版本，BF16 为稠密算力），取近似值。 GPUS = { # SM 数, BF16 稠密 FLOP/s, HBM 带宽（字节/s）, HBM 容量（GB） "A100": (108, 312e12, 2.0e12, 80), "H100": (132, 989.5e12, 3.35e12, 80), "B200": (148, 2.25e15, 8.0e12, 180), } def ridge_point(peak_flops, bandwidth): """ Roofline 的"拐点"：算术强度达到多少，才能从内存受限变成计算受限。 【公式】平衡点 = 峰值算力 / 内存带宽（单位：FLOP/字节） """ return peak_flops / bandwidth def roofline(intensity, peak_flops, bandwidth): """ Roofline 模型预测的可达算力：min(峰值算力, 算术强度 × 带宽)。 【参数】intensity : 算术强度（FLOP/字节），> 0 【返回】FLOP/s """ return min(peak_flops, intensity * bandwidth) def demo_specs(): """演示 1：硬件规格与平衡点。""" print_header("演示 1：GPU 规格与 roofline 的拐点") print(f" {'GPU':<6}{'SM 数':>6}{'BF16 算力':>14}{'HBM 带宽':>12}{'显存':>8}{'平衡点':>18}") for name, (sms, flops, bw, mem) in GPUS.items(): print(f" {name:<6}{sms:>6}{flops / 1e12:>10,.0f} TF/s{bw / 1e12:>9.2f} TB/s{mem:>6} GB" f"{ridge_point(flops, bw):>10.0f} FLOP/字节") print(""" 平衡点约 150~300：每从显存读一个字节，要做一两百次运算，才能让计算单元不闲着！ 而且代际之间算力涨得比带宽快（A100 → B200：算力 ×7，带宽 ×4），"喂不饱"的问题越来越严重 —— 这就是 CS336 课上说的"计算的增长快于内存的增长"（内存墙）。""") # ============================================================================= # 第 2 部分：常见运算的算术强度 # ============================================================================= def intensity_table(n=4096, bytes_per=2): """ 常见运算的 (FLOPs, 字节数, 算术强度)。n 是向量长度 / 矩阵边长，bytes_per 是每个数的字节数。 【各运算的计数】（每个输入读一次、输出写一次） ReLU（逐元素） ：n 次运算；读 n、写 n 个数 点积 x·y ：2n 次运算；读 2n 个数，写 1 个 矩阵 × 向量（n×n @ n） ：2n² 次运算；读 n² + n、写 n 个数 矩阵 × 矩阵（n×n @ n×n）：2n³ 次运算；读 2n²、写 n² 个数 """ b = bytes_per ops = { "ReLU（逐元素）": (n, 2 * n * b), "点积": (2 * n, (2 * n + 1) * b), "矩阵×向量": (2 * n * n, (n * n + 2 * n) * b), "矩阵×矩阵": (2 * n ** 3, 3 * n * n * b), } return {k: (f, by, f / by) for k, (f, by) in ops.items()} def demo_intensity(): """演示 2：哪些运算在 H100 上是内存受限的？""" print_header("演示 2：常见运算的算术强度（BF16，n = 4096，H100）") _, peak, bw, _ = GPUS["H100"] ridge = ridge_point(peak, bw) print(f" {'运算':<16}{'算术强度':>12}{'roofline 预测可达算力':>22}{'占峰值':>8} 结论") for name, (f, by, ai) in intensity_table().items(): achievable = roofline(ai, peak, bw) kind = "计算受限" if ai >= ridge else "内存受限" print(f" {name:<16}{ai:>10.2f}{achievable / 1e12:>18.1f} TF/s{achievable / peak:>9.1%} {kind}") print(f""" （H100 平衡点 ≈ {ridge:.0f} FLOP/字节） - 逐元素运算、点积、矩阵×向量的强度都小于 1，只能用到峰值算力的零头。 - 矩阵×矩阵的强度约为 n/3，n 足够大时远超平衡点 —— 这是 GPU 真正擅长的事。 所以让 GPU 跑得快的秘诀就是：**尽量把工作变成大矩阵乘法，并减少其他运算搬运的字节数**。""") def demo_decode_vs_train(): """演示 3：推理解码 vs 训练 —— 批次大小决定了算术强度。""" print_header("演示 3：为什么推理解码是内存受限的？") _, peak, bw, _ = GPUS["H100"] d = 8192 print(f" 一个 d×d（d = {d}）的 BF16 权重矩阵乘以 B 个 token 的激活：") print(f" {'批次 B（token 数）':>18}{'算术强度':>10}{'可达算力占峰值':>16}") for B in [1, 8, 64, 256, 1024, 4096]: flops = 2 * B * d * d bytes_moved = 2 * (d * d + 2 * B * d) # 读权重 + 读输入 + 写输出 ai = flops / bytes_moved print(f" {B:>18}{ai:>10.1f}{roofline(ai, peak, bw) / peak:>16.1%}") print(""" 算术强度约等于批次中的 token 数 B：权重读一次，被 B 个 token 重复使用。 - 训练 / 预填充（prefill）：一次处理成千上万个 token，轻松计算受限； - 逐个生成 token 的解码（decode）：每个请求每步只有 1 个 token，必须把很多请求凑成大批次， 否则大部分时间都在从显存搬权重（和 KV 缓存）。这是第 10 讲推理优化的出发点。""") # ============================================================================= # 第 3 部分：你自己电脑的 roofline # ============================================================================= def get_device(): """优先使用 GPU（CUDA / Apple MPS），否则用 CPU。""" if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def sync(device): """等待设备上的运算完成（GPU 是异步的，计时前必须同步）。""" if device.type == "cuda": torch.cuda.synchronize() elif device.type == "mps": torch.mps.synchronize() def bench(fn, device, repeats=10): """热身后计时，返回每次平均秒数。""" for _ in range(3): fn() sync(device) t0 = time.perf_counter() for _ in range(repeats): fn() sync(device) return (time.perf_counter() - t0) / repeats def demo_my_roofline(): """演示 4：在你的电脑上实测 roofline。""" device = get_device() print_header(f"演示 4：你的电脑的 roofline（设备：{device}）") # 带宽：复制一个大数组（读 n 个 + 写 n 个 float32），几乎不做运算 n = 1 << 25 x = torch.randn(n, device=device) t = bench(lambda: x.clone(), device) bw = 2 * n * 4 / t # 峰值算力：大矩阵乘法 m = 4096 a, b = torch.randn(m, m, device=device), torch.randn(m, m, device=device) t = bench(lambda: a @ b, device, repeats=5) peak = 2 * m ** 3 / t ridge = ridge_point(peak, bw) print(f" 实测内存带宽 ≈ {bw / 1e9:.0f} GB/s，实测矩阵乘法算力 ≈ {peak / 1e12:.2f} TFLOP/s（FP32）") print(f" → 你的电脑的平衡点 ≈ {ridge:.0f} FLOP/字节（H100 是 295）") print(f"\n {'运算':<26}{'算术强度':>10}{'实测算力':>14}{'roofline 预测':>14}") tests = [] # 逐元素：y = x * 2（1 次运算，读 4 字节、写 4 字节） tests.append(("逐元素 x*2", 1 / 8, lambda: x * 2, n)) # 逐元素：y = x * 2 + 1（2 次运算）。如果"一次读、一次写"完成，强度是 2/8 tests.append(("逐元素 x*2+1", 2 / 8, lambda: x * 2 + 1, 2 * n)) # 矩阵乘法：不同大小 for s in [64, 256, 1024, 4096]: A, B = torch.randn(s, s, device=device), torch.randn(s, s, device=device) tests.append((f"矩阵乘法 {s}×{s}", (2 * s ** 3) / (3 * s * s * 4), (lambda A=A, B=B: A @ B), 2 * s ** 3)) for name, ai, fn, flops in tests: t = bench(fn, device, repeats=5) pred = roofline(ai, peak, bw) print(f" {name:<26}{ai:>10.2f}{flops / t / 1e9:>10.0f} GF/s{pred / 1e9:>10.0f} GF/s") print(""" - x*2 的实测算力与 roofline 的预测基本吻合：完全被带宽卡住。 - x*2+1 只达到预测的一半左右：PyTorch 把它分成两步执行（先算 x*2 写回内存，再读出来 +1）， 实际搬运的字节数翻了一倍，有效强度只有 2/16。把多个运算"融合"成一次读写， 正是 03 号文件和第 6 讲要讲的算子融合。 - 小矩阵乘法远低于 roofline 预测：除了带宽，还有启动开销、线程调度等"固定成本" （roofline 只是上界，不是保证）。 - 大矩阵乘法接近峰值。 （具体数字取决于你的电脑。笔记本的内存带宽一般是每秒几十到几百 GB， 远低于 H100 的 3350 GB/s；算力也低得多，所以平衡点通常比 H100 小。）""") def main(): torch.manual_seed(0) demo_specs() demo_intensity() demo_decode_vs_train() demo_my_roofline() print("\n全部演示运行完毕。") if __name__ == "__main__": main()