一篇读懂训练稳定性：Loss Spike、梯度爆炸与数值问题的诊断与修复

一次 Loss Spike 的完整病历

现象：一条平静曲线上的垂直悬崖

先把现象定义清楚。正常的预训练 loss 曲线在双对数坐标下是一条平滑下降的曲线，带 learning rate schedule 造成的轻微起伏。Loss spike 是这条曲线上突然出现的垂直跳升——典型形态是在几十个 step 内 loss 从 2.0 跳到 6.0 甚至更高，随后要么自然恢复，要么从此失稳发散。它不是小数点后的抖动，而是训练事故级的事件：spike 前后的 checkpoint 即使 loss 表面恢复，下游能力的恢复也常常滞后数亿 token，更坏的情况是模型悄悄学坏，loss 看不出异常但某个能力维度（如数学）永久性降级。

公开文献里最著名的病历来自两处。PaLM 论文（Chowdhery 等人，2022）披露 540B 模型训练中经历了两次显著 spike，处理方式是从 spike 前的 checkpoint 重启、跳过可疑时间窗内正在使用的数据段继续训练，两次都成功恢复；论文同时披露他们持续使用 z-loss（logit 平方的正则项，系数 1e-4）作为稳定性手段。OPT-175B 的训练 logbook（Meta 公开的 SCROLLS 报告附录）更是一部惨烈实录：训练全程记录了数十次 loss 尖峰与梯度爆炸事件，团队尝试了从丢弃 bad batch、重启、调整 clip 阈值到替换 Adam beta2 的全部手段。这两份公开材料的共同信息是：在千亿级训练里，spike 是常态而非事故，稳定性工程是预训练的必修课而不是可选保险。

监控曲线：三个比 loss 更早报警的信号

发现 spike 时往往已经晚了，真正的工程能力体现在提前几个 hundred step 看到前兆。三个先行指标。其一，梯度范数（grad norm 全局 L2）：健康的训练里它在数量级上稳定，spike 前常见 2 到 10 倍的持续抬升——这是最灵敏的领先指标。其二，logit 的最大绝对值：注意力与输出层 logits 的峰值缓慢爬升，往往先于 loss 异常几百个 step。其三，梯度的分层数值分布：某一层（常见是嵌入层或第一层 FFN）的梯度范数与其他层比例失衡，提示问题在局部结构。生产管线应把这三个指标按 step 粒度写入监控，并配阈值告警——阈值不是拍脑袋，取过去 1000 个 step 中位数的 5 倍动态计算比固定值更抗漂移。

# 训练循环里的稳定性监控骨架
total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# clip_grad_norm_ 返回裁剪前的全局范数，务必记录
metrics.log("grad_norm", total_norm)

if total_norm > SPIKE_FACTOR * rolling_median(grad_norms):
    alarm("grad_norm 异常抬升")          # 先行告警
if not torch.isfinite(total_norm):
    skip_batch_and_rollback_check()       # NaN/Inf 立即处置

五类常见根因

数据批次异常

最常见也最易被忽视的根因：某个 batch 里的脏数据触发异常梯度。典型形态是超长重复文本（爬虫残留的模板页）、乱码（编码错误的二进制被当文本读入）、以及「不可能样本」——比如 tokenizer 里数万个 UNK 连片出现。这类 spike 的特征是一次性且可恢复：跳过该 batch 后训练自然回到正常轨迹。诊断手段是把 spike 时间窗内消费的数据 batch 记录下来（W&B 或自建的 batch 级溯源日志），事后逐条检查；预防手段是数据管线的启发式过滤（重复串长度、非文本字符占比、UNK 率阈值），以及训练侧的输入断言。OPT 的 logbook 显示他们的多次事故与数据异常直接相关，这是跳批策略（后面细讲）成为标配的直接原因。

fp16 的数值边界

