""" ============================================================================= 第 8 讲 · 代码 1：roofline——一个运算在 GPU 上花多久，受什么限制 ============================================================================= 对应讲义：lectures/lecture08_gpu_flashattention.md 第 2、3、4、6 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- GPU 的说明书上写着"312 TFLOP/s"。为什么真实的训练只跑到它的一半？哪些运算慢，慢在哪里？ 演示 1：几代加速器的规格与 roofline 的"拐点"（峰值算力 ÷ 显存带宽）。 演示 2：常见运算的算术强度（FLOPs / 字节）：大矩阵乘法、小矩阵乘法、生成时的矩阵 × 向量、 逐元素运算、LayerNorm、softmax、单个注意力头——谁受算力限制，谁受带宽限制。 演示 3：把一个 Transformer 层拆成十几个运算，逐个记账：FLOPs 占比 vs 时间占比； 融合（fusion）+ FlashAttention 之后能省多少。 演示 4：MFU 与 HFU：按 PaLM 论文附录 B 的公式重算 46.2% 和 57.8%，并给几个大模型的训练效率记账。 全部是"纸上计算"（roofline 模型 + 论文里的数字），没有计时，结果是确定的。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 01_roofline.py # 约 1 秒；需要 numpy（只用到本目录的 flash_tools.py 和第 1 讲的 console.py） 【Python 小知识】 - 列表里放元组 (名字, 类别, flops, 字节数)，再用 for name, cat, f, b in ops: 一次拆开，叫"解包"（unpacking）。 - dict.get(key, 默认值)：键不存在时返回默认值，不报错。 - f"{x:.1%}"：把 0.462 格式化成 "46.2%"。 """ import math import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import pad, print_header # noqa: E402 from flash_tools import (HARDWARE, flash_block_sizes, flash_io, fmt_si, matmul_cost, # noqa: E402 ridge_point, roofline_time) A100 = HARDWARE["A100 40GB (2020)"] BYTES = 2 # fp16 / bf16：每个数 2 个字节 # ============================================================================= # 演示 1：规格与拐点 # ============================================================================= def demo1_specs(): print_header("演示 1：几代加速器的规格，以及 roofline 的拐点（峰值算力 ÷ 显存带宽）") print(f"{pad('加速器', 18)}{pad('16 位矩阵峰值', 14, '>')}{pad('普通运算峰值', 14, '>')}" f"{pad('显存带宽', 12, '>')}{pad('显存', 8, '>')}{pad('拐点(矩阵)', 12, '>')}{pad('拐点(普通)', 12, '>')}") for name, hw in HARDWARE.items(): simt = fmt_si(hw["peak_simt"], "FLOP/s") if hw["peak_simt"] else "—" r_tc = ridge_point(hw) r_simt = ridge_point(hw, use_tensor_core=False) print(f"{pad(name, 18)}{pad(fmt_si(hw['peak_tc'], 'FLOP/s'), 14, '>')}{pad(simt, 14, '>')}" f"{pad(fmt_si(hw['hbm_bw'], 'B/s'), 12, '>')}{pad(str(hw['hbm_gb']) + ' GB', 8, '>')}" f"{pad(f'{r_tc:.0f}', 12, '>')}{pad(f'{r_simt:.1f}' if r_simt else '—', 12, '>')}") v100, h100 = HARDWARE["V100 (2017)"], HARDWARE["H100 SXM (2022)"] print(f"\nV100 → H100：矩阵峰值 ×{h100['peak_tc'] / v100['peak_tc']:.1f}，显存带宽只 ×{h100['hbm_bw'] / v100['hbm_bw']:.1f}" f"——拐点从 {ridge_point(v100):.0f} 升到 {ridge_point(h100):.0f}：算力涨得比带宽快，越来越多的运算受带宽限制。") sram_total = A100["sram_kb"] * A100["n_sm"] / 1024 print(f"A100 的片上 SRAM：{A100['n_sm']} 个 SM × {A100['sram_kb']} KB = {sram_total:.2f} MB，" f"是 40 GB 显存的 1/{40 * 1024 / sram_total:.0f}；带宽约 19 TB/s，是显存的 {19 / 1.555:.0f} 倍左右。") print("\n观察 Observations：") print(" 1. 拐点 ≈ 200 FLOPs/字节（A100）：一个 16 位运算如果每读 1 个字节做不到约 200 次运算，") print(" tensor core 就会闲着等数据。逐元素运算每个数只做几次运算，离 200 差两个数量级。") print(" 2. 不走 tensor core 的普通运算峰值只有矩阵峰值的 1/16（A100：19.5 vs 312）——") print(" 一个模型 99% 的 FLOPs 是矩阵乘法，并不意味着 99% 的时间花在矩阵乘法上（演示 3）。") print(" 3. 每一代硬件的拐点都在升高：'少搬数据'比'少算'越来越重要——这是 FlashAttention 的大背景。") # ============================================================================= # 演示 2：算术强度 # ============================================================================= def demo2_intensity(): print_header("演示 2：常见运算的算术强度（A100，16 位，每个数 2 字节；拐点约 200）") rows = [] for n in [64, 512, 4096, 16384]: f, b = matmul_cost(n, n, n) rows.append((f"方阵乘法 n = {n}", f, b, True)) d = 12288 # GPT-3 175B 的宽度 for T in [1, 32, 256, 2048]: f, b = matmul_cost(T, d, 4 * d) what = "生成：矩阵 × 向量" if T == 1 else f"{T} 个 token" rows.append((f"GPT-3 MLP 第一层，{what}", f, b, True)) n_el = 2048 * 12288 # 一个 2048 token 的序列、宽度 12288 的激活 rows.append(("逐元素：加偏置 / 残差相加", 1 * n_el, BYTES * 2 * n_el, False)) rows.append(("逐元素：GELU（tanh 近似）", 8 * n_el, BYTES * 2 * n_el, False)) rows.append(("LayerNorm（均值、方差、缩放）", 8 * n_el, BYTES * 2 * n_el, False)) N = 2048 rows.append(("softmax，N×N 的注意力分数", 5 * N * N, BYTES * 2 * N * N, False)) for dh in [64, 128]: f, b = matmul_cost(N, dh, N) rows.append((f"一个注意力头 QKᵀ，N = {N}，d = {dh}", f, b, True)) print(f"{pad('运算', 40)}{pad('FLOPs', 12, '>')}{pad('读写字节', 12, '>')}{pad('强度', 10, '>')}" f"{pad('受限于', 10, '>')}{pad('最多达到峰值的', 16, '>')}") for name, f, b, tc in rows: r = roofline_time(f, b, A100, use_tensor_core=tc) frac = (f / r["t"]) / (A100["peak_tc"] if tc else A100["peak_simt"]) bound = "算力" if r["bound"] == "compute" else "带宽" intensity = f"{r['intensity']:.1f}" # 先算成字符串，避免 f-string 里再嵌套同种引号（Python ≤ 3.11 不允许） print(f"{pad(name, 40)}{pad(fmt_si(f), 12, '>')}{pad(fmt_si(b, 'B'), 12, '>')}" f"{pad(intensity, 10, '>')}{pad(bound, 10, '>')}{pad(f'{frac:.1%}', 16, '>')}") print("\n观察 Observations：") print(" 1. 方阵乘法的强度 = n/3：n = 64 只有 21，最多用到 tensor core 的 1/10；n = 512 也还差一点（171 < 200）；") print(" n = 4096 起远超拐点。'把矩阵做大'本身就是提高硬件利用率的办法（大 batch、宽模型）。") print(" 2. 同一个 GPT-3 的 MLP 层：训练时一次喂 2048 个 token，强度过千；生成时一次只算 1 个 token，") print(" 强度约 1——只能用到峰值的 0.5%。这就是为什么推理要攒 batch、要量化（少读权重字节），") print(" 以及为什么'生成一个 token 的时间 ≈ 把整个模型从显存读一遍的时间'。") print(" 3. 逐元素运算、LayerNorm、softmax 的强度只有 0.2–2：它们的时间几乎全花在读写显存上。") print(" 一个注意力头的 QKᵀ 虽然是矩阵乘法，但 d 只有 64–128、输出却是 N×N，强度只有 60–114，") print(" 也低于拐点——标准注意力是受带宽限制的，这正是 FlashAttention 下手的地方。") # ============================================================================= # 演示 3：一个 Transformer 层的记账 # ============================================================================= def layer_ops(B, N, d, H, fused=False, sram_elements=None): """ 一个 pre-LN Transformer 层（前向）拆成的运算清单。 【思路】未融合（fused=False）：PyTorch 朴素实现里每个运算是一个独立的 kernel， 每个 kernel 都要从显存读输入、把输出写回显存。 融合（fused=True）：注意力的 QKᵀ、缩放、掩码、softmax、PV 换成一个 FlashAttention kernel （读写量用 flash_io 的公式）；加偏置 + GELU、残差相加并进前一个矩阵乘法的"收尾"（epilogue）， 不再单独读写一遍。 【参数】B：batch 里的序列数；N：序列长度；d：模型宽度；H：头数；sram_elements：FlashAttention 的 M 【返回】列表 [(名字, 类别, flops, 字节数, 是否走 tensor core)] 【边界】只算前向；反向的结构类似（每个运算约 2 倍 FLOPs），占比结论相同。不算 dropout 和 embedding。 """ T, dh = B * N, d // H att_el = B * H * N * N ops = [("LayerNorm 1", "归一化", 8 * T * d, BYTES * 2 * T * d, False), ("QKV 投影", "线性层", *matmul_cost(T, d, 3 * d), True)] if not fused: f, b = matmul_cost(N, dh, N) ops += [("QKᵀ（每个头）", "注意力矩阵乘", B * H * f, B * H * b, True), ("缩放 ÷√d", "逐元素", att_el, BYTES * 2 * att_el, False), ("加因果掩码", "逐元素", att_el, BYTES * 2 * att_el, False), ("softmax", "softmax", 5 * att_el, BYTES * 2 * att_el, False)] f, b = matmul_cost(N, N, dh) ops += [("P·V（每个头）", "注意力矩阵乘", B * H * f, B * H * b, True), ("输出投影", "线性层", *matmul_cost(T, d, d), True), ("残差相加", "逐元素", T * d, BYTES * 3 * T * d, False)] else: Br, Bc = flash_block_sizes(sram_elements, dh, N) io = flash_io(N, dh, Br, Bc)["forward"] ops += [("FlashAttention（融合）", "注意力矩阵乘", B * H * (4 * N * N * dh + 5 * N * N), BYTES * B * H * io, True)] f, b = matmul_cost(T, d, d) ops += [("输出投影 + 残差（融合）", "线性层", f + T * d, b + BYTES * T * d, True)] ops += [("LayerNorm 2", "归一化", 8 * T * d, BYTES * 2 * T * d, False)] f, b = matmul_cost(T, d, 4 * d) if not fused: ops += [("MLP 第一层", "线性层", f, b, True), ("GELU", "逐元素", 8 * 4 * T * d, BYTES * 2 * 4 * T * d, False)] else: ops += [("MLP 第一层 + GELU（融合）", "线性层", f + 8 * 4 * T * d, b, True)] f, b = matmul_cost(T, 4 * d, d) if not fused: ops += [("MLP 第二层", "线性层", f, b, True), ("残差相加", "逐元素", T * d, BYTES * 3 * T * d, False)] else: ops += [("MLP 第二层 + 残差（融合）", "线性层", f + T * d, b + BYTES * T * d, True)] return ops def print_ledger(title, ops, hw): """打印一张运算账本，返回总时间（秒）。""" total_f = sum(o[2] for o in ops) times = [roofline_time(f, b, hw, use_tensor_core=tc) for _, _, f, b, tc in ops] total_t = sum(t["t"] for t in times) print(f"\n{title}") print(f"{pad('运算', 30)}{pad('FLOPs 占比', 12, '>')}{pad('时间占比', 10, '>')}{pad('强度', 8, '>')}{pad('受限于', 8, '>')}") for (name, _, f, _, _), t in zip(ops, times): share_f, share_t = f"{f / total_f:.2%}", f"{t['t'] / total_t:.1%}" intensity = f"{t['intensity']:.0f}" bound = "算力" if t["bound"] == "compute" else "带宽" print(f"{pad(name, 30)}{pad(share_f, 12, '>')}{pad(share_t, 10, '>')}{pad(intensity, 8, '>')}{pad(bound, 8, '>')}") by_cat = {} for (name, cat, f, _, _), t in zip(ops, times): cf, ct = by_cat.get(cat, (0, 0)) by_cat[cat] = (cf + f, ct + t["t"]) summary = "，".join(f"{cat} {cf / total_f:.1%} 的 FLOPs / {ct / total_t:.1%} 的时间" for cat, (cf, ct) in by_cat.items()) print(f"按类别：{summary}") peak_t = total_f / hw["peak_tc"] print(f"理论用时下界 {total_t * 1e3:.3f} ms（若全部按 tensor core 峰值算只要 {peak_t * 1e3:.3f} ms → " f"最多达到峰值的 {peak_t / total_t:.0%}）") return total_t def demo3_layer_ledger(): print_header("演示 3：一个 Transformer 层的运算账本（A100，前向，roofline 估计）") M = A100["sram_kb"] * 1024 // BYTES # 一个 SM 的 SRAM 能放多少个 16 位数 cases = [("GPT-2 medium 的一层（d = 1024，16 头，N = 1024，batch 8）", 8, 1024, 1024, 16), ("GPT-3 175B 的一层（d = 12288，96 头，N = 2048，batch 1）", 1, 2048, 12288, 96)] for title, B, N, d, H in cases: t0 = print_ledger("【未融合】" + title, layer_ops(B, N, d, H), A100) t1 = print_ledger("【融合 + FlashAttention】" + title, layer_ops(B, N, d, H, fused=True, sram_elements=M), A100) print(f"→ 融合后理论用时 {t0 * 1e3:.3f} → {t1 * 1e3:.3f} ms，快 {t0 / t1:.2f} 倍") print("\n观察 Observations：") print(" 1. GPT-2 medium 这一层（未融合）：线性层占 85% 的 FLOPs，却只占 29% 的时间；") print(" 逐元素运算和 softmax 合起来只占 0.5% 的 FLOPs，却占了 52% 的时间——全花在搬数据上。") print(" （Ivanov et al. 2021 对 BERT 的实测结论相同：非矩阵乘法的运算只占约 0.2% 的 FLOPs，") print(" 但跑得比矩阵乘法慢几百倍，占了可观的运行时间。）") print(" 2. 注意力的 N×N 中间结果是大头：缩放、掩码、softmax 三趟读写 N×N 矩阵，各占 15% 的时间。") print(" FlashAttention + 简单的融合让这一层的理论用时快了 2.7 倍，最多能达到峰值的比例从 34% 升到 93%；") print(" FlashAttention 自己的强度约 171（GPT-3 那层 d = 128，约 114），仍低于拐点——读写量砍掉了大半，但没有消除。") print(" 3. GPT-3 那一层宽度 12288，线性层占 97% 的 FLOPs、83% 的时间，融合只快 1.14 倍：") print(" 模型越宽，注意力和逐元素运算越不重要；序列越长，它们越重要（第 8 讲 §6 长上下文）。") print(" 4. 这只是 roofline 下界。真实 kernel 达不到峰值，但'哪些运算是瓶颈'这个结论是可靠的。") # ============================================================================= # 演示 4：MFU 与 HFU # ============================================================================= def demo4_mfu(): print_header("演示 4：MFU 与 HFU——按 PaLM 论文附录 B 的公式重算") N, L, H, Q, T = 540.35e9, 118, 48, 256, 2048 # PaLM 540B：参数、层数、头数、头维度、序列长度 tokens_per_s, chips, peak = 238.3e3, 6144, 275e12 dense = 6 * N # 每个 token 的矩阵乘法 FLOPs（前向 2N + 反向 4N） attn = 12 * L * H * Q * T # 注意力的额外 FLOPs：6LH(2QT) print(f"每个 token：6N = {dense:.4g} FLOPs，注意力 12LHQT = {attn:.4g}，合计 {dense + attn:.4g}" f"（论文表 22：3.28 TFLOPs / token）") for label, f in [("MFU（不算注意力）", dense), ("MFU（算注意力）", dense + attn)]: print(f"{pad(label, 22)} = {tokens_per_s:.4g} × {f:.4g} / ({chips} × {peak:.3g}) = " f"{tokens_per_s * f / (chips * peak):.1%}") remat = 4.10e12 # 论文表 22：算上重算（rematerialization）每 token 4.10 TFLOPs print(f"{pad('HFU（算上重算）', 22)} = {tokens_per_s:.4g} × {remat:.4g} / ({chips} × {peak:.3g}) = " f"{tokens_per_s * remat / (chips * peak):.1%} （论文：57.8%）") ideal_days = 780e9 / tokens_per_s / 86400 chip_days = (6144 * 1200 + 3072 * 336) / 24 print(f"\n780B token ÷ {tokens_per_s:.4g} token/s = {ideal_days:.1f} 天（按 6144 块芯片、实测吞吐不间断地算）") print(f"论文报告：6144 块芯片 1200 小时 + 3072 块芯片 336 小时（含停机和重复的步数）= " f"{chip_days:,.0f} 芯片·天，折合 6144 块芯片 {chip_days / 6144:.1f} 天") print(f"→ 实际花掉的芯片时间是'理想不间断'的 {chip_days / 6144 / ideal_days:.2f} 倍" f"（batch 预热阶段吞吐更低、停机、损失尖峰回滚——第 10 讲）") print("\n几个大模型的训练效率（MFU 或'每块芯片实测 FLOP/s ÷ 峰值'；各论文的记账方式不完全相同）：") rows = [("GPT-3 175B（V100，2020）", "PaLM 论文表 3", 0.213), ("Gopher 280B（TPU v3，2021）", "PaLM 论文表 3", 0.325), ("Megatron-Turing NLG 530B（A100，2021）", "PaLM 论文表 3", 0.302), ("PaLM 540B（TPU v4，2022）", "PaLM 论文表 3", 0.462), ("Megatron-LM 1T 基准（A100，2021）", "163 / 312 TFLOP/s", 163 / 312), ("OPT-175B（A100，2022）", "147 / 312 TFLOP/s", 147 / 312), ("BLOOM 176B（A100，2022）", "156 / 312 TFLOP/s", 156 / 312), ("LLaMA-65B（A100，2023）", "380 token/s/GPU × 6N / 312T", 380 * 6 * 65.2e9 / 312e12), ("GPT 训练 + FlashAttention-2（A100，2023）", "225 / 312 TFLOP/s", 225 / 312), ("Llama 3 405B（H100，2024）", "论文：38%–43%", 0.43)] for name, src, v in rows: print(f" {pad(name, 42)}{pad(src, 30)}{v:6.1%}") print("\n观察 Observations：") print(" 1. 公式 6N + 12LHQT 和论文给的吞吐，正好重算出 45.7% / 46.2% / 57.8%：MFU 只数'模型必须做的'运算，") print(" HFU 还数上了为省显存而重算的前向（激活重算，第 9 讲）——同一次训练，HFU 永远 ≥ MFU。") print(" 2. 2020 年 GPT-3 只用到峰值的约 1/5；2022 年的 PaLM、OPT、BLOOM 在 45%–50%；") print(" FlashAttention-2 让一般的 GPT 训练到了 72%。效率翻一倍 = 同样的模型少花一半的钱和时间。") print(" 3. 实际花掉的芯片时间比'按实测吞吐不间断'多出一半：硬件利用率之外，还有'有效训练时间'这一层（第 10 讲）。") if __name__ == "__main__": demo1_specs() demo2_intensity() demo3_layer_ledger() demo4_mfu()