Transformer 与 LLM：结构、实现与算量

01 · Transformer 长什么样：只有五种运算

结论：embedding 查表 → attention（唯一让 token 互相看的地方） → FFN（逐 token 的非线性，知识在这，占一层 2/3）→ 残差 + LayerNorm → lm_head；attention 是集合运算不知道顺序，位置必须显式给。

%%{init: {"flowchart": {"wrappingWidth": 330}}}%%
flowchart LR
    IN["token 编号<br/>[464, 3797, 3332, 319, 262]"] --> EMB["embedding 查表<br/>50257 × 768 → 5 个 768 维向量<br/>+ 位置表"]
    EMB --> B
    subgraph B["× 12 个 block"]
        direction TB
        A["LayerNorm → causal self-attention → 加回"]
        F["LayerNorm → FFN 768→3072→768 → 加回"]
        A --> F
    end
    B --> HEAD["LayerNorm → lm_head 768 × 50257<br/>softmax → p(mat) = 0.31"]

一个 attention 头的六步手算（d = 4、T = 3）

• (\text{Attention}(Q,K,V) = \text{softmax}(QK^\top/\sqrt d + M)\,V)；t2 的权重 (0.27, 0.27, 0.45)，与 PyTorch 对拍差 (6 \times 10^{-8})
• 随机 (q \cdot k) 的标准差是 (\sqrt d)：不除，d = 128 时 softmax 一个 token 独占

GPT-2 真实的头在看什么

• 多头不增加算量：h 个小乘法 = 一个大乘法，只多一个 (W_O)
• 换序实验：把输入打乱，attention 输出跟着换位——所以位置必须另给（GPT-2 查表 / Llama RoPE）

02 · 一个 token 的旅程：训练侧与推理侧

结论：causal mask + teacher forcing 让一句话 = T 个训练样本；推理时旧 token 的 K、V 不变，算一次存下来就是 KV cache——训练 / prefill / decode 是三种形态。

prefill 与 decode：为什么前面的 token 不用重算

数字

• 有 / 无 KV cache 输出逐 token 一致；prompt 256 生成 256 快 7.9 倍
• 代价：Llama-3-8B 128 KiB / token = 32 层 × 2 × 8 个 KV 头 × 128 × 2 B
• KV cache 不存 Q——旧 token 的 Q 用过即弃
• 训练长上下文比推理多一份与 B × T 成正比的激活值

03 · 手搓 GPT（上）：nanoGPT model.py 逐行

结论：330 行、6 个类，结构本身（LayerNorm、CausalSelfAttention、MLP、Block）不到 90 行；与 HF GPT-2 对拍相对差 (9 \times 10^{-5})。

class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # pre-norm：先 norm 再子层，残差直连
        x = x + self.mlp(self.ln_2(x))
        return x

# CausalSelfAttention.forward 的骨架
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)          # 一次 GEMM 算 QKV
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # 头换到第 1 维
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # 融合 kernel
y = y.transpose(1, 2).contiguous().view(B, T, C)              # 拼回去

每个类落在结构的哪个方框

• Llama 相对 GPT-2 只改五处：RMSNorm、RoPE、SwiGLU、GQA、去 bias（lm_head 不共享）
• 三个独立 Linear 与一个 c_attn 数学等价，合并只为一次 GEMM

04 · 手搓 GPT（下）：train.py 与训一个会续写的模型

结论：1.1 MB 莎士比亚、0.8M 参数、2000 步 7 分钟，loss 从 ln 65 = 4.17 → 1.66；层是串行的——参数与耗时线性、loss 收益递减。

训练循环每一步在做什么

%%{init: {"flowchart": {"wrappingWidth": 170}}}%%
flowchart LR
    LR["设 lr<br/>warmup + cosine"] --> EV["到点评估<br/>存 checkpoint 五样"] --> ACC["k 个 micro-batch<br/>loss ÷ k 累积"] --> CL["unscale<br/>+ 梯度裁剪"] --> ST["optimizer.step()<br/>zero_grad()"] --> LR

