""" ============================================================================= 第 8 讲 · 代码 2：在你自己的机器上测一测——算力、带宽、开销与融合 ============================================================================= 对应讲义：lectures/lecture08_gpu_flashattention.md 第 2、3、4 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 代码 01 是"纸上的 roofline"。这里在你正在用的这台电脑上（CPU，PyTorch 单线程）实际测一遍： 三种"受限"——算力、带宽、开销——在 CPU 上同样存在，只是数字不同。 演示 1：矩阵乘法的速度随矩阵大小怎样变化（小矩阵 → 开销受限；大矩阵 → 接近这台机器的矩阵峰值）。 演示 2：逐元素运算的速度（每秒搬多少字节）——这台机器的"显存带宽"。 演示 3：融合（fusion）：a*b+c vs torch.addcmul、手写 softmax vs F.softmax、手写 GELU vs F.gelu、 x.cos().cos() vs x.cos()。哪些融合有用，为什么。 演示 4：第 1 讲的 base model 训练一步：有多少个运算（kernel）、多少 FLOPs、batch 变大时时间怎样变。 演示 5：同样的训练循环，把模型变宽：每秒完成的 FLOPs 怎样变化——"矩阵越大，硬件用得越满"。 【关于计时】 计时类的数字每次运行都会略有不同，也和你的机器有关（作者的机器是 Apple M4 Pro）。 所有带计时数字的行都以 ⏱ 开头；其余行（运算个数、FLOPs、读写字节）是确定的。 每个计时取 5 次重复里最快的一次（最快的那次最接近"没有被其他程序打扰"的情况）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 02_measure_this_cpu.py # 约 30 秒；需要 PyTorch 【PyTorch 小知识】 - time.perf_counter()：高精度的计时器（秒）。测很短的运算时要重复很多次再平均，否则测到的都是噪声。 - x.mul_(c)：带下划线的方法是"原地"（in-place）运算，直接改 x，不新分配内存。 - torch.addcmul(c, a, b)：一次算出 c + a*b——一个 kernel 读 3 个数组、写 1 个； 写成 a*b+c 则是两个 kernel：先算 a*b 写到一个临时数组，再读回来加 c。 - FlopCounterMode：PyTorch 自带的 FLOP 计数器（第 4 讲用过），在 with 块里数矩阵乘法等运算的 FLOPs。 - TorchDispatchMode：PyTorch 的"分发钩子"，每个底层运算（aten 算子）执行前都会经过它； 我们用它数"训练一步一共调用了多少个运算"。它在 torch.utils._python_dispatch 里，名字带下划线， 说明是半内部的接口，以后的 PyTorch 版本可能会变。 """ import math import os import sys import time import torch import torch.nn.functional as F from torch.utils._python_dispatch import TorchDispatchMode from torch.utils.flop_counter import FlopCounterMode HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import pad, print_header # noqa: E402 from tiny_gpt import BASE_CONFIG, GPTConfig, count_parameters, make_model # noqa: E402 （导入时会设置单线程） def bench(fn, repeats=5, min_seconds=0.05): """ 测一个函数平均每次调用要多久（秒）。 【思路】先调用一次"热身"（第一次调用常常要分配内存、加载代码）；再找一个循环次数 n， 让 n 次调用至少花 min_seconds 秒（太短的测量全是噪声）；最后测 repeats 轮，取最快的一轮 ÷ n。 【参数】fn：不带参数的函数；repeats：测几轮；min_seconds：每轮至少多长 【返回】每次调用的秒数（浮点数） 【边界】fn 有副作用（例如原地修改数组）时，多次调用会累积——演示 2 用乘以一个接近 1 的数，避免数值跑飞。 """ fn() n = 1 while True: t = time.perf_counter() for _ in range(n): fn() if time.perf_counter() - t >= min_seconds / 5: break n *= 2 best = float("inf") for _ in range(repeats): t = time.perf_counter() for _ in range(n): fn() best = min(best, (time.perf_counter() - t) / n) return best class OpCounter(TorchDispatchMode): """ 数一段代码一共调用了多少个底层运算（aten 算子），并按名字分类。 【思路】在 GPU 上，每个算子通常就是一次 kernel 启动；每次启动都有固定的开销（几微秒）， 所以"一步训练调用多少个算子"决定了开销的下限。 【用法】with OpCounter() as oc: ...；之后看 oc.n（总数）和 oc.by_name（dict：算子名 → 次数） """ def __init__(self): super().__init__() self.n = 0 self.by_name = {} def __torch_dispatch__(self, func, types, args=(), kwargs=None): self.n += 1 name = func.overloadpacket.__name__ # 例如 "mm"、"add"、"_softmax" self.by_name[name] = self.by_name.get(name, 0) + 1 return func(*args, **(kwargs or {})) # ============================================================================= # 演示 1：矩阵乘法 # ============================================================================= def demo1_matmul(gen): print_header("演示 1：fp32 矩阵乘法 (n×n)·(n×n) 的速度（单线程）") print(f"{pad('n', 8, '>')}{pad('FLOPs', 14, '>')}{pad('强度 FLOP/字节', 16, '>')}{pad('每次用时', 14, '>')}{pad('GFLOP/s', 12, '>')}") best = 0.0 for n in [8, 16, 32, 64, 128, 256, 512, 1024]: a, b = torch.randn(n, n, generator=gen), torch.randn(n, n, generator=gen) t = bench(lambda: a @ b) flops = 2 * n ** 3 intensity = flops / (4 * 3 * n * n) # fp32：每个数 4 字节，读 2 个矩阵写 1 个 best = max(best, flops / t) print(f"⏱{pad(n, 7, '>')}{pad(f'{flops:,}', 14, '>')}{pad(f'{intensity:.1f}', 16, '>')}" f"{pad(f'用时 {t * 1e6:.1f} µs', 14, '>')}{pad(f'{flops / t / 1e9:.1f}', 12, '>')}") print("\n观察 Observations：") print(" 1. 小矩阵（n ≤ 32）每次调用都要零点几微秒到一微秒，几乎和 n 无关——这是'开销'（调用 PyTorch、") print(" 检查形状、选 kernel）。这时 GFLOP/s 很低：机器不是算不动，而是在忙别的。") print(" 2. n 增大，每秒完成的 FLOPs 增长几百倍，最后接近这台机器的矩阵峰值。") print(" 在作者的 Apple M4 Pro 上，单线程矩阵乘法能到约 3 TFLOP/s——比普通运算快两个数量级，") print(" 因为 macOS 的数学库（Accelerate）用了芯片里专门做矩阵乘法的单元，正像 GPU 的 tensor core。") print(" 在别的 CPU 上，单线程矩阵乘法通常是几十到一两百 GFLOP/s；规律相同，数字不同。") return best # ============================================================================= # 演示 2：逐元素运算的带宽 # ============================================================================= def demo2_bandwidth(gen): print_header("演示 2：原地乘一个常数 x.mul_(c) 的速度——这台机器每秒能搬多少字节") print(f"{pad('元素个数', 14, '>')}{pad('数组大小', 12, '>')}{pad('每次用时', 16, '>')}{pad('读+写 GB/s', 14, '>')}") bw = 0.0 for p in [10, 14, 18, 22, 24, 26]: n = 1 << p x = torch.randn(n, generator=gen) t = bench(lambda: x.mul_(1.0000001)) gbs = 8 * n / t / 1e9 # 每个元素读 4 字节、写 4 字节 if p >= 24: bw = max(bw, gbs * 1e9) size = f"{4 * n / 2 ** 20:.3g} MB" print(f"⏱{pad(f'{n:,}', 13, '>')}{pad(size, 12, '>')}{pad(f'用时 {t * 1e6:.1f} µs', 16, '>')}{pad(f'{gbs:.1f}', 14, '>')}") print("\n观察 Observations：") print(" 1. 很小的数组（1024 个数）每次也要约一微秒：又是开销，带宽远没用上。") print(" 2. 数组变大后，速度稳定在每秒一百多 GB（作者的机器）——这就是这台机器的'显存带宽'（对 CPU 是内存带宽）。") print(" x.mul_(c) 每个数只做 1 次乘法、搬 8 个字节，强度 0.125：它的速度完全由带宽决定。") return bw # ============================================================================= # 演示 3：融合 # ============================================================================= def gelu_manual(x): """手写的 GELU（tanh 近似）：7 个 PyTorch 运算，每个都把整个数组读一遍、写一遍。""" return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x * x * x))) def softmax_manual(s): """手写的安全 softmax：max、减、exp、sum、除，5 个运算。""" m = s.max(dim=-1, keepdim=True).values e = (s - m).exp() return e / e.sum(dim=-1, keepdim=True) def demo3_fusion(gen): print_header("演示 3：融合（fusion）——把几个运算并成一个 kernel，少读写几遍内存") n = 1 << 24 # 1677 万个数，64 MB，远大于 CPU 缓存 a, b, c, x = (torch.randn(n, generator=gen) for _ in range(4)) s = torch.randn(32 * 4 * 96, 96, generator=gen) # 第 1 讲 base model 一个 batch 的注意力分数的形状 cases = [ ("a*b+c vs addcmul（一个 kernel）", lambda: a * b + c, lambda: torch.addcmul(c, a, b), "6 → 4 趟"), ("手写 softmax（5 个运算）vs F.softmax", lambda: softmax_manual(s), lambda: F.softmax(s, dim=-1), "约 9 → 2 趟"), ("手写 GELU（7 个运算）vs F.gelu", lambda: gelu_manual(x), lambda: F.gelu(x, approximate="tanh"), "约 14 → 2 趟"), ("x.cos().cos() vs x.cos()", lambda: x.cos().cos(), lambda: x.cos(), "4 → 2 趟"), ] err = (gelu_manual(x) - F.gelu(x, approximate="tanh")).abs().max().item() err_sm = (softmax_manual(s) - F.softmax(s, dim=-1)).abs().max().item() print(f"（数值核对：手写 GELU 与 F.gelu 最大差 {err:.1e}，手写 softmax 与 F.softmax 最大差 {err_sm:.1e}——同一个函数）\n") print(f"{pad('比较', 38)}{pad('读写数组的趟数', 16)}{pad('前者用时', 14, '>')}{pad('后者用时', 14, '>')}{pad('前者 ÷ 后者', 12, '>')}") for name, f1, f2, passes in cases: t1, t2 = bench(f1), bench(f2) print(f"⏱{pad(name, 37)}{pad(passes, 16)}{pad(f'用时 {t1 * 1e3:.2f} ms', 14, '>')}" f"{pad(f'{t2 * 1e3:.2f} ms', 14, '>')}{pad(f'{t1 / t2:.2f}', 12, '>')}") print("\n观察 Observations：") print(" 1. a*b+c 要读写 6 趟（a、b 读，临时数组写；临时数组、c 读，结果写），addcmul 只要 4 趟：") print(" 这种受带宽限制的运算，融合后的加速比接近'趟数之比' 1.5（作者的机器上实测约 1.5）。") print(" 2. softmax 拆成 5 个运算要多读写好几遍，融合的 F.softmax 明显更快（作者的机器上约 1.6 倍）。") print(" 3. 但 GELU 融合几乎没用（约 1.1 倍），x.cos().cos() 也正好是 x.cos() 的 2 倍：") print(" 在这台 CPU 上 tanh、cos 本身很贵，这些运算是受'算力'限制的，少读写几遍省不了多少。") print(" GPU 上恰恰相反：算 cos 几乎不花时间，时间全在读写显存，所以 Horace He 的文章说") print(" 融合后的 x.cos().cos() 和 x.cos() 一样快。同一个运算在不同硬件上可以落在 roofline 的不同一侧——") print(" 优化之前先判断'受什么限制'，这正是 roofline 的用处。") # ============================================================================= # 演示 4：base model 训练一步 # ============================================================================= def make_step(model, B, T, gen): """ 返回一个"训练一步"的函数（前向 → 交叉熵 → 反向 → AdamW 更新），输入是固定的随机字符。 【思路】这里只关心速度，不关心学到什么，所以用随机输入，省掉切 batch 的开销。 【参数】model：TinyGPT；B, T：batch 大小和序列长度；gen：torch.Generator 【返回】不带参数的函数 step() """ V = model.cfg.vocab_size x = torch.randint(0, V, (B, T), generator=gen) y = torch.randint(0, V, (B, T), generator=gen) opt = torch.optim.AdamW(model.parameters(), lr=1e-4) def step(): loss = F.cross_entropy(model(x).view(-1, V), y.view(-1)) opt.zero_grad() loss.backward() opt.step() return step def count_step(model, B, T, gen): """数训练一步（不含 AdamW 的第一次状态初始化）的算子个数和 FLOPs。返回 (算子个数, 按名字的计数, FLOPs)。""" step = make_step(model, B, T, gen) step() # 第一次调用会初始化 AdamW 的状态，不计入 with OpCounter() as oc: step() with FlopCounterMode(display=False) as fc: step() return oc.n, oc.by_name, fc.get_total_flops() def demo4_base_model(gen, matmul_peak): print_header("演示 4：第 1 讲的 base model（2 层，d = 64，上下文 96）训练一步") model = make_model(BASE_CONFIG, seed=0) n_ops, by_name, _ = count_step(model, 32, 96, gen) top = sorted(by_name.items(), key=lambda kv: -kv[1])[:8] print(f"训练一步一共调用 {n_ops} 个底层运算；最多的几种：" + "，".join(f"{k} ×{v}" for k, v in top)) print(f"\n{pad('batch', 8, '>')}{pad('运算个数', 10, '>')}{pad('每步 FLOPs', 16, '>')}{pad('每步用时', 16, '>')}" f"{pad('GFLOP/s', 10, '>')}{pad('÷ 矩阵峰值', 12, '>')}") prev = None for B in [1, 2, 4, 8, 16, 32, 64, 128]: model = make_model(BASE_CONFIG, seed=0) n_ops, _, flops = count_step(model, B, 96, gen) t = bench(make_step(model, B, 96, gen), repeats=5, min_seconds=0.1) print(f"{pad(B, 8, '>')}{pad(n_ops, 10, '>')}{pad(f'{flops:,}', 16, '>')}") grow = f"（batch ×2，时间 ×{t / prev:.2f}）" if prev else "" print(f"⏱{pad('', 33)}{pad(f'用时 {t * 1e3:.2f} ms', 16, '>')}{pad(f'{flops / t / 1e9:.0f}', 10, '>')}" f"{pad(f'{flops / t / matmul_peak:.1%}', 12, '>')} {grow}") prev = t print("\n观察 Observations：") print(" 1. 训练一步要调用约 600 个底层运算，个数几乎不随 batch 变：前向、反向、AdamW 对每个参数张量的更新……") print(" 每个运算都有固定的开销（在作者的机器上，batch = 1 时一步约 1.2 ms，平均每个运算约 2 µs）。") print(" 2. batch 从 1 到 2，工作量翻倍，时间只增加三四成——这就是 Horace He 说的'开销受限'的信号；") print(" batch ≥ 8 以后，时间基本和 batch 成正比，开销被摊薄了。") print(" 3. 但即使在大 batch，base model 也只用到这台机器矩阵峰值的约 5%：") print(" 它的矩阵只有 64 宽（演示 1：n = 64 的矩阵乘法本身就很慢），而且 LayerNorm、GELU、softmax 等") print(" 逐元素运算占了很大一部分时间。小模型天生'用不满'硬件。") # ============================================================================= # 演示 5：模型变宽 # ============================================================================= def demo5_width(gen, matmul_peak): print_header("演示 5：同样 2 层、batch 32、上下文 96，只改宽度 d——每秒完成的 FLOPs") print(f"{pad('d', 6, '>')}{pad('参数量', 14, '>')}{pad('每步 FLOPs', 18, '>')}{pad('每步用时', 16, '>')}" f"{pad('GFLOP/s', 10, '>')}{pad('÷ 矩阵峰值', 12, '>')}") for d in [32, 64, 128, 256, 512, 1024]: cfg = GPTConfig(vocab_size=44, block_size=96, n_layer=2, n_head=4, d_model=d) model = make_model(cfg, seed=0) _, _, flops = count_step(model, 32, 96, gen) t = bench(make_step(model, 32, 96, gen), repeats=3, min_seconds=0.1) print(f"{pad(d, 6, '>')}{pad(f'{count_parameters(model):,}', 14, '>')}{pad(f'{flops:,}', 18, '>')}") print(f"⏱{pad('', 37)}{pad(f'用时 {t * 1e3:.1f} ms', 16, '>')}{pad(f'{flops / t / 1e9:.0f}', 10, '>')}" f"{pad(f'{flops / t / matmul_peak:.1%}', 12, '>')}") print("\n观察 Observations：") print(" 1. 宽度从 32 到 1024，参数量涨了约 850 倍，每步时间只涨了几十倍：每秒完成的 FLOPs 涨了二十多倍，") print(" 从矩阵峰值的约 2% 升到一半左右（作者的机器）。矩阵越大，强度越高、开销占比越低。") print(" 2. 这就是第 4 讲 §3 的 6ND 只是'需要的计算量'，而'花的时间'还要除以硬件利用率的原因；") print(" 也是大模型训练能做到 40%–50% MFU、而我们的玩具模型只有几个百分点的原因。") if __name__ == "__main__": gen = torch.Generator().manual_seed(0) peak = demo1_matmul(gen) bw = demo2_bandwidth(gen) print(f"\n⏱ 这台机器（单线程）：矩阵峰值约 {peak / 1e9:.0f} GFLOP/s，带宽约 {bw / 1e9:.0f} GB/s，" f"拐点约 {peak / bw:.0f} FLOP/字节（fp32；作为对照，A100 的 16 位拐点约 200）") demo3_fusion(gen) demo4_base_model(gen, peak) demo5_width(gen, peak)