第 2 章 Transformer 架构：工程师视角

模型参数（Llama 2 7B）:
- vocab_size = 32000       # 词表大小
- hidden_dim = 4096        # 隐藏层维度
- n_layers = 32            # Transformer 层数
- n_heads = 32             # Attention 头数
- head_dim = 128           # 每个头的维度 (4096 / 32)
- intermediate_dim = 11008 # FFN 中间层维度

输入文本: "什么是 KV Cache"
    │
    ▼
┌──────────────────────────────────────────────────────────┐
│ Step 1: Tokenize                                         │
│ "什么是 KV Cache" → [20345, 12876, 476, 8067, 28747, 5765] │
│ 输出 shape: [6]  (6 个 token ID)                          │
└──────────────────────────────────────────────────────────┘
    │
    ▼
┌──────────────────────────────────────────────────────────┐
│ Step 2: Token Embedding + Positional Encoding            │
│ Token Embedding：token ID → 稠密向量的查表步骤            │
│ Positional Encoding (PE)：位置编码，把"第几个 token"      │
│   也变成向量信号注入进去，不然 Attention 看不出顺序        │
│ 每个 token ID 查表得到一个 4096 维向量                     │
│ 再叠加位置编码（Llama 2 用 RoPE，旋转位置编码）            │
│ 输出 shape: [6, 4096]                                    │
└──────────────────────────────────────────────────────────┘
    │
    ▼
┌──────────────────────────────────────────────────────────┐
│ Step 3: 32 × Transformer Block (重复 32 次)               │
│ Transformer Block = Attention + FFN + 两次 Norm + 两次   │
│ Residual 的固定积木，整个模型靠它堆叠出来                  │
│                                                          │
│   ┌────────────────────────────────────────────┐         │
│   │ 3a. RMSNorm（归一化层，2.1 末尾会展开）       │         │
│   │     [6, 4096] → [6, 4096]                  │         │
│   ├────────────────────────────────────────────┤         │
│   │ 3b. Multi-Head Self-Attention (MHA)        │         │
│   │     即多头自注意力，把 Attention 拆成多个独立  │         │
│   │     的"头"并行计算，再拼回来                 │         │
│   │     Q (Query) / K (Key) / V (Value)：       │         │
│   │     attention 的三个投影矩阵，类比一次       │         │
│   │     带权 lookup：Q 是查询、K 是键、V 是值     │         │
│   │     Q: [6, 4096] → [6, 32, 128]           │         │
│   │     K: [6, 4096] → [6, 32, 128]           │         │
│   │     V: [6, 4096] → [6, 32, 128]           │         │
│   │     Attention 计算后: [6, 32, 128]          │         │
│   │     投影回: [6, 4096]                       │         │
│   │     注：Llama 2 7B 用 MHA（32 个 KV 头）；   │         │
│   │     Llama 3 / Llama 2 70B 改用 GQA          │         │
│   │     （Grouped Query Attention，分组查询注意 │         │
│   │     力；8 个 KV 头，K/V 变成 [6, 8, 128]）  │         │
│   ├────────────────────────────────────────────┤         │
│   │ 3c. Residual Connection（残差连接）         │         │
│   │     x = x + attention_output               │         │
│   │     把输入"短路"加到输出上，让梯度能穿过深层   │         │
│   ├────────────────────────────────────────────┤         │
│   │ 3d. RMSNorm                                │         │
│   │     [6, 4096] → [6, 4096]                  │         │
│   ├────────────────────────────────────────────┤         │
│   │ 3e. FFN (SwiGLU)                           │         │
│   │     SwiGLU 是 Llama 系列 FFN 用的激活函数，  │         │
│   │     可看作 ReLU 的升级版                    │         │
│   │     [6, 4096] → [6, 11008] → [6, 4096]    │         │
│   ├────────────────────────────────────────────┤         │
│   │ 3f. Residual Connection                    │         │
│   │     x = x + ffn_output                     │         │
│   └────────────────────────────────────────────┘         │
│                                                          │
│ 输出 shape: [6, 4096]                                    │
└──────────────────────────────────────────────────────────┘
    │
    ▼
