Transformer 是当前自然语言处理领域最核心的模型结构但很多人学的时候只停留在调包阶段真要用 PyTorch 从零写一遍才发现自注意力、QKV、位置编码这些概念是断开的每个模块都知道大概合在一起就不知道怎么接。这篇文章不绕弯子直接按一条完整实操路线走从 token 化开始把嵌入层、位置编码、自注意力、多头注意力、残差连接、前馈网络和训练循环全部手写最后跑一个最小可用的文本生成 demo。适合两类读者一类是 PyTorch 用过但没完整写过 Transformer 的人另一类是公式看了不少、但不知道张量在每个阶段长什么样的人。最值得关注的不是某一行代码而是整条数据流怎么从 token 一路走到损失函数。下面按落地顺序拆一遍。1. 先想清楚手撕 Transformer 要撕的到底是什么1.1 你至少需要什么基础“手撕”不是从零造数学理论也不是复刻 GPT-4。一个最小可用的 Transformer核心就是把下面这几件事做出来输入是一段 token 序列。模型先把每个 token 变成向量。通过自注意力让每个位置能看到其他位置的信息。用位置编码补上顺序关系。用前馈网络做非线性变换。最后输出下一个 token 的概率分布。这段流程拆开后非常固定。我认为新手最怕的不是公式而是不知道张量在每一步变成什么形状。所以这篇文章会一直强调形状变化代码里也尽量把 shape 写清楚。基础要求其实不高会用 Python会写函数和循环用过 nn.Module 和 tensor 基本运算知道矩阵乘法。线性代数不需要懂得很深能把(batch, seq_len, d_model)这类形状在脑子里过一遍就够。1.2 一条主线串起所有组件先记住整条数据流原始文本 → token 序列 → embedding 向量 → 位置编码 → 自注意力(QKV) → 多头拼接 → 残差/层归一化 → 前馈网络 → 线性输出 → softmax后面每个章节都会围绕这条主线展开。代码从零写不调用 PyTorch 自带的 Transformer 封装这样才能真正看清内部结构。1.3 每个阶段张量长什么样我先给出一张形状对照表后面写代码时会反复对应回来张量形状含义输入 x(batch, seq_len)token ID 序列嵌入后(batch, seq_len, d_model)每个 token 的向量表示Q / K / V(batch, seq_len, d_head)注意力的三个投影注意力分数(batch, seq_len, seq_len)两两位置之间的相关程度注意力输出(batch, seq_len, d_head)按权重聚合后的结果模型最终输出(batch, seq_len, vocab_size)每个位置下一个 token 的分布这张表建议收藏。排错的时候第一反应应该是现在这步输出形状是什么下一步想要的是什么形状。2. 环境准备与最小数据管线先让 Token 真正跑起来2.1 PyTorch 环境怎么选先确认机器有没有可用 GPU。我第一次跑这种手写 demo 时其实是在 CPU 上完成的完全能跑。先执行这一句python -c import torch; print(torch.cuda.is_available())输出 True 说明 GPU 可用False 也没关系。小 demo 用 CPU 就能跑只是要把 d_model、层数、seq_len 调小一点。安装 PyTorch 建议直接去官方站点选对应系统的命令不要凭记忆敲。Windows 下最容易踩的坑是 CUDA 版本不匹配。判断标准先看nvidia-smi里显示的 CUDA 版本再选不超过它的 PyTorch CUDA 版本。CPU 环境只装 CPU 版就够了硬装 CUDA 版反而会引入一堆 driver 报错。2.2 从文本到 Token 序列这篇文章的 demo 用字符级 tokenizer目的就是少踩 OOV未登录词的坑。词表小、实现简单几行代码就能完成。import torch import torch.nn as nn import torch.nn.functional as F # 原始语料 text hello world this is a tiny transformer demo # 建立字符级词表 chars sorted(set(text)) stoi {c: i for i, c in enumerate(chars)} # 字符 - 索引 itos {i: c for c, i in stoi.items()} # 索引 - 字符 vocab_size len(chars) # 整段文本变成 token 序列 ids [stoi[c] for c in text] print(ids[:20]) print(vocab_size:, vocab_size)为什么用字符级因为词表通常只有几十个字符训练时不会出现未知词。缺点也很明显模型要同时学拼写和语法。但我们的目标是理解 Transformer 结构不是拿这个模型去打榜所以字符级是最合适的选择。如果你以后换到真实任务可以把这段替换成 BPE 或 SentencePiece但外层结构完全不用动。2.3 训练样本怎么切Transformer 训练不能把整段文本一次丢进去要把长文本切成固定长度的输入块。我一般切完后让输入和标签错开一位这样每个位置都在预测下一个 token。def get_batch(ids, batch_size16, seq_len16): starts torch.randint(0, len(ids) - seq_len - 1, (batch_size,)) x torch.stack([torch.tensor(ids[i:iseq_len]) for i in starts]) y torch.stack([torch.tensor(ids[i1:iseq_len1]) for i in starts]) return x, y这里 x 是输入y 是目标输出。x 的第 t 个 token 预测的是 y 的第 t 个 token也就是原文本中的下一个字符。这个“错位一位”的操作是整个语言模型训练的核心逻辑。3. 自注意力与 QKV别背公式直接看张量形状3.1 Q、K、V 从哪里来自注意力的输入是同一个序列 x但会通过三个独立的线性层投影成 Q、K、V。为什么不能直接用 x 算相似度因为模型需要根据任务学出“用哪些信息去查询、按哪些特征匹配”。如果直接用 x每个位置的查询和键完全一样模型就没有自由度了。class SelfAttention(nn.Module): def __init__(self, d_model, d_head, dropout0.1): super().__init__() self.d_head d_head self.q nn.Linear(d_model, d_head) self.k nn.Linear(d_model, d_head) self.v nn.Linear(d_model, d_head) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): Q self.q(x) # (batch, seq_len, d_head) K self.k(x) # (batch, seq_len, d_head) V self.v(x) # (batch, seq_len, d_head) scores Q K.transpose(-2, -1) / (self.d_head ** 0.5) # (batch, seq_len, seq_len) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) weights torch.softmax(scores, dim-1) weights self.dropout(weights) out weights V # (batch, seq_len, d_head) return out这 30 行不到的代码就是单头自注意力的全部。关键动作就两个Q 和 K 做点积得到两两位置的相关度。softmax 把相关度变成权重再用权重去加权 V。Q 决定“我在找什么”K 决定“我是什么内容”V 决定“我最终提供什么信息”。三者分工明确。3.2 为什么除根号 d_head如果不除以根号 d_headd_head 越大点积的数值会越大softmax 之后容易进入饱和区梯度变得非常小。除一个根号 d_head 是为了让分数方差保持稳定。这行代码很小的优化但对训练稳定性影响很大。很多新手直接照抄公式从不调这一项结果模型训练时 loss 不稳定然后到处怀疑数据问题。3.3 因果掩码为什么训练 decoder 要遮住未来语言模型生成时只能看当前时刻之前的内容不能偷看未来 token。训练阶段也必须模拟这个限制否则模型会学到“作弊”模式。实现方式叫因果掩码。def causal_mask(seq_len): return torch.tril(torch.ones(seq_len, seq_len, dtypetorch.bool))这个函数生成一个下三角矩阵当前位置允许看到自己和之前的位置不能看到之后的位置。之后在注意力分数计算时把这个掩码传进去mask causal_mask(seq_len).unsqueeze(0) # (1, seq_len, seq_len) scores scores.masked_fill(mask 0, float(-inf))把未来位置对应的分数变成负无穷softmax 之后权重就趋近于 0。这样每个位置只能聚合到当前位置和之前位置的信息。这里要特别注意掩码的形状。如果 batch size 不为 1通常要扩展成(batch, 1, seq_len, seq_len)或者(1, seq_len, seq_len)让它广播千万不要直接拿(seq_len, seq_len)去和(batch, seq_len, seq_len)的 scores 做 masked_fill会报形状错误。4. 位置编码与多头注意力补上顺序和多视角4.1 为什么注意力本身没有顺序概念自注意力是集合运算。把 token 序列乱序打乱如果不加位置信息输出排列也不会按语义变化。这对语言任务来说是致命的I love you和you love I在集合意义上看起来一样但语义完全不同。所以必须在 embedding 阶段把位置信息加进去。最直接的方式是准备一个同样维度的位置向量和 token embedding 相加。4.2 正弦编码还是可学习编码两种方案都常见可学习位置编码用nn.Embedding(max_len, d_model)实现简单短序列任务表现稳定。正弦位置编码用固定公式生成不依赖训练理论上对任意长度都能给出位置向量外推能力更强。小 demo 我建议直接用可学习位置编码代码少不容易错。如果以后要做长度外推或者长文本场景再看正弦编码。def sinusoidal_encoding(max_len, d_model): pe torch.zeros(max_len, d_model) pos torch.arange(max_len).unsqueeze(1).float() div torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model)) pe[:, 0::2] torch.sin(pos * div) pe[:, 1::2] torch.cos(pos * div) return pe.unsqueeze(0) # (1, max_len, d_model)这个公式本身不需要背但要理解它的作用偶数维用 sin奇数维用 cos不同频率组合起来让模型可以区分不同位置。4.3 多头注意力多个子空间并行单头注意力只能关注一种关系。多头注意力的思路是用多个独立的注意力头每个头在不同的特征子空间里计算相关性最后把所有头的结果拼起来再做一次线性变换。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads 0, d_model 必须能被 num_heads 整除 self.num_heads num_heads self.d_head d_model // num_heads self.heads nn.ModuleList( [SelfAttention(d_model, self.d_head) for _ in range(num_heads)] ) self.proj nn.Linear(d_model, d_model) def forward(self, x, maskNone): heads [h(x, mask) for h in self.heads] out torch.cat(heads, dim-1) # (batch, seq_len, d_model) return self.proj(out)为什么需要多个头因为一个头只能学一种权重分布。有些头可能负责相邻词搭配有些头负责长距离指代有些头负责某个具体句式。多个头相当于多个视角同时看问题。实际实现里主流做法是把 Q、K、V 一次投影到整个 d_model再 reshape 拆成多个头。我这里为了可读性用了 ModuleList每个头独立投影。代码直观一些不过计算效率不如高端写法。理解阶段不需要过度优化。4.4 常见错误d_model 不能被 num_heads 整除每个头的维度是 d_model / num_heads所以 d_model 必须能被 num_heads 整除。新手最容易这里报错。我一般开头加一行 assert至少报错信息清晰。5. 完整 Block残差、层归一化、前馈网络5.1 一个 Transformer Block 的标准结构到了这一节前面写好的模块就开始组合了。一个 decoder-only 的 Transformer Block 通常包含LayerNorm 归一化。多头自注意力。残差连接。LayerNorm 归一化。前馈网络FFN。残差连接。class TransformerBlock(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout0.1): super().__init__() self.attn MultiHeadAttention(d_model, num_heads) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) self.ln1 nn.LayerNorm(d_model) self.ln2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): x x self.dropout(self.attn(self.ln1(x), mask)) x x self.dropout(self.ffn(self.ln2(x))) return x这段代码几乎是所有 Transformer 大模型的骨架。GPT 系列、Llama 系列的 block 虽然细节有改动但主结构都能看到这个影子。5.2 为什么先 LayerNorm 再进注意力我写的是 Pre-LN先归一化再计算注意力。这和原始 Transformer 论文里的 Post-LN 顺序不同。原始论文是先注意力、后归一化训练起来对学习率更敏感调参稍不注意就发散。Pre-LN 的好处是训练更稳定适合新手。你可以简单理解成每层输入先被拉回一个正常分布再进注意力或前馈网络这样梯度不容易爆炸。5.3 前馈网络不是可有可无注意力本质是 token 之间的信息交换和加权求和整体计算偏线性。前馈网络提供逐位置的非线性变换让模型表达能力更强。d_ff 通常取 d_model 的 2 到 4 倍。我这个小 demo 里 d_model64d_ff128已经够用。你如果调大 d_ff模型容量会上升但训练时间和显存占用也会明显增加。6. 从零训练一个小模型文本生成 Demo6.1 完整模型定义把前面的模块拼起来就是一个极小的 decoder-only Transformerclass TinyTransformer(nn.Module): def __init__(self, vocab_size, d_model64, num_heads4, d_ff128, num_layers2, max_len128, dropout0.1): super().__init__() self.token_embedding nn.Embedding(vocab_size, d_model) self.position_embedding nn.Embedding(max_len, d_model) self.blocks nn.ModuleList([ TransformerBlock(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.ln nn.LayerNorm(d_model) self.head nn.Linear(d_model, vocab_size) def forward(self, x, maskNone): B, T x.shape pos torch.arange(T, devicex.device).unsqueeze(0) x self.token_embedding(x) self.position_embedding(pos) for block in self.blocks: x block(x, mask) return self.head(self.ln(x))这里的位置编码用nn.Embedding做的可学习位置编码。每次 forward 时根据当前序列长度 T 生成位置索引。6.2 训练循环怎么写训练用交叉熵损失。注意 logits 和目标形状要对应logits 是(batch, seq_len, vocab_size)目标 y 是(batch, seq_len)。CrossEntropyLoss 默认输入是(N, C)所以要把前两维展平。model TinyTransformer(vocab_size) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) loss_fn nn.CrossEntropyLoss() for step in range(500): x, y get_batch(ids, batch_size16, seq_len16) mask causal_mask(16) logits model(x, mask) loss loss_fn(logits.reshape(-1, vocab_size), y.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if step % 100 0: print(fstep {step}: loss {loss.item():.4f})这个循环里新手最常忽略的是optimizer.zero_grad()。如果不每次清零梯度会在多个 step 之间累加loss 会表现得非常诡异。6.3 学习率怎么调1e-3 对这个 5 万参数级别的小模型通常能跑。如果 loss 发散优先降到 3e-4 或 1e-4。不要一上来就加 batch size先把学习率降下来往往更有效。这里也提醒一句小数据任务不要开大步长训练否则模型很快过拟合到语料上loss 降了但生成时会一直重复。6.4 生成阶段采样而不是贪心训练完成后生成的基本思路是把已经生成的 token 放回模型预测下一个 token然后拼回去继续预测。def generate(model, start_text, max_new50, temperature0.8): model.eval() idx torch.tensor([[stoi[c] for c in start_text]]) with torch.no_grad(): for _ in range(max_new): x idx[:, -32:] mask causal_mask(x.size(1)) logits model(x, mask)[0, -1] probs torch.softmax(logits / temperature, dim-1) next_id torch.multinomial(probs, 1).item() idx torch.cat([idx, torch.tensor([[next_id]])], dim1) return .join(itos[i] for i in idx[0].tolist())temperature 控制随机性。温度越低输出越确定但容易重复温度越高输出越随机。小模型上我一般先试 0.6 到 0.9。直接选概率最大的 token贪心也不是不行但输出会非常单调。print(generate(model, start_texthello))6.5 怎样算训练成功不要指望这个迷你模型生成通顺长句。判断标准很简单loss 明显下降不振荡过大生成的字符串里开始出现训练语料的高频字符组合。比如语料里有 hello模型能输出 h、e、l 这类常见顺序就说明已经学到东西了。我的建议是先把单条训练 pass 跑通再去看生成质量。输出不对的时候不要急着改模型结构先确认 loss 有没有降、采样代码有没有问题。7. 常见报错与排查顺序新手最容易卡在哪7.1 形状不匹配先打印 shape最典型的报错是mat1 and mat2 shapes cannot be multiplied。原因通常是某个 Linear 层输入维度不是期望的 d_model。排查顺序很固定先看报错发生在哪一行。在那一层之前加一行print(x.shape)。对照第三节的张量形状表确认当前输出和下一层期望输入是否一致。我几乎每次都靠打印 shape 解决问题。模型结构越复杂越不要靠肉眼猜。常见报错对应关系现象常见原因排查方向mat1 and mat2 形状不匹配某层输入维度不对打印每层入口 shapemask 广播失败掩码形状不是 (1, T, T)打印 mask.shape必要时加 batch 维度device mismatch模型在 GPU、输入在 CPU模型和输入统一 .to(device)loss 变成 nan学习率太大或输入异常降低学习率检查 tokenizer7.2 训练不收敛先看数据再改参数loss 一直不降不要急着认为模型写得有问题。优先看三件事输入 token 序列是否为空词表是否建对。输入和标签是否错位正确。如果 x 和 y 完全相同模型就是在学复制粘贴loss 不会按预期下降。学习率是不是太大导致梯度震荡。loss 变成 nan最常见原因是学习率太大。小模型直接用 1e-3 一般没事但如果数据量很小还是降到 3e-4 更稳。7.3 显存、内存和任务规模怎么匹配低配置能跑通 demo不代表适合批量训练。不同任务规模的参数选择可以参考这张表任务类型d_modellayers数据规模预期学习 demo32-641-2几 KBCPU 分钟级真实小语料128-2564-6MB 级CPU 较慢建议 GPU大规模训练512 以上6-12GB 级需要多卡和分布式优化如果显存不够不要无脑加 batch size。优先缩小 seq_len或者减小 d_model。很多时候 seq_len 对显存的影响比 batch size 更明显因为注意力矩阵是 seq_len 的平方。7.4 生成重复调整温度而不是改模型小模型生成时容易陷入重复循环。这时候不要急着增大网络先试以下几种方式降低 temperature让采样更接近高概率 token。生成时限制最大长度避免无限循环。检查训练 loss 是不是已经降到很低如果过拟合会加剧重复。如果你希望后续做更深入的调试还可以在注意力层把 weights 返回出来画一个注意力热力图看看模型到底在关注哪些位置。新手阶段不用急先把主链路跑通。如果把这套代码完整跑一遍你会发现 Transformer 真正难的不是某个公式而是数据、形状、训练和生成之间的衔接。我个人更建议先把单条训练 pass 跑稳再调参数。每一步都确认张量形状和 loss 变化比一次性把所有组件堆完要快得多。下次再有人问“能不能手撕 Transformer”你可以直接把这篇文章发给他让 TA 从 token 开始跑一遍自己的小模型。