Transformer从原理到代码实现：LLM训练推理与故障排查-开发者社区

Transformer 这套结构从 2017 年被提出来到现在，已经成了 LLM 绕不过去的地基。我带过几个刚入行的同学，他们能把 self-attention 的公式背得一字不差，但真让他们解释为什么要除以根号 d_k、为什么多头不是简单地把一个大头拆成几个小头、训练到一半 loss 突然飞了该从哪里查，往往就卡壳了。公式会背和真能跑起来之间，隔着大量只有亲手写过、调过、炸过的人才清楚的细节。这篇长文我打算把 Transformer 从结构原理到代码实现、再到训练推理阶段的坑，按我自己踩过的顺序完整捋一遍，重点讲清楚每个设计背后的“为什么”，而不是复述论文。

内容会覆盖自注意力的计算本质、多头机制的真正作用、位置编码的几种主流方案取舍、残差与归一化的位置之争、从零手写一个能跑通的 Block、显存与算力的估算方法、KV Cache 的实现思路，以及训练推理阶段最常见的几类故障排查。适合已经看过结构图但没真正动手实现过的读者，也适合实现过但训练总出问题的读者。我不假设你数学功底多强，但假设你至少写过 PyTorch，知道张量维度是怎么回事。

1. 为什么 Transformer 能取代 RNN 成为 LLM 的地基

1.1 从循环结构的两个硬伤说起

在 Transformer 出现之前，处理序列的主流方案是 RNN 及其变体 LSTM、GRU。这类结构按时间步展开，第 t 个位置的隐藏状态依赖第 t-1 个位置，天然带着顺序性。问题也正出在这个顺序性上。第一个硬伤是没法并行：训练时一个长度为 512 的序列要串行算 512 步，GPU 上那几千个核心大部分时间在等前一步的结果，利用率低得可怜。第二个硬伤是长依赖衰减：信息要从位置 1 传到位置 500，中间经过几百次非线性变换，梯度在反向传播时要么消失要么爆炸，LSTM 的门控机制只是缓解，没有根治。

我早期做过一个文本分类的小项目，序列长度大概 300 左右，LSTM 训练一个 epoch 要将近四十分钟，换到 Transformer 之后同样的数据量降到了七八分钟。这个差距不是模型本身算得快，而是并行度完全不同——整个序列一次性喂进去，注意力矩阵一次矩阵乘法就能算完，GPU 的吞吐被吃满了。这也是为什么后来所有的大模型都建立在 Transformer 之上，规模越大，并行带来的收益越明显。

1.2 自注意力到底在算什么

自注意力的核心动作可以用一句话概括：序列里每个位置都去“看”其他所有位置，根据相关程度加权汇总信息。具体到计算，输入序列的每个 token 向量会通过三个独立的线性变换投影成 Query、Key、Value 三组向量。Query 代表“我在找什么”，Key 代表“我能提供什么”，Value 代表“我实际携带的信息”。然后用 Query 和所有 Key 做点积得到相关性分数，softmax 归一化成权重，再对 Value 加权求和。

这个过程里最关键的一点是：任意两个位置之间的路径长度是 1。位置 1 和位置 500 之间不需要经过中间的 499 步，直接在注意力矩阵里就对应一个元素。这就是长依赖问题被根治的原因——梯度回传的路径短了，衰减自然就弱了。当然代价是计算量：注意力矩阵的大小是序列长度的平方，序列翻倍，计算量和显存都涨四倍。这个平方复杂度后来衍生出一大堆稀疏注意力、线性注意力的研究，但对大多数中等长度场景，标准注意力仍然是性价比最高的选择。

1.3 一个常被忽略的取舍：归纳偏置的减少

RNN 自带“顺序”这个强归纳偏置，CNN 自带“局部性”和“平移不变性”。Transformer 几乎不带任何结构性的先验，位置信息要靠额外编码注入，局部性也不做假设。这种“少假设”的设计在数据量小的时候是劣势——模型学得慢、容易过拟合；但在数据量极大的时候反而成了优势，模型可以自由地学习任意位置之间的关系，不被人为假设束缚。

这也是为什么 Transformer 先在机器翻译这种数据量还算大的任务上证明自己，然后才随着数据规模的增长在 LLM 上彻底爆发。如果你手上的数据只有几万条，硬上 Transformer 未必比一个调好的 CNN 或 LSTM 效果好，这一点在选型时要有清醒认识。我个人建议是小数据场景先用轻量结构打底，确认数据管线没问题，再换 Transformer 做对比，不然很容易把数据问题误判成模型问题。

2. 拆开结构：核心组件逐个过一遍

2.1 Self-Attention 的计算过程与维度推演

先约定维度。输入张量 X 的形状是[batch, seq_len, d_model]，比如[32, 128, 512]。三个投影矩阵 W_Q、W_K、W_V 的形状都是[d_model, d_k]，其中 d_k 是每个头的维度，通常取 d_model 除以头数。投影之后 Q、K、V 的形状都是[batch, seq_len, d_k]。

接下来算分数矩阵：Q 乘以 K 的转置，得到[batch, seq_len, seq_len]，这个矩阵的第 i 行第 j 列就代表位置 i 对位置 j 的关注程度。这里有个必须做的操作——除以根号 d_k。原因很直接：点积的结果随着 d_k 增大而增大，维度是 64 的时候点积的方差大概是 64 量级，直接送进 softmax 会让分布极度尖锐，几乎全部分到某一个位置上，梯度接近于零，训练直接卡死。除以根号 d_k 相当于把方差拉回到 1 附近，softmax 的输出分布就平缓得多。

然后 softmax 按行归一化，得到注意力权重，再乘以 V，输出形状[batch, seq_len, d_k]。如果要做多头，就把多个头的输出在最后一维拼接回[batch, seq_len, d_model]，再乘一个输出投影 W_O 做一次融合。整个过程的计算量，主要的开销在那次seq_len × seq_len × d_k的矩阵乘法上，序列长度是主导因素。

注意：实现时一定要给 softmax 前的分数加上 padding mask，把补齐位置的分量设成一个极大的负数（比如 -1e9），否则 softmax 会给这些无效位置分到权重，模型会学到一堆噪声。这个坑我见过不止一次，症状是 loss 能降但验证集效果始终差一截。

2.2 Multi-Head 不是简单地把大头拆小头

很多人对多头的理解停留在“把 512 维拆成 8 个 64 维并行算”，觉得跟单头 512 维差不多。实际上两者的表达能力差距很明显。单头注意力的 softmax 只能产生一个关注分布，模型在一个位置上只能突出一种关系。多头则允许每个头学习不同的关注模式——有的头盯相邻词，有的头盯句法依存，有的头盯远距离指代。这些模式在训练中是自发分化出来的，不需要人工指定。

我做过一个粗糙的验证实验：在一个小型语言模型上把 8 个头换成一个 512 维的单头，参数量基本对齐，困惑度差了将近 15%。这个差距在训练初期不明显，到中期才拉开，说明多头带来的不是参数量的优势，而是表示空间的多样性。另外一个实操细节是，头数不是越多越好。头数太多会导致每个头的维度太小，比如 d_model=512、16 个头，每头只有 32 维，表达容量不够，效果反而下降。一般每头维度保持在 64 左右是比较稳妥的经验值。

2.3 位置编码：正余弦、可学习与旋转三种方案对比

自注意力本身对顺序是无感知的——把输入序列打乱，输出只是跟着打乱，内容不变。所以必须额外注入位置信息。最早的方案是正余弦编码，用不同频率的正弦和余弦函数生成每个位置的向量，直接加到词嵌入上。这种方案的好处是能外推到训练时没见过的长度，坏处是表达能力有限，且与词嵌入简单相加会互相干扰。

第二种是可学习的位置嵌入，给每个位置分配一个可训练向量，实现简单，效果在大模型上也不错，缺点是不能外推，超过最大长度就没法处理。第三种是旋转位置编码，它对 Query 和 Key 按位置做旋转操作，把相对位置信息编码进点积里，既保留了相对位置的表达能力，又支持一定程度的外推。目前主流的中长上下文模型基本都转向了这种方案。

方案	外推能力	实现复杂度	常见适用场景
正余弦编码	较强	低	早期模型、教学实现
可学习嵌入	无	低	固定长度训练
旋转位置编码	中等偏强	中	长上下文模型

提示：如果你用的是可学习位置嵌入，做长度外推时不要直接把超长序列截断就完事，截断会丢掉尾部信息。比较稳妥的做法是先在训练时就把最大长度设得比实际需求富裕一些，比如实际最长 1024，就把上限设成 2048，留出余量。

2.4 残差、LayerNorm 与 FFN 的分工

每个 Transformer Block 里除了注意力，还有两个不起眼但极其关键的部件。第一个是残差连接，把子层的输入直接加到输出上。这个设计解决的是深层网络的梯度回传问题——有了残差，梯度可以沿着捷径直接传回去，不用穿过每一层的非线性变换。没有残差，堆到十几层以上训练基本就崩了。

第二个是 LayerNorm。注意它是在特征维度上做归一化，而不是在 batch 维度上，这一点跟 BatchNorm 有本质区别。序列任务里每个样本的长度可能不同，batch 统计量不稳定，LayerNorm 按单个样本的特征做归一化，不受 batch 组成影响，训练和推理行为一致。至于 LayerNorm 放在残差前还是残差后，也就是常说的 Pre-LN 和 Post-LN，实践下来 Pre-LN 更稳，训练不需要太精细的学习率 warmup 也能收敛，Post-LN 理论上表达力略强但对超参敏感，现在主流实现基本都用 Pre-LN。

FFN 也就是前馈网络，结构是两个线性层中间夹一个激活函数，中间维度通常是 d_model 的 4 倍。它的作用是对每个位置独立地做非线性变换，不涉及位置间交互。参数量上 FFN 占了 Block 里的大头，大约三分之二。这里有个容易忽略的点：中间的扩展倍数不是必须为 4，小模型上可以适当调小以减小参数，大模型上也有用门控线性单元替代普通 FFN 的做法，效果通常更好。

3. 从零实现一个能跑通的 Transformer Block

3.1 环境准备与张量维度约定

我建议先用 PyTorch 在一个小任务上把单层 Block 跑通，再往上堆。环境上没什么特殊要求，装好 PyTorch 就行，CPU 也能跑，只是慢。关键是先把维度约定固定下来，避免后面到处是形状不匹配的报错。我的习惯是统一用[batch, seq_len, d_model]这个顺序，所有中间变量都按这个约定推导。

先定义超参：d_model 取 128，头数取 4，每头维度 32，FFN 中间维度取 512，也就是 4 倍。序列长度取 64，batch 取 16，词表大小取 1000。这些数字很小，单卡甚至 CPU 都能在几秒内跑完一个 batch，方便快速验证逻辑。

import torch import torch.nn as nn d_model = 128 n_heads = 4 d_k = d_model // n_heads d_ff = 512 seq_len = 64 batch = 16 vocab = 1000

维度约定清楚之后，每一步的形状变化都可以在纸上推出来，调试的时候对着打印的形状一看就知道哪一步错了。这个习惯看起来笨，但比在几百行代码里瞎找要快得多。

3.2 手写注意力并做数值校验

先实现单头注意力，验证逻辑正确后再扩展成多头。核心就是四次矩阵运算加一次 softmax。写完之后一定要做数值校验：随机生成一个输入，和一个现成的官方实现对比输出，误差应该在 1e-5 以内。这一步不能省，很多人自己写的版本在 mask 处理或者缩放上出了细微错误，训练时表现为收敛慢但不算崩，极难排查。

def attention(q, k, v, mask=None): scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k ** 0.5) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) weights = torch.softmax(scores, dim=-1) return torch.matmul(weights, v), weights

校验的时候除了对比输出，还要把注意力权重矩阵打印出来看一眼。正常的权重每一行加起来应该是 1，而且分布不应该全部集中在一个位置，也不应该均匀得像噪声。如果发现某几行全是 0，大概率是 mask 加错了，把有效位置也屏蔽了。

3.3 组装 Block 与第一个训练任务

把多头注意力、FFN、残差、LayerNorm 按 Pre-LN 的顺序拼起来。Pre-LN 的顺序是：输入先过 LayerNorm，再进注意力，输出加残差；然后再过 LayerNorm，进 FFN，输出再加残差。这个顺序跟原始论文的 Post-LN 是反的，一定要注意别混着写。

组装完之后接一个简单的语言建模头，也就是一个线性层把 d_model 映射回词表大小，用交叉熵损失训练一个“预测下一个 token”的任务。数据可以自己造，比如用一段重复的文本或者简单的序列规律，先确认模型能把这个规律学会，loss 能降到接近 0，再换真实数据。这个“先过拟合一个小样本”的验证方法非常有效，能快速区分是模型结构有问题还是数据或训练配置有问题。

3.4 参数初始化与训练的坑

初始化对 Transformer 的训练稳定性影响很大。默认的初始化方差可能偏大，深层网络容易在训练初期就发散。比较稳妥的做法是对线性层用较小的方差初始化，比如正态分布标准差取 0.02，偏置全部置零。有些实现还会对残差分支的输出投影做额外的缩放，进一步稳定深层训练。

注意：残差分支里的最后一个线性层，如果初始化方差过大，残差累加之后数值会逐层放大，到十几层就溢出成 inf。我遇到过一次 loss 直接变 nan，查了两天才发现是某个自定义层的初始化没处理好。建议每加一层就把中间激活的均值和方差打印出来看，正常应该在 1 附近波动，逐层放大就说明有问题。

4. 训练与推理阶段的实操经验

4.1 显存估算与批大小选择

显存不够是训练时最常见的拦路虎。粗略估算，训练时显存主要花在四块：模型参数、梯度、优化器状态、激活值。用 Adam 优化器、混合精度训练的话，每个参数大约占 16 到 20 字节，激活值另算，跟 batch 大小和序列长度成正比。一个 1 亿参数的模型，光参数和优化器状态就要 1.6G 到 2G，加上激活值，batch 稍微大一点 8G 卡就撑不住了。

实操上我一般先用一个很小的 batch 跑起来，用工具看实际显存占用，然后按比例往上推。如果显存卡在临界点，优先考虑梯度累积——用小 batch 跑几步再更新一次，等效于大 batch，显存占用不变。其次是激活重计算，用计算换显存，代价是训练速度慢百分之二三十。这两个手段组合起来，通常能让原本跑不动的配置跑起来。

4.2 学习率 warmup 与梯度裁剪

Transformer 对学习率比较敏感，尤其是训练初期。直接用大学习率，前几百步很容易发散。标准做法是加 warmup，前若干步学习率从很小线性增长到目标值，之后再按余弦或者反平方根衰减。warmup 步数一般取总步数的百分之几，小数据集上几千步就够，大数据集上可能要上万步。

梯度裁剪也是标配，把梯度的范数裁到一个阈值以内，比如 1.0。这个操作能防止个别异常 batch 把参数带偏。我习惯在裁剪前后都记录一下梯度范数，如果发现频繁触发裁剪，说明学习率可能偏大或者数据里有异常样本，这时候要回头查而不是无脑调大裁剪阈值。

配置项	常用取值	说明
warmup 比例	1% 到 5%	数据越大比例可越小
峰值学习率	1e-4 到 3e-4	小模型可适当调大
梯度裁剪阈值	1.0	频繁触发需排查
权重衰减	0.01 到 0.1	与 LayerNorm 参数分开处理

4.3 KV Cache 是怎么把推理加速的

推理阶段跟训练有个本质区别：生成是自回归的，每次只出一个 token。如果每生成一个 token 都把整个序列重新算一遍注意力，那计算量会随生成长度平方增长。KV Cache 的思路很直接——已经算过的 Key 和 Value 缓存起来，新 token 只需要算自己的 Query，然后跟缓存的 K、V 做注意力。这样每步的计算量从跟序列长度平方相关降到线性相关。

代价是显存。缓存的大小跟层数、头数、序列长度、batch 都成正比。一个 7B 级别的模型，上下文开到几千，batch 开到几十，KV Cache 吃掉十几 G 显存是常事。所以实际部署时 batch 大小往往不是被模型本身限制的，而是被 KV Cache 限制的。后来出现的一些压缩方案，比如把 K、V 量化到低精度、或者分组共享，都是为了缓解这个瓶颈。

提示：做流式生成的时候，一定要确认缓存的更新逻辑是对的——每步追加新的 K、V，位置编码的偏移也要跟着走。我见过缓存没清空导致两轮对话串味的 bug，症状是第二轮回答里混进了第一轮的内容，排查时还以为是模型本身的问题。

4.4 精度选择：FP16 还是 BF16

混合精度训练基本是现在的默认操作，但 FP16 和 BF16 的差别值得说清楚。FP16 的尾数位多、精度高，但表示范围小，训练中梯度稍微大一点就溢出成 inf，需要配合损失缩放来动态调整。BF16 的尾数位少、精度低一些，但表示范围跟 FP32 一样大，基本不会溢出，训练更省心。

我的经验是，新项目直接用 BF16，省去调损失缩放的麻烦，虽然精度略低，但对最终效果影响很小。只有在硬件不支持 BF16、或者对数值精度极其敏感的场景才退回 FP16 加损失缩放。另外要注意，LayerNorm 和 softmax 这些对数值敏感的操作，最好保持在 FP32 下算，只把矩阵乘法放到低精度，这样能避免不少莫名其妙的 nan。

5. 常见问题与排查技巧实录

5.1 损失不降或者剧烈震荡

这是最高频的问题，原因可能有很多，但我一般按固定顺序排查。先看数据管线：把 batch 里的样本打印出来，确认输入和标签是对齐的。标签错位是新手最容易犯的错，症状是 loss 卡在一个比随机猜略好的水平，怎么调都下不去。再看学习率：把学习率降一个数量级重跑，如果 loss 开始稳定下降，那就是学习率太大。

如果数据和超参都没问题，就要怀疑结构实现。常见的隐性错误包括 mask 方向搞反、位置编码没加、LayerNorm 归一化的维度搞错。这几个错误都不会直接报错，只会让模型表现变差。我的做法是写一组单元测试，把每个子模块单独拿出来跟参考实现对比，一次性把这些低级错误排干净。

5.2 注意力权重坍缩

训练一段时间后，如果发现注意力权重几乎全部集中在一个位置，或者反过来均匀得像没训练，那说明注意力机制没起作用。坍缩到单点通常是缩放没做对或者初始化方差太大，导致 softmax 进入饱和区。均匀分布则常见于梯度消失，注意力层的参数几乎没更新。

排查方法是把每一层的注意力权重矩阵可视化出来看，横轴是 Query 位置，纵轴是 Key 位置。健康的模型应该能看到一些有规律的对角线或者块状结构。如果所有层都长得一模一样，说明表示没有分化，可能是层数太少或者训练不充分。这个可视化手段非常直观，比盯着 loss 曲线猜要有效得多。

5.3 位置信息失效

模型能学会词与词的搭配，但对顺序不敏感，典型表现是打乱句子里的词序，输出几乎不变。这基本可以断定是位置编码没生效。检查点包括：位置编码有没有真的加到输入上、可学习的位置嵌入有没有被优化器更新、旋转编码的旋转角度有没有算对。

还有一个隐蔽的情况是位置编码的取值范围跟词嵌入差太多，加完之后被词嵌入完全淹没。这种情况可以在训练初期监控位置编码那部分的梯度范数，如果长期接近零，说明它没在起作用。解决办法是调整位置编码的缩放，或者改用相对位置方案。

5.4 推理输出重复或乱码

生成阶段的问题跟训练阶段不太一样。输出反复重复同一句话，通常是解码策略的问题，贪心解码容易陷入循环，加一点随机采样或者重复惩罚一般能缓解。如果输出的是完全无意义的字符，那要考虑是不是词表和分词器对不上，或者生成时的起始 token 用错了。

还有一个容易被忽略的点是训练和推理的长度分布不一致。训练时序列都补齐到固定长度，推理时序列是逐渐变长的，如果 mask 的处理在两种模式下不一致，模型看到的东西就跟训练时不一样，输出自然会乱。这个问题的排查方法很简单：用训练时的完整序列做一次前向，再模拟逐步生成的过程，对比中间每一步的输出，哪里开始分叉就查哪里。

故障现象	优先排查方向	快速验证方法
损失卡住不降	标签对齐、学习率	小样本过拟合测试
损失震荡发散	学习率、梯度裁剪	降学习率重跑
注意力权重异常	缩放、mask、初始化	权重矩阵可视化
输出重复循环	解码策略	加采样或重复惩罚
训练推理不一致	mask 处理、位置偏移	单步生成对比