跳转至

Transformer 训练稳定性:Loss Spike、梯度爆炸与混合精度训练深度解析

Transformer 的训练常伴随着极端的数值不稳定性。Loss Spike、梯度爆炸、混合精度下溢等问题不单影响模型收敛,更可能让持续数天的训练功亏一篑。系统性理解这些不稳定现象的成因与应对策略,是大模型训练的必修课。

什么是 Loss Spike?在训练 Transformer 时,可能由哪些原因引起?

Loss Spike(损失尖峰)指的是训练过程中损失值在某个时刻突然剧烈上升,可以达到正常值的数倍甚至数十倍,形成一个尖锐的波峰。这种现象在 Transformer 的大规模训练中相当常见,轻则拖慢收敛,重则导致梯度爆炸和 NaN,使训练中断。

Loss Spike 的根源是某些批次的梯度异常巨大,导致参数更新量远超正常范围,破坏了已经学到的表示。具体诱因可以分为以下几个层面:

数据层面:训练数据中混入了极端样本或噪声。例如一段全是乱码的文本、一张全白或全黑的图像、或者文本与图像完全不匹配的图文对。这些样本的损失值天然就大,产生的梯度也异常大。特别是当多个这样的“坏样本”恰好出现在同一个批次中时,尖峰就出现了。数据集的局部偏差也会导致尖峰。比如某个特定领域(医学、法律)的文本在预训练语料中占比极小,模型对它的分布非常陌生,一旦采样到这一小簇数据,梯度就会剧烈波动。

优化器层面:Adam 类优化器的自适应学习率在某些情况下会放大问题。当某个参数的梯度连续几个步骤都很小时,其二阶矩估计(v)也变得很小,这会导致该参数的有效学习率(除以根号 v)变得极大。如果此时恰好来一个稍大的梯度,就会被这个放大的学习率推到一个极端的参数值,产生连锁反应。这就是为何 Adam 的训练有时会突然崩溃。

模型架构层面:Transformer 的深层残差结构在梯度反向传播时,各层的梯度尺度可能差异巨大。如果某些层的参数初始值不佳(比如 FFN 的第二层权重初始化得太大),早期训练中这些层就会产生过大的激活值,经过残差累积,在深层形成数值爆炸。此外,注意力机制的 softmax 操作在输入值过大时会趋于饱和,梯度趋近于零,但这往往伴随着前向数值的溢出。

数值精度层面:使用 FP16 混合精度时,梯度的表示范围极窄。正常范围的梯度在经过多次链式求导后,绝对值可能变得极小或极大。过小的梯度下溢为零,过大的梯度上溢为 Inf,直接引发尖峰。甚至在 FP32 下,如果初始化不当,深层 Transformer 的注意力 logits 也可能溢出。

遇到 Loss Spike 后,通常应该怎么处理?跳过更新、回滚检查点还是调整超参数?

面对 Loss Spike,处理策略需要根据尖峰的严重程度和发生频率分层次决策。

第一层:自动跳过异常更新。 这是最轻量、最自动化的处理方式。在训练循环中监控损失值或梯度的范数。如果当前 batch 的损失超过历史移动平均的若干倍(比如 5-10 倍),或者梯度的 L2 范数超过预设阈值,则自动跳过本次参数更新,不将这批梯度应用到优化器。这在混合精度训练中尤为常见——当动态损失缩放检测到梯度溢出,会自动减小缩放因子并重试。这种方式能过滤掉偶然的数据噪声,但无法解决系统性的训练不稳定。

第二层:从检查点回滚并调整学习率。 如果尖峰导致模型参数严重破坏,损失持续高位震荡甚至 NaN,需要立即回滚到尖峰前的检查点。恢复后,将当前学习率减半(或更大幅度),并从这批数据重新开始训练。通常还需要将 Warmup 重新开启一小段,让模型平稳过渡。这是最稳妥的恢复方式,代价是丢弃了尖峰后的一段训练。

