Transformer 是当前大语言模型和许多深度学习任务的核心架构，但很多人学它时卡在“概念听懂了、代码写不出来”这一步。网上讲解自注意力、QKV、位置编码的文章很多，真正能让人从 Token 开始一步步把模型写出来、跑起来、看到损失下降的材料却不多。这篇文章就用 PyTorch 手写一个极小的 Transformer 字符级语言模型，覆盖 Token 化、Embedding、位置编码、多头自注意力、Transformer Block、训练循环和生成验证。学完以后，你既能理解 QKV 和自注意力的计算过程，也能在自己机器上跑通一个可复现的最小项目，后续再去看大模型源码会轻松很多。
1. 先拆清 Transformer 的输入输出，再动手写代码
1.1 Token 是什么？为什么要先把文本变成数字
在 Transformer 眼里，文本不是文字，而是一串离散编号。这里的“离散编号”就是 Token ID。Token 可以是一个字符、一个词语、一个子词，不同分词策略对应不同的词表大小。
举例来说，字符串"abc"如果按字符切分，可以切成['a', 'b', 'c']三个 Token；如果按词语切分，可能只有一个 Token"abc"。无论哪种切法，最终都要通过一个词表把 Token 映射成整数。模型内部计算只能处理数字，不能直接处理字符。
这一步决定了整个任务的边界：
- 词表太小，模型很难表达复杂语义。
- 词表太大，Embedding 参数会变多，训练成本上升。
- 分词不一致，会让同一句话在不同场景下变成完全不同的序列。
在本文的最小项目中，我会先用“字符级分词”演示完整流程。它简单、透明、容易复现，适合零基础理解；真实项目里再换成 BPE 等子词分词方式。
1.2 Transformer 看到的输入和输出
假设我们有一个 batch 里面有batch_size句话，每句话固定长度为seq_len，那么 Transformer 的输入通常是一个形状为(batch_size, seq_len)的整型张量。每一个元素都是 Token ID。
模型内部会做大量矩阵运算，最终输出形状为(batch_size, seq_len, vocab_size)的浮点张量，称为 logits。logits[b, t, v]表示第b个样本、第t个位置上，下一 Token 是词表第v个词语的得分。
这种“输入序列、输出序列”的结构非常通用：
- 文本分类：把输出序列压缩成一个向量，再映射到类别。
- 机器翻译：Encoder 读入源语言，Decoder 输出目标语言。
- 生成式语言模型：每一个位置都预测下一个 Token。
本文采用的是生成式语言模型任务：给定前面seq_len个字符，让模型预测下一个字符。这样最小项目既简单，又能直接看出训练效果。
1.3 最小模型包含的模块清单
一个完整的 Transformer 生成模型，主要包含下面几个模块：
| 模块 | 作用 | 本文中的代码 |
|---|---|---|
| Token Embedding | 把 Token ID 映射成稠密向量 | nn.Embedding |
| 位置编码 | 给序列元素注入顺序信息 | PositionalEncoding |
| 多头自注意力 | 计算每个 Token 与其他 Token 的关系 | MultiHeadAttention |
| 残差连接 | 让深层网络梯度更容易回传 | x + attn(x) |
| LayerNorm | 稳定训练过程 | nn.LayerNorm |
| FFN 前馈网络 | 对每个位置做非线性变换 | Linear + GELU + Linear |
| 输出投射层 | 把向量映射回词表维度 | nn.Linear |
很多人一上来就写完整代码，结果每一行都似懂非懂。建议按照模块逐个实现，最后再组装成完整模型。
1.4 学习环境和生产环境的要求分别是什么
零基础手写 Transformer 时，目标是“跑通逻辑”，不是“训练大模型”。学习和生产的要求差异很大：
| 维度 | 学习环境 | 生产环境 |
|---|---|---|
| 数据规模 | 几十个字符即可 | 几十 GB 甚至更大 |
| 模型规模 | d_model=64，num_layers=2 | 几亿到几千亿参数 |
| GPU | 没有也能跑 | 需要多卡训练 |
| 分词 | 字符级即可 | BPE、SentencePiece |
| 训练技巧 | 最简单交叉熵 | 混合精度、分布式、学习率调度 |
| 部署 | 无需考虑 | 推理优化、量化、服务化 |
这个表格不是让大家跳过生产问题，而是要区分“理解原理”和“工程落地”。本文先完成理解原理这一步。
2. 准备 PyTorch 环境并搭好项目骨架
2.1 用 Anaconda 创建 PyTorch 环境
推荐用 Anaconda 管理环境，避免污染系统 Python。下面是创建一个名为transformer的 Python 3.10 环境，并安装 CPU 版 PyTorch 的命令：
conda create -n transformer python=3.10 conda activate transformer conda install pytorch torchvision torchaudio cpuonly -c pytorch如果机器有 NVIDIA GPU，并且已经安装好 CUDA 驱动，可以去 PyTorch 官网选择对应的安装命令。常见写法如下，但具体版本以官网为准：
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia学习本文的字符级模型时，CPU 版本的 PyTorch 已经足够。训练时间通常只需要几十秒，不需要纠结 GPU 版本。
2.2 验证 PyTorch 安装是否成功
创建环境并安装后，在命令行进入 Python，执行下面代码：
import torch print(torch.__version__) print(torch.cuda.is_available())如果输出类似于：
2.1.0 False说明 CPU 版安装成功。如果torch.cuda.is_available()输出True，说明 GPU 可用，后续可以把模型和数据放到 GPU 上训练。
2.3 项目目录结构
本文为了演示方便，会把代码写在 Jupyter Notebook 或者单个 Python 文件里。实际工程项目中，建议把不同模块拆到不同文件：
transformer_handwrite/ ├── data.py # 分词、数据集、DataLoader ├── model.py # 位置编码、注意力、Transformer Block、完整模型 ├── train.py # 训练循环 ├── generate.py # 生成函数 └── config.py # 超参数配置零基础学习时，先在一个文件里跑通，再拆分成模块。目录结构的意义是让代码可维护，不是一开始就追求完美架构。
2.4 依赖版本核对清单
虽然代码只用到了 PyTorch 标准库，但版本差异会影响 API 和运行结果。建议运行前核对：
| 依赖 | 作用 | 建议 |
|---|---|---|
| Python | 解释器版本 | 3.10 或 3.11 |
| PyTorch | 深度学习框架 | 2.0 以上 |
| NumPy | 数值计算 | PyTorch 自带依赖 |
如果使用 PyTorch 2.x，后续还可以尝试用torch.nn.functional.scaled_dot_product_attention替换手写注意力，速度更快、内存更少。不过那一步放在后面进阶，手写实现是为了看清细节。
3. 从 Token 到 Embedding：先把文本变成模型能吃的数据
3.1 用字符级词表演示 Token 化
我们用一段短文本作为训练数据：
text = "the quick brown fox jumps over the lazy dog the early bird catches the worm"字符级分词就是把text里面的每个字符取出来，去重后建立“字符到编号”和“编号到字符”的映射：
chars = sorted(list(set(text))) vocab_size = len(chars) stoi = {ch: i for i, ch in enumerate(chars)} itos = {i: ch for i, ch in enumerate(chars)} def encode(s): return [stoi[c] for c in s] def decode(ids): return "".join([itos[i] for i in ids])执行encode("the")会得到类似[8, 3, 4]的整数列表，执行decode([8, 3, 4])会还原成"the"。这里不要直接对字符做计算，因为这个编号只是索引，不代表数值大小。
3.2 Embedding：为什么整数 ID 不能直接进入网络
一个常见误区是：把 Token ID 当作普通数值，直接塞给全连接层。这样做的问题是，ID 1 和 ID 2 本身没有大小关系，整数 ID 的数值差异会误导模型。
正确做法是把 Token ID 映射成稠密向量。在 PyTorch 里，nn.Embedding(vocab_size, d_model)就是一张可学习的查表：
import torch import torch.nn as nn vocab_size = len(chars) d_model = 64 emb = nn.Embedding(vocab_size, d_model) token_ids = torch.tensor(encode("the")) token_vectors = emb(token_ids) # shape: (3, d_model)每个 Token ID 对应一行向量。模型训练时，这些向量会不断更新，最终学到具有一定语义的分布式表示。
3.3 构造序列数据：输入与标签错开一位
语言模型的任务是预测下一个 Token。给定一段字符序列，我们需要把数据切成很多长度为seq_len的小窗口，每个窗口的输入是前seq_len个字符，标签是向右移动一位的后seq_len个字符。
from torch.utils.data import Dataset, DataLoader seq_len = 8 data = torch.tensor(encode(text), dtype=torch.long) class CharSeqDataset(Dataset): def __init__(self, data, seq_len): self.data = data self.seq_len = seq_len def __len__(self): return len(self.data) - self.seq_len def __getitem__(self, idx): x = self.data[idx: idx + self.seq_len] y = self.data[idx + 1: idx + self.seq_len + 1] return x, y train_dataset = CharSeqDataset(data, seq_len) train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)因为所有样本长度都等于seq_len，所以 DataLoader 的默认collate_fn就能自动堆叠成(batch_size, seq_len)的批量张量。
3.4 批次长度不一致时：Padding 和 Attention Mask
上面的数据集固定长度，不需要 Padding。但真实项目中文本长度往往不一致，常见做法是把短序列补到当前 batch 的最大长度，再生成一个 Mask 告诉注意力机制：哪些位置是无效的。
def pad_sequences(sequences, pad_idx=0): max_len = max(len(s) for s in sequences) padded = torch.full((len(sequences), max_len), pad_idx, dtype=torch.long) for i, seq in enumerate(sequences): padded[i, :len(seq)] = torch.tensor(seq, dtype=torch.long) return padded # 有效位置为 True，Padding 位置为 False padding_mask = (padded != pad_idx).unsqueeze(1).unsqueeze(2)padding_mask的形状是(batch_size, 1, 1, seq_len)，后面可以直接广播到注意力分数矩阵上。这一点在实现多头注意力时会用到。
4. 位置编码：让自注意力感知先后顺序
4.1 自注意力天然是集合操作，而不是序列操作
如果把(batch_size, seq_len, d_model)的序列输入到自注意力层，把 Token 之间互换位置，注意力分数会完全相同。这是因为自注意力对每个 Token 做的是“两两交互”，没有默认的顺序概念。
对语言来说，顺序极其重要。"we eat chicken"和"chicken eat we"的语义完全不同。因此模型必须额外把“位置信息”注入到每个 Token 的向量里，这就是位置编码。
4.2 正余弦位置编码的公式和实现
论文《Attention Is All You Need》使用了一组正余弦函数：
[ PE(pos, 2i) = \sin(pos / 10000^{2i/d_model}) ]
[ PE(pos, 2i+1) = \cos(pos / 10000^{2i/d_model}) ]
其中pos是序列位置，i是向量维度索引。用不同频率的正余弦波，可以让模型在向量空间中区分不同位置。
对应 PyTorch 实现如下：
import math def sinusoidal_position_encoding(seq_len, d_model): pe = torch.zeros(seq_len, d_model) position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp( torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model) ) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) return pe.unsqueeze(0) # shape: (1, seq_len, d_model)注意这里要求d_model是偶数，因为偶数和奇数维度各占一半。实际模型中d_model几乎都是偶数。
4.3 可学习位置编码与正余弦位置编码怎么选
正余弦位置编码是固定生成的，不参与训练；另一种常见方案是让位置编码参与训练，例如把位置向量定义为可学习的nn.Parameter。
| 方案 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| 正余弦位置编码 | 不增加训练参数，可外推到更长序列 | 表达方式固定，可能不够灵活 | 教学演示、经典 Transformer |
| 可学习位置编码 | 训练时能自适应数据 | 无法自然外推到训练长度之外 | BERT、GPT 系列早期实现 |
| RoPE 等相对位置编码 | 泛化能力强，适合长文本 | 实现稍复杂 | 当前主流大模型 |
本文用正余弦位置编码，原因是最简单、最容易理解。理解之后，替换成可学习位置编码只需要几行代码。
4.4 位置编码的常见误区
第一个误区是忘记截断。模型初始化时通常会生成max_len的位置编码，但实际输入可能只有seq_len，需要切片：self.pe[:, :x.size(1)]。
第二个误区是让位置编码变成训练参数。固定编码应该用register_buffer注册，而不是nn.Parameter，否则model.parameters()会认为它在参与训练，增加不必要的优化压力。
第三个误区是用float位置编码去加int型 Embedding。Embedding 输出的已经是浮点张量，直接相加即可。
5. 手写多头自注意力：QKV 一次讲透
5.1 Query、Key、Value 的类比和技术含义
把自注意力想象成一个检索过程：你在图书馆里找一本书，手里拿着纸条，纸条上的问题就是 Query；书架上的书名都是 Key；你要找的具体内容是 Value。模型对所有 Key 计算与 Query 的相关性，再把 Value 按相关性加权汇总。
在 Transformer 中，每个 Token 有三个向量：
- Query：当前 Token 想从其他 Token 那里获取什么信息。
- Key：当前 Token 能提供什么信息，用来和 Query 匹配。
- Value：当前 Token 真正携带的信息内容。
这些向量都来自输入x经过不同的线性变换得到，即Q = xW_q，K = xW_k，V = xW_v。
5.2 缩放点积注意力公式和最小代码
缩放点积注意力的计算过程可以用公式概括：
[ Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V ]
其中d_k是每个注意力头的向量维度。除以sqrt(d_k)是为了防止点积过大导致 softmax 进入饱和区，梯度消失。
import torch.nn.functional as F def scaled_dot_product_attention(q, k, v, mask=None): d_k = q.size(-1) scores = q @ k.transpose(-2, -1) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, float("-inf")) attn = torch.softmax(scores, dim=-1) return attn @ v这段代码里的q @ k.transpose(-2, -1)是在计算所有 Token 两两之间的相关分数。masked_fill会把无效位置替换成负无穷，这样 softmax 之后权重为 0。
5.3 多头注意力：为什么 head 要分开
单个注意力头只能学到一种“关注模式”。多头注意力把向量切分成多个子空间，每个头可以关注不同关系：
- 一个头关注语法依赖。
- 一个头关注相邻 Token。
- 一个头关注远距离指代。
实现时需要把d_model维向量切分成n_head份，每份维度是d_k = d_model / n_head。下面是完整实现：
class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_head, dropout=0.1): super().__init__() assert d_model % n_head == 0 self.d_model = d_model self.n_head = n_head self.d_k = d_model // n_head 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_size, seq_len, _ = x.shape q = self.w_q(x).view(batch_size, seq_len, self.n_head, self.d_k).transpose(1, 2) k = self.w_k(x).view(batch_size, seq_len, self.n_head, self.d_k).transpose(1, 2) v = self.w_v(x).view(batch_size, seq_len, self.n_head, self.d_k).transpose(1, 2) scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k) if mask is not None: if mask.dim() == 2: mask = mask.unsqueeze(0).unsqueeze(1) elif mask.dim() == 3: mask = mask.unsqueeze(1) scores = scores.masked_fill(mask == 0, float("-inf")) attn = torch.softmax(scores, dim=-1) attn = self.dropout(attn) context = attn @ v context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) return self.out_proj(context)关键点在于view和transpose。x本来形状是(batch, seq_len, d_model)，先view成(batch, seq_len, n_head, d_k)，再transpose(1, 2)变成(batch, n_head, seq_len, d_k)。这样每个头就能独立计算注意力。
5.4 因果遮罩：生成式模型为什么只能看过去
生成式语言模型必须满足一个约束：预测第t个位置时，不能看到第t+1及之后的位置。否则训练时模型“偷看”了正确答案，生成时又无法偷看，训练和推理就不一致。
因果遮罩可以用一个下三角矩阵表示：
def causal_mask(seq_len): return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))当mask[i][j]为True时，表示位置i可以关注位置j；为False时，表示要屏蔽。在MultiHeadAttention中，这个二维矩阵会广播到(batch, n_head, seq_len, seq_len)，所以不需要每个样本都存一份。
5.5 QKV 形状变化速查表
以下表格以输入x形状为(batch_size, seq_len, d_model)为例：
| 步骤 | 操作 | 形状 |
|---|---|---|
| 输入 | Token Embedding + 位置编码 | (batch, seq_len, d_model) |
| 生成 Q/K/V | Linear 投射 | (batch, seq_len, d_model) |
| 切分多头 | view + transpose | (batch, n_head, seq_len, d_k) |
| 计算注意力分数 | Q @ K^T | (batch, n_head, seq_len, seq_len) |
| 加权 Value | attn @ V | (batch, n_head, seq_len, d_k) |
| 合并多头 | transpose + view | (batch, seq_len, d_model) |
| 输出投影 | Linear | (batch, seq_len, d_model) |
这个表格是调试维度问题最直接的参考。任何一步形状不匹配，都能顺着它找到原因。
6. 组装一个可训练的最小 Transformer Block
6.1 Block 的标准组成：注意力、残差、LayerNorm、FFN
单独的注意力只能做“信息交换”，还不能完成复杂的非线性变换。Transformer 的每个 Block 通常包含：
- 多头自注意力。
- 残差连接。
- LayerNorm。
- 前馈网络 FFN。
- 第二个残差连接和 LayerNorm。
关于归一化的位置，有两种常见写法：
- Post-LN：先加残差，再归一化。原始论文使用这种方式，但深层训练容易不稳定。
- Pre-LN：先归一化，再做注意力或 FFN，最后加残差。现代实现更常用，训练更稳定。
本文使用 Pre-LN，原因是对小模型和大模型都更友好。
6.2 用 PyTorch 实现 TransformerBlock
class TransformerBlock(nn.Module): def __init__(self, d_model, n_head, d_ff, dropout=0.1): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = MultiHeadAttention(d_model, n_head, dropout) self.ln2 = nn.LayerNorm(d_model) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), nn.Dropout(dropout), ) def forward(self, x, mask=None): x = x + self.attn(self.ln1(x), mask) x = x + self.ffn(self.ln2(x)) return x注意x + self.attn(...)的残差结构。这样即使中间层效果不好，梯度也能直接流过捷径，避免深层网络退化。
6.3 组装 TinyTransformer
把前面所有模块拼接起来，就是一个最小可用模型：
class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=512): super().__init__() self.register_buffer("pe", sinusoidal_position_encoding(max_len, d_model)) def forward(self, x): return x + self.pe[:, :x.size(1)] class TinyTransformer(nn.Module): def __init__(self, vocab_size, d_model=64, n_head=4, d_ff=128, num_layers=2, max_len=128, dropout=0.1): super().__init__() self.tok_emb = nn.Embedding(vocab_size, d_model) self.pos_enc = PositionalEncoding(d_model, max_len) self.blocks = nn.ModuleList([ TransformerBlock(d_model, n_head, d_ff, dropout) for _ in range(num_layers) ]) self.ln_f = nn.LayerNorm(d_model) self.head = nn.Linear(d_model, vocab_size) self._init_weights() def _init_weights(self): for m in self.modules(): if isinstance(m, nn.Linear): nn.init.normal_(m.weight, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, std=0.02) def forward(self, idx, mask=None): x = self.tok_emb(idx) x = self.pos_enc(x) for block in self.blocks: x = block(x, mask) x = self.ln_f(x) logits = self.head(x) return logits_init_weights的作用是让参数初始值更稳。Transformer 对初始化比较敏感，直接用默认初始化也能跑，但固定标准差为0.02是常见实践。
6.4 参数量估算
以vocab_size=30, d_model=64, n_head=4, d_ff=128, num_layers=2为例：
| 模块 | 计算方式 | 大致参数量 |
|---|---|---|
| Token |