• get_batch：memmap + 随机窗口 + stack，y 是 x 右移一位，没有 epoch
• 不除累积步数 ≠ 只是尺度不同：等价于学习率放大 k 倍；DDP 只在最后一个 micro-batch 同步
• checkpoint 五样：model、optimizer、model_args、iter_num、config——续训不带优化器状态，前几百步抖

05 · Transformer 解剖：从 config.json 算出 8.03B

结论：dense Transformer 没有隐藏参数——每层四个 attention 矩阵 + 三个 SwiGLU 矩阵，乘层数加词表，Llama-3-8B 算出 8,030,261,248，精确到个位。

[ N = L\big[d(2d + 2d_{kv}) + 3d \cdot d_{ff} + 2d\big] + 2Vd + d,\qquad d_{kv} = n_{kv}\,d_{head} ]

常见的算错方式，以及第一次分裂

• 词表占比随规模下降：8B 13%、70B 3%、405B 约 1%
• DeepSeek-V3：骨架相同，attention 换 MLA 六个矩阵、FFN 换 257 个专家——总参数 671B、每 token 激活 37B，「参数量」第一次不再单独对应成本

06 · Attention 变体与 KV cache：MHA、GQA、MQA、MLA

结论：KV cache 每 token = (2\,L\,n_{kv}\,d_{head} \cdot) bytes/elem，公式里没有 (n_h)——V3 128 头 61 层的 KV（68.6 KiB）比 8B 32 头（128 KiB）还小。

MLA：只有 c^KV 与 k^R 进 cache

%%{init: {"flowchart": {"wrappingWidth": 260}}}%%
flowchart LR
    H["h_t (7168)"] --> DKV["W_DKV + W_KR<br/>7168 → 576"]
    DKV --> CKV["c^KV_t (512)<br/>进 cache"]
    DKV --> KR["k^R_t (64)，RoPE<br/>128 个 head 共享，进 cache"]
    CKV --> UKV["W_UK / W_UV 升维<br/>512 → 128 head × 128"]
    UKV --> ATT["attention<br/>q·[k^C ; k^R] / √192"]
    KR --> ATT
    H --> Q["W_DQ → W_UQ + W_QR<br/>q^C 128 维 + q^R 64 维"] --> ATT
    ATT --> WO["W_O 16384 → 7168"]
    classDef cache fill:#fde68a,stroke:#b45309,stroke-width:2px,stroke-dasharray:5 3
    class CKV,KR cache

• 解耦 RoPE：位置相关的旋转吸不进与位置无关的升维矩阵，所以单独留 64 维
• decode 把 (W_{UK})、(W_{UV}) 吸收进 (W_Q)、(W_O)：等价于 128 头共享一个 576 维 KV 头的 MQA，强度约 242；代价 attention FLOPs 约 3.4 倍，所以 prefill 走非吸收路径

物化的 S = QKᵀ 有多大，FlashAttention 省了什么

• 8K 上下文、32 头：每层 (s \times s) 的 logits 4 GiB（BF16），写出 HBM 再读回
• FlashAttention：分块 + online softmax，从不物化 S，HBM 流量降到 (O(s^2 d^2 / M))；算量不变
• 8B 单卡 8K 的最大并发约 59、128K 只有 3：KV 定并发
• MLA 不是免费的：多算 3.4 倍 attention FLOPs、维护两条等价路径、专用 kernel、不能从 MHA checkpoint 直接转换

07 · 位置编码与长上下文：RoPE 的波长

结论：RoPE 把 (d_{head}) 维向量看成 64 对复数，每对以自己的波长旋转；8K 训练时 14 对低频维度没转完一圈——推到 32K 出现从未见过的相位。不是装不下，是没见过。

长上下文的账：线性项、二次项，与交叉点

