第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 的每个零件：

要求	用到的组件
输出第 1 个数字时必须"看"输入的最后一个数字	Cross-Attention 必须学会对齐
要知道每个数字在第几位	位置编码不能写错，也不能被淹没
输入长度 3~12 不等	Padding Mask 必须正确
生成时不能偷看后面的答案	Causal Mask 必须正确
有明确的标准答案	可以用"整句完全正确率"量化效果

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

20.1.2 整体架构

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章）的实测环境：

项目	版本
Python	3.12
PyTorch	2.7.0+cu126（CUDA 12.6）
GPU	NVIDIA RTX 4090 24GB

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 要乘 $dmodel\sqrt{d_{model}}$ ：原论文 3.4 节的做法，目的是让 Token 信号和位置编码处在相近的数值范围。
为什么要手动初始化 Embedding：这是从其他框架移植时最容易踩的坑。PyTorch 的 nn.Embedding 默认按 $N(0,1)\mathcal{N}(0, 1)$ 初始化，乘以 $128≈11.3\sqrt{128} ≈ 11.3$ 之后标准差约为 11，而位置编码只在 $[- 1, 1]$ 之间——位置信号只剩不到十分之一，被淹没了。按 $N(0,d−0.5)\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 为例）：

步骤	张量	Shape
输入	`q`	[64, 14, 128]
线性投影	`W_q(q)`	[64, 14, 128]
拆头 `_split`	`Q`	[64, 4, 14, 32]
打分	`Q @ K.transpose(-2, -1)`	[64, 4, 14, 14]
加权求和	`weights @ V`	[64, 4, 14, 32]
合并 `_merge`	—	[64, 14, 128]
输出投影	`W_o(...)`	[64, 14, 128]

两处 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 不一样：

位置	Q 来自	K/V 来自	需要屏蔽	Mask Shape
Encoder Self-Attention	src	src	源端 PAD	[B, 1, 1, Ls]
Decoder Self-Attention	tgt	tgt	未来位置 + 目标端 PAD	[B, 1, Lt, Lt]
Decoder Cross-Attention	tgt	memory	源端 PAD	[B, 1, 1, Ls]

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=5a=\sqrt5$ ），方差比 Xavier 小。nn.Transformer 内部、The Annotated Transformer 都会改用 Xavier，本章也一样。
encode 和 decode 拆开：推理时 Encoder 只需要跑一次，Decoder 要循环跑很多次（20.11 节），拆开才能复用 memory。

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

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

因果性：改动 Decoder 输入中第 3 个位置之后的 Token，前 3 个位置的输出必须完全不变。
Padding 不变性：源序列末尾多补几个 PAD，输出必须完全不变。

# 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 学习率

$d_{model}^{-0.5} \cdot \min\left(step^{-0.5},\ step \cdot warmup^{-1.5}\right)$

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

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 写法相比的几个要点：

要点	做法	原因
忽略 PAD	`F.cross_entropy(..., ignore_index=PAD_ID)`	一行代替手写 Mask 求平均；否则模型只要学会输出 PAD 就能拿到很低的 loss
bf16 混合精度	`torch.autocast(dtype=torch.bfloat16)`	4090 的 bf16 Tensor Core 吞吐远高于 FP32；bf16 的指数位和 FP32 一样，不需要 fp16 的 `GradScaler`
loss 用 FP32 算	`logits.float()`	Softmax + log 对精度敏感，交叉熵放在 FP32 下计算更稳
每步三件套	`zero_grad` → `backward` → `step`	PyTorch 的梯度默认累加，不清零就会和上一步的梯度混在一起
`set_to_none=True`	把梯度设为 None 而不是填 0	少一次显存写入，PyTorch 2.x 的默认行为
训练/评估模式	`model.train()` / `model.eval()`	控制 Dropout 开关；`greedy_decode` 里切到 eval，回到训练循环时要切回 train

20.11 推理：贪心解码

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

# 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。

排查思路：

单元测试通过 → Mask 没问题，排除"偷看答案"或"PAD 没屏蔽"。
loss 在稳定下降 → 学习率、梯度方向没问题。
剩下的差异只可能来自框架默认值 → 逐个对比 TensorFlow 和 PyTorch 各层的默认初始化：

层	Keras 默认	PyTorch 默认	本章的问题
Embedding	均匀分布 ±0.05	$N(0,1)\mathcal{N}(0, 1)$	乘 $128\sqrt{128}$ 后标准差 ≈ 11，位置编码被淹没
Dense / Linear	Glorot（Xavier）均匀	Kaiming 均匀（ $a=5a=\sqrt5$ ）	方差偏小，影响较轻

翻转任务完全依赖位置信息，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]

指标	修复前	修复后
epoch 1 token_acc	0.349	0.795
epoch 5 整句正确率	0.166	0.964
epoch 10 整句正确率	0.604	0.996

读数要点：

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

实现	耗时	峰值显存
手写 Attention	2.97 ms	712 MB
SDPA（FlashAttention）	0.25 ms	79 MB
差距	快约 12 倍	约 1/9

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

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