第三层:根本性的超参数调整。 如果 Loss Spike 频繁发生,说明训练配置存在系统性问题。需要审视:

  • 学习率是否过高:降低峰值学习率,或延长 Warmup 步数。

  • 梯度裁剪阈值是否合适:将裁剪阈值设为历史梯度范数的 90 分位数或一个保守常数(如 1.0)。

  • 初始化方案:检查是否所有层使用了正确的初始化和缩放。

  • 数据质量:检查引发尖峰的批次数据,剔除明显的噪声样本。

  • 优化器参数:适当增大 Adam 的 β2β2(如从 0.999 到 0.9999),让二阶矩更平滑,避免有效学习率骤然放大。

  • 混合精度设置:如果使用 FP16,确保开启了动态损失缩放且初始缩放因子不过大;如果条件允许,切换到 BF16。

为什么深层的 Transformer 在训练初期容易出现梯度爆炸或消失?初始化如何帮助?

深层 Transformer 由多个残差块堆叠而成,每个残差块内包含自注意力和 FFN 两个子层。在训练初期,权重是随机初始化的,信号的传播极不稳定。

梯度爆炸和消失的根源:在反向传播时,梯度需要从最后一层逐层回传。每经过一个子层,梯度都会被该层的雅可比矩阵相乘。如果大部分层的雅可比矩阵的谱半径(最大奇异值)大于 1,梯度就会指数级放大(爆炸);如果小于 1,梯度就会指数级衰减(消失)。在标准 Transformer 的随机初始状态下,不同层的谱半径分布不均,既可能大于 1 也可能小于 1,这使得深层网络在初期极易出现极端的梯度值。

初始化如何帮助:初始化的目标是在训练初期让每一层的前向激活方差和反向梯度方差都保持恒定,避免信号在层间放大或缩小。Xavier 初始化假设激活函数是线性或 Tanh,将权重初始化为方差与输入输出维度之和成反比的分布。Kaiming 初始化针对 ReLU 类激活的修正版本。在 Transformer 中,由于残差连接的存在,信号传播的动力学有所不同——残差分支本身是恒等映射,不需要特殊初始化,因此重点是控制子层内部各线性变换的方差,使得子层的输出相对于其输入不会过大或过小。现代 Transformer 通常对 QKV 投影和 FFN 的第一层采用标准初始化的修正版,而对残差分支的最后一层(如注意力输出投影和 FFN 的第二层)采用较小的初始化,使得初始阶段每个子层的输出很小,网络行为接近恒等映射。DeepNet 等方案更是将这种思想系统化,通过理论推导为每层设置缩放因子,保证信号的平稳传播。

请解释 Xavier 初始化和 Kaiming 初始化的原理,Transformer 各层通常使用哪种初始化?

Xavier 初始化(Glorot 初始化)的出发点是:对于没有非线性或使用 Tanh 的网络,希望前向传播时每层激活值的方差保持一致,反向传播时梯度的方差也保持一致。为此,权重应从均值为 0、方差为 2fan_in+fan_outfan_in+fan_out2 的分布中采样,其中 fan_in 是输入维度,fan_out 是输出维度。这个方差使得在前向和反向传播中信号的尺度保持平衡。

Kaiming 初始化(He 初始化)专门针对 ReLU 类激活函数。ReLU 会将负半轴的输入清零,相当于丢弃了一半的方差。为了补偿,Kaiming 初始化将权重的方差扩大为 2fan_infan_in2(仅考虑前向)或 2fan_outfan_out2(仅考虑反向)。对于标准 Transformer 的 FFN,中间层的激活函数(GELU、ReLU 等)与 ReLU 类似,因此 FFN 的第一层常采用 Kaiming 初始化或其变体(如对 GELU 校正的版本)。

Transformer 各层的实践:由于 Transformer 有残差连接,子层的输出并非直接传递到下一层,而是与原输入相加。因此,子层内部的初始化策略需要保证该子层的输出方差相对于输入较小,而不是恰好等于 1。对于注意力层的 QKV 投影和 FFN 的第一层,通常采用标准初始化(Xavier 或 Kaiming,具体取决于激活函数),这些层的输出方差被设为 1 左右。而对于残差分支最后的线性层(注意力输出投影和 FFN 第二层),通常使用较小的初始化,比如从均值为零、标准差为 12⋅layers⋅fan_in2⋅layers⋅fan_in1 的分布中采样,甚至初始化为零。这确保训练开始时模型接近恒等映射,极大地提升了训练稳定性。GPT 系列和 BERT 等模型的具体初始化参数经过精心调优,本质思想相同。