┌──────────────────────────────────────────────────────────┐
│ Step 4: Final RMSNorm + Linear Head                      │
│ Linear Head（也叫 LM Head，Language Model Head，         │
│ 语言模型头）：一层全连接，把 hidden 向量映射到词表大小      │
│ [6, 4096] → RMSNorm → [6, 4096] → Linear → [6, 32000]  │
│ 取最后一个位置的 logits: [32000]                           │
│ （logits 是词表上每个 token 的"原始分数"，可正可负，         │
│   未经 softmax 之前不是概率分布）                           │
└──────────────────────────────────────────────────────────┘
    │
    ▼
┌──────────────────────────────────────────────────────────┐
│ Step 5: Sampling（采样，按概率从 logits 里挑下一个 token） │
│ 对 [32000] 维 logits 做 softmax（把任意实数向量压成概率   │
│ 分布的函数）+ temperature scaling（温度缩放，调节生成的    │
│ "随机度"，越高越散）                                       │
│ 然后 top-p/top-k 采样（截断式采样：top-k 只从概率最高的 k │
│ 个里挑、top-p 从累积概率达到 p 的最小集合里挑），           │
│ 输出 1 个 token ID                                       │
│ 比如: 15043 → decode → "是"                              │
└──────────────────────────────────────────────────────────┘

初始:  l o w e r    (5个字符)
第1轮: lo w e r     (lo 合并，因为 l+o 最频繁)
第2轮: low e r      (lo+w 合并)
第3轮: lower        (low+er 合并)

输入: "大语言模型的推理优化"

tiktoken (GPT-4):     ["大语言", "模型", "的", "推理", "优化"]          → 5 tokens
SentencePiece (Llama): ["大", "语言", "模型", "的", "推", "理", "优化"] → 7 tokens

Q = X @ Wq    # [n, d] @ [d, d] → [n, d]
K = X @ Wk    # [n, d] @ [d, d] → [n, d]
V = X @ Wv    # [n, d] @ [d, d] → [n, d]

Scores = Q @ K^T     # [n, d] @ [d, n] → [n, n]   ← 这步是 O(n²)
Scores = Scores / sqrt(d)
Weights = softmax(Scores)  # [n, n]
Output = Weights @ V       # [n, n] @ [n, d] → [n, d]

MHA (Multi-Head Attention，多头注意力，原始 Transformer 的标准做法):
  Q: 32 heads    K: 32 heads    V: 32 heads

GQA (Grouped Query Attention, 分组查询注意力, 8 KV heads):
  Q: 32 heads    K: 8 heads     V: 8 heads
  每 4 个 Q head 共享 1 个 KV head

MQA (Multi-Query Attention，多查询注意力，GQA 的极端版):
  Q: 32 heads    K: 1 head      V: 1 head
  所有 Q head 共享 1 个 KV head

第 1 步: 输入 [A, B, C]         → 计算全部 Attention → 生成 D
第 2 步: 输入 [A, B, C, D]     → 重新计算全部 Attention → 生成 E
第 3 步: 输入 [A, B, C, D, E] → 又重新计算全部 Attention → 生成 F

第 1 步: 输入 [A, B, C]     → 计算 K_ABC, V_ABC 并缓存 → 生成 D
第 2 步: 只输入 [D]          → 计算 K_D, V_D，和缓存的 KV 拼接 → 生成 E
第 3 步: 只输入 [E]          → 计算 K_E, V_E，和缓存的 KV 拼接 → 生成 F

KV Cache 大小 = 2 × n_layers × n_kv_heads × head_dim × seq_len × batch_size × bytes_per_param

2 × 32 layers × 32 heads × 128 dim × 2048 seq × 1 batch × 2 bytes(FP16)
= 2 × 32 × 32 × 128 × 2048 × 2
= 1,073,741,824 bytes
≈ 1 GB

Llama 2 7B (32 KV heads): seq_len=2048, batch=1 → 1 GB
Llama 3 8B (8 KV heads):  seq_len=2048, batch=1 → 256 MB

First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

First Citizen:
You are all resolved rather to die than to famish?