混合精度是 spike 的第二大来源，机理要讲透。fp16 的可表示范围是约 ±65504，最小规约数约 6e-5，动态范围极窄。训练里两处最容易越界：注意力 logits（未缩放的 dot product 在长序列或异常 token 上轻松超过 65504，直接溢出为 Inf）；损失函数的中间项（交叉熵前的 log_softmax 对极大 logit 的计算）。溢出的表现是梯度变 Inf/NaN，一次失败后优化器状态（Adam 的动量与二阶矩）被污染，之后的每个 step 都在传播 NaN——这就是 fp16 训练里「一次溢出毁掉整场训练」的机理。bf16 用 8 位指数换来约 1e38 的动态范围，从结构上消灭了溢出问题（代价是精度位数少，对累加敏感的场景要靠 fp32 主权重与 Kahan 求和补偿）。所以当代预训练的默认配置是 bf16 计算 + fp32 主权重 + fp32 优化器状态；只有在硬件不支持 bf16 时才用 fp16，且必须配动态 loss scaling 与 Inf 检测回退。

注意力 logits 爆炸与优化器状态损坏

第三类根因在注意力本身：QK^T 的数值随训练演化可能整体增大（attention logits 的「低温化」现象，Gemma 2 论文对此有系统讨论），logit 过大导致 softmax 退化为 one-hot，梯度出现极端值。对策是两个已被多家实验室验证的手段：logit 软截断（soft-capping，把注意力 logits 限幅到 ±30 左右但不硬截断，保留梯度通路，Gemma 2 采用）、以及qk-layernorm（对 Q 与 K 投影后做 RMSNorm，从源头控制 logits 尺度，PaLM 与多家后续模型采用）。第四类根因是优化器状态损坏：数值上表现为 loss 恢复但 grad norm 持续偏高，或某个 shard 的动量出现 Inf。诊断靠优化器状态的校验和与逐 shard 数值检查，修复靠从干净 checkpoint 重启。第五类是学习率与 warmup 的次生效应：warmup 不充分时早期 grad norm 大起大落，decay 末段学习率过小时 spike 恢复力下降——这些不制造 spike 但决定 spike 的伤害大小。

flowchart TD
    A[检测到 loss spike] --> B{grad norm 是否 Inf 或 NaN}
    B -- 是 --> C[fp16 溢出或状态损坏]
    C --> D[检查注意力 logits 与 loss scaling]
    B -- 否 --> E{spike 后是否自然回落}
    E -- 是 --> F[疑似脏 batch]
    F --> G[调出时间窗内 batch 逐条检查]
    G --> H[补数据过滤规则]
    E -- 否 --> I[持续失稳]
    I --> J[回滚到 spike 前 checkpoint]
    J --> K[跳过可疑数据段与参数扰动重启]
    L[监控 grad norm 与 logits 峰值] -.前兆告警.-> A

修复的标准操作：回滚、跳批与恢复

为什么要跳数据而不是只回滚

标准处置流程分三步。第一步，止损：确认 spike 后立即暂停，把 spike 时间窗的消费水位记录下来。第二步，回滚：从 spike 发生前（grad norm 开始抬升前更好）的最近 checkpoint 恢复参数与优化器状态——只回滚参数不回滚优化器状态是常见错误，被污染的动量会把 spike 「记忆」回来。第三步，跳批：把从 checkpoint 到 spike 点之间消费过的数据段标记为隔离区，重启后跳过这段直接消费新数据。为什么不原样重放？因为如果根因是数据，重放等于把同一颗雷再踩一遍；如果根因是数值的边际状态，同样的数据加同样的参数序列会精确复现同样的轨迹。PaLM 与 OPT 的实践都采用跳过策略，PaLM 论文明确描述了「重启时跳过可疑数据窗」的操作。配套细节：跳批导致数据顺序变化，对这条数据流的 shuffle 种子要能按水位重算；隔离区数据别扔，换 batch 组合方式后通常还能安全回流。

稳定性正则：z-loss 与它的同事

预防手段里最值得一文不值地展开的是 z-loss：在交叉熵之外加一项 λ·(Σ_i e^{z_i})² / (Σ_i e^{z_i})²（对 logit 指数和的平方做正则，PaLM 用 λ=1e-4）。它压的是 logit 的整体绝对尺度，机理与软截断同向但更温和——不改变 softmax 的形状假设，只是持续施压防止 logits「漂移变大」。与它配合的手段：梯度裁剪（max_norm=1.0 是最常见配置，注意裁剪是止损不是治愈，频繁触发裁剪本身就是告警信号）、qk-layernorm、以及 embedding 层的输出缩放。这些手段的共同哲学：预训练稳定性问题的 80% 是 logit 尺度管理问题，把所有中间表示的尺度按住，绝大多数 spike 根本不会发生。代价要心里有数：z-loss 与软截断对最终能力的影响在多数报告里是中性或轻微正面的，但它们确实改变了有效学习率与梯度流，超参需要重新微调而不是直接沿用旧配置。

