第20章 从零实现 Transformer（PyTorch）

第六篇 实战

第20章 从零实现 Transformer（PyTorch）

本章目标：

不调用 nn.Transformer / nn.MultiheadAttention，用 PyTorch 从零写出一个完整的 Encoder-Decoder Transformer，在 RTX 4090 上 17 秒训练到 99.6% 的整句准确率，并亲眼看到 Cross-Attention 学出来的对齐矩阵。

本章目录

• 20.1 本章做什么 —— 任务选择、整体架构、项目结构与环境
• 20.2 输入层：Embedding + 位置编码 —— 对应第3章
• 20.3 Multi-Head Attention —— 对应第6、7章，逐行标注 Shape
• 20.4 Feed Forward Network —— 对应第8章
• 20.5 Encoder —— Pre-Norm 与最后那一层 LayerNorm
• 20.6 Decoder —— Masked Self-Attention + Cross-Attention
• 20.7 三种 Mask 与完整模型 —— 对应第13章，一次理清
• 20.8 先测再训：两个单元测试
• 20.9 数据：序列翻转任务
• 20.10 训练 —— Teacher Forcing、Warmup、bf16 混合精度
• 20.11 推理：贪心解码
• 20.12 运行结果：一次"代码全对、模型却学不会"的排查
• 20.13 实验：可视化 Cross-Attention
• 20.14 实验：手写 Attention vs SDPA（FlashAttention）
• 20.15 Hugging Face 源码定位
• 20.16 常见坑

20.1 本章做什么

20.1.1 为什么用"序列翻转"做任务？

真实的机器翻译需要下载语料、训练分词器、跑几个小时，大部分精力会花在数据工程上。本章选一个合成任务：把一串数字倒过来。

输入: 7 3 0 9 5 <eos>
输出: <bos> 5 9 0 3 7 <eos>

这个任务看似简单，却恰好能检验 Transformer 的每个零件：

任何一个组件出问题，模型都学不好——20.12 节就会看到一个真实的例子。

20.1.2 整体架构

#mermaid-svg-QjkyVt3CTxqgwa7R{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-QjkyVt3CTxqgwa7R .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-QjkyVt3CTxqgwa7R .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-QjkyVt3CTxqgwa7R .error-icon{fill:#552222;}#mermaid-svg-QjkyVt3CTxqgwa7R .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-QjkyVt3CTxqgwa7R .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-QjkyVt3CTxqgwa7R .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-QjkyVt3CTxqgwa7R .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-QjkyVt3CTxqgwa7R .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-QjkyVt3CTxqgwa7R .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-QjkyVt3CTxqgwa7R .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-QjkyVt3CTxqgwa7R .marker{fill:#333333;stroke:#333333;}#mermaid-svg-QjkyVt3CTxqgwa7R .marker.cross{stroke:#333333;}#mermaid-svg-QjkyVt3CTxqgwa7R svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-QjkyVt3CTxqgwa7R p{margin:0;}#mermaid-svg-QjkyVt3CTxqgwa7R .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-QjkyVt3CTxqgwa7R .cluster-label text{fill:#333;}#mermaid-svg-QjkyVt3CTxqgwa7R .cluster-label span{color:#333;}#mermaid-svg-QjkyVt3CTxqgwa7R .cluster-label span p{background-color:transparent;}#mermaid-svg-QjkyVt3CTxqgwa7R .label text,#mermaid-svg-QjkyVt3CTxqgwa7R span{fill:#333;color:#333;}#mermaid-svg-QjkyVt3CTxqgwa7R .node rect,#mermaid-svg-QjkyVt3CTxqgwa7R .node circle,#mermaid-svg-QjkyVt3CTxqgwa7R .node ellipse,#mermaid-svg-QjkyVt3CTxqgwa7R .node polygon,#mermaid-svg-QjkyVt3CTxqgwa7R .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-QjkyVt3CTxqgwa7R .rough-node .label text,#mermaid-svg-QjkyVt3CTxqgwa7R .node .label text,#mermaid-svg-QjkyVt3CTxqgwa7R .image-shape .label,#mermaid-svg-QjkyVt3CTxqgwa7R .icon-shape .label{text-anchor:middle;}#mermaid-svg-QjkyVt3CTxqgwa7R .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-QjkyVt3CTxqgwa7R .rough-node .label,#mermaid-svg-QjkyVt3CTxqgwa7R .node .label,#mermaid-svg-QjkyVt3CTxqgwa7R .image-shape .label,#mermaid-svg-QjkyVt3CTxqgwa7R .icon-shape .label{text-align:center;}#mermaid-svg-QjkyVt3CTxqgwa7R .node.clickable{cursor:pointer;}#mermaid-svg-QjkyVt3CTxqgwa7R .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-QjkyVt3CTxqgwa7R .arrowheadPath{fill:#333333;}#mermaid-svg-QjkyVt3CTxqgwa7R .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-QjkyVt3CTxqgwa7R .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-QjkyVt3CTxqgwa7R .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-QjkyVt3CTxqgwa7R .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-QjkyVt3CTxqgwa7R .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-QjkyVt3CTxqgwa7R .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-QjkyVt3CTxqgwa7R .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-QjkyVt3CTxqgwa7R .cluster text{fill:#333;}#mermaid-svg-QjkyVt3CTxqgwa7R .cluster span{color:#333;}#mermaid-svg-QjkyVt3CTxqgwa7R div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-QjkyVt3CTxqgwa7R .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-QjkyVt3CTxqgwa7R rect.text{fill:none;stroke-width:0;}#mermaid-svg-QjkyVt3CTxqgwa7R .icon-shape,#mermaid-svg-QjkyVt3CTxqgwa7R .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-QjkyVt3CTxqgwa7R .icon-shape p,#mermaid-svg-QjkyVt3CTxqgwa7R .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-QjkyVt3CTxqgwa7R .icon-shape .label rect,#mermaid-svg-QjkyVt3CTxqgwa7R .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-QjkyVt3CTxqgwa7R .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-QjkyVt3CTxqgwa7R .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-QjkyVt3CTxqgwa7R :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

源序列 src[B, Ls]

Token Embedding × √d+ 位置编码

Encoder Layer × NSelf-Attn + FFN

memory[B, Ls, d_model]

目标输入 tgt_in[B, Lt]

Token Embedding × √d+ 位置编码

Decoder Layer × NMasked Self-AttnCross-AttnFFN

Linear[B, Lt, vocab]

logits

20.1.3 项目结构与环境

ch20_transformer/
├── model.py # 全部模型组件（20.2 ~ 20.7）
├── test_model.py # Mask 单元测试（20.8）
├── train.py # 数据、训练、推理（20.9 ~ 20.11）
├── viz.py # Cross-Attention 可视化（20.13）
└── mha_vs_sdpa.py # 手写 Attention vs SDPA（20.14）

本章（以及第21、22章）的实测环境：

pip install numpy==2.2.6 matplotlib==3.10.3 # PyTorch 按官网命令单独安装对应 CUDA 版本

代码会自动检测设备：没有 GPU 时退回 CPU，本章的小模型在 CPU 上也能在几分钟内训完。下文所有输出都来自真实运行。

20.2 输入层：Embedding + 位置编码

第3章讲过：Transformer 自身不认识顺序，必须把位置信息加到 Token 向量上。

# ch20_transformer/model.py
"""第20章 从零实现 Transformer（PyTorch 版）—— 模型

Mask 约定（全文件统一）：bool 张量，True = 屏蔽（不许看），False = 可见。
注意：F.scaled_dot_product_attention 的 bool mask 语义正好相反（True = 可见），
本章为了拿到注意力权重做可视化，手写 Attention，不调用 SDPA。
"""
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

PAD_ID = 0