Pre-Norm 和 Post-Norm 对训练稳定性的影响机制是什么?

Pre-Norm 和 Post-Norm 是 Transformer 中两种层归一化的放置方式,对训练稳定性影响深刻。

Post-Norm(原始 Transformer):每个子层的输出与输入相加后,再通过 LayerNorm。结构是 x + Sublayer(x) → LN。这种设计的问题在于,残差路径的梯度必须经过 LayerNorm 才能回传。LayerNorm 的梯度包含对输入的缩放因子,当层数加深时,多层的缩放因子相乘会指数级放大或缩小梯度,极易导致梯度爆炸或消失。因此,Post-Norm 对学习率和初始化极其敏感,训练深层网络非常困难。DeepNorm 等改进正是针对这一问题,通过特殊的权重缩放来抑制 LayerNorm 的梯度缩放效应。

Pre-Norm(现代标准):LayerNorm 放置在子层之前。结构是 x + Sublayer(LN(x))。残差路径完全绕过了 LayerNorm,梯度可以沿着残差连接无衰减地流向浅层。LayerNorm 只用于调整进入子层的输入分布,使它更稳定,但不影响梯度的跨层传播。这就好比给每层安装了一个输入调节器,同时保留了一条畅通的梯度高速公路。Pre-Norm 让深层 Transformer 的训练稳定性大幅提升,对学习率的容忍范围也宽得多,因此几乎所有现代大模型都采用 Pre-Norm。

深层机制:从初始化角度看,Pre-Norm 下,每个子层看到的输入是经过归一化的,方差被控制在 1 左右。而 Post-Norm 下,子层的输入是未归一化的残差累积,其方差随层数增长,子层需要不断适应漂移的输入分布,训练更困难。

随机深度(Stochastic Depth)是什么?它是如何作为一种正则化手段提高稳定性的?

随机深度是一种在训练时随机丢弃整个 Transformer 层的技术。对每个训练样本,以一定概率 pp 将某些层的输出完全替换为输入(即该层被跳过,退化为恒等映射)。被跳过的层不参与前向和反向计算。通常对较浅的层使用较小的丢弃概率,对深层使用较大的丢弃概率,呈线性递增。

正则化效果:随机深度迫使模型不依赖于任何单层,因为任意一层都可能被丢弃。这类似于 Dropout,但作用在层级别而非神经元级别,能有效防止过拟合。此外,它训练了不同深度的子网络组合,产生了集成学习的效果。在推理时,所有层被激活,但需要根据训练时的存活概率调整每层的输出(或直接使用,因为平均效应已在权重中隐式体现)。

提高训练稳定性:通过随机跳过一些层,训练时的有效网络深度变浅,梯度传播路径缩短,减少了梯度爆炸和消失的风险。这对于极深 Transformer 尤其有益。它还可以看作是一种隐式的深度课程学习:早期训练时网络实际上较浅,随着训练进行,所有层逐渐协同工作。

Dropout 在 Transformer 中通常加在哪些位置?过大的 Dropout 率对训练有什么影响?

Transformer 中的 Dropout 通常加在以下位置:

  • 注意力权重上:在计算 softmax 之后,对注意力概率矩阵施加 Dropout,随机丢弃某些 token 间的注意力连接。这迫使模型不依赖单一路径的注意力。

  • 残差连接之前:在每个子层(注意力或 FFN)的输出上,在加到残差流之前施加 Dropout。这是最标准的做法,用于防止子层的输出过拟合。

  • FFN 隐层激活上:在 FFN 的两层之间,激活函数之后施加 Dropout。

  • 嵌入层上:在词嵌入和位置嵌入相加之后,有时也会施加 Dropout。