20.15 Hugging Face 源码定位

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

本章组件	Hugging Face 对应位置	差异
`MultiHeadAttention`	`models/bart/modeling_bart.py` → `BartAttention`	支持 KV Cache，可切换 eager / SDPA / FlashAttention-2 后端
`EncoderLayer`	`models/bart/modeling_bart.py` → `BartEncoderLayer`	BART 是 Post-Norm
`DecoderLayer`	`models/bart/modeling_bart.py` → `BartDecoderLayer`	同上，含 `encoder_attn`（Cross-Attention）
正弦位置编码	`models/marian/modeling_marian.py` → `MarianSinusoidalPositionalEmbedding`	BART 本身用可学习位置编码
`greedy_decode`	`generation/utils.py` → `GenerationMixin.generate`	统一支持 Greedy / Beam / Sampling
PyTorch 官方实现	`torch/nn/modules/transformer.py` → `nn.Transformer`	`norm_first=True` 即 Pre-Norm；`_reset_parameters` 里用 Xavier 初始化

20.16 常见坑

现象	原因	解决
单元测试全过，训练却慢得反常	`nn.Embedding` 默认 $N(0,1)\mathcal{N}(0,1)$ ，乘 $d\sqrt{d}$ 后淹没位置编码	按 $N(0,d−0.5)\mathcal{N}(0, d^{-0.5})$ 初始化（20.12）
`RuntimeError: view size is not compatible...`	`transpose` 后内存不连续，不能用 `view`	改用 `reshape`，或先 `.contiguous()`
`Expected all tensors to be on the same device`	forward 里新建的 Mask 默认在 CPU	创建时传 `device=x.device`
位置编码没跟着模型上 GPU	用普通属性 `self.pe = ...` 保存	用 `register_buffer`
子层参数没被训练	用普通 `list` 装子模块	用 `nn.ModuleList`
改成 SDPA 后结果全错	SDPA 的 bool mask 是 True = 可见，与 `masked_fill` 相反	传 `~mask`，或用 `is_causal=True`
loss 降得很快，生成时却全是乱码	Causal Mask 方向写反或没传，训练时偷看了答案	跑 20.8 的因果性测试
短句效果好、长句效果差	只屏蔽了源端 PAD，没屏蔽目标端 PAD	Decoder 自注意力用 `causal \| padding`
输出整体错一位	`tgt_in` / `tgt_out` 切片写反	`tgt_in = tgt[:, :-1]`，`tgt_out = tgt[:, 1:]`
bf16/fp16 下出现 NaN	屏蔽值 `-1e9` 超出 fp16 范围	用 `torch.finfo(scores.dtype).min`
推理结果每次都不一样	忘了 `model.eval()`，Dropout 还开着	推理前 `model.eval()`

本章练习与实验

复现初始化问题：删掉 TokenAndPosition 里的 nn.init.normal_(...) 那几行，重新训练，复现 20.12.1 的慢收敛。再试试保留默认初始化、但去掉 * math.sqrt(self.d_model)，效果如何？
去掉位置编码：把 + self.pe[:L] 删掉重新训练，exact_match 会怎样？为什么翻转任务离不开位置信息？
只用 1 个 Head：num_heads=1，对比收敛速度和最终准确率。
Post-Norm 版本：把 Encoder/Decoder Layer 改成 LN(x + Sublayer(x))，去掉 ln_final，在 num_layers=6 时与 Pre-Norm 对比；再把 warmup 改成 1 看看两者谁更容易崩。
换任务：把目标改成"排序"（np.sort(digits)）。Cross-Attention 矩阵还会是反对角线吗？
SwiGLU：参照第8章，把 FeedForward 换成 SwiGLU，并保持参数量大致相同（隐藏维度取 $23×4d\frac{2}{3} \times 4d$ ）。
与 nn.Transformer 对比：用 nn.Transformer(d_model=128, nhead=4, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=512, batch_first=True, norm_first=True) 替换本章的 Encoder/Decoder，训练结果是否一致？

本章总结

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

本章思考题

推理时 Encoder 只跑一次，Decoder 要跑多次。如果源序列很长、目标序列很短（如文本摘要），计算瓶颈在哪一侧？
Cross-Attention 的 Mask 为什么只需要源端 Padding，不需要 Causal Mask？
为什么 val_exact_match 会比 token_acc 低？两者的差距来自哪里？（提示：Exposure Bias）
20.12 节中，Embedding 的数值比位置编码大 10 倍时，为什么模型"方向对、但分不清相邻位置"？从点积注意力的角度解释。
本章的位置编码是固定的正弦函数。如果测试时输入长度超过训练时的最大长度 12，模型还能正确翻转吗？试着做实验验证（注意 max_len 也要调大）。

下一章预告

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

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

转载自 CSDN-专业IT技术社区

原文链接：https://blog.csdn.net/richard_first/article/details/168775264

第20章从零实现 Transformer（PyTorch）

第六篇实战