def sinusoidal_position_encoding(max_len: int, d_model: int) –> torch.Tensor:
"""正弦位置编码，返回 [max_len, d_model]"""
pos = torch.arange(max_len, dtype=torch.float32).unsqueeze(1) # [L, 1]
div = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32)
* (–math.log(10000.0) / d_model)) # [d/2]
pe = torch.zeros(max_len, d_model)
pe[:, 0::2] = torch.sin(pos * div) # 偶数维 sin
pe[:, 1::2] = torch.cos(pos * div) # 奇数维 cos
return pe

class TokenAndPosition(nn.Module):
"""Token Embedding × sqrt(d_model) + 位置编码"""
def __init__(self, vocab_size, d_model, max_len, dropout=0.1):
super().__init__()
self.d_model = d_model
self.emb = nn.Embedding(vocab_size, d_model, padding_idx=PAD_ID)
# PyTorch 的 Embedding 默认按 N(0, 1) 初始化，乘 sqrt(d_model) 后标准差≈11，
# 而位置编码只在 [-1, 1]，位置信号会被淹没。按 N(0, d^-0.5) 初始化，乘完后标准差≈1。
nn.init.normal_(self.emb.weight, mean=0.0, std=d_model ** –0.5)
with torch.no_grad():
self.emb.weight[PAD_ID].zero_()
# register_buffer：随模型 .to(device)，但不是可训练参数，也会存进 state_dict
self.register_buffer("pe", sinusoidal_position_encoding(max_len, d_model), persistent=False)
self.drop = nn.Dropout(dropout)

def forward(self, ids): # ids: [B, L]
L = ids.size(1)
x = self.emb(ids) * math.sqrt(self.d_model) + self.pe[:L]
return self.drop(x)

三个值得注意的细节：

• register_buffer：位置编码不是可训练参数，但要跟着 model.to("cuda") 一起搬到 GPU。注册成 buffer 就能做到；persistent=False 表示不存进 state_dict（它可以随时重新算出来）。
• 为什么 Embedding 要乘 $\sqrt{d_{model}}$：原论文 3.4 节的做法，目的是让 Token 信号和位置编码处在相近的数值范围。
• 为什么要手动初始化 Embedding：这是从其他框架移植时最容易踩的坑。PyTorch 的 nn.Embedding 默认按 $\mathcal{N}(0,1)$ 初始化，乘以 $\sqrt{128}\approx 11.3$ 之后标准差约为 11，而位置编码只在 $[-1,1]$ 之间——位置信号只剩不到十分之一，被淹没了。按 $\mathcal{N}(0,d^{-0.5})$ 初始化后，乘完正好是标准差 1。20.12 节会看到不改这一行时，模型学得有多慢。

20.3 Multi-Head Attention

这是第6、7章的代码化。forward 同时返回了注意力权重，20.13 节要用它画图。

# ch20_transformer/model.py（续）
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads, dropout=0.1):
super().__init__()
assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"
self.h, self.d_k = num_heads, d_model // num_heads
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
self.attn_drop = nn.Dropout(dropout)

def _split(self, x): # [B, L, d_model] -> [B, h, L, d_k]
B, L, _ = x.shape
return x.view(B, L, self.h, self.d_k).transpose(1, 2)

def _merge(self, x): # [B, h, L, d_k] -> [B, L, d_model]
B, _, L, _ = x.shape
return x.transpose(1, 2).reshape(B, L, self.h * self.d_k)

def forward(self, q, k, v, attn_mask=None):
Q, K, V = self._split(self.W_q(q)), self._split(self.W_k(k)), self._split(self.W_v(v))
scores = Q @ K.transpose(–2, –1) / math.sqrt(self.d_k) # [B, h, Lq, Lk]
if attn_mask is not None: # True = 屏蔽
# 用 dtype 的最小值而不是 -1e9：bf16/fp16 下 -1e9 会溢出成 -inf 或失真
scores = scores.masked_fill(attn_mask, torch.finfo(scores.dtype).min)
weights = scores.softmax(dim=–1) # [B, h, Lq, Lk]
out = self.attn_drop(weights) @ V # [B, h, Lq, d_k]
return self.W_o(self._merge(out)), weights

Shape 变化一览（以本章配置 B=64, L=14, d_model=128, h=4 为例）：

两处 PyTorch 特有的写法：

• view vs reshape：_split 里用 view，因为线性层的输出是连续内存；_merge 里必须用 reshape，因为 transpose 之后内存不再连续，view 会报错。
• 屏蔽值用 torch.finfo(dtype).min 而不是 -1e9：本章训练用 bf16 混合精度（20.10 节），不同精度下"足够小的数"不一样，直接取当前 dtype 能表示的最小值最稳妥。

Mask 约定：本章所有 Mask 都是 bool 张量，True = 屏蔽，与 masked_fill 的语义一致。要特别小心的是，PyTorch 内置的 F.scaled_dot_product_attention 的 bool mask 恰好相反（True = 可以看），Hugging Face 的 attention_mask 则是整数 1 = 可以看。三种约定混用是最常见的 bug 来源。

20.4 Feed Forward Network

# ch20_transformer/model.py（续）
class FeedForward(nn.Module):
def __init__(self, d_model, dff, dropout=0.1):
super().__init__()
self.net = nn.Sequential(
nn.Linear(d_model, dff), nn.GELU(),
nn.Linear(dff, d_model), nn.Dropout(dropout),
)

def forward(self, x):
return self.net(x)

dff = 4 × d_model（第8章），激活函数用 GELU。想换成 SwiGLU 可以作为练习（见章末）。

20.5 Encoder

本章采用现代 LLM 通用的 Pre-Norm：x + Sublayer(LayerNorm(x))（第9、10章）。

# ch20_transformer/model.py（续）
class EncoderLayer(nn.Module):
"""Pre-Norm：x + Sublayer(LN(x))"""
def __init__(self, d_model, num_heads, dff, dropout=0.1):
super().__init__()
self.ln1, self.ln2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
self.mha = MultiHeadAttention(d_model, num_heads, dropout)
self.ffn = FeedForward(d_model, dff, dropout)
self.drop = nn.Dropout(dropout)

def forward(self, x, pad_mask):
h = self.ln1(x)
a, w = self.mha(h, h, h, attn_mask=pad_mask)
x = x + self.drop(a)
x = x + self.ffn(self.ln2(x))
return x, w

# ch20_transformer/model.py（续）
class Encoder(nn.Module):
def __init__(self, vocab_size, d_model, num_heads, dff, num_layers, max_len, dropout=0.1):
super().__init__()
self.embed = TokenAndPosition(vocab_size, d_model, max_len, dropout)
self.layers = nn.ModuleList(EncoderLayer(d_model, num_heads, dff, dropout) for _ in range(num_layers))
self.ln_final = nn.LayerNorm(d_model) # Pre-Norm 需要在出口补一个 LN

def forward(self, src, pad_mask):
x = self.embed(src)
for layer in self.layers:
x, _ = layer(x, pad_mask)
return self.ln_final(x)

为什么 Pre-Norm 最后要多一个 ln_final？ Pre-Norm 中，主干 x 一路只做加法，从未被归一化过，层数越多数值越大。所以要在整个 Stack 的出口补一个 LayerNorm，GPT-2（ln_f）、Llama（norm）都有这一层。Post-Norm 每层出口已经归一化过，不需要。

nn.ModuleList 不能换成普通的 Python list：只有注册进 ModuleList 的子模块，参数才会出现在 model.parameters() 里、才会被优化器更新、才会跟着 .to(device) 搬到 GPU。

20.6 Decoder

Decoder Layer 比 Encoder Layer 多一个 Cross-Attention：Q 来自 Decoder 自己，K/V 来自 Encoder 的输出 memory（第13章）。