• 滑窗把 KV 与算量从 O(s) 变 O(W)（Mistral W = 4096）但丢信息；sink + 滑窗保留开头 4 个 token
• 交错局部 / 全局把系数变 1/k；MLA 减系数不减阶——两者正交可叠加
• 对 Infra：KV 定并发、二次项定 TTFT、chunked prefill 防一条 128K 请求独占 GPU 11 s

08 · MoE：三个「参数量」分开算

结论：总参数定显存（671B，FP8 也放不进 8 卡 640 GB）、激活参数定 FLOPs（37B，是 70B 的一半）、每步实际读取的参数定 decode 带宽——中等 batch 下几乎读全部专家，稀疏省了算量没省访存。

MoE 层：router 打分、top-k 选择，只算选中的

%%{init: {"flowchart": {"wrappingWidth": 240}}}%%
flowchart LR
    TOK["token x（d 维）"] --> RT["router W_r x<br/>E 个分数 → top-k（k = 8）"]
    RT -->|"g_1"| E1["FFN_1 选中"]
    RT -.->|"g_2 = 0"| E2["FFN_2 不计算"]
    RT -->|"g_E"| EE["… FFN_E 选中"]
    TOK --> SH["共享专家<br/>不经 router"]
    E1 & EE & SH --> SUM["y = Σ g_i · FFN_i(x) + FFN_shared(x)<br/>算了 k + 1 个，持有 E + 1 个"]
    classDef active fill:#dff0d8,stroke:#3c763d
    classDef idle fill:#f5f5f5,stroke:#aaa,color:#888
    classDef sh fill:#fcf8e3,stroke:#8a6d3b
    class E1,EE active
    class E2 idle
    class SH sh

部署为什么难：EP 与 all-to-all

• 出路是专家并行（EP32 每卡约 37 GB、EP320 每卡 19.6 GB，简化模型），代价：
  • 每层两次 all-to-all：每 token 每层 dispatch FP8 56 KiB + combine BF16 112 KiB，是 TP-8 all-reduce 的 6.7 倍；4096 token 的 prompt 58 层共 38 GiB
  • 每专家 GEMM 只有 (Tk/E) 行，强度比 dense 低 E/k = 32 倍，过 ridge 需一层里约 9600 个 token
  • 最慢的卡决定全层时间 → 节点受限路由（每 token 最多 4 个节点）、aux-loss-free 均衡、冗余专家
• TP-8 不行：把 2048 宽的专家切成 256 列太瘦，且不减少每卡读的专家数

09 · MTP：改训练目标而不改主干

结论：每个位置额外预测 (t_{i+2})，信号密度 ×(1 + D)，主干被逼编码更远的未来；模块顺序喂真实 (t_{i+1}) 保持因果链（teacher forcing 的延伸）；推理时丢弃或当投机 draft。

%%{init: {"flowchart": {"wrappingWidth": 220}}}%%
flowchart LR
    T["主干 Transformer<br/>h_i"] --> H0["lm_head → 预测 t_{i+1}<br/>L_main"]
    T --> C["RMSNorm(h_i) ‖ RMSNorm(emb(t_{i+1}))<br/>投影 2d → d"]
    E["真实 t_{i+1}"] --> C
    C --> BLK["一个 Transformer block"] --> H1["共享 lm_head → 预测 t_{i+2}<br/>λ · L_MTP"]
    classDef mtp fill:#e0e7ff,stroke:#4338ca
    class C,BLK,H1 mtp

多花了什么、多得了什么

• 「MTP 头 loss 更低所以更强」——条件不同；「小模型没收益所以无效」——规模依赖
• 顺序模块 vs 并行头是两条路线：并行头跳过中间 token 会干扰主干

10 · 前向的算量与访存量：一条 Roofline

结论：每参数每 token 2 FLOPs；decode 每步把 16 GB 权重读一遍、算术强度在数值上就等于 B，ridge 是 989 / 3.35 ≈ 295——这就是「decode 是 memory-bound」的全部含义。