过大的 Dropout 率会严重损害训练:

  • 信号不足:过度丢弃导致可用的信息过少,模型难以学习到有效的特征,训练损失下降缓慢甚至停滞。

  • 训练-推理不一致加剧:Dropout 在训练时关闭部分神经元,推理时全部激活,过高的丢弃率使得推理时的激活分布与训练时差异巨大,导致性能显著下降。

  • 欠拟合:模型容量被过度限制,无法拟合训练数据,即使训练足够久也无法达到理想的准确率。

通常 Transformer 的 Dropout 率设置在 0.1 左右。对于大规模数据和模型,Dropout 的作用降低,有时甚至不用(如 GPT-3 仅在嵌入层使用 0.1 的 Dropout,其他层不使用)。

训练不收敛时,如何从学习率、初始化、归一化位置等方面系统性排查?

系统性排查应遵循“从外到内、从粗到细”的顺序:

第一步:检查数据。确保输入没有异常值、NaN 或 Inf。验证分词器工作正常,序列长度分布合理。不收敛最常见的原因往往是数据预处理出现 bug。

第二步:检查学习率。学习率过高是导致不收敛的头号元凶。尝试大幅降低学习率(如 10 倍),观察 loss 是否开始下降。使用学习率查找器(LR Finder)来确定安全的学习率范围。如果 loss 一开始就 NaN,几乎肯定是学习率过大或初始化问题。

第三步:检查初始化。如果 loss 在几个步骤后突然 NaN,很可能是初始化不当导致的前向数值溢出。验证是否所有线性层都使用了正确的初始化方案(如 Kaiming 或 Xavier 的适当变体)。检查残差分支的最后一层是否用了较小的初始化。可以尝试从一个已经验证能收敛的小模型配置开始,逐步增大。

第四步:检查归一化位置。确认使用的是 Pre-Norm 还是 Post-Norm。如果使用 Post-Norm,训练不稳定是常态,需要配合更小的学习率和特殊的初始化(如 DeepNorm)。如果使用 Pre-Norm,确保 LayerNorm 在子层之前,残差路径未被错误地插入归一化。

第五步:检查梯度裁剪。确认梯度裁剪已开启且阈值合理(通常 1.0)。观察梯度的 L2 范数,如果长期过高,说明模型配置需要调整。

第六步:检查混合精度设置。如果使用 FP16,确认 Loss Scaling 已开启且为动态模式,缩放因子在正常范围内(如 2^16 左右)。可以尝试切换到 BF16 或 FP32 以排除精度问题。

第七步:监控训练日志。定期记录每层的梯度范数、参数更新范数、激活值均值和方差。可视化这些指标的趋势,定位突然变化的层级,有助于精确诊断。

为什么大模型训练需要使用数据并行和模型并行,有时还会遇到通信死锁?如何避免?

大模型的参数量巨大,远超单卡显存,因此必须使用分布式训练将模型切分到多卡。

数据并行:每张卡持有一份完整的模型副本,数据被均分到各卡。每张卡独立前向、反向得到梯度,然后通过 AllReduce 操作将所有卡的梯度求和(或平均),最后每张卡用聚合后的梯度更新自己的副本。当模型本身单卡可以容纳时,这是最直接高效的并行方式。

模型并行:模型本身被切分到不同卡上。分为张量并行(将单个矩阵乘法切分到多卡)和流水线并行(将不同层放在不同卡)。这些方式会在卡之间引入通信,比如张量并行在每次前向和反向时需要做 AllReduce 或 AllGather。

通信死锁的成因:分布式训练中,不同卡上的操作需要协调一致。如果通信原语不匹配(例如一张卡在发送数据,另一张卡却在等待接收不同大小的数据),或者由于代码 bug 导致一张卡跳过了某个通信步骤,就会发生死锁——所有卡互相等待,训练卡死。常见触发条件包括:

  • 条件分支导致不同卡执行了不同的通信操作(比如 if-else 中某张卡进入了另一分支)。

  • 数据加载不均衡导致某张卡提前结束了数据迭代,而其他卡还在等待聚合梯度。

  • 使用 NCCL 等通信库时,配置不当(如环境变量设置错误)。

