用PyTorch从零手写Transformer并跑通训练全流程

很多人学 Transformer 都有一个共同的卡点：论文看了好几遍，公式抄了又抄，QKV、多头注意力、位置编码这些名字都能脱口而出，但一落到 PyTorch 代码里，就不知道从哪里下手。要么是照着开源仓库抄一遍但完全看不懂，要么是只会在 HuggingFace 里调model = BertModel.from_pretrained(...)，真让你自己写一个能训练的 Transformer，大脑一片空白。

这篇文章想做的事情很直接：用 PyTorch 从零手写一个极简但完整的 Transformer，并让它真正跑通一次训练。不调nn.Transformer封装好的模块，Attention、QKV、位置编码、残差连接、层归一化，全部自己实现。代码不做任何花哨包装，能跑、能训练、能生成文本，让你完整看到数据从一个字符到最终输出的全过程。

如果你正准备面试、刚接触 NLP、或者想真正理解 GPT 这类模型为什么能工作，这篇文章非常适合你。读完你会有两个收获：一是对 Transformer 每个模块的输入输出形状和数学含义不再含糊；二是拿到一段可以继续改造和学习的极简 PyTorch 实现。

1. 这篇文章真正要解决的问题

先说说为什么值得手撕一遍 Transformer。现在学习资源非常多，光是讲解 Attention 的博客就数以千计，但大部分学习者仍然停留在“看得懂、写不出”的状态。原因在于，Transformer 不是一个孤立的模型结构，它把分词、嵌入、位置编码、注意力计算、归一化、前馈网络、训练循环串成了一条完整的链路。任何一个环节只懂概念不懂实现，整条链路在脑子里就是断裂的。

另一个现实问题是：面试和工程实战对“理解深度”的要求越来越高。面试官问“为什么 Transformer 要加位置编码”“QKV 到底是什么形状”“多头注意力为什么有效”，你如果只回答“位置编码是给模型提供顺序信息”“QKV 是查询、键、值”，大概率会被追问到细节。这些问题靠背答案是撑不住的，只有亲手写过一遍，才会真正知道答案在哪里。

这篇文章解决的具体痛点有三个：

1. 概念与代码对不上：知道 QKV 三个字母，但不知道它们在 PyTorch 里应该创建什么Tensor、怎么变换形状、怎么计算。
2. 只懂前向不懂训练：很多教程讲模型结构讲得很细，但训练循环、损失函数、梯度回传语焉不详。本文会带你把训练部分也跑通。
3. 不知道如何验证：模型写完了，怎么判断它真的学到了东西？本文会给出一个可观察的验证方式。

这里先给出一个核心判断：手写 Transformer 的价值，不在于你写出来的模型性能有多好，而在于你亲手建立了一条从数据到梯度回传的完整认知链。这条认知链，是读十篇论文都给不了的。

2. 核心概念：从 Token 到 QKV 再到多头注意力

在写代码之前，我们需要把几个基础概念讲透。这些概念不是用来背的，而是后续代码里每一个Tensor形状变换的根据。

2.1 Token 与词嵌入：文本怎么变成数字

神经网络无法直接处理字符或单词，必须把它们变成数字向量。第一步是“分词”：把一段文本切分成最小的语义单元。最简单的做法是字符级分词，比如"hello"被切成['h','e','l','l','o']。实际项目中常用的是 BPE 等子词分词，但原理一致。

每个 Token 对应一个整数 ID，这就是“索引”。然后我们构建一个嵌入矩阵nn.Embedding(vocab_size, d_model)，通过 Token ID 从矩阵中取出一行，得到该 Token 的稠密向量。这个向量就是“词嵌入”。它相当于把“猫”和“狗”这样语义接近的词，映射到向量空间中距离较近的位置。

2.2 位置编码：为什么要加，怎么加

Transformer 的核心是 Self-Attention，它天然对序列顺序不敏感。在 Attention 计算时，query会和序列里所有位置计算相似度，换个顺序得到的结果几乎一样。如果不加位置信息，模型会把“我爱你”和“你爱我”当成同一个输入来对待，这显然是错误的。

位置编码的作用就是给每个位置注入一种“绝对或者相对的顺序信号”。经典的正余弦位置编码公式长这样：