# ch20_transformer/model.py（续）
class DecoderLayer(nn.Module):
def __init__(self, d_model, num_heads, dff, dropout=0.1):
super().__init__()
self.ln1, self.ln2, self.ln3 = nn.LayerNorm(d_model), nn.LayerNorm(d_model), nn.LayerNorm(d_model)
self.self_mha = MultiHeadAttention(d_model, num_heads, dropout)
self.cross_mha = MultiHeadAttention(d_model, num_heads, dropout)
self.ffn = FeedForward(d_model, dff, dropout)
self.drop1, self.drop2 = nn.Dropout(dropout), nn.Dropout(dropout)

def forward(self, x, memory, self_mask, cross_mask):
# ① Masked Self-Attention：只能看自己和之前的 Token
h = self.ln1(x)
a, _ = self.self_mha(h, h, h, attn_mask=self_mask)
x = x + self.drop1(a)
# ② Cross-Attention：Q 来自 Decoder，K/V 来自 Encoder 输出
h = self.ln2(x)
a, cross_w = self.cross_mha(h, memory, memory, attn_mask=cross_mask)
x = x + self.drop2(a)
# ③ FFN
x = x + self.ffn(self.ln3(x))
return x, cross_w

# ch20_transformer/model.py（续）
class Decoder(nn.Module):
def __init__(self, vocab_size, d_model, num_heads, dff, num_layers, max_len, dropout=0.1):
super().__init__()
self.embed = TokenAndPosition(vocab_size, d_model, max_len, dropout)
self.layers = nn.ModuleList(DecoderLayer(d_model, num_heads, dff, dropout) for _ in range(num_layers))
self.ln_final = nn.LayerNorm(d_model)

def forward(self, tgt, memory, self_mask, cross_mask):
x = self.embed(tgt)
cross_weights = []
for layer in self.layers:
x, w = layer(x, memory, self_mask, cross_mask)
cross_weights.append(w)
return self.ln_final(x), cross_weights

Decoder 收集了每一层的 Cross-Attention 权重，后面可视化时直接取用。

20.7 三种 Mask 与完整模型

Encoder-Decoder 里一共有三处 Attention，各自需要的 Mask 不一样：

#mermaid-svg-1DQiVraax83oQ4oV{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-1DQiVraax83oQ4oV .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-1DQiVraax83oQ4oV .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-1DQiVraax83oQ4oV .error-icon{fill:#552222;}#mermaid-svg-1DQiVraax83oQ4oV .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-1DQiVraax83oQ4oV .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-1DQiVraax83oQ4oV .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-1DQiVraax83oQ4oV .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-1DQiVraax83oQ4oV .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-1DQiVraax83oQ4oV .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-1DQiVraax83oQ4oV .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-1DQiVraax83oQ4oV .marker{fill:#333333;stroke:#333333;}#mermaid-svg-1DQiVraax83oQ4oV .marker.cross{stroke:#333333;}#mermaid-svg-1DQiVraax83oQ4oV svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-1DQiVraax83oQ4oV p{margin:0;}#mermaid-svg-1DQiVraax83oQ4oV .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-1DQiVraax83oQ4oV .cluster-label text{fill:#333;}#mermaid-svg-1DQiVraax83oQ4oV .cluster-label span{color:#333;}#mermaid-svg-1DQiVraax83oQ4oV .cluster-label span p{background-color:transparent;}#mermaid-svg-1DQiVraax83oQ4oV .label text,#mermaid-svg-1DQiVraax83oQ4oV span{fill:#333;color:#333;}#mermaid-svg-1DQiVraax83oQ4oV .node rect,#mermaid-svg-1DQiVraax83oQ4oV .node circle,#mermaid-svg-1DQiVraax83oQ4oV .node ellipse,#mermaid-svg-1DQiVraax83oQ4oV .node polygon,#mermaid-svg-1DQiVraax83oQ4oV .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-1DQiVraax83oQ4oV .rough-node .label text,#mermaid-svg-1DQiVraax83oQ4oV .node .label text,#mermaid-svg-1DQiVraax83oQ4oV .image-shape .label,#mermaid-svg-1DQiVraax83oQ4oV .icon-shape .label{text-anchor:middle;}#mermaid-svg-1DQiVraax83oQ4oV .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-1DQiVraax83oQ4oV .rough-node .label,#mermaid-svg-1DQiVraax83oQ4oV .node .label,#mermaid-svg-1DQiVraax83oQ4oV .image-shape .label,#mermaid-svg-1DQiVraax83oQ4oV .icon-shape .label{text-align:center;}#mermaid-svg-1DQiVraax83oQ4oV .node.clickable{cursor:pointer;}#mermaid-svg-1DQiVraax83oQ4oV .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-1DQiVraax83oQ4oV .arrowheadPath{fill:#333333;}#mermaid-svg-1DQiVraax83oQ4oV .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-1DQiVraax83oQ4oV .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-1DQiVraax83oQ4oV .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-1DQiVraax83oQ4oV .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-1DQiVraax83oQ4oV .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-1DQiVraax83oQ4oV .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-1DQiVraax83oQ4oV .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-1DQiVraax83oQ4oV .cluster text{fill:#333;}#mermaid-svg-1DQiVraax83oQ4oV .cluster span{color:#333;}#mermaid-svg-1DQiVraax83oQ4oV div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-1DQiVraax83oQ4oV .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-1DQiVraax83oQ4oV rect.text{fill:none;stroke-width:0;}#mermaid-svg-1DQiVraax83oQ4oV .icon-shape,#mermaid-svg-1DQiVraax83oQ4oV .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-1DQiVraax83oQ4oV .icon-shape p,#mermaid-svg-1DQiVraax83oQ4oV .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-1DQiVraax83oQ4oV .icon-shape .label rect,#mermaid-svg-1DQiVraax83oQ4oV .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-1DQiVraax83oQ4oV .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-1DQiVraax83oQ4oV .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-1DQiVraax83oQ4oV :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

Cross-Attn

padding_mask(src)和 Encoder 复用

Decoder Self-Attn

causal_mask(Lt)

按位或 OR

padding_mask(tgt_in)

Encoder Self-Attn

padding_mask(src)

Shape 里的 1 都靠广播展开：[B, 1, 1, Ls] 广播到 [B, h, Lq, Ls]，意思是"所有 Head、所有 Query 都不许看这几个 PAD Key"。

# ch20_transformer/model.py（续）
def make_padding_mask(ids):
"""[B, L] -> [B, 1, 1, L]，PAD 位置为 True"""
return (ids == PAD_ID)[:, None, None, :]

def make_causal_mask(L, device):
"""[1, 1, L, L]，上三角（未来位置）为 True"""
return torch.triu(torch.ones(L, L, dtype=torch.bool, device=device), diagonal=1)[None, None]

class Transformer(nn.Module):
def __init__(self, src_vocab, tgt_vocab, d_model=128, num_heads=4,
dff=512, num_layers=2, max_len=64, dropout=0.1):
super().__init__()
self.encoder = Encoder(src_vocab, d_model, num_heads, dff, num_layers, max_len, dropout)
self.decoder = Decoder(tgt_vocab, d_model, num_heads, dff, num_layers, max_len, dropout)
self.out_proj = nn.Linear(d_model, tgt_vocab)
# 线性层用 Xavier 初始化（nn.Transformer、The Annotated Transformer、Keras Dense 默认都是 Glorot/Xavier）；
# PyTorch nn.Linear 默认的 Kaiming-uniform(a=√5) 方差更小
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
nn.init.zeros_(m.bias)

def encode(self, src):
src_mask = make_padding_mask(src)
return self.encoder(src, src_mask), src_mask