数字：8B 在一张 H100 上

• 训练 6ND（激活重算 8ND）；激活 (sbh(34 + 5as/h))，8K 时每层 11 GiB 其中 10 GiB 是 s² 项
• 对 decode，「模型多大」的度量是字节不是 FLOPs：FFN 砍一半下界不变，INT4 才降到 1.3 ms

11 · 浮点格式：指数位定范围，尾数位定精度

结论：前向反向要范围、权重更新要精度——所以 BF16 计算 + FP32 master weights；(\Delta w / w \sim 10^{-4}) 低于 BF16 的单位舍入 (2^{-8} = 0.004)，1.0 + 0.001 会被舍回 1.0。

各格式能表示的范围，与训练里梯度的分布

• 训练状态 16 B / 参数（bf16 权重 2 + 梯度 2 + fp32 主权重 4 + Adam m、v 8）：8B 128 GB、70B 1.1 TB、671B 10.7 TB
• FP8 分工：E4M3 存权重与激活、E5M2 存梯度；V3 每 128 项提升到 FP32 累加，与 1 × 128 / 128 × 128 的分块 scale 对齐

数值丢失的四个位置，以及什么才是 bug

• 两个 kernel 的差异在 ε 到 ε√k 之间是噪声（BF16 GEMM 相对 FP32 参考 10⁻³–10⁻²）；大几个数量级或有系统性符号才是 bug
• QK-norm 把 attention logit 上界压到 (\sqrt{d_{head}}\,g_q g_k \approx 11.3\,g_q g_k)

12 · 量化、投机解码与 LoRA：三种方法，同一条 Roofline

结论：都不改结构，各改一个变量——量化改 (W_{bytes})、投机改每步的 m、LoRA 改训练时的 N；前两者都在兑现 memory-bound 区间里空转的算力，过 ridge 收益同时消失。

三种方法各改哪个变量

• GPTQ 用 Hessian 把误差补偿到未量化的列；AWQ 保护激活幅度最大的 1% 通道；SmoothQuant 把激活的离群迁到权重
• W8A8 字节减半且算力翻倍，对 prefill 也有效；实测 decode 约 3 倍而非 4 倍（lm_head 保留、KV 读取不减）

13 · 多模态：一张图等于多少 token

结论：图片贵的不是 encoder（一次性、compute-bound、12 ms），是它变成的 token 在 decoder 里占的 KV——与同长文本同价、活到请求结束；1024² 在 Qwen2-VL 里 = 1369 个 token，一段无法被 tokenizer 压短的 system prompt。

三段各一笔账（70B 规格，1024² 一张图）

• 比值 (2Ln_{kv}d_{head}/d_{model})：70B 20、8B 16、Qwen2-VL-7B 8
• cross-attention 注入（Llama 3.2 Vision）用 0.5B 参数换序列长度，图片 KV 800 → 200 MiB
• 一分钟 720p 1 fps 视频 35,880 token；按像素预算，不按张数

第三段共用的成本模型

%%{init: {"flowchart": {"wrappingWidth": 260}}}%%
flowchart TB
    S["结构：L、d、d_ff、n_kv、E、k<br/>（01、03、05）"] --> N["参数量 N、激活参数量（05、08）"]
    S --> K["KV / token = 2 L n_kv d_head · bytes/elem（06）"]
    D["数值：bytes/param、bytes/elem（11、12）"] --> W["权重字节 = N · bytes/param"]
    N --> W
    D --> K
    N --> F["FLOPs / token ≈ 2N + 4dLs（10、07）"]
    R["运行点：B、s、prefill / decode（02、07、08、10）"] --> F
    R --> H["每步字节 = 权重字节 + B · s · KV/token"]
    K --> H
    W --> H
    F --> I["算术强度 I = FLOPs / 字节"]
    H --> I
    I --> T["与 ridge ≈ 295 比较 → memory / compute-bound → 时间下界"]

五条贯穿线