Transformer 是当前大语言模型和许多深度学习任务的核心架构但很多人学它时卡在“概念听懂了、代码写不出来”这一步。网上讲解自注意力、QKV、位置编码的文章很多真正能让人从 Token 开始一步步把模型写出来、跑起来、看到损失下降的材料却不多。这篇文章就用 PyTorch 手写一个极小的 Transformer 字符级语言模型覆盖 Token 化、Embedding、位置编码、多头自注意力、Transformer Block、训练循环和生成验证。学完以后你既能理解 QKV 和自注意力的计算过程也能在自己机器上跑通一个可复现的最小项目后续再去看大模型源码会轻松很多。1. 先拆清 Transformer 的输入输出再动手写代码1.1 Token 是什么为什么要先把文本变成数字在 Transformer 眼里文本不是文字而是一串离散编号。这里的“离散编号”就是 Token ID。Token 可以是一个字符、一个词语、一个子词不同分词策略对应不同的词表大小。举例来说字符串abc如果按字符切分可以切成[a, b, c]三个 Token如果按词语切分可能只有一个 Tokenabc。无论哪种切法最终都要通过一个词表把 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.LayerNormFFN 前馈网络对每个位置做非线性变换Linear GELU Linear输出投射层把向量映射回词表维度nn.Linear很多人一上来就写完整代码结果每一行都似懂非懂。建议按照模块逐个实现最后再组装成完整模型。1.4 学习环境和生产环境的要求分别是什么零基础手写 Transformer 时目标是“跑通逻辑”不是“训练大模型”。学习和生产的要求差异很大维度学习环境生产环境数据规模几十个字符即可几十 GB 甚至更大模型规模d_model64num_layers2几亿到几千亿参数GPU没有也能跑需要多卡训练分词字符级即可BPE、SentencePiece训练技巧最简单交叉熵混合精度、分布式、学习率调度部署无需考虑推理优化、量化、服务化这个表格不是让大家跳过生产问题而是要区分“理解原理”和“工程落地”。本文先完成理解原理这一步。2. 准备 PyTorch 环境并搭好项目骨架2.1 用 Anaconda 创建 PyTorch 环境推荐用 Anaconda 管理环境避免污染系统 Python。下面是创建一个名为transformer的 Python 3.10 环境并安装 CPU 版 PyTorch 的命令conda create -n transformer python3.10 conda activate transformer conda install pytorch torchvision torchaudio cpuonly -c pytorch如果机器有 NVIDIA GPU并且已经安装好 CUDA 驱动可以去 PyTorch 官网选择对应的安装命令。常见写法如下但具体版本以官网为准conda install pytorch torchvision torchaudio pytorch-cuda12.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.11PyTorch深度学习框架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), dtypetorch.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_size16, shuffleTrue)因为所有样本长度都等于seq_len所以 DataLoader 的默认collate_fn就能自动堆叠成(batch_size, seq_len)的批量张量。3.4 批次长度不一致时Padding 和 Attention Mask上面的数据集固定长度不需要 Padding。但真实项目中文本长度往往不一致常见做法是把短序列补到当前 batch 的最大长度再生成一个 Mask 告诉注意力机制哪些位置是无效的。def pad_sequences(sequences, pad_idx0): max_len max(len(s) for s in sequences) padded torch.full((len(sequences), max_len), pad_idx, dtypetorch.long) for i, seq in enumerate(sequences): padded[i, :len(seq)] torch.tensor(seq, dtypetorch.long) return padded # 有效位置为 TruePadding 位置为 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, 2i1) \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, dtypetorch.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_qK xW_kV 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, maskNone): 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, dropout0.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, maskNone): 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个位置时不能看到第t1及之后的位置。否则训练时模型“偷看”了正确答案生成时又无法偷看训练和推理就不一致。因果遮罩可以用一个下三角矩阵表示def causal_mask(seq_len): return torch.tril(torch.ones(seq_len, seq_len, dtypetorch.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/VLinear 投射(batch, seq_len, d_model)切分多头view transpose(batch, n_head, seq_len, d_k)计算注意力分数Q K^T(batch, n_head, seq_len, seq_len)加权 Valueattn V(batch, n_head, seq_len, d_k)合并多头transpose view(batch, seq_len, d_model)输出投影Linear(batch, seq_len, d_model)这个表格是调试维度问题最直接的参考。任何一步形状不匹配都能顺着它找到原因。6. 组装一个可训练的最小 Transformer Block6.1 Block 的标准组成注意力、残差、LayerNorm、FFN单独的注意力只能做“信息交换”还不能完成复杂的非线性变换。Transformer 的每个 Block 通常包含多头自注意力。残差连接。LayerNorm。前馈网络 FFN。第二个残差连接和 LayerNorm。关于归一化的位置有两种常见写法Post-LN先加残差再归一化。原始论文使用这种方式但深层训练容易不稳定。Pre-LN先归一化再做注意力或 FFN最后加残差。现代实现更常用训练更稳定。本文使用 Pre-LN原因是对小模型和大模型都更友好。6.2 用 PyTorch 实现 TransformerBlockclass TransformerBlock(nn.Module): def __init__(self, d_model, n_head, d_ff, dropout0.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, maskNone): 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_len512): 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_model64, n_head4, d_ff128, num_layers2, max_len128, dropout0.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, std0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, std0.02) def forward(self, idx, maskNone): 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_size30, d_model64, n_head4, d_ff128, num_layers2为例模块计算方式大致参数量Token