def decode(self, tgt_in, memory, src_mask):
L = tgt_in.size(1)
# Decoder 自注意力：未来位置 或 目标端 PAD，任一成立就屏蔽
self_mask = make_causal_mask(L, tgt_in.device) | make_padding_mask(tgt_in)
h, cross_w = self.decoder(tgt_in, memory, self_mask, src_mask)
return self.out_proj(h), cross_w

def forward(self, src, tgt_in):
memory, src_mask = self.encode(src)
logits, _ = self.decode(tgt_in, memory, src_mask)
return logits # [B, Lt, tgt_vocab]

几点说明：

• 两个 bool Mask 用 | 合并：任一条件成立就屏蔽。
• make_causal_mask 要传 device：Mask 在 forward 里临时创建，必须和数据在同一块 GPU 上，否则会报 Expected all tensors to be on the same device。
• Xavier 初始化：PyTorch nn.Linear 默认是 Kaiming-uniform（$a=\sqrt{5}$），方差比 Xavier 小。nn.Transformer 内部、The Annotated Transformer 都会改用 Xavier，本章也一样。
• encode 和 decode 拆开：推理时 Encoder 只需要跑一次，Decoder 要循环跑很多次（20.11 节），拆开才能复用 memory。

20.8 先测再训：两个单元测试

Mask 写错时模型照样能跑、loss 照样下降，只是学到了"作弊"或"噪声"。所以训练前先用两个性质测一下：

# ch20_transformer/test_model.py
"""训练前先验证 Mask：因果性 + Padding 不变性"""
import torch

from model import Transformer

torch.manual_seed(0)
m = Transformer(13, 13, max_len=14, dropout=0.0).eval() # 关掉 Dropout，两次前向才可比

src = torch.tensor([[3, 4, 5, 2, 0, 0]])
tgt = torch.tensor([[1, 5, 4, 3, 2]])
with torch.no_grad():
a = m(src, tgt)

# 1) 因果性：改动第 3 个位置之后的 Token，前 3 个位置输出必须不变
b = m(src, torch.tensor([[1, 5, 4, 9, 9]]))
print("causal ok:", torch.allclose(a[:, :3], b[:, :3], atol=1e-5),
"后面变了:", not torch.allclose(a[:, 3:], b[:, 3:]))

# 2) Padding：源端多补几个 PAD，输出必须不变
c = m(torch.tensor([[3, 4, 5, 2, 0, 0, 0, 0, 0]]), tgt)
print("padding ok:", torch.allclose(a, c, atol=1e-5))

causal ok: True 后面变了: True
padding ok: True

注意 .eval()：测试时要关掉 Dropout，否则两次前向本来就不一样。把 make_causal_mask 里的 diagonal=1 故意改成 diagonal=0（连自己都看不到）或去掉 triu，第一个测试会立刻失败——建议亲手试一次。

20.9 数据：序列翻转任务

词表只有 13 个 Token：PAD=0, BOS=1, EOS=2，数字 0~9 映射到 3~12。

# ch20_transformer/train.py
"""第20章：序列翻转任务 —— 数据、训练、贪心解码

运行：python train.py
"""
import time

import numpy as np
import torch
import torch.nn.functional as F

from model import PAD_ID, Transformer

# —————- 设备与精度 —————-
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
USE_BF16 = DEVICE == "cuda" and torch.cuda.is_bf16_supported() # RTX 4090 支持 bf16
torch.backends.cuda.matmul.allow_tf32 = True # Ampere 及以上：FP32 矩阵乘走 TF32
torch.backends.cudnn.allow_tf32 = True

# —————- 数据：序列翻转任务 —————-
BOS_ID, EOS_ID = 1, 2
NUM_OFFSET = 3 # 数字 0~9 映射为 id 3~12
VOCAB = NUM_OFFSET + 10 # 13
MIN_LEN, MAX_LEN = 3, 12
SEQ_LEN = MAX_LEN + 2 # 加上 BOS / EOS

def make_dataset(n, seed):
rng = np.random.default_rng(seed)
src = np.zeros((n, SEQ_LEN), dtype=np.int64)
tgt = np.zeros((n, SEQ_LEN), dtype=np.int64)
for i in range(n):
L = rng.integers(MIN_LEN, MAX_LEN + 1)
digits = rng.integers(0, 10, L) + NUM_OFFSET
src[i, :L] = digits
src[i, L] = EOS_ID
tgt[i, 0] = BOS_ID
tgt[i, 1:L + 1] = digits[::–1] # 目标 = 源序列翻转
tgt[i, L + 1] = EOS_ID
return torch.from_numpy(src), torch.from_numpy(tgt)

一条样本（数字 7 3 0 9 5）在内存里的样子：

src : [10 6 3 12 8 2 0 0 0 0 0 0 0 0] ← 7 3 0 9 5 <eos> <pad>…
tgt : [ 1 8 12 3 6 10 2 0 0 0 0 0 0 0] ← <bos> 5 9 0 3 7 <eos> <pad>…

文件开头的两行 allow_tf32 = True：Ampere 及之后的 GPU（包括 4090）可以用 TF32 格式计算 FP32 矩阵乘法，精度损失很小、速度快数倍。

20.10 训练

20.10.1 Teacher Forcing：输入和标签错开一位

训练时 Decoder 的输入是正确答案去掉最后一位，标签是正确答案去掉第一位：

tgt : <bos> 5 9 0 3 7 <eos>
tgt_in : <bos> 5 9 0 3 7 ← 喂给 Decoder
tgt_out : 5 9 0 3 7 <eos> ← 每个位置要预测的下一个 Token

每个位置都"看着正确的前文、预测下一个"，再配合 Causal Mask，一次前向就能并行算出所有位置的 loss（第14章 14.4.2）。

20.10.2 Warmup 学习率

$$lr=d_{model}^{-0.5}\cdot \min \left(ste{p}^{-0.5},step\cdot warmu{p}^{-1.5}\right)$$

前 warmup 步线性升高，之后按 $1/\sqrt{step}$ 衰减。训练初期参数是随机的，梯度方向不可靠，Adam 的二阶矩估计也还不准，小学习率起步更稳。

