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_model512、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 Block3.1 环境准备与张量维度约定我建议先用 PyTorch 在一个小任务上把单层 Block 跑通再往上堆。环境上没什么特殊要求装好 PyTorch 就行CPU 也能跑只是慢。关键是先把维度约定固定下来避免后面到处是形状不匹配的报错。我的习惯是统一用[batch, seq_len, d_model]这个顺序所有中间变量都按这个约定推导。先定义超参d_model 取 128头数取 4每头维度 32FFN 中间维度取 512也就是 4 倍。序列长度取 64batch 取 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, maskNone): 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 处理、位置偏移单步生成对比这些排查顺序是我这几年攒下来的基本上八九成的故障都能在前两步定位到。真正难查的是那种“每个模块单独测都对拼起来就不对”的问题这种往往出在模块之间的接口约定上比如维度顺序、mask 的形状、归一化的位置建议在写代码时就把这些约定写在注释里省得后面自己都忘了当初怎么设计的。最后分享一个我个人的习惯每次改结构或者改训练配置只改一个变量跑完记下结果再改下一个。Transformer 的可调项太多同时改三四个地方出了问题根本不知道是哪个引起的。这个笨办法救过我很多次尤其是在赶进度、想一口气调好几个参数的时候克制住这种冲动反而走得最快。