class MiniGPT(nn.Module):
    def __init__(self, vocab_size, n_embd, n_head, n_layer, block_size):
        super().__init__()
        # 1. Token Embedding：把 token ID 映射成 n_embd 维向量
        self.token_embedding = nn.Embedding(vocab_size, n_embd)
        # 2. Position Embedding：每个位置（0..block_size-1）也学一个 n_embd 向量
        self.position_embedding = nn.Embedding(block_size, n_embd)
        # 3. 堆 n_layer 个 Transformer Block
        self.blocks = nn.Sequential(*[
            TransformerBlock(n_embd, n_head, block_size) for _ in range(n_layer)
        ])
        # 4. 最后一层 LayerNorm + 输出投影到词表大小
        self.ln_f = nn.LayerNorm(n_embd)
        self.lm_head = nn.Linear(n_embd, vocab_size)
 
    def forward(self, idx, targets=None):
        B, T = idx.shape                                     # [B, T] 一批 token ID 序列
        tok_emb = self.token_embedding(idx)                  # [B, T, n_embd]
        pos_emb = self.position_embedding(torch.arange(T))   # [T, n_embd]
        x = tok_emb + pos_emb                                # 广播相加 → [B, T, n_embd]
        x = self.blocks(x)                                   # 经过 n_layer 个 Block
        x = self.ln_f(x)
        logits = self.lm_head(x)                             # [B, T, vocab_size]
 
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

class TransformerBlock(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = MultiHeadAttention(n_embd, n_head, block_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ffn = FeedForward(n_embd)
 
    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # Pre-norm 后过 Attention，再做 Residual
        x = x + self.ffn(self.ln2(x))    # Pre-norm 后过 FFN，再做 Residual
        return x

class SelfAttention(nn.Module):
    def __init__(self, n_embd, head_dim, block_size):
        super().__init__()
        # Q/K/V 三个线性投影，把 n_embd 维向量分别投到 head_dim 维
        self.query = nn.Linear(n_embd, head_dim, bias=False)
        self.key   = nn.Linear(n_embd, head_dim, bias=False)
        self.value = nn.Linear(n_embd, head_dim, bias=False)
        # Causal Mask（因果掩码，强制每个 token 只能看见自己和左侧 token，
        # 不能"偷看"右边的未来 token，这是 decoder-only 模型的核心约束）：
        # 下三角矩阵，"已知/未知"的标记。register_buffer 表示
        # 这是模型的一部分（会跟 .to(device) 一起搬），但不是可训练参数
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))
 
    def forward(self, x):
        B, T, C = x.shape                                # [batch, seq_len, embed_dim]
        q, k, v = self.query(x), self.key(x), self.value(x)   # 各 [B, T, head_dim]
 
        # Attention scores = Q @ K^T / sqrt(d)
        scores = q @ k.transpose(-2, -1) * (k.shape[-1] ** -0.5)   # [B, T, T]
 
        # 把未来位置（上三角）填成 -inf，softmax 后它们的权重就是 0
        scores = scores.masked_fill(self.mask[:T, :T] == 0, float("-inf"))
 
        weights = F.softmax(scores, dim=-1)              # [B, T, T] 每行加起来等于 1
        return weights @ v                               # [B, T, head_dim]

1 0 0 0          位置 0 只能看自己
1 1 0 0          位置 1 能看位置 0、1
1 1 1 0          位置 2 能看位置 0、1、2
1 1 1 1          位置 3 能看全部前面的

device = "cuda" if torch.cuda.is_available() else "cpu"
model = MiniGPT(vocab_size, n_embd=128, n_head=4, n_layer=4, block_size=128).to(device)
# AdamW：Adam 优化器加权重衰减（weight decay）的变体，目前是训练 Transformer 事实上的默认优化器
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
 
for step in range(max_steps):
    # 1. 从数据集随机切一个 batch：xb 是输入 [B, T]，yb 是 xb 整体右移一位的 target
    xb, yb = dataset.get_batch(batch_size=32, device=device)
 
    # 2. Forward：模型同时返回 logits 和 loss（targets 不为 None 时）
    logits, loss = model(xb, yb)
 
    # 3. Backward + Update（训练三件套）
    optimizer.zero_grad()           # 清空上一步的梯度
    loss.backward()                  # 反向传播，把梯度填到每个参数的 .grad 上
    optimizer.step()                 # AdamW 根据 .grad 更新参数

ROMEO:
Rast thire the hey an to hame prechere lope,
To the rodwld thath pawath thou cand thavemes bod
To the thy this here and sose nong mand lambres wither
the pathe the on thy rown,
Now hey inteterion shi