#mermaid-svg-yVVHZkAsMRxEq70K{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-yVVHZkAsMRxEq70K .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-yVVHZkAsMRxEq70K .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-yVVHZkAsMRxEq70K .error-icon{fill:#552222;}#mermaid-svg-yVVHZkAsMRxEq70K .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-yVVHZkAsMRxEq70K .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-yVVHZkAsMRxEq70K .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-yVVHZkAsMRxEq70K .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-yVVHZkAsMRxEq70K .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-yVVHZkAsMRxEq70K .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-yVVHZkAsMRxEq70K .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-yVVHZkAsMRxEq70K .marker{fill:#333333;stroke:#333333;}#mermaid-svg-yVVHZkAsMRxEq70K .marker.cross{stroke:#333333;}#mermaid-svg-yVVHZkAsMRxEq70K svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-yVVHZkAsMRxEq70K p{margin:0;}#mermaid-svg-yVVHZkAsMRxEq70K .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-yVVHZkAsMRxEq70K .cluster-label text{fill:#333;}#mermaid-svg-yVVHZkAsMRxEq70K .cluster-label span{color:#333;}#mermaid-svg-yVVHZkAsMRxEq70K .cluster-label span p{background-color:transparent;}#mermaid-svg-yVVHZkAsMRxEq70K .label text,#mermaid-svg-yVVHZkAsMRxEq70K span{fill:#333;color:#333;}#mermaid-svg-yVVHZkAsMRxEq70K .node rect,#mermaid-svg-yVVHZkAsMRxEq70K .node circle,#mermaid-svg-yVVHZkAsMRxEq70K .node ellipse,#mermaid-svg-yVVHZkAsMRxEq70K .node polygon,#mermaid-svg-yVVHZkAsMRxEq70K .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-yVVHZkAsMRxEq70K .rough-node .label text,#mermaid-svg-yVVHZkAsMRxEq70K .node .label text,#mermaid-svg-yVVHZkAsMRxEq70K .image-shape .label,#mermaid-svg-yVVHZkAsMRxEq70K .icon-shape .label{text-anchor:middle;}#mermaid-svg-yVVHZkAsMRxEq70K .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-yVVHZkAsMRxEq70K .rough-node .label,#mermaid-svg-yVVHZkAsMRxEq70K .node .label,#mermaid-svg-yVVHZkAsMRxEq70K .image-shape .label,#mermaid-svg-yVVHZkAsMRxEq70K .icon-shape .label{text-align:center;}#mermaid-svg-yVVHZkAsMRxEq70K .node.clickable{cursor:pointer;}#mermaid-svg-yVVHZkAsMRxEq70K .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-yVVHZkAsMRxEq70K .arrowheadPath{fill:#333333;}#mermaid-svg-yVVHZkAsMRxEq70K .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-yVVHZkAsMRxEq70K .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-yVVHZkAsMRxEq70K .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-yVVHZkAsMRxEq70K .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-yVVHZkAsMRxEq70K .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-yVVHZkAsMRxEq70K .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-yVVHZkAsMRxEq70K .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-yVVHZkAsMRxEq70K .cluster text{fill:#333;}#mermaid-svg-yVVHZkAsMRxEq70K .cluster span{color:#333;}#mermaid-svg-yVVHZkAsMRxEq70K div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-yVVHZkAsMRxEq70K .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-yVVHZkAsMRxEq70K rect.text{fill:none;stroke-width:0;}#mermaid-svg-yVVHZkAsMRxEq70K .icon-shape,#mermaid-svg-yVVHZkAsMRxEq70K .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-yVVHZkAsMRxEq70K .icon-shape p,#mermaid-svg-yVVHZkAsMRxEq70K .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-yVVHZkAsMRxEq70K .icon-shape .label rect,#mermaid-svg-yVVHZkAsMRxEq70K .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-yVVHZkAsMRxEq70K .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-yVVHZkAsMRxEq70K .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-yVVHZkAsMRxEq70K :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

step 0lr≈0

Warmup线性升高

step=warmuplr 峰值

之后∝ 1/√step 衰减

PyTorch 里用 LambdaLR 实现：它把学习率设为 base_lr × f(step)，所以优化器的 lr 设成 1.0，f 直接返回上面公式的值。

20.10.3 训练代码

# ch20_transformer/train.py（续）
def noam_lambda(d_model, warmup):
"""lr = d^-0.5 · min(step^-0.5, step · warmup^-1.5)，配合 base lr = 1.0 使用"""
def f(step):
step = step + 1 # LambdaLR 从 step=0 开始，避免除 0
return d_model ** –0.5 * min(step ** –0.5, step * warmup ** –1.5)
return f

def main():
torch.manual_seed(0)
src_train, tgt_train = make_dataset(20000, seed=0)
src_val, tgt_val = make_dataset(1000, seed=1)
loader = torch.utils.data.DataLoader(
torch.utils.data.TensorDataset(src_train, tgt_train),
batch_size=64, shuffle=True, drop_last=True)

D_MODEL, EPOCHS = 128, 10
model = Transformer(VOCAB, VOCAB, d_model=D_MODEL, num_heads=4,
dff=512, num_layers=2, max_len=SEQ_LEN).to(DEVICE)
print(f"device={DEVICE} bf16={USE_BF16} 参数量={sum(p.numel() for p in model.parameters()):,}")

optimizer = torch.optim.Adam(model.parameters(), lr=1.0, betas=(0.9, 0.98), eps=1e-9)
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, noam_lambda(D_MODEL, warmup=1000))

t0 = time.time()
for epoch in range(1, EPOCHS + 1):
model.train()
for src, tgt in loader:
src, tgt = src.to(DEVICE, non_blocking=True), tgt.to(DEVICE, non_blocking=True)
tgt_in, tgt_out = tgt[:, :–1], tgt[:, 1:] # Teacher Forcing：右移一位
with torch.autocast(device_type=DEVICE, dtype=torch.bfloat16, enabled=USE_BF16):
logits = model(src, tgt_in)
# ignore_index=PAD：PAD 位置不计入 loss（等价于 TF 版的 masked_loss）
loss = F.cross_entropy(logits.reshape(–1, VOCAB).float(),
tgt_out.reshape(–1), ignore_index=PAD_ID)
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
scheduler.step()

with torch.no_grad(): # 最后一个 batch 的逐 Token 准确率
mask = tgt_out != PAD_ID
acc = ((logits.argmax(–1) == tgt_out) & mask).sum().item() / mask.sum().item()
em = exact_match(model, src_val[:500], tgt_val[:500])
print(f"epoch {epoch:2d} loss={loss.item():.4f} token_acc={acc:.4f} "
f"val_exact_match={em:.3f} lr={scheduler.get_last_lr()[0]:.2e} {time.time() – t0:.0f}s")

# 单条推理演示
demo = [7, 3, 0, 9, 5]
s = torch.zeros(1, SEQ_LEN, dtype=torch.long)
s[0, :len(demo)] = torch.tensor(demo) + NUM_OFFSET
s[0, len(demo)] = EOS_ID
p = greedy_decode(model, s.to(DEVICE))[0].tolist()
print("输入:", demo, "-> 输出:", [i – NUM_OFFSET for i in p[1:] if i >= NUM_OFFSET])
torch.save(model.state_dict(), "reverse_transformer.pt")

if __name__ == "__main__":
main()

与 TensorFlow 写法相比的几个要点：

20.11 推理：贪心解码

推理时没有正确答案可以"喂"，只能自己生成一个、拼回去、再生成下一个：

#mermaid-svg-Q29LWbEJJsraQtIO{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-Q29LWbEJJsraQtIO .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-Q29LWbEJJsraQtIO .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-Q29LWbEJJsraQtIO .error-icon{fill:#552222;}#mermaid-svg-Q29LWbEJJsraQtIO .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-Q29LWbEJJsraQtIO .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-Q29LWbEJJsraQtIO .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-Q29LWbEJJsraQtIO .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-Q29LWbEJJsraQtIO .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-Q29LWbEJJsraQtIO .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-Q29LWbEJJsraQtIO .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-Q29LWbEJJsraQtIO .marker{fill:#333333;stroke:#333333;}#mermaid-svg-Q29LWbEJJsraQtIO .marker.cross{stroke:#333333;}#mermaid-svg-Q29LWbEJJsraQtIO svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-Q29LWbEJJsraQtIO p{margin:0;}#mermaid-svg-Q29LWbEJJsraQtIO .label{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;color:#333;}#mermaid-svg-Q29LWbEJJsraQtIO .cluster-label text{fill:#333;}#mermaid-svg-Q29LWbEJJsraQtIO .cluster-label span{color:#333;}#mermaid-svg-Q29LWbEJJsraQtIO .cluster-label span p{background-color:transparent;}#mermaid-svg-Q29LWbEJJsraQtIO .label text,#mermaid-svg-Q29LWbEJJsraQtIO span{fill:#333;color:#333;}#mermaid-svg-Q29LWbEJJsraQtIO .node rect,#mermaid-svg-Q29LWbEJJsraQtIO .node circle,#mermaid-svg-Q29LWbEJJsraQtIO .node ellipse,#mermaid-svg-Q29LWbEJJsraQtIO .node polygon,#mermaid-svg-Q29LWbEJJsraQtIO .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-Q29LWbEJJsraQtIO .rough-node .label text,#mermaid-svg-Q29LWbEJJsraQtIO .node .label text,#mermaid-svg-Q29LWbEJJsraQtIO .image-shape .label,#mermaid-svg-Q29LWbEJJsraQtIO .icon-shape .label{text-anchor:middle;}#mermaid-svg-Q29LWbEJJsraQtIO .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-Q29LWbEJJsraQtIO .rough-node .label,#mermaid-svg-Q29LWbEJJsraQtIO .node .label,#mermaid-svg-Q29LWbEJJsraQtIO .image-shape .label,#mermaid-svg-Q29LWbEJJsraQtIO .icon-shape .label{text-align:center;}#mermaid-svg-Q29LWbEJJsraQtIO .node.clickable{cursor:pointer;}#mermaid-svg-Q29LWbEJJsraQtIO .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-Q29LWbEJJsraQtIO .arrowheadPath{fill:#333333;}#mermaid-svg-Q29LWbEJJsraQtIO .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-Q29LWbEJJsraQtIO .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-Q29LWbEJJsraQtIO .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-Q29LWbEJJsraQtIO .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-Q29LWbEJJsraQtIO .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-Q29LWbEJJsraQtIO .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-Q29LWbEJJsraQtIO .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-Q29LWbEJJsraQtIO .cluster text{fill:#333;}#mermaid-svg-Q29LWbEJJsraQtIO .cluster span{color:#333;}#mermaid-svg-Q29LWbEJJsraQtIO div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-Q29LWbEJJsraQtIO .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-Q29LWbEJJsraQtIO rect.text{fill:none;stroke-width:0;}#mermaid-svg-Q29LWbEJJsraQtIO .icon-shape,#mermaid-svg-Q29LWbEJJsraQtIO .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-Q29LWbEJJsraQtIO .icon-shape p,#mermaid-svg-Q29LWbEJJsraQtIO .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-Q29LWbEJJsraQtIO .icon-shape .label rect,#mermaid-svg-Q29LWbEJJsraQtIO .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-Q29LWbEJJsraQtIO .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-Q29LWbEJJsraQtIO .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-Q29LWbEJJsraQtIO :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

No

Yes

src

Encoder 只跑一次得到 memory

out = [BOS]

Decoder(out, memory)

取最后一个位置的 argmax

全部样本都输出了 EOS?

out = out + [next]

结束

# ch20_transformer/train.py（续）
# —————- 推理：贪心解码 —————-
@torch.no_grad()
def greedy_decode(model, src, max_len=SEQ_LEN):
"""src: [B, Ls] -> [B, <=max_len]"""
model.eval()
memory, src_mask = model.encode(src)
B = src.size(0)
out = torch.full((B, 1), BOS_ID, dtype=torch.long, device=src.device)
finished = torch.zeros(B, dtype=torch.bool, device=src.device)
for _ in range(max_len – 1):
logits, _ = model.decode(out, memory, src_mask)
nxt = logits[:, –1].argmax(–1)
nxt = nxt.masked_fill(finished, PAD_ID) # 已结束的样本补 PAD
out = torch.cat([out, nxt[:, None]], dim=1)
finished |= nxt == EOS_ID
if finished.all():
break
return out

def exact_match(model, src, tgt):
pred = greedy_decode(model, src.to(DEVICE)).cpu().tolist()
ok = 0
for p, t in zip(pred, tgt.tolist()):
t = t[1:t.index(EOS_ID) + 1] # 去掉 BOS，保留到 EOS
ok += p[1:len(t) + 1] == t
return ok / len(src)

@torch.no_grad() 让这段代码不记录计算图：推理不需要反向传播，省显存也更快。

这里每生成一个 Token 都把整个 out 重新送进 Decoder，前面位置的 K/V 被反复计算。这正是第15章 KV Cache 要解决的问题，下一章会在 GPT 上亲手实现它。

20.12 运行结果：一次"代码全对、模型却学不会"的排查

python test_model.py # 先过单元测试
python train.py

20.12.1 第一次运行：慢得反常

笔者第一版代码没有 20.2 节那几行 Embedding 初始化，也没有 20.7 节的 Xavier 初始化（用的是 PyTorch 默认值）。单元测试全部通过，训练结果却是这样：

device=cuda bf16=True 参数量=931,213
epoch 1 loss=1.6167 token_acc=0.3487 val_exact_match=0.020 lr=8.75e-04 2s
epoch 2 loss=1.3009 token_acc=0.4825 val_exact_match=0.046 lr=1.75e-03 4s
epoch 3 loss=1.2627 token_acc=0.4906 val_exact_match=0.040 lr=2.62e-03 5s
epoch 4 loss=1.0233 token_acc=0.5657 val_exact_match=0.118 lr=2.50e-03 7s
epoch 5 loss=0.9464 token_acc=0.5835 val_exact_match=0.166 lr=2.24e-03 9s
epoch 6 loss=0.8655 token_acc=0.6466 val_exact_match=0.186 lr=2.04e-03 10s
epoch 7 loss=0.6120 token_acc=0.7418 val_exact_match=0.370 lr=1.89e-03 12s
epoch 8 loss=0.5351 token_acc=0.8015 val_exact_match=0.482 lr=1.77e-03 14s
epoch 9 loss=0.4389 token_acc=0.8445 val_exact_match=0.660 lr=1.67e-03 15s
epoch 10 loss=0.3849 token_acc=0.8509 val_exact_match=0.604 lr=1.58e-03 17s

模型在学，但学得很慢：10 个 epoch 后整句正确率只有 0.60。同样的结构和超参数，用 TensorFlow 实现时 5 个 epoch 就到了 0.97。

排查思路：

翻转任务完全依赖位置信息，Token 信号比位置信号大 10 倍，模型就很难分清"这个 7 在第几位"。20.13 节的注意力矩阵也印证了这一点：第一版模型的 Cross-Attention 每行都分散在 2~3 个位置上。

20.12.2 修复后

加上 20.2 节和 20.7 节的初始化，其他一行不改：

device=cuda bf16=True 参数量=931,213
epoch 1 loss=0.5494 token_acc=0.7948 val_exact_match=0.386 lr=8.75e-04 2s
epoch 2 loss=0.1301 token_acc=0.9545 val_exact_match=0.876 lr=1.75e-03 4s
epoch 3 loss=0.1097 token_acc=0.9639 val_exact_match=0.872 lr=2.62e-03 5s
epoch 4 loss=0.0552 token_acc=0.9841 val_exact_match=0.946 lr=2.50e-03 7s
epoch 5 loss=0.0449 token_acc=0.9819 val_exact_match=0.964 lr=2.24e-03 9s
epoch 6 loss=0.0376 token_acc=0.9893 val_exact_match=0.992 lr=2.04e-03 10s
epoch 7 loss=0.0299 token_acc=0.9926 val_exact_match=0.988 lr=1.89e-03 12s
epoch 8 loss=0.0185 token_acc=0.9904 val_exact_match=0.996 lr=1.77e-03 14s
epoch 9 loss=0.0300 token_acc=0.9910 val_exact_match=0.996 lr=1.67e-03 15s
epoch 10 loss=0.0211 token_acc=0.9907 val_exact_match=0.996 lr=1.58e-03 17s
输入: [7, 3, 0, 9, 5] -> 输出: [5, 9, 0, 3, 7]

读数要点：

• token_acc 是 Teacher Forcing 下的逐 Token 准确率（最后一个 batch），val_exact_match 是自回归生成时整句完全正确的比例，后者更严格：一句 12 个数字错 1 个就算错。
• epoch 2→3 的整句正确率略降（0.876 → 0.872）：1 个 epoch ≈ 312 步，warmup=1000，第 3 个 epoch 结束时学习率才升到峰值 2.6e-3，这时出现小幅波动是正常的。
• 模型约 93 万参数，4090 上 10 个 epoch 共 17 秒。

教训：跨框架移植模型时，“代码逻辑一样"不等于"模型一样”——默认初始化、默认 epsilon、默认激活函数近似方式都可能不同。单元测试能保证逻辑正确，但保证不了训练动态一致。

20.13 实验：可视化 Cross-Attention

模型真的学会了"输出第 i 个数时去看输入的倒数第 i 个"吗？直接看 Cross-Attention 权重。

# ch20_transformer/viz.py
"""可视化 Cross-Attention：模型是否学会了"倒着对齐"？

先运行 python train.py 生成 reverse_transformer.pt，再运行 python viz.py
"""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch

from model import Transformer
from train import DEVICE, EOS_ID, NUM_OFFSET, SEQ_LEN, VOCAB, greedy_decode

model = Transformer(VOCAB, VOCAB, d_model=128, num_heads=4, dff=512,
num_layers=2, max_len=SEQ_LEN).to(DEVICE)
model.load_state_dict(torch.load("reverse_transformer.pt", map_location=DEVICE))
model.eval()

demo = [7, 3, 0, 9, 5, 1, 8]
src = torch.zeros(1, SEQ_LEN, dtype=torch.long)
src[0, :len(demo)] = torch.tensor(demo) + NUM_OFFSET
src[0, len(demo)] = EOS_ID
src = src.to(DEVICE)

with torch.no_grad():
pred = greedy_decode(model, src)
memory, src_mask = model.encode(src)
_, cross_w = model.decode(pred[:, :–1], memory, src_mask)

w = cross_w[–1][0].mean(dim=0).float().cpu() # 最后一层，所有 Head 平均 [Lt, Ls]
n = len(demo) + 1
w = w[:n, :n]
torch.set_printoptions(precision=2, sci_mode=False)
print(w)

xl = [str(d) for d in demo] + ["<eos>"]
yl = ["<bos>"] + [str(d) for d in demo[::–1]]
plt.imshow(w.numpy(), cmap="Blues")
plt.xticks(range(len(xl)), xl); plt.yticks(range(len(yl)), yl)
plt.xlabel("Source (Encoder)"); plt.ylabel("Decoder input"); plt.colorbar()
plt.title("Cross-Attention (last layer, mean of heads)")
plt.savefig("cross_attn.png", dpi=120, bbox_inches="tight")
print("已保存 cross_attn.png")

实测输出（行 = Decoder 当前输入，列 = 源序列 7 3 0 9 5 1 8 <eos>；行列标签为便于阅读手工添加）：

7 3 0 9 5 1 8 <eos>
<bos> [[0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 1.00, 0.00],
8 [0.01, 0.01, 0.00, 0.00, 0.00, 0.97, 0.00, 0.00],
1 [0.00, 0.00, 0.00, 0.00, 1.00, 0.00, 0.00, 0.00],
5 [0.00, 0.00, 0.00, 0.96, 0.04, 0.00, 0.00, 0.00],
9 [0.00, 0.00, 0.99, 0.01, 0.00, 0.00, 0.00, 0.00],
0 [0.00, 0.94, 0.02, 0.00, 0.00, 0.00, 0.04, 0.00],
3 [0.99, 0.01, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00],
7 [0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 1.00]]

一条清晰的反对角线：

• 第 1 行，Decoder 输入是 <bos>，要预测第一个输出 8，注意力 100% 落在源序列的 8 上；
• 第 2 行，输入是 8，要预测 1，注意力落在源序列的 1 上；
• ……依次往左移，每行的最大权重都在 0.94 以上。
• 最后一行要预测 <eos>，注意力 100% 落在源序列的 <eos> 上——模型学到了"源序列读完了，该结束了"。

没有任何人告诉模型"要倒着对齐"，它是从 loss 里自己学出来的。这就是第5章说的"Attention = 软对齐"。

作为对照，20.12.1 节初始化有问题的模型，同一个输入的注意力矩阵是这样的：

<bos> [[0.00, 0.00, 0.00, 0.00, 0.00, 0.12, 0.88, 0.00],
8 [0.00, 0.00, 0.00, 0.00, 0.18, 0.66, 0.15, 0.00],
1 [0.00, 0.00, 0.03, 0.09, 0.77, 0.06, 0.05, 0.00],
5 [0.00, 0.01, 0.29, 0.63, 0.06, 0.02, 0.01, 0.00],
9 [0.02, 0.28, 0.57, 0.08, 0.01, 0.00, 0.00, 0.03],
0 [0.21, 0.51, 0.03, 0.01, 0.00, 0.00, 0.00, 0.24],
3 [0.55, 0.04, 0.00, 0.00, 0.00, 0.00, 0.00, 0.40],
7 [0.36, 0.11, 0.00, 0.00, 0.00, 0.00, 0.00, 0.53]]

方向是对的，但每行都"糊"在相邻的 2~3 个位置上——模型大致知道该往哪看，却分不清相邻的两个位置。这正是位置信号被淹没的表现。

20.14 实验：手写 Attention vs SDPA（FlashAttention）

本章手写 Attention 是为了拿到注意力权重。实际工程中，PyTorch 2.x 提供了 F.scaled_dot_product_attention（SDPA），会根据硬件自动选择 FlashAttention、Memory-Efficient Attention 等高效内核（第15章 15.6 节）。两者在数学上完全等价，下面验证这一点，并比较速度和显存：

# ch20_transformer/mha_vs_sdpa.py
"""对照：同一个 Multi-Head Attention，手写版 vs PyTorch 内置 SDPA

运行：python mha_vs_sdpa.py
（第20章用手写版是为了拿到注意力权重；第21章换成 SDPA 是为了速度）
"""
import time

import torch
import torch.nn.functional as F

from model import MultiHeadAttention, make_causal_mask

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(0)
B, L, D, H = 8, 1024, 512, 8
mha = MultiHeadAttention(D, H, dropout=0.0).to(DEVICE).eval()
x = torch.randn(B, L, D, device=DEVICE)

def sdpa_forward(x):
sp = lambda t: t.view(B, L, H, D // H).transpose(1, 2)
y = F.scaled_dot_product_attention(sp(mha.W_q(x)), sp(mha.W_k(x)), sp(mha.W_v(x)), is_causal=True)
return mha.W_o(y.transpose(1, 2).reshape(B, L, D))

with torch.no_grad():
torch.backends.cuda.matmul.allow_tf32 = False
ours, _ = mha(x, x, x, attn_mask=make_causal_mask(L, DEVICE))
print("手写版与 SDPA 一致:", torch.allclose(ours, sdpa_forward(x), atol=1e-4))

if DEVICE == "cuda":
for name, fn in [("手写版", lambda: mha(x, x, x, attn_mask=make_causal_mask(L, DEVICE))),
("SDPA ", lambda: sdpa_forward(x))]:
with torch.autocast("cuda", dtype=torch.bfloat16):
fn(); torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
t = time.perf_counter()
for _ in range(20):
fn()
torch.cuda.synchronize()
print(f"{name} bf16: {(time.perf_counter() – t) / 20 * 1000:.2f} ms/次, "
f"峰值显存 {torch.cuda.max_memory_allocated() / 2**20:.0f} MB")

实测（RTX 4090，B=8, L=1024, d_model=512, h=8，bf16）：

手写版与 SDPA 一致: True
手写版 bf16: 2.97 ms/次, 峰值显存 712 MB
SDPA bf16: 0.25 ms/次, 峰值显存 79 MB

差距从哪来？手写版必须把 [8, 8, 1024, 1024] 的注意力矩阵（6700 万个元素）完整写进显存，再读出来做 Softmax、再读出来乘 V。FlashAttention 把这个矩阵切成小块放在 GPU 片上的 SRAM 里算完，从不把完整矩阵写回显存（附录A 推导了分块 Softmax 为什么精确等价）。序列越长，差距越大。

所以第21章的 GPT 会直接用 SDPA。

20.15 Hugging Face 源码定位

读完本章代码，可以直接对照 transformers 库的实现（路径相对于 src/transformers/）：

20.16 常见坑

本章练习与实验

本章总结

#mermaid-svg-utoh6k5YdpNDfHrG{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-utoh6k5YdpNDfHrG .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-utoh6k5YdpNDfHrG .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-utoh6k5YdpNDfHrG .error-icon{fill:#552222;}#mermaid-svg-utoh6k5YdpNDfHrG .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-utoh6k5YdpNDfHrG .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-utoh6k5YdpNDfHrG .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-utoh6k5YdpNDfHrG .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-utoh6k5YdpNDfHrG .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-utoh6k5YdpNDfHrG .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-utoh6k5YdpNDfHrG .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-utoh6k5YdpNDfHrG .marker{fill:#333333;stroke:#333333;}#mermaid-svg-utoh6k5YdpNDfHrG .marker.cross{stroke:#333333;}#mermaid-svg-utoh6k5YdpNDfHrG svg{font-family:\”trebuchet ms\”,verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-utoh6k5YdpNDfHrG p{margin:0;}#mermaid-svg-utoh6k5YdpNDfHrG .edge{stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .section–1 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section–1 path,#mermaid-svg-utoh6k5YdpNDfHrG .section–1 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section–1 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section–1 path{fill:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section–1 text{fill:#ffffff;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon–1{font-size:40px;color:#ffffff;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge–1{stroke:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth–1{stroke-width:17;}#mermaid-svg-utoh6k5YdpNDfHrG .section–1 line{stroke:hsl(60, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-0 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-0 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-0 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-0 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-0 path{fill:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-0 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-0{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-0{stroke:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-0{stroke-width:14;}#mermaid-svg-utoh6k5YdpNDfHrG .section-0 line{stroke:hsl(240, 100%, 83.5294117647%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-1 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-1 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-1 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-1 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-1 path{fill:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-1 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-1{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-1{stroke:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-1{stroke-width:11;}#mermaid-svg-utoh6k5YdpNDfHrG .section-1 line{stroke:hsl(260, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-2 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-2 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-2 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-2 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-2 path{fill:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-2 text{fill:#ffffff;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-2{font-size:40px;color:#ffffff;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-2{stroke:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-2{stroke-width:8;}#mermaid-svg-utoh6k5YdpNDfHrG .section-2 line{stroke:hsl(90, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-3 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-3 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-3 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-3 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-3 path{fill:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-3 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-3{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-3{stroke:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-3{stroke-width:5;}#mermaid-svg-utoh6k5YdpNDfHrG .section-3 line{stroke:hsl(120, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-4 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-4 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-4 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-4 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-4 path{fill:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-4 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-4{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-4{stroke:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-4{stroke-width:2;}#mermaid-svg-utoh6k5YdpNDfHrG .section-4 line{stroke:hsl(150, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-5 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-5 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-5 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-5 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-5 path{fill:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-5 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-5{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-5{stroke:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-5{stroke-width:-1;}#mermaid-svg-utoh6k5YdpNDfHrG .section-5 line{stroke:hsl(180, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-6 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-6 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-6 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-6 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-6 path{fill:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-6 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-6{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-6{stroke:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-6{stroke-width:-4;}#mermaid-svg-utoh6k5YdpNDfHrG .section-6 line{stroke:hsl(210, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-7 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-7 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-7 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-7 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-7 path{fill:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-7 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-7{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-7{stroke:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-7{stroke-width:-7;}#mermaid-svg-utoh6k5YdpNDfHrG .section-7 line{stroke:hsl(270, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-8 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-8 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-8 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-8 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-8 path{fill:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-8 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-8{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-8{stroke:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-8{stroke-width:-10;}#mermaid-svg-utoh6k5YdpNDfHrG .section-8 line{stroke:hsl(330, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-9 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-9 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-9 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-9 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-9 path{fill:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-9 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-9{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-9{stroke:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-9{stroke-width:-13;}#mermaid-svg-utoh6k5YdpNDfHrG .section-9 line{stroke:hsl(0, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-10 rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-10 path,#mermaid-svg-utoh6k5YdpNDfHrG .section-10 circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-10 polygon,#mermaid-svg-utoh6k5YdpNDfHrG .section-10 path{fill:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-10 text{fill:black;}#mermaid-svg-utoh6k5YdpNDfHrG .node-icon-10{font-size:40px;color:black;}#mermaid-svg-utoh6k5YdpNDfHrG .section-edge-10{stroke:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .edge-depth-10{stroke-width:-16;}#mermaid-svg-utoh6k5YdpNDfHrG .section-10 line{stroke:hsl(30, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled,#mermaid-svg-utoh6k5YdpNDfHrG .disabled circle,#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:lightgray;}#mermaid-svg-utoh6k5YdpNDfHrG .disabled text{fill:#efefef;}#mermaid-svg-utoh6k5YdpNDfHrG .section-root rect,#mermaid-svg-utoh6k5YdpNDfHrG .section-root path,#mermaid-svg-utoh6k5YdpNDfHrG .section-root circle,#mermaid-svg-utoh6k5YdpNDfHrG .section-root polygon{fill:hsl(240, 100%, 46.2745098039%);}#mermaid-svg-utoh6k5YdpNDfHrG .section-root text{fill:#ffffff;}#mermaid-svg-utoh6k5YdpNDfHrG .section-root span{color:#ffffff;}#mermaid-svg-utoh6k5YdpNDfHrG .section-2 span{color:#ffffff;}#mermaid-svg-utoh6k5YdpNDfHrG .icon-container{height:100%;display:flex;justify-content:center;align-items:center;}#mermaid-svg-utoh6k5YdpNDfHrG .edge{fill:none;}#mermaid-svg-utoh6k5YdpNDfHrG .mindmap-node-label{dy:1em;alignment-baseline:middle;text-anchor:middle;dominant-baseline:middle;text-align:center;}#mermaid-svg-utoh6k5YdpNDfHrG :root{–mermaid-font-family:\”trebuchet ms\”,verdana,arial,sans-serif;}

从零实现 Transformer

组件

位置编码 register_buffer

Multi-Head Attention

FFN

Pre-Norm 与 ln_final

Mask

Encoder 源端 Padding

Decoder 因果 或 目标端 Padding

Cross 源端 Padding

训练

Teacher Forcing

Warmup LambdaLR

ignore_index 忽略 PAD

bf16 autocast

推理

Encoder 跑一次

Decoder 循环生成

验证

单元测试 因果性 与 Padding

Cross-Attention 反对角线

初始化决定收敛速度

一句话总结：Transformer 的代码量不大，真正容易出错的是 Mask、数据错位和初始化；先用单元测试验证性质，再用可解释的小任务（看注意力矩阵）验证效果，比盯着 loss 曲线可靠得多。

本章思考题

下一章预告

下一章去掉 Encoder 和 Cross-Attention，只保留 Decoder，训练一个字符级 GPT：

• 在莎士比亚全集上做 Next Token Prediction，观察过拟合并用最佳 checkpoint 早停
• 用 SDPA（FlashAttention）+ bf16 + torch.compile 在 4090 上 90 秒训完 1000 万参数的模型
• 实现 Temperature / Top-k 采样
• 亲手实现 KV Cache，验证它与朴素生成逐字一致，并找出它真正发挥作用的条件