Transformer 这几年的热度不用多说从 NLP 一路干到 CV、语音、多模态几乎所有主流模型都有它的影子。很多朋友看完《Attention Is All You Need》之后觉得懂了但一打开源码就懵了尤其是 QKV 矩阵、多头切分、mask 掩码这些细节每行代码都认识串起来就不明白为什么要这么写。这篇文章不聊宏观趋势直接带着你把 Transformer 的每个模块一步步拆开看配合 PyTorch 代码和实际踩坑经验讲清楚每个设计背后的逻辑。适合想彻底搞懂 Transformer 架构、准备手撕代码或者做模型训练的读者读完之后你会对整套结构有一种“原来如此”的通透感。1. 输入表示Embedding 与位置编码是怎么把文字变成张量的Transformer 本质上做的是“序列到序列”的变换但它不像 RNN 那样按时间步一个一个吃输入而是一次性把整个序列灌进去。这就带来一个核心问题模型必须通过某种方式把“词的含义”和“词的位置”同时编码成向量。这一步如果做得不对后面所有模块都白搭。1.1 Token Embedding词嵌入层的维度设计与直觉Embedding 层的作用很直接把离散的 token ID 映射成稠密的连续向量。假设词表大小是vocab_size 30000嵌入维度d_model 512那么 Embedding 层就是一个[30000, 512]的查找表。输入一个形状为[batch, seq_len]的索引矩阵查表后得到[batch, seq_len, 512]的张量。这里的核心问题是为什么嵌入维度通常是 512 而不是 128 或者 4096这里有一个工程和效果的权衡。维度太小语义表达能力不足词与词之间的区分度不够维度太大参数量爆炸训练成本和显存占用都会翻倍。512 这个数值在 2017 年的论文里是一个平衡点后来许多模型沿用或者在此基础上做缩放。比如 ViT 的 patch embedding 用的也是这个逻辑只不过输入从 token 换成了图像 patch本质没变。实际写代码时有个很容易被忽略的细节Embedding 层的初始化方式。标准做法是使用均值为 0、标准差为d_model ** -0.5的正态分布初始化这样做的目的是控制初始嵌入向量的范数在一个合理范围内避免激活值过大或过小。如果直接用 PyTorch 默认的初始化早期训练 loss 会偏高收敛速度也会变慢。我测试过不同初始化方式对训练曲线的影响差距在 5% 到 10% 之间在资源有限的情况下这个优化是稳赚不赔的。import torch import torch.nn as nn import math class TokenEmbedding(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.d_model d_model def forward(self, x): # x shape: [batch, seq_len] return self.embedding(x) * math.sqrt(self.d_model)注意第 11 行这里乘了sqrt(d_model)。论文里没有特别强调但在原版实现中是有的。乘法的作用是在后续加位置编码时保持位置编码的相对影响力。嵌入向量的数值通常在[-1, 1]之间乘以sqrt(512) ≈ 22.6之后嵌入向量占据主导地位位置编码只是微调这样模型初期可以更专注于学习词本身的信息。1.2 位置编码为什么 Transformer 必须另起炉灶设计位置信息RNN 天生是按顺序处理输入的第 3 个词就是第 3 个时间步位置信息隐含在结构里。Transformer 是并行处理整个序列输入张量同时包含所有词如果不加位置信息模型看到“我爱你”和“你爱我”是完全一样的因为没有顺序概念。论文用了三角函数位置编码PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))这里pos是 token 在序列中的位置i是维度索引。为什么要用三角函数而不是直接用整数[0, 1, 2, ...]呢有三个原因。第一周期函数的值域固定在[-1, 1]不会因为序列很长导致数值爆炸第二相对位置可以通过线性变换表示数学上有优雅的推导第三不同频率的三角函数覆盖不同尺度的位置关系低频表整体、高频表局部这种多分辨率特性和人类理解位置的直觉一致。不过实际工程中越来越多的模型选择了可学习位置编码比如 BERT 直接初始化一个[max_len, d_model]的矩阵去训练。这样做的好处是能从数据中自适应学到位置模式坏处是外推性差超过训练时最大长度就会出现奇怪行为。这里有一个非常典型的坑训练长度和推理长度不一致。如果你用可学习位置编码训练时最大长度是 512推理时碰到 600 长度的输入位置编码矩阵直接越界模型立刻崩掉。如果模型有这种场景需求要么用三角函数编码对任意长度有外推性要么在训练时做长度采样让模型看到不同长度的序列。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() self.dropout nn.Dropout(pdropout) pe torch.zeros(max_len, d_model) position torch.arange(0, max_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) pe pe.unsqueeze(0) # [1, max_len, d_model] self.register_buffer(pe, pe) def forward(self, x): x x self.pe[:, :x.size(1)] return self.dropout(x)这段代码里div_term的构造等价于1 / 10000^(2i/d_model)用exp和log组合是为了避免幂运算的数值不稳定。register_buffer将pe注册为模型的持久缓冲区这样参数保存、.to(device)都会自动带上它不会被当作需要梯度更新的参数。1.3 输入模块的组合顺序与实际调试经验整体组合是token embedding → 乘 sqrt(d_model) → 加位置编码 → Dropout。这个顺序不要随便调换。如果先加再乘位置编码被放大后噪声太大如果不做 Dropout深层模型在小数据集上很容易过拟合。我在训练过程中体验很深刻位置编码的 Dropout 值设成 0.1 就比较合适过大会导致位置信息被大量抹掉模型收敛非常慢过小则在小数据集上 loss 降不下去模型会把位置当成硬编码来记。另外PyTorch 的nn.Embedding默认是不做缩放初始化的如果发现初始训练曲线不太对劲第一个要检查的就是这个输入模块。2. 多头注意力Transformer 的灵魂引擎注意力机制是 Transformer 的核心其他所有模块都是围绕它构建的。理解多头注意力的关键在于把三个问题想透自注意力在做什么、为什么除以sqrt(d_k)、多头分别学到了什么。这三个问题想透了代码只是换个表达的事情。2.1 QKV 三件套自注意力的完整计算流程先看单个头的自注意力计算。输入是经过位置编码的向量序列X [x_1, x_2, ..., x_n]每个x_i是d_model维向量。通过三个不同的权重矩阵W_Q、W_K、W_V分别映射出 Query、Key、ValueQ X W_Q K X W_K V X W_V用生活化的类比来理解Query 是你在心中问的问题——“我该关注谁”Key 是每个候选词的自我介绍——“我包含什么信息”Value 是候选词的实质内容。注意力计算就是“拿你的问题去和所有候选词的自我介绍做匹配根据匹配程度加权提取内容”。具体计算是 $Attention(Q,K,V) softmax(\frac{QK^T}{\sqrt{d_k}})V$。QK^T得到注意力分数矩阵[seq_len, seq_len]第 i 行第 j 列表示第 i 个 Query 与第 j 个 Key 的匹配得分。除以sqrt(d_k)后做 softmax得到归一化的注意力权重最后乘V得到加权汇总。def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.size(-1) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, -1e9) attn_weights torch.softmax(scores, dim-1) output torch.matmul(attn_weights, V) return output, attn_weights这段代码里有几个关键操作要留意。K.transpose(-2, -1)是交换最后两个维度把[batch, heads, seq_len, d_k]转成[batch, heads, d_k, seq_len]才能和 Q 做矩阵乘法。masked_fill把 mask 中为 0 的位置填成-1e9经过 softmax 后这些位置的概率趋近于 0相当于完全忽略这些位置。这里填-1e9而不是0是因为 softmax 对 0 值会分配概率只有填一个极大的负数才能让概率归零。2.2 缩放因子为什么一定要除以根号 d_k这是新手最容易忽略的细节。注意力分数QK^T的每个元素是d_k个乘积的和。如果d_k很大比如 64那么分数的方差也会变大分布会变得非常尖锐softmax 之后几乎变成了 one-hot 分布梯度消失模型学不动。举个例子假设q和k的每个维度均值 0、方差 1那么一个点积的方差是d_k标准差是sqrt(d_k)。除以sqrt(d_k)后方差重新回到 1softmax 的输入分布不再尖锐梯度可以顺畅回传。这个设计是理论推导加实验验证的结果论文里专门用了一段话解释这一点。实际测试中如果用d_model 512且不除sqrt(64) 8训练到前几个 step 就会看到 logits 指数爆炸loss 直接 NaN。即使勉强训练曲线的收敛速度也会明显变慢。所以这个缩放因子不是可选项而是必须项。2.3 多头切分每个头到底学到了什么东西多头注意力的设计思路是与其用一个注意力头去捕获所有关系不如用多个头分别关注不同类型的关系。比如一个头关注词法上的相邻关系一个头关注跨距离的指代关系一个头关注语法角色。具体实现方式将d_model 512拆成 8 个d_k 64的头每个头独立做注意力计算最后拼接到一起再通过输出矩阵W_O融合。这样做的额外好处是计算效率高——8 个头的矩阵计算可以并行一次性完成。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0, d_model must be divisible by num_heads self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads 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.W_O nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, query, key, value, maskNone): batch_size query.size(0) # 1. 线性投影后拆头 Q self.W_Q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K self.W_K(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V self.W_V(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 2. 缩放点积注意力 attn_output, _ scaled_dot_product_attention(Q, K, V, mask) # 3. 拼接所有头过输出矩阵 attn_output attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) return self.W_O(attn_output)注意contiguous()这行转置后的张量在内存中不是连续排布的直接view会报错必须调用contiguous()让内存连续化。这是 PyTorch 新手最常见的报错点之一。多头注意力还有一个不容易想到的好处每个头的梯度路径是独立的相当于集成了多个弱分类器具备一定的鲁棒性。实践中如果某个头学到的东西退化了其他头还能补上模型的整体表现不会崩塌。这也解释了为什么多头数量太小时模型能力不足、太多时边际收益递减还增加显存消耗——8 到 16 头是一个经验平衡区。2.4 注意力 mask 的两种形态和实现陷阱attention mask 有两种典型应用Padding Mask 和 Look-Ahead Mask。Padding Mask 用于遮蔽 padding 位置避免模型把注意力放在无意义的补零 token 上。Look-Ahead Mask 用于 Decoder 的自注意力确保第 i 个位置只能看到前 i-1 个位置的信息防止未来信息泄露。Look-Ahead Mask 是一个上三角全 1 矩阵对角线以下为 0。实现时通过torch.tril(torch.ones(seq_len, seq_len))生成然后用masked_fill把 0 位置替换成-1e9。这里有一个非常隐蔽的坑padding mask 和 look-ahead mask 需要叠加使用。Decoder 的输入同时有 padding 和未来位置两种信息需要屏蔽。很多人分开实现没问题合到一起就逻辑混乱。正确做法是把两个 mask 做逻辑与运算然后统一传给注意力函数。我在实战中见过不少模型在推理时出现“训练正常、生成乱序”的情况最后定位都是这个 mask 叠加逻辑写错了。3. 残差连接与层归一化让训练稳定的隐藏功臣注意力层和前馈网络本身并不复杂真正让 Transformer 在深层也能稳定训练的是包裹在每个子层外面的残差连接和层归一化。这两个组件常被人一笔带过实际踩坑时才意识到它们的重要性。3.1 残差连接给梯度修高速公路残差连接最早在 ResNet 中被证明能有效解决深层网络的梯度消失问题。在 Transformer 中每个子层输出会与自己输入相加output LayerNorm(x Sublayer(x))这样做的意义在于即使某个子层学到了非常复杂的映射梯度也能通过绕行的“高速公路”直接回传到更前面的层。如果没有残差连接梯度要在数十层的矩阵乘法中反复相乘中后期层的梯度会指数级衰减深层参数几乎学不动。观察实现细节真正常见的做法是在子层操作之后、残差相加之前做 Dropout。也就是说实际计算是x dropout(sublayer(x))。这个顺序是有讲究的对子层输出做 Dropout相当于给残差路径增加噪声起到正则化作用降低过拟合风险如果对相加结果做 Dropout效果不如前者明显因为残差路径中原本没噪声的信息也会被扰动。3.2 LayerNorm 与 BatchNorm 的核心差异层归一化在 Transformer 中是不可或缺的。BatchNorm 在 CV 中很常见但在序列模型中不适用因为序列长度不固定batch 内长度不同会导致统计量不稳定。而 LayerNorm 是对每一个样本的每一个 token 在特征维度上做归一化不依赖于 batch 内的其他样本天然适配变长序列。class LayerNorm(nn.Module): def __init__(self, d_model, eps1e-6): super().__init__() self.gamma nn.Parameter(torch.ones(d_model)) self.beta nn.Parameter(torch.zeros(d_model)) self.eps eps def forward(self, x): mean x.mean(dim-1, keepdimTrue) std x.std(dim-1, keepdimTrue) return self.gamma * (x - mean) / (std self.eps) self.betaLayerNorm 内部维护两个可学习参数gamma和beta。gamma初始化为 1beta初始化为 0模型可以学习到是否对归一化后的结果做缩放和平移。eps防止分母为 0一般取1e-6到1e-5之间太大会让标准化后的数值方差偏大影响稳定性和表达能力。3.3 Pre-Norm 与 Post-Norm 的实战区别原版 Transformer 使用 Post-Norm 结构先做子层计算再残差相加最后 LayerNorm。而现代实现GPT、BERT 等大多使用 Pre-Norm 结构先 LayerNorm再子层计算最后残差相加。为什么会有这个转变Post-Norm 在训练深层模型时梯度不稳定随着层数增加很容易崩。Pre-Norm 则不同因为它将 LayerNorm 放在子层之前相当于在梯度回传路径上加了一个尺度变换无论网络多深残差路径上的恒等映射始终存在梯度回传更稳定。我用一个 12 层的 Transformer 分别用两种结构做过对比Post-Norm 在 8 层以内问题不大到 12 层时 loss 振荡明显需要更小的学习率Pre-Norm 则一路平稳下降唯一的缺点是最终收敛精度略低一点点但换来的稳定性收益远超这个损失。class TransformerBlock(nn.Module): def __init__(self, d_model, num_heads, dim_ff, dropout0.1): super().__init__() self.attn MultiHeadAttention(d_model, num_heads, dropout) self.ffn FeedForward(d_model, dim_ff, dropout) self.norm1 LayerNorm(d_model) self.norm2 LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): # Pre-Norm 结构 x x self.dropout(self.attn(self.norm1(x), self.norm1(x), self.norm1(x), mask)) x x self.dropout(self.ffn(self.norm2(x))) return x这段代码中有个容易踩坑的写法Pre-Norm 结构中Q、K、V 都传入了self.norm1(x)也就是同一个归一化后的张量。有人会误解成应该对原始 x 做x后面才做 norm。其实 Pre-Norm 的核心理念就是把 norm 放在子层内部的第一位三处共享同一个归一化结果完全没问题。如果写成self.attn(x, x, x, mask)再在输出后做 norm就退化成了 Post-Norm失去了稳定性优势。4. 前馈网络 FFN被严重低估的“记忆层”很多讲解 Transformer 的文章会把 FFNFeed-Forward Network一笔带过说它就是“两个全连接层夹一个激活函数”。但实际上FFN 是 Transformer 参数量最大的模块也是模型记忆能力的主要来源。没有它注意力模块再花哨也学不动复杂函数。4.1 FFN 的结构与激活函数选择标准 FFN 由两个线性变换和中间一个激活函数组成FFN(x) max(0, x W_1 b_1) W_2 b_2中间隐藏维度dim_ff通常是d_model的 4 倍。对于d_model 512dim_ff 2048参数量大约是2 * 512 * 2048 ≈ 2M参数相比多头注意力模块的4 * 512 * 512 ≈ 1M参数FFN 的参数量占了总参数的近三分之二。为什么中间维度要放得这么大因为注意力层负责“聚合信息”FFN 负责“加工信息”。d_model维度的空间对单层线性变换来说表达能力有限需要把数据投影到更高维空间做非线性变换再压缩回来。这与核方法的思想有些类似——高维空间中更容易找到决策边界。激活函数原版使用的是 ReLU后来 GPT 系列换成了 GELU效果略好。GELU 是 ReLU 的平滑版本在负半轴不是完全截断而是保留了一部分梯度有助于深层模型的梯度流动。如果做分类任务或者中小规模模型ReLU 够用如果训练大规模模型我建议直接用 GELU训练曲线更平滑。class FeedForward(nn.Module): def __init__(self, d_model, dim_ff, dropout0.1): super().__init__() self.fc1 nn.Linear(d_model, dim_ff) self.fc2 nn.Linear(dim_ff, d_model) self.dropout nn.Dropout(dropout) self.activation nn.GELU() def forward(self, x): return self.fc2(self.dropout(self.activation(self.fc1(x))))4.2 FFN 与注意力模块的分工逻辑要理解 Transformer 为什么有效必须理解这两个子层之间“分治”的关系。注意力层的作用是 token 之间的信息交流——某个 token 需要看哪些其他 token把信息聚合过来。FFN 的作用是对聚合后的信息做独立加工——每个 token 在自己的位置上进行一次非线性变换提取更高层的语义特征。这样交替堆叠相当于模型经历了“交流 → 思考 → 交流 → 思考”的循环。注意力层是“社交环节”FFN 是“独自消化时间”。如果去掉 FFN模型对每个 token 的表示只能做线性变换表达能力急剧下降如果去掉注意力层每个 token 永远只能看到自己无法整合上下文信息。两者缺一不可交替堆叠是经过反复验证的最优排列方式。4.3 FFN 的参数量计算与加速技巧以d_model512, dim_ff2048为例fc1权重是[512, 2048]偏置[2048]fc2权重是[2048, 512]偏置[512]。合计(512*2048 2048 2048*512 512) ≈ 2.1M参数。如果总共有 12 层 Transformer仅 FFN 就有 25M 参数。这里就引出一个优化技巧FFN 的 Dropout 应该比注意力层设置得更大一些通常在 0.1 到 0.2 之间因为 FFN 参数量大、容量高更容易过拟合。前向计算时FFN 是纯矩阵乘法GPU 利用率很高一般不需要特殊优化。但在 CPU 推理时需要考虑 Batch 合并尽量把多个样本的 token 拼成一个大的矩阵一次性计算避免逐样本循环调用nn.Linear那样 CPU 上的效率会差好几倍。5. Decoder 模块与输出层从序列到序列的完整闭环如果只做理解类任务比如 BERT 式的编码器到 FFN 这一层就已经够了。但要做生成类任务机器翻译、文本生成、股票预测等就必须理解 Decoder 的设计。Decoder 和 Encoder 的区别主要集中在三个地方Masked Self-Attention、Cross-Attention 和输出层处理。5.1 Masked Self-Attention 与因果掩码Decoder 的第一个子层也是自注意力但和 Encoder 不同的是它必须添加一个因果掩码Causal Mask。为什么要这么做因为在生成第 t 个词的时候模型不应该看到第 t1 个及之后的词。如果能看到未来信息这个问题就变成了“抄答案”推理阶段根本无法实现。因果掩码的实现很直接一个[seq_len, seq_len]的上三角矩阵对角线以下为 1可以看到对角线以上为 0要被 mask 掉。用torch.tril(torch.ones(seq_len, seq_len))可以生成下三角全 1 矩阵再配合masked_fill把 0 位置填成-1e9。训练时Transformer 采用 Teacher Forcing 策略——一次性输入完整的目标序列通过因果掩码保证每个位置的输出只依赖之前的位置。这和 RNN 逐时间步生成的方式完全不同是 Transformer 训练速度优势的重要来源。但这是否意味着训练和推理完全一致呢并非如此。训练时每步都能看到真实的前文推理时每一步都用自己的上一步输出作为输入这种“训练-生成差异”被称为 exposure bias应对方法包括计划采样Scheduled Sampling和强化学习微调这属于进阶话题了。5.2 Cross-AttentionDecoder 如何利用 Encoder 的信息Decoder 的第二个子层是 Cross-Attention交叉注意力这是 Encoder-Decoder 架构的独特之处。Q 来自 Decoder 上一层的输出K 和 V 来自 Encoder 的最终输出。这样设计好理解Decoder 每生成一个 token都去 Encoder 的上下文里“查资料”——问题Q来自我已经生成的内容资料库K、V来自原始输入。交叉注意力的实现代码和多头注意力几乎一致区别就在forward传入的key和value不是同一个张量而是 Encoder 的输出。如果只看 PyTorch 代码很多人在MultiHeadAttention的forward中看到 query、key、value 三个参数都觉得多余到 Cross-Attention 这一步才真正理解为什么要把三者区分开来。# Decoder 中 Cross-Attention 的调用方式 attn_output self.cross_attn( querydecoder_output, # 来自 decoder 自注意力层 keyencoder_output, # 来自 encoder 最后一层 valueencoder_output # 同上 )注意这里 encoder 的输出是否需要 mask需要但 mask 的逻辑和 Decoder 自注意力完全不同。Cross-Attention 需要屏蔽的是 Encoder 输入中的 padding 位置即 padding mask。Decoder 的因果 mask 只在自注意力层使用不能用在 Cross-Attention 上因为解码器生成第 t 个 token 时理应是能看到 Encoder 完整输入信息的。5.3 输出层与 Softmax 温度参数Decoder 最后一层输出[batch, seq_len, d_model]要通过一个线性层映射回词表大小的 logits然后做 softmax 得到概率分布。这个线性层的权重通常和 Embedding 层共享用nn.Linear(d_model, vocab_size, biasFalse)同时token_embedding.weight也绑定到这个权重上。这样做能大幅减少参数量并且实验表明共享权重有助于提高生成质量。推理阶段还有一个重要细节温度参数。直接使用 softmax 的原始概率分布做采样容易出现两个问题——分布太平坦导致文本缺乏确定性或者分布太尖锐导致文本过于重复。通过温度系数调整 logitslogits logits / temperature probs torch.softmax(logits, dim-1)温度大于 1 时分布更平滑输出更多样温度小于 1 时分布更尖锐输出更确定。做序列生成任务时温度一般设 0.8 到 1.2 之间。我在实际生成场景中测试过温度过低时模型会陷入重复序列的循环温度过高则输出变得无意义。6. 训练实战中的常见问题与排查技巧理论拆完了最后分享一些实际训练 Transformer 模型时遇到的典型问题。大部分问题和模型架构本身无关而是操作细节不到位导致的但这些坑几乎每个初次上手的人都会踩。6.1 训练 Loss 不下降的排查思路如果模型训练了几个 epochloss 还是纹丝不动除了学习率设置不当最常见的两个原因分别是Embedding 层未缩放和注意力 mask 错误。未缩放时嵌入向量的值域和位置编码不匹配位置信息干扰过大模型初始阶段会混乱mask 错误时如果是训练阶段 mask 没生效模型能够看到未来信息loss 会直接降到非常低但一测试就立刻崩掉。排查方法是先打印单个 batch 的前向结果手动验证 mask 的结构。用一个小例子 [2, 3] 的输入打印 mask 矩阵看看需要对哪些位置屏蔽、实际屏蔽的是哪些位置。这个小动作能节省数小时的排查时间。6.2 显存不足的优化手段训练 Transformer 最大的痛点之一就是显存开销。显存消耗主要分布在激活值存储、梯度、优化器状态和参数四个部分。在 12GB 显存的显卡上训练一个d_model512, num_heads8, batch_size16, seq_len128的模型基本就是在崩溃边缘试探。几个非常有效的显存优化技巧使用混合精度训练AMP显存几乎减半梯度累积Gradient Accumulation用更小的 batch 分步累积梯度激活检查点Activation Checkpointing重计算前向激活值来换取显存。优先级排序是 AMP 效果最明显且无副作用梯度累积适合大 batch 场景激活检查点空间换时间代价较大最后考虑。6.3 常见错误速查表问题表现可能原因排查方法loss 长时间不变学习率过大/过小用学习率扫描器找到合理区间loss 变成 NaN注意力分数未缩放或 logits 过大检查是否除以 sqrt(d_k)显存 OOMbatch 太大或序列过长开启 AMP减小 batch训练正常测试崩坏mask 未正确应用打印 mask 矩阵验证推理结果全是乱码位置编码外推失败换三角函数编码或加长训练长度收敛速度慢未使用 Pre-Norm 结构改用 Pre-Norm 的 TransformerBlock6.4 训练策略细节Transformer 对学习率非常敏感尤其是 Adam 优化器配合 warmup 策略几乎是标配。原版论文使用的是 Noam 学习率调度先线性增长到峰值再按平方根倒数衰减。class NoamSchedule: def __init__(self, optimizer, d_model, warmup_steps4000): self.optimizer optimizer self.d_model d_model self.warmup_steps warmup_steps self.step_num 0 def step(self): self.step_num 1 lr self.d_model ** (-0.5) * min(self.step_num ** (-0.5), self.step_num * self.warmup_steps ** (-1.5)) for param_group in self.optimizer.param_groups: param_group[lr] lr self.optimizer.step()Warmup 阶段学习率从 0 线性增长到峰值主要作用是让模型在初始阶段“熟悉”参数的梯度尺度避免一开始就把训练带偏后面的衰减阶段则逐步收敛到更精细的局部最优。Warmup steps 在 4000 左右是一个标准起点如果是小数据集可以适当减少大数据集可能需要 10000 步以上的 warmup。最后分享一个小感受Transformer 拆解完之后你会发现它其实就是一个“注意力模块 前馈模块”反复堆叠的结构每个模块单个看都不算复杂难点在于理解它们组合在一起时各自承担什么角色。我在最开始接触的时候一直在纠结 QKV 的物理意义后来发现先把计算流程跑通、再回头体会设计意图是效率更高的路径。如果你也在学习这个架构建议把代码从零开始手写一遍把每个矩阵的形状都打印出来看一下比看十遍论文都管用。