分布式训练特有的失稳源

单卡训练的数值问题到了千卡集群上会被放大并出现新的形态。第一类是非确定性漂移：数据并行下，梯度 all-reduce 的求和顺序取决于通信拓扑与分块策略，浮点加法不满足结合律，同一份数据在不同卡数下的训练轨迹会分叉——这不是 bug，但它意味着「精确复现 spike」在分布式下几乎不可能，回滚调试时要有心理预期。第二类是通信精度的坑：all-reduce 如果用 fp16 做梯度聚合，百万卡时的累积舍入误差会把梯度整体拉偏，主流框架默认 fp32 或 bf16 通信正是为此；自建训练栈时这是必须显式确认的配置项。第三类是rank 间数值发散：混合并行（流水线 + 张量并行）下，某个 rank 的算子降级（比如某台机器不支持某种 bf16 内核而静默回退 fp32）会让各 rank 的中间表示出现细微差异，几千步后累积成异常梯度——这类问题的诊断手段是在固定 step 打各 rank 的梯度校验和并对比，差异出现的 rank 就是有算子降级的机器。

spike 恢复后的能力暗伤验证

最后一个容易被跳过的环节：loss 曲线恢复平坦后，训练并不能直接继续——loss 恢复不等于能力恢复。spike 期间的异常梯度可能对特定能力维度的参数造成选择性损伤，语言建模损失（它衡量的是通用 next-token 预测）对此并不敏感。标准做法是维护一组固定的下游 probe 任务（数学、代码、多语言、长文本各取一两个小评测集），在回滚恢复后立即跑一遍与 spike 前的 checkpoint 对比；若某个维度掉了超过噪声水平，就要从更早的 checkpoint 回滚，或者对该维度做短期强化数据补训。PaLM 论文在两次 spike 后都做了这类下游验证——这是「表面恢复」与「真恢复」的分水岭，跳过它的团队会在数周后的评测里发现模型悄悄变笨，那时已无法定位原因。

一张预防清单

把前文收敛成上线前的检查表：精度选 bf16（fp32 主权重 + 优化器状态），没有 bf16 硬件时 fp16 必须配动态 loss scaling；注意力 logits 上软截断或 qk-layernorm 至少其一；z-loss 1e-4 起步观察；梯度裁剪 1.0 且记录裁剪前范数；warmup 占总步数 0.5% 到 2%，长训练取低值；数据管线有 UNK 率、重复率、长度分布的在线断言；checkpoint 频率按「损失 30 分钟内进度」设定（spike 越早回滚伤害越小）；监控面板有 grad norm、logit 峰值、分层数值分布三个先行指标加动态阈值告警；演练过一次完整的回滚加跳批流程。这份清单的成本是一次性的，而千亿训练里一次未妥善处置的 spike 的代价是数天算力加永久性能力损伤，性价比不言自明。

小结

Loss spike 在大规模训练里是常态：PaLM 540B 走过两次、OPT-175B 记录了数十次，稳定性工程的定位是「让 spike 少发生、发生了快速恢复、恢复后不留暗伤」。诊断的核心是三个先行指标——grad norm、logit 峰值、分层数值分布——它们比 loss 早几百步报警；根因的五类归纳里，脏数据 batch 一次可恢复、fp16 溢出会污染优化器状态持续恶化、注意力 logits 漂移需要 z-loss 与软截断这类尺度管理手段从源头压制。修复的标准动作是「回滚到 spike 前的完整状态（含优化器）+ 隔离跳过可疑数据窗 + 重启」，只回滚参数是常见错误。而最好的稳定性投资永远在预防侧：bf16、qk-layernorm、z-loss、动态阈值告警加一次演练过的回滚预案——spike 管理做得好的团队，差别不在反应速度，而在 spike 根本没机会被看见。