PE(pos, 2i) = sin(pos / 10000^(2i / d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

其中pos是位置下标，i是向量维度下标。这个公式的含义是：不同维度使用不同的波长，这样每个位置的编码向量都是唯一的，而且相邻位置的编码差异是平滑的。正余弦编码是固定值，不需要学习。

更严谨地说，这个公式还有一个数学上的好处：对位置pos + k的编码可以近似表示成pos位置编码的线性变换。这有利于模型学习相对位置关系。

在实现上，位置编码和词嵌入的计算方式是直接相加：

x = token_embedding + position_encoding

相加之后，向量仍然携带语义信息和位置信息，后续 Attention 就可以同时利用这两类信息。

2.3 QKV 与注意力权重：自注意力的数学本质

很多初学者在 QKV 这里被劝退，它其实没有想象中复杂。一句话解释：Q（Query）是“我要找什么”，K（Key）是“我有什么”，V（Value）是“我给出的内容”。

可以把这个过程类比成搜索引擎。你输入一个查询词（Q），搜索引擎遍历所有网页标题（K），计算查询词和网页标题的匹配度，然后按匹配度加权返回网页内容（V）。Attention 做的就是这件事，只不过整个过程是连续的、可微的。

具体到代码里，输入序列的向量x会分别经过三个线性层，得到:

Q = x @ W_q K = x @ W_k V = x @ W_v

这里的W_q、W_k、W_v是形状为(d_model, d_model)的可学习权重矩阵。注意 Q、K、V 都来自同一个输入x，这就是“Self-Attention（自注意力）”名字的由来——序列和自己在做注意力。

然后计算注意力分数并加权求和：

scores = Q @ K^T / sqrt(d_k) weights = softmax(scores, dim=-1) output = weights @ V

除以sqrt(d_k)是为了防止点积结果过大导致 softmax 落在梯度饱和区。d_k是每个头的维度，之所以要开放大维度，是因为如果向量的维度很大，点积的方差也会变大，softmax 的输出会过于接近 one-hot，导致梯度消失。

2.4 多头注意力：多角度看序列

单一注意力只能关注一种关系模式，这限制了模型的表达力。多头注意力的想法很简单：把 Q、K、V 分别拆分到多个“头”中，每个头独立计算注意力，最后拼接在一起。

比如输入向量维度是 64，我们设置 8 个头，那么每个头的维度是 8。每个头内的 QKV 计算是独立的，可以并行完成。在 PyTorch 中，我们通常把形状从(batch, seq_len, d_model)变换成(batch, seq_len, num_heads, head_dim)，然后转置成(batch, num_heads, seq_len, head_dim)进行批量矩阵乘法。

多头注意力的意义在于：每个头可以关注不同类型的相关性。比如一个头关注语法关系，另一个头关注指代关系，第三个头关注相邻词。这种多角度的信息提取能力，是 Transformer 表达力的核心来源之一。

2.5 残差连接、层归一化与 FFN

Attention 输出之后并不直接进入下一层，而是经过一系列标准化和变换。这里有两个关键设计。

残差连接：把输入直接加到输出上，即x = x + sublayer(x)。它解决了深层网络的梯度消失问题，让信息可以跨层流动。Transformer 的每一层几乎都用了残差结构。

层归一化（LayerNorm）：对每个样本的每个位置的向量做标准化，让均值接近 0、方差接近 1。它让数据分布更稳定，减少训练过程中的内部协变量偏移。注意它和 BatchNorm 的区别：LayerNorm 是对某一层内所有神经元做归一化，与 batch 大小无关，因此特别适合序列模型。

前馈网络（FFN）：Attention 层是线性变换加注意力加权，本质上还是在做“信息的收集”。要增强模型的非线性拟合能力，还需要一个带有激活函数的两层 MLP：

FFN(x) = GELU(x @ W1 + b1) @ W2 + b2

它先升维到4 * d_model，再降维回d_model。每个位置共享同一组参数，对序列逐位置进行非线性变换。

到这里，核心概念已经全部介绍完。下一章开始搭建环境并进入代码实现。

3. PyTorch 环境准备与版本建议

手写 Transformer 只需要 PyTorch 基础组件，不需要额外安装 Transformer 库。环境准备是这个项目里最容易卡住的环节，尤其是 GPU 版本安装。

建议版本组合如下（以实际环境为准，本文代码不依赖特定主版本）：

• Python 3.9 或 3.10
• PyTorch 2.0 或以上，CPU 版本即可运行本文示例
• NumPy（PyTorch 会携带依赖，通常无需单独安装）

创建虚拟环境并安装 CPU 版 PyTorch 的示例命令：

conda create -n transformer python=3.10 -y conda activate transformer # CPU 版，适用于学习环境 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu # 如果已有 CUDA 环境，也可以安装 GPU 版 # pip install torch --index-url https://download.pytorch.org/whl/cu118

安装完成后，验证一下 PyTorch 是否可用：

python -c "import torch; print(torch.__version__)"

如果输出类似2.1.2的版本号，说明安装成功。需要提醒的是，如果电脑没有 NVIDIA GPU 或者没有配置好 CUDA，直接用 CPU 版训练本文的极简模型完全足够，训练时间通常只有几十秒到几分钟。

4. 数据准备：从字符到 Batch

为了让代码能在 CPU 上快速跑起来，我选用一个字符级语言模型作为示例任务。训练目标很简单：给定前几个字符，预测下一个字符。这个任务虽然小，但足以完整展示 Transformer 的训练流程。

先准备一份小文本语料。这里直接内嵌一段英文文本，避免下载数据集的麻烦：

# 文件路径：prepare_data.py import torch # 示例训练语料（可替换为任何 txt 文件内容） text = ( "the quick brown fox jumps over the lazy dog. " "machine learning is powerful and fascinating. " "transformers are changing the world of artificial intelligence. " "attention is all you need, but you also need practice. " "practice makes perfect if you always code with patience." ) # 字符级分词：收集所有不重复字符 chars = sorted(list(set(text))) vocab_size = len(chars) # 建立字符到 ID 的映射 char_to_id = {ch: i for i, ch in enumerate(chars)} id_to_char = {i: ch for i, ch in enumerate(chars)} print(f"语料长度: {len(text)}") print(f"字符种类数: {vocab_size}") print(f"字符表: {chars}")

在字符级分词中，每个字符就是一个 Token。这样的好处是词表很小，训练速度快。如果我们换成英文单词级分词，词表可能有几万个，但这在原理上和字符级完全一致。

接着把文本转成 ID 序列，并按固定窗口长度构造训练样本：

# 文件路径：prepare_data.py（续） seq_len = 32 # 每个训练样本的长度 # 转成 ID 序列 ids = [char_to_id[ch] for ch in text] # 用滑动窗口构造 (input, target) 对 # input: 连续 seq_len 个字符的 ID # target: input 右移一位，即每个位置的下一个字符 ID inputs = [] targets = [] for i in range(len(ids) - seq_len): input_ids = ids[i : i + seq_len] target_ids = ids[i + 1 : i + seq_len + 1] inputs.append(input_ids) targets.append(target_ids) # 转成 Tensor inputs = torch.tensor(inputs, dtype=torch.long) targets = torch.tensor(targets, dtype=torch.long) batch_size, _ = inputs.shape print(f"训练样本数: {batch_size}") print(f"输入形状: {inputs.shape}") # (样本数, 序列长度) print(f"目标形状: {targets.shape}")

这里有一个细节值得注意：target并不是input的下一整句，而是右移一位的整个窗口。对语言模型来说，我们希望在看到句子的第t个字符时，能预测第t+1个字符。窗口滑动让每个样本内部都包含多个预测任务，最大化数据利用率。

在实际项目中，数据量远大于这段小文本，我们通常用DataLoader配合Dataset组织数据。这里为了保持代码简洁，直接用 Tensor 切片做训练。

5. Transformer 核心模块逐行实现

现在进入本文的核心部分。我会按照从底向上的顺序实现模块：先实现位置编码，再实现注意力，然后组合成多头注意力、FFN、Encoder 层，最后构造完整的 Transformer 语言模型。

5.1 正余弦位置编码

先写位置编码类。这个类不包含可学习参数，前向时返回一个形状为(seq_len, d_model)的位置编码矩阵：

# 文件路径：transformer_model.py import torch import torch.nn as nn import math class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=512): super().__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp( torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model) ) # 偶数维度用 sin pe[:, 0::2] = torch.sin(position * div_term) # 奇数维度用 cos pe[:, 1::2] = torch.cos(position * div_term) # 注册为 buffer，不参与梯度更新，但会随模型迁移到 GPU self.register_buffer('pe', pe.unsqueeze(0)) def forward(self, x): # x 形状: (batch, seq_len, d_model) seq_len = x.size(1) return x + self.pe[:, :seq_len, :]

这里div_term的计算方式等价于1 / (10000^(2i/d_model))，但用指数形式计算更稳定。register_buffer是 PyTorch 里的一个重要细节：buffer 会随模型的to('cuda')一起迁移设备，但又不会在反向传播时更新梯度，非常适合保存位置编码这类固定数据。

5.2 单头注意力：完整实现 QKV 计算

接下来实现自注意力模块。为方便理解，我先写一个不拆分多头的基础版本，然后在其基础上扩展到多头。

# 文件路径：transformer_model.py（续） class SelfAttention(nn.Module): def __init__(self, d_model, dropout=0.1): super().__init__() self.d_model = d_model self.d_k = d_model # 单头时 head_dim 等于 d_model 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.out_proj = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def attention(self, q, k, v, mask=None): scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) weights = torch.softmax(scores, dim=-1) weights = self.dropout(weights) return torch.matmul(weights, v) def forward(self, x, mask=None): # x 形状: (batch, seq_len, d_model) q = self.W_q(x) k = self.W_k(x) v = self.W_v(x) output = self.attention(q, k, v, mask) return self.out_proj(output)

这就是自注意力的全部核心逻辑。scores的形状是(batch, seq_len, seq_len)，它记录的是序列中任意两个位置之间的相关程度。mask参数用于控制哪些位置不能互相看见，这一点在 GPT 类的因果语言模型中特别重要。

5.3 因果掩码：为什么 GPT 不能看见未来 token

在语言模型训练中，我们预测第t+1个 Token 时，模型只能看到第 1 到第t个 Token，绝不能看到第t+1个以及后面的 Token。如果允许看到，模型就可以直接“抄答案”，训练就失去了意义。

这时需要一个上三角掩码矩阵。矩阵中第i行第j列如果是 0，表示位置i不能看位置j；如果是 1，表示可以看。对每个位置i，能让它看的位置是j <= i。

def create_causal_mask(seq_len): # 返回形状: (1, 1, seq_len, seq_len) # 下三角为 1，上三角为 0 mask = torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len) return mask

上面代码返回的 mask 可以直接传入scores.masked_fill(mask == 0, float('-inf'))。被掩码的位置在 softmax 前被设为负无穷，softmax 后权重就会变成 0，表示这个位置的 token 对当前输出没有任何贡献。

5.4 多头注意力：拆分、计算、拼接

现在把单头注意力升级成多头。拆分的核心思路是：不创建多个独立的 QKV 线性层，而是一次性计算得到3 * d_model维度的结果，再按头数切分。这是工程实现中最常见的做法，效率更高。

# 文件路径：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.d_model = d_model self.num_heads = num_heads self.d_k = 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.out_proj = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): batch, seq_len, _ = x.shape # 分别计算 Q、K、V q = self.W_q(x) # (batch, seq_len, d_model) k = self.W_k(x) v = self.W_v(x) # 拆分成多头形状: (batch, seq_len, num_heads, d_k) q = q.view(batch, seq_len, self.num_heads, self.d_k) k = k.view(batch, seq_len, self.num_heads, self.d_k) v = v.view(batch, seq_len, self.num_heads, self.d_k) # 转置为: (batch, num_heads, seq_len, d_k) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) # 点积注意力 scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: # mask 的形状要能广播到 scores: (batch, 1, seq_len, seq_len) scores = scores.masked_fill(mask == 0, float('-inf')) weights = torch.softmax(scores, dim=-1) weights = self.dropout(weights) context = torch.matmul(weights, v) # 重新拼接各头: (batch, num_heads, seq_len, d_k) context = context.transpose(1, 2).contiguous() # 合并头: (batch, seq_len, d_model) context = context.view(batch, seq_len, self.d_model) return self.out_proj(context)

这个实现里有几个形状变换的细节值得反复品味：

1. view把(batch, seq_len, d_model)拆成(batch, seq_len, num_heads, d_k)，其中d_model = num_heads * d_k。
2. transpose(1, 2)让num_heads维度提前，这样batch和num_heads可以合并成两个批次维度，方便后续批量矩阵乘法。
3. 最后contiguous().view()的作用是先把 Transpose 后不连续的内存变成连续内存，再合并维度。如果少了contiguous()，某些情况下会报错。

多头注意力的核心意义：不同头学到的注意力模式是不同的，有的头更关注相邻词，有的头关注远距离词。多个头并行工作，相当于一个团队里不同成员各自关注不同的线索，最后把线索汇总起来。

5.5 前馈网络与 Transformer 层

FFN 实现很简单，但它是 Transformer 中参数量最大的模块之一。一般认为 Attention 负责“收集信息”，FFN 负责“加工信息”。

# 文件路径：transformer_model.py（续） class FeedForward(nn.Module): def __init__(self, d_model, d_ff=2048, dropout=0.1): super().__init__() self.linear1 = nn.Linear(d_model, d_ff) self.linear2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): x = self.linear1(x) x = torch.relu(x) x = self.dropout(x) x = self.linear2(x) return x

然后组装一个完整的 Transformer 层。一个层由多头注意力和 FFN 两大部分组成，每个部分外面都套着“残差 + LayerNorm”的结构。

# 文件路径：transformer_model.py（续） class TransformerBlock(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.attention = MultiHeadAttention(d_model, num_heads, dropout) self.norm1 = nn.LayerNorm(d_model) self.ffn = FeedForward(d_model, d_ff, dropout) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): # 子层 1：多头注意力 + 残差 attn_out = self.attention(x, mask) x = x + self.dropout(attn_out) x = self.norm1(x) # 子层 2：前馈网络 + 残差 ffn_out = self.ffn(x) x = x + self.dropout(ffn_out) x = self.norm2(x) return x

有一点需要说明：Transformer 原论文中采用的是 Post-Norm（先残差再 Norm），而 GPT 等现代模型更常用 Pre-Norm（先 Norm 再进入子层）。这里为了代码清晰，使用的是先残差再 LayerNorm 的简化写法。在极小模型上两者差异不大，但在大型模型中，Pre-Norm 往往训练更稳定。

5.6 组装完整 Transformer 语言模型

最后把前面所有模块组装成一个完整的 decoder-only 语言模型。这个模型接收 Token ID 序列，输出每个位置的下一个 Token 的概率分布。

# 文件路径：transformer_model.py（续） class TransformerLM(nn.Module): def __init__(self, vocab_size, d_model=64, num_heads=4, num_layers=2, d_ff=128, max_len=128, dropout=0.1): super().__init__() self.token_embedding = nn.Embedding(vocab_size, d_model) self.positional_encoding = PositionalEncoding(d_model, max_len) self.blocks = nn.ModuleList([ TransformerBlock(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.norm = nn.LayerNorm(d_model) self.lm_head = nn.Linear(d_model, vocab_size) def forward(self, token_ids, mask=None): # token_ids 形状: (batch, seq_len) x = self.token_embedding(token_ids) # (batch, seq_len, d_model) x = self.positional_encoding(x) for block in self.blocks: x = block(x, mask) x = self.norm(x) logits = self.lm_head(x) # (batch, seq_len, vocab_size) return logits

lm_head把向量映射回词表大小，输出的是 logits——也就是未经过 softmax 的概率分数。训练时我们用交叉熵损失函数直接吃 logits，内部会自动计算 softmax。

到这里，模型核心代码已经全部完成。下面准备训练脚本。

6. 训练循环：让模型真正跑起来

训练一个语言模型的标准流程是：将输入 ID 序列传入模型得到 logits，把 logits 和 target ID 序列计算交叉熵损失，然后反向传播更新参数。关键点是CrossEntropyLoss要求的目标形状是(batch * seq_len)，logits 形状是(batch, seq_len, vocab_size)，需要对 logits 做 view。

# 文件路径：train.py import torch import torch.nn as nn from torch.optim import AdamW from transformer_model import TransformerLM from prepare_data import inputs, targets, vocab_size device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # 超参数 batch_size = 16 learning_rate = 1e-3 epochs = 100 # 创建模型 model = TransformerLM( vocab_size=vocab_size, d_model=64, num_heads=4, num_layers=2, d_ff=128, max_len=64, dropout=0.1 ).to(device) # 因果掩码 seq_len = inputs.size(1) causal_mask = torch.tril(torch.ones(seq_len, seq_len)).view(1, 1, seq_len, seq_len).to(device) optimizer = AdamW(model.parameters(), lr=learning_rate) criterion = nn.CrossEntropyLoss() # 训练 model.train() total_steps = len(inputs) // batch_size for epoch in range(epochs): total_loss = 0.0 # 打乱数据顺序 perm = torch.randperm(len(inputs)) inputs = inputs[perm] targets = targets[perm] for step in range(total_steps): batch_x = inputs[step * batch_size : (step + 1) * batch_size].to(device) batch_y = targets[step * batch_size : (step + 1) * batch_size].to(device) logits = model(batch_x, mask=causal_mask) logits = logits.view(-1, vocab_size) batch_y = batch_y.view(-1) loss = criterion(logits, batch_y) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() avg_loss = total_loss / total_steps if epoch % 20 == 0 or epoch == epochs - 1: print(f"Epoch {epoch:3d} | Loss: {avg_loss:.4f}") # 保存模型权重 torch.save(model.state_dict(), "transformer_lm.pt") print("训练完成，模型已保存到 transformer_lm.pt")

这段代码有几个训练细节值得说明：

1. 交叉熵损失的形状处理：logits.view(-1, vocab_size)把批量和序列长度合并，对应地对 target 做同样的操作。这样每个位置上的预测都参与损失计算，模型学的是“所有位置的平均预测能力”。
2. 损失下降的判断：如果初始语料有 50 个不同的字符，随机猜测的交叉熵大约在log(50) ≈ 3.91。训练后 loss 下降到 1 以下，说明模型已经学会了基本的字符共现规律。
3. AdamW 优化器：目前的大模型训练事实标准，它在 Adam 的基础上做了权重衰减解耦，比普通 Adam 在 Transformer 上效果更稳定。

如果你在 Windows 或 Mac 上使用 CPU 训练，这个例子大概只需要 1 到 3 分钟。

7. 运行结果与效果验证

训练完模型之后，我们需要一种直观的方法来验证模型是否真的学到了东西。验证方法就是让模型自己“续写文本”：给它一个开头的字符序列，让它逐字预测并生成后面的内容。

写一个简单的文本生成函数：

# 文件路径：generate.py import torch from transformer_model import TransformerLM from prepare_data import char_to_id, id_to_char, vocab_size device = torch.device("cuda" if torch.cuda.is_available() else "cpu") seq_len = 32 model = TransformerLM( vocab_size=vocab_size, d_model=64, num_heads=4, num_layers=2, d_ff=128, max_len=128, dropout=0.1 ).to(device) model.load_state_dict(torch.load("transformer_lm.pt", map_location=device)) model.eval() def generate(prompt, max_new_tokens=40, temperature=0.8): model.eval() # 将 prompt 转成 ID 列表 ids = [char_to_id[ch] for ch in prompt][-seq_len:] input_ids = torch.tensor([ids], dtype=torch.long, device=device) with torch.no_grad(): for _ in range(max_new_tokens): logits = model(input_ids) next_token_logits = logits[0, -1, :] / temperature probs = torch.softmax(next_token_logits, dim=-1) next_token = torch.multinomial(probs, num_samples=1).item() ids.append(next_token) input_ids = torch.tensor([ids[-seq_len:]], dtype=torch.long, device=device) return ''.join(id_to_char[i] for i in ids) print(generate("the quick brown"))

运行后你可能会看到类似这样的输出：

the quick brown fox jumps over the lazy dog. machine

不同训练随机种子下结果会有差异，但只要 loss 下降到一定程度，模型应该能生成一些语法上像英文片段的文本。

如果生成的结果完全不可读，按以下顺序排查：

1. Loss 是否下降：训练后 loss 如果还在 3.0 以上，说明模型没有学到有效信息，先检查数据集是否太短、epoch 是否太少。
2. 温度参数是否合理：temperature过大会导致输出随机性太强，文本看起来像乱码。建议在 0.5 到 1.0 之间调节。
3. 矩阵形状是否匹配：生成时多次调用了模型，输入长度变化不能超过max_len。这里用[-seq_len:]做窗口截断就是为了保证长度固定。

8. 常见问题与排查思路

手写 Transformer 的过程中，大部分人都会遇到类似的报错和诡异问题。我整理了一份高频问题排查表，建议直接收藏。

还有一个非常隐蔽的坑：位置编码的 buffer 和设备迁移问题。如果你在模型已to(device)之后才创建位置编码矩阵，那么矩阵还在 CPU 上，前向时会报 “Expected all tensors to be on the same device”。因此我们在PositionalEncoding中使用register_buffer，并在模型初始化时定义好所有模块，这样model.to(device)会把所有参数和 buffer 一并迁移。

9. 最佳实践与工程建议

手写模型和在生产环境中实际使用 Transformer，是两件不同的事。下面整理一些从“能跑”到“好用”的工程建议。

9.1 从最小实验开始

永远不要一开始就尝试复现 GPT-3 或 BERT-Large。推荐的路线是：先用极小的d_model=64、num_heads=4、num_layers=2跑通训练流程，然后逐步扩大规模。这样每个阶段的问题都可控，调试成本极低。

9.2 使用内置 API 加速开发

当你理解了手写实现之后，在实际项目中可以直接使用 PyTorch 提供的高效实现nn.MultiheadAttention和nn.TransformerEncoderLayer。这些内部实现经过高性能优化并支持 FlashAttention 等新特性，训练速度和显存利用率远高于我们手写的基础版本。

但要注意：手写版是为了理解，内置版是为了生产。二者并不冲突，可以先用自己写的模型做一个小实验验证思路，确定无误后再迁移到内置版。

9.3 训练时的梯度监控

训练 Transformer 时，梯度异常通常是最早出现的故障信号。建议在训练循环中定期打印梯度的 L2 范数：

total_norm = 0.0 for p in model.parameters(): if p.grad is not None: param_norm = p.grad.detach().norm(2) total_norm += param_norm.item() ** 2 total_norm = total_norm ** 0.5 print(f"grad norm: {total_norm:.4f}")

如果梯度范数爆炸（比如超过 10 甚至 100），可以在优化器更新之前做梯度裁剪：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

这是 GPT 训练中常用的技巧，能有效避免训练发散。

9.4 数据质量优先于模型结构

在真实业务中，模型结构的改进带来的收益往往不如数据清洗和预处理带来的收益大。训练语言模型前，先统计字符分布、检查文本编码、确认没有乱码和异常空格，这些基础工作的重要性不亚于调参。

9.5 版本兼容与日志

训练脚本要养成记录超参数的习惯。保存模型权重时，建议同时保存超参数和词表映射：

import json checkpoint = { "model_state_dict": model.state_dict(), "vocab_size": vocab_size, "char_to_id": char_to_id, "d_model": 64, "num_heads": 4, "num_layers": 2, } torch.save(checkpoint, "transformer_checkpoint.pt")

这样以后加载模型时不需要依赖训练脚本里的全局变量，方便复用和部署。

10. 总结与后续学习路线

到这一步，你已经亲手实现了一个完整的 Transformer 语言模型。不妨回顾一下本文真正讲清楚的几个点：Token 是如何从字符变成 ID 再从 ID 变成向量的；位置编码为什么加、怎么加；QKV 的形状变换和注意力计算过程；多头注意力如何拆分和拼接；因果掩码为什么必不可少；一个训练循环里损失函数、梯度回传和优化器更新是怎么串起来的。

建议下一步从这几个方向继续深入：

1. 把编码器结构也实现一遍：本文实现的是 decoder-only 结构，也就是 GPT 系列的基础。你可以对比学习 BERT 使用的 encoder-only 结构，理解两者在 mask、训练目标和应用场景上的差异。
2. 阅读原始论文 Attention Is All You Need：带着代码经验去读论文，很多之前看不懂的细节会突然变得清晰。
3. 尝试训练一个更大规模的模型：把语料换成莎士比亚全集或中文小说，把d_model提升到 256 甚至 512，在 GPU 上观察模型生成能力的显著提升。
4. 尝试使用 PyTorch 内置模块重构：用nn.TransformerDecoderLayer替换手写模块，对比训练速度和代码量差异，这会让你更深刻地理解框架为什么提供这些封装。

最后给一个实用建议：把这份代码保存在自己的 Git 仓库里，经常回来改一改、加一个功能、调一组参数。Transformer 的核心知识在动手写代码的过程中会真正内化成你自己的直觉。下一次再看到新的注意力变体论文，你就不会觉得它是一堆陌生的公式，而只是在你熟悉的主干上做了一些改动。