""" ============================================================================= CS349E · 第 2 讲 · 03_transformer_accounting.py —— Transformer 的账本：逐个矩阵数参数、FLOPs 和字节，再和公式对账 ============================================================================= 【这个脚本要回答的问题】 第 1 讲用三条近似公式估计了 LLM 推理（参数 ≈ 矩阵、FLOPs/token ≈ 2 × 参数、KV/token = 2·L·G·dh·字节）。 这里用一个**真的在跑**的模型（transformer_lab）把每一项数出来，和公式逐项对账，然后把（验证过的）公式用到 Llama-3.1 这种跑不动的大模型上。 演示 1：参数量：小模型逐个数组数 vs 公式；Llama-3.1 的参数都花在哪（注意力 / MLP / 词表）。 演示 2：一次前向的账：prefill 和 decode 各数一遍（FLOPs、权重字节、KV 字节、激活字节），和第 1 讲 llm_math 对账。 演示 3：einsum 记号：把每一步写成 einsum，用"2 × 所有维度之积"的规则算 FLOPs（How to Scale Your Model 第 4 章）。 演示 4：注意力的运算什么时候不能忽略：上下文多长时，注意力核心的 FLOPs 追上矩阵乘法。 演示 5：prefill 和 decode 的"计算图"：每一类运算各自是 compute-bound 还是 memory-bound（Llama-3.1-8B，H100）。 【怎么读输出】全部是确定性的计算，每次运行结果相同。字节一律按 BF16（每个数 2 字节）记账。 【运行】python3 03_transformer_accounting.py （约 1 秒） 【依赖】numpy；同目录的 transformer_lab；从 ../lecture01 导入 console、hw_specs、llm_math。 """ import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_num, fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 from llm_math import LLAMA31_8B, LLAMA31_70B, LLAMA31_405B, decode_step, prefill_time # noqa: E402 from transformer_lab import (ATTN_CORE_OPS, ATTN_PROJ_OPS, MATMUL_OPS, MLP_OPS, OP_GROUPS, TINY, # noqa: E402 Counter, KVCache, analytic_ops, count_params, einsum_flops, forward, init_params, sum_rows) # --------------------------------------------------------------------------- # 演示 1：参数量 # --------------------------------------------------------------------------- def demo1_params(): print_header("演示 1：参数量——逐个数组数 vs 公式；大模型的参数花在哪") cfg = TINY p = init_params(cfg, seed=0) mc = cfg.to_model_config() lp = p["layers"][0] counted_attn = sum(lp[k].size for k in ("wq", "wk", "wv", "wo")) counted_mlp = sum(lp[k].size for k in ("w_gate", "w_up", "w_down")) rows = [ ["每层注意力 Wq+Wk+Wv+Wo", f"{counted_attn:,}", f"{mc.attn_params_per_layer():,}", "d·H·dh + 2·d·G·dh + H·dh·d"], ["每层 MLP gate+up+down", f"{counted_mlp:,}", f"{mc.mlp_params_per_layer():,}", "3·d·f"], ["全部参数", f"{count_params(p):,}", f"{mc.n_params():,}", "嵌入 + L·(注意力+MLP+2d) + d + lm_head"], ["参与矩阵乘法的参数", f"{count_params(p) - p['embed'].size - (2 * cfg.n_layers + 1) * cfg.d_model:,}", f"{mc.matmul_params():,}", "不含词嵌入和 RMSNorm"], ] print(f"小模型 {cfg.name}：L = {cfg.n_layers}，d = {cfg.d_model}，H = {cfg.n_heads}，G = {cfg.n_kv_heads}，" f"dh = {cfg.d_head}，f = {cfg.d_ff}，V = {cfg.vocab}") print_table(["项目", "逐个数组数出来", "llm_math 公式", "公式"], rows) rows = [] for m in (LLAMA31_8B, LLAMA31_70B, LLAMA31_405B): tot = m.n_params() attn = m.n_layers * m.attn_params_per_layer() mlp = m.n_layers * m.mlp_params_per_layer() emb = m.vocab * m.d_model rows.append([m.name, fmt_num(tot, 2), f"{attn / tot * 100:.1f}%", f"{mlp / tot * 100:.1f}%", f"{emb / tot * 100:.1f}%", f"{emb / tot * 100:.1f}%"]) print() print_table(["模型", "参数", "注意力", "MLP", "词嵌入", "lm_head"], rows) print_observations([ "逐个数组数出来的参数和第 1 讲的公式**完全相等**：公式里每一项都对应模型里实实在在的一个矩阵。", "MLP 占了大头（8B 是 70%，405B 是 80%）；注意力只占约 17%：GQA 把 K、V 投影压到了 Q 投影的 1/4（8B）到 1/16（405B）。", "模型越小，词表占比越大：8B 的词嵌入 + lm_head 占 13%，405B 只占 1%。小模型更喜欢共享（tied）这两个矩阵。", ]) # --------------------------------------------------------------------------- # 演示 2：一次前向的账 # --------------------------------------------------------------------------- def run_counted(cfg, params, batch, n_new, pos0, seed=0): """ 在小模型上真的跑一次前向：先不记账地写入 pos0 个 token 的 KV cache，再记账地算 n_new 个新 token。 【返回】Counter """ rng = np.random.default_rng(seed) cache = KVCache(cfg, batch) if pos0: forward(params, cfg, rng.integers(0, cfg.vocab, size=(batch, pos0)), cache=cache) c = Counter(bytes_per_elem=2) forward(params, cfg, rng.integers(0, cfg.vocab, size=(batch, n_new)), cache=cache, counter=c, last_only=(n_new == 1)) return c def demo2_one_forward(): print_header("演示 2：一次前向的账——prefill 与 decode，数出来的 vs 第 1 讲的公式") cfg = TINY p = init_params(cfg, seed=0) mc = cfg.to_model_config() cases = [("prefill：B=1, T=128", 1, 128, 0), ("decode：B=8, 已有 127 个 token", 8, 1, 127)] for label, B, T, pos0 in cases: c = run_counted(cfg, p, B, T, pos0) merged = c.merged(OP_GROUPS) rows = [] for g, r in merged.items(): rows.append([g, fmt_num(r["flops"]), fmt_num(r["useful"]), fmt_bytes(r["w"]), fmt_bytes(r["kv"]), fmt_bytes(r["act"])]) print(f"\n{label}") print_table(["运算", "执行的 FLOPs", "有用的 FLOPs", "权重字节", "KV 字节", "激活字节"], rows) rows_a = analytic_ops(cfg, B, T, pos0, last_only=(T == 1)) same = all(abs(c.rows[op][k] - r[k]) < 1e-6 for op, r in rows_a.items() for k in ("flops", "useful", "w", "kv", "act")) print(f" 逐项和解析公式 analytic_ops 比较：{'完全相等' if same else '不相等！'}") # 和第 1 讲 llm_math 对账 print("\n和第 1 讲 llm_math 的公式对账：") rows = [] c = run_counted(cfg, p, 1, 128, 0) lin_counted = c.total("flops", MATMUL_OPS) rows.append(["prefill 矩阵乘法 FLOPs", fmt_num(lin_counted, 3), fmt_num(128 * mc.linear_flops_per_token(), 3), "T × 2·matmul_params"]) pf = prefill_time(mc, H100, 128) attn_useful = c.total("useful", ATTN_CORE_OPS) rows.append(["prefill 注意力（有用的）", fmt_num(attn_useful, 3), fmt_num(pf["flops"] - 128 * mc.linear_flops_per_token(), 3), "attn(1)·T(T+1)/2"]) rows.append(["prefill 注意力（执行的）", fmt_num(c.total("flops", ATTN_CORE_OPS), 3), "-", "朴素实现算满 T×T"]) c = run_counted(cfg, p, 8, 1, 127) st = decode_step(mc, H100, batch=8, ctx=128) rows.append(["decode 读权重", fmt_bytes(c.total("w", MATMUL_OPS)), fmt_bytes(st["weight_bytes"]), "matmul_params × 2 B"]) rows.append(["decode 读 KV", fmt_bytes(c.total("kv", ATTN_CORE_OPS)), fmt_bytes(st["kv_bytes"]), "B × ctx × KV/token（ctx = 127 + 1）"]) rows.append(["decode 写 KV", fmt_bytes(c.total("kv", ["kv_write"])), "-", "B × 1 × KV/token（公式忽略）"]) rows.append(["decode FLOPs（矩阵乘法 + 注意力）", fmt_num(c.total("flops", MATMUL_OPS + ATTN_CORE_OPS), 3), fmt_num(st["flops"], 3), "B × flops_per_token(ctx)"]) print_table(["项目", "数出来的", "llm_math", "公式 / 说明"], rows) print_observations([ "数出来的矩阵乘法 FLOPs、注意力的有用 FLOPs、decode 读权重和读 KV 的字节，和第 1 讲的公式**分毫不差**" "（decode 时公式里的 ctx 要取'包括新 token 在内'的 128）。第 1 讲的账本可以放心用。", "prefill 时朴素实现'执行的'注意力 FLOPs 约是'有用的'两倍：先把 T×T 的分数全算出来再掩掉一半。" "FlashAttention 一类的内核会跳过整块被掩掉的部分（第 11 讲）。", "decode 时'读权重'被批里 8 个 token 共享，'读 KV'却是每个序列各读各的——第 1 讲 §4.3 的结论在这里变成了两行数字。", "公式忽略了激活字节和写 KV。prefill 时激活字节很大（每个运算都读写一遍 [T, d] 级别的张量），" "这部分要靠**算子融合**省掉（第 11、12 讲）；decode 时它们相对权重很小。", ]) # --------------------------------------------------------------------------- # 演示 3：einsum 记号 # --------------------------------------------------------------------------- def demo3_einsum(): """ How to Scale Your Model 用单个字母表示维度：B 批、T query 长度、S key 长度、D 模型维、F MLP 中间维、 N query 头、K KV 头、G 每组的 query 头数（N = K·G）、H 每头维度、V 词表。 规则：一次缩并的 FLOPs = 2 × 所有出现过的维度之积（每个字母只算一次）。 """ print_header("演示 3：einsum 记号——每一步写成一个 einsum，FLOPs = 2 × 所有维度之积") m = LLAMA31_8B B, T, S = 1, 2048, 2048 D, F, N, K, H, V = m.d_model, m.d_ff, m.n_heads, m.n_kv_heads, m.d_head, m.vocab G = N // K ops = [ ("q_proj", "btd,dnh->btnh", (B, T, D), (D, N, H)), ("k_proj", "btd,dkh->btkh", (B, T, D), (D, K, H)), ("v_proj", "btd,dkh->btkh", (B, T, D), (D, K, H)), ("QKᵀ", "btkgh,bskh->btskg", (B, T, K, G, H), (B, S, K, H)), ("P·V", "btskg,bskh->btkgh", (B, T, S, K, G), (B, S, K, H)), ("o_proj", "btnh,nhd->btd", (B, T, N, H), (N, H, D)), ("mlp_gate", "btd,df->btf", (B, T, D), (D, F)), ("mlp_up", "btd,df->btf", (B, T, D), (D, F)), ("mlp_down", "btf,fd->btd", (B, T, F), (F, D)), ] rows = [] total = 0 for name, spec, s1, s2 in ops: f = einsum_flops(spec, s1, s2) total += f rows.append([name, spec, fmt_num(f, 2)]) print(f"Llama-3.1-8B 的一层，B = {B}、T = S = {T}（不计因果掩码）：") print_table(["运算", "einsum", "FLOPs"], rows) a = analytic_ops(m, B, T, 0, with_cache=False) per_layer = (sum_rows(a, "flops", ATTN_PROJ_OPS + MLP_OPS + ATTN_CORE_OPS)) / m.n_layers print(f"一层合计：einsum 规则 {fmt_num(total, 3)}，analytic_ops（执行的 FLOPs）{fmt_num(per_layer, 3)}") # 用 numpy 的 einsum 真算一个小例子，确认规则和 transformer_lab.attention_core 的写法一致 rng = np.random.default_rng(0) q = rng.standard_normal((1, 4, 2, 3, 8)) k = rng.standard_normal((1, 5, 2, 8)) s1 = np.einsum("btkgh,bskh->btskg", q, k) s2 = np.einsum("btgrd,bsgd->bgrts", q, k).transpose(0, 3, 4, 1, 2) print(f"小例子：书上的 'btkgh,bskh->btskg' 与 transformer_lab 的 'btgrd,bsgd->bgrts' 最大差 {np.max(np.abs(s1 - s2)):.1e}") print_observations([ "einsum 字符串本身就说明了一切：哪些维度相乘、哪些被求和（在输入里出现、输出里没有的字母）、结果是什么形状。" "FLOPs 规则只看字母，不用画图。", "QKᵀ 和 P·V 里 B、K 出现在两个输入里，是'批维'（batching）；它们都含 T 和 S，所以运算量 ∝ T·S——注意力的平方项。", "GQA 在记号上就是把 N 拆成 K·G：K 组 KV 头，每组服务 G 个 query 头。K、V 的形状里只有 K 没有 G，" "这就是它省 KV cache 的原因。", "第 4 讲的切分记号（比如 D 被切到 Y 轴上写成 D_Y）就是在这些字母上加下标——先熟悉 einsum，切分就只是'给字母加下标'。", ]) # --------------------------------------------------------------------------- # 演示 4：注意力的运算什么时候不能忽略 # --------------------------------------------------------------------------- def demo4_attention_share(): print_header("演示 4：注意力核心的 FLOPs 什么时候追上矩阵乘法") rows = [] for m in (LLAMA31_8B, LLAMA31_70B, LLAMA31_405B): lin = m.linear_flops_per_token() cells = [m.name] for ctx in (1024, 8192, 32768, 131072): a = m.attn_flops_per_token(ctx) cells.append(f"{a / (a + lin) * 100:.0f}%") ctx_star = lin / m.attn_flops_per_token(1) cells.append(f"{ctx_star / 1000:.0f} K") cells.append(f"{2 * ctx_star / 1000:.0f} K") rows.append(cells) print("decode 一个 token（上下文 ctx）时注意力核心占总 FLOPs 的比例；以及两者相等的长度：") print_table(["模型", "ctx=1K", "8K", "32K", "128K", "decode 相等点 ctx*", "prefill 相等点 ≈ 2·ctx*"], rows) print_observations([ "decode 时：ctx* = 2·matmul_params / (4·L·H·dh)。8B 约 29K、70B 约 53K、405B 约 98K——模型越大，" "矩阵乘法越重，注意力越晚成为运算的主要部分。How to Scale Your Model 对 MHA、F = 4D 的模型给出 T > 8D 的经验式，同一个道理。", "prefill 时第 i 个 token 只看前 i 个，平均上下文是 n/2，所以相等点约是 2·ctx*。", "但**运算不是 decode 的瓶颈**：哪怕注意力的 FLOPs 很少，读 KV cache 的字节也可能远超读权重（第 1 讲 §4.3）。" "看 FLOPs 占比会低估长上下文的代价，要看字节（演示 5）。", ]) # --------------------------------------------------------------------------- # 演示 5：prefill 和 decode 的计算图 # --------------------------------------------------------------------------- def roofline_rows(model, chip, B, T, pos0, last_only): """ 对每一类运算单独套 roofline：时间下界 = max(FLOPs / 峰值, 字节 / 带宽)，字节 = 权重 + KV + 激活。 【假设】每一类运算是独立的内核、彼此不重叠（所以总时间 = 各项之和，比"整体一个 roofline"更悲观也更接近真实的 未融合实现）；注意力用有用的 FLOPs（假设内核会跳过被掩掉的块）；BF16。 【返回】(表格行, 各项时间之和, 总 FLOPs, 总字节) """ a = analytic_ops(model, B, T, pos0, last_only=last_only) peak, bw = chip.peak_flops("bf16"), chip.mem_bw rows, t_sum, f_sum, b_sum = [], 0.0, 0.0, 0.0 for g, names in OP_GROUPS.items(): fl = sum_rows(a, "useful", names) by = sum_rows(a, "w", names) + sum_rows(a, "kv", names) + sum_rows(a, "act", names) if by == 0: continue t = max(fl / peak, by / bw) inten = fl / by rows.append([g, fmt_num(fl), fmt_bytes(by), f"{inten:.1f}", "compute" if fl / peak >= by / bw else "memory", fmt_time(t)]) t_sum += t f_sum += fl b_sum += by return rows, t_sum, f_sum, b_sum def demo5_prefill_vs_decode(): print_header("演示 5：prefill 与 decode 的'计算图'——每一类运算的强度和瓶颈（Llama-3.1-8B，H100，BF16）") m = LLAMA31_8B print(f"H100 的屋脊点：{H100.critical_intensity('bf16'):.0f} FLOP/byte；每一类运算单独套 roofline，再把时间加起来") cases = [("prefill：1 个序列，提示 2048 个 token", 1, 2048, 0, True), ("decode：B = 1，上下文 2048", 1, 1, 2047, True), ("decode：B = 64，上下文 2048", 64, 1, 2047, True)] summary = [] for label, B, T, pos0, last in cases: rows, t_sum, f_sum, b_sum = roofline_rows(m, H100, B, T, pos0, last) print(f"\n{label}") print_table(["运算", "FLOPs", "字节", "强度", "瓶颈", "时间下界"], rows) if T > 1: ref = prefill_time(m, H100, T) else: ref = decode_step(m, H100, batch=B, ctx=pos0 + 1) summary.append([label, fmt_time(t_sum), fmt_time(ref["t_step"]), f"{t_sum / ref['t_step']:.2f}×"]) print("\n逐项相加 vs 第 1 讲的'整体一个 roofline'（不计激活字节、计算和搬运完全重叠）：") print_table(["场景", "逐项相加", "llm_math 下界", "比值"], summary) print_observations([ "prefill：矩阵乘法的强度上千（远高于屋脊点 295），是 compute-bound；但 RMSNorm、RoPE、SwiGLU 这些逐元素运算" "强度只有 1–2，单独执行时是 memory-bound 的，一共要 3.7 ms，相当于矩阵乘法时间（约 30 ms）的 12%。" "融合它们是编译器和内核的工作（第 11–13 讲）。lm_head 这里只算最后一个 token，所以是 memory-bound。", "decode B = 1：所有运算的强度都 ≈ 1，**整张计算图都是 memory-bound**，时间几乎全是读权重。", "decode B = 64：投影和 MLP 的强度涨到约 60，仍低于 295；注意力核心的强度停在 ≈ G = 4（GQA 每组 4 个 query 头）——" "批再大它也不会变：每个序列读自己的 KV，只被自己的 4 个 query 头用（How to Scale Your Model 第 4 章问题 4）。" "上下文才 2048，读 KV 的注意力核心（5.1 ms）就已经比整个 MLP（3.4 ms）还慢了。", "decode 的逐项相加和整体下界几乎相等（1.00–1.02 倍，时间都在搬字节）；prefill 逐项相加多出 6%，主要来自逐元素运算——" "**整体 roofline 是下界，逐项 roofline 更接近未融合的真实实现**。", ]) if __name__ == "__main__": demo1_params() demo2_one_forward() demo3_einsum() demo4_attention_share() demo5_prefill_vs_decode()