避免方法:保持所有卡上的代码逻辑严格一致,避免任何可能产生分支的操作(包括 Python 的 if 语句、assert 等);使用框架提供的分布式训练 API(如 PyTorch DDP、DeepSpeed),它们已经处理了这些细节;在遇到死锁时,通过检查 NCCL 日志、设置超时和添加调试信息来定位。

如何通过监控梯度的 L2 范数和参数更新的范数来判断训练是否健康?

定期记录每一层的梯度和参数更新的统计量,是诊断训练健康度的强力手段。健康的训练通常呈现以下特征:

  • 梯度范数适度且稳定:梯度的 L2 范数在训练初期可能较大,随后应稳定在一个合理范围(如 0.01 到 1 之间),没有持续上升或下降的趋势。如果梯度范数持续增大,预示着即将发生梯度爆炸。如果持续减小到极低水平(如 1e-6),说明梯度消失或已经收敛。

  • 梯度范数的比例:查看梯度范数与参数范数的比值。如果这个比值过大,意味着单步更新可能剧烈改变模型。通常这个比值在 0.01 到 0.1 量级。

  • 更新范数与参数范数成比例:参数更新量(即优化器步骤后的参数变化)的范数应与当前参数范数保持合理比例。如果更新范数远大于参数范数,模型可能被推向未知区域;如果极小,则学习率可能过低。

  • 不同层的梯度差异:浅层和深层的梯度范数不应有数量级的差异(在 Pre-Norm 下,差异通常不大)。如果某层梯度总是为零或极大,该层可能存在问题。

通过可视化这些范数的曲线,可以及早发现异常并调整训练策略,而不是等到 Loss Spike 或 NaN 出现才被动应对。

在混合精度训练中,梯度下溢(Underflow)是什么?如何检测和解决?

梯度下溢指的是在反向传播时,某些梯度值太小,小于当前浮点格式能表示的最小正规数,导致在 FP16 下被截断为 0。当梯度变为 0 后,相应的参数就无法更新,模型学习停滞或偏向次优解。FP16 能表示的最小正规数约为 6×10⁻⁸,而训练中许多梯度(尤其是深层网络传播到浅层的梯度)可能远小于此值。

检测:通过梯度直方图查看梯度值分布。如果大量梯度值恰好为 0,而对应的 FP32 主副本中的梯度不为 0,即可确认下溢。也可以检查损失是否停止下降而 FP32 的梯度仍然健康。

解决:

  • Loss Scaling:这是最主要的解决方案。将损失放大,梯度随之放大,从而被推入 FP16 的可表示范围。动态损失缩放自动完成这一过程。

  • 切换到 BF16:BF16 的动态范围与 FP32 相同,几乎不会有下溢问题,是现代大模型训练的首选。

  • 关键操作保持 FP32:对易出现下溢的操作(如 softmax、LayerNorm)强制使用 FP32。

为什么有时候训练到一个 epoch 后期,损失反而上升?可能的原因有哪些?

训练后期损失上升,通常不是正常的过拟合表现,而是存在更深层的问题:

  • 学习率衰减不当:学习率衰减得过早或过快,导致模型在尚未收敛到最优区域时就停止了有效更新。后续训练只能在有限的参数空间内震荡,损失上升。如果衰减太快,模型可能无法继续沿着梯度方向下降。

  • 过拟合:在训练集上损失下降,但在验证集上损失上升。如果在训练集上损失也上升,那通常不是单纯的过拟合,而更可能是其他问题。

  • 数据问题:训练数据在 epoch 后期出现了质量较差的样本(如果数据未随机打乱,某些特定类型的样本集中在 epoch 末尾)。

  • 优化器状态不稳定:Adam 的二阶矩在训练后期可能因为累积了历史上的大梯度而变得异常大,导致有效学习率被过度缩小,模型更新量几乎为零,损失由于数值波动而微幅上升。

  • 梯度噪声:在接近收敛时,梯度的信噪比下降,随机的参数更新可能使损失小幅波动,但不会持续大幅上升。

排查时,应首先区分是否在训练集上上升。如果是,检查学习率衰减策略和数据顺序,并监控优化器状态。如果不是,则是过拟合,应增加正则化。