预训练稳定性工程：loss spike 检测与断点重启策略

说明：本篇为综述/无第一手实验数据，采用公开实践与工程经验路线。

实验目标

大模型预训练以"天"为单位的运行中，最大的敌人不是慢，而是不稳定：loss 突然跳高（spike）后不回落、某个节点 OOM 导致整批 checkpoint 失效、梯度爆炸把训练参数打飞。本篇目标：建立预训练稳定性的工程闭环——实时监控（loss/梯度范数）→ spike 自动检测与归因 → 自动保存与断点重启 → 超参修复。读完能对"训练挂了怎么救、要不要回滚、怎么判断是坏数据还是坏超参"给出可操作答案。

代码逐段拆解

第一段：训练循环里的稳定性监控

import torch, numpy as np

EMA_ALPHA = 0.99
loss_ema = None
grad_norm_ema = None
spike_ratio = 3.0        # 触发判定：当前 loss / EMA 基线

for step, batch in enumerate(train_loader):
    loss, grad_norm = train_step(batch)   # train_step 内部返回 loss 与梯度范数

    # 指数滑动平均：建立"正常基线"
    loss_ema = loss if loss_ema is None else EMA_ALPHA * loss_ema + (1 - EMA_ALPHA) * loss
    grad_norm_ema = grad_norm if grad_norm_ema is None else \
        EMA_ALPHA * grad_norm_ema + (1 - EMA_ALPHA) * grad_norm

    # spike 检测：loss 相对基线跳升超过阈值
    if loss > spike_ratio * loss_ema and step > 500:   # 前 500 步是 warmup，不判
        print(f"[SPIKE] step={step} loss={loss:.3f} ema={loss_ema:.3f} "
              f"grad_norm={grad_norm:.2f} ema={grad_norm_ema:.2f}")
        save_marker(step, loss, grad_norm)   # 记录 spike 现场，供归因

    # 梯度爆炸护栏：超过阈值直接裁剪并降学习率（详见下）
    if grad_norm > 10 * grad_norm_ema:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        lr_scheduler.eta *= 0.5              # 半衰式降 LR，帮助恢复

要点：

1. EMA 是"基线"而非固定阈值：loss 曲线在预训练中整体下行，固定阈值会失效；EMA 跟随趋势，loss > 3×EMA 才是"偏离趋势的异常"。
2. warmup 阶段不判 spike：前几百步 loss 从初值快速下降，本身波动大，判定会误报。
3. 梯度范数是第二信号：loss spike 前往往先有梯度范数飙升。双信号（loss + grad_norm）能区分"坏 batch"（loss 突跳但梯度正常）和"训练发散"（梯度先爆）。
4. 护栏要"温和干预"而非"硬切"：clip 梯度 + 降 LR（半衰）让训练自己恢复，比直接回滚 checkpoint 少损失进度。

第二段：自动保存与断点重启

# 双保险：周期 Checkpoint + 异常 Checkpoint
save_every = 1000                       # 周期保存（配合 04 章存储设计）
marker_hits = 0

if step % save_every == 0:
    save_checkpoint(step, tag=f'periodic-{step}')
if loss > spike_ratio * loss_ema:       # spike 现场也存一份
    marker_hits += 1
    if marker_hits <= 3:                # 只保留前 3 次 spike 现场（防爆盘）
        save_checkpoint(step, tag=f'spike-marker-{step}')

# 异常重启：连续 N 步 loss 不回落后回滚
if loss > spike_ratio * loss_ema for last 200 steps:
    restore_checkpoint(nearest_before_spike)   # 回滚到 spike 前最近周期点
    lr_scheduler.eta *= 0.5
    print("[RECOVER] rolled back to pre-spike checkpoint, LR halved")

设计逻辑：

• 周期保存管"正常回滚"：每 1000 步一份，spike 后回滚的损失上限是 1000 步。
• spike 现场保存管"归因"：spike 那几步的权重快照用于事后分析（是不是某个 batch 有毒、是不是某层参数飞了）。
• 回滚判定要"连续"而非"单次"：单次 spike 可能自己恢复；连续 200 步不回落才回滚，避免误伤正常波动。

原理讲解

spike 的五类归因

类别	特征	处理
坏样本/脏数据	loss 单点跳、梯度正常	记录样本 ID，数据管线过滤
梯度爆炸	grad_norm 先飙升	clip + 降 LR；检查 LR 峰值是否过高
学习率过高	周期性小 spike、loss 震荡	降 LR 或换 warmup 调度
数值不稳定	特定层 NaN/Inf 前兆	检查 AMP 的 loss scale、LayerNorm eps
硬件/通信抖动	伴随节点报错/网络告警	回到 02 章 DCGM 与网络排查

判别顺序：先看梯度范数（爆 → 超参/数值问题），再看是否伴随硬件告警（是 → 环境问题），最后才怀疑数据——这个顺序避免在错误方向花时间。

为什么"自动重启"必须配合 Checkpoint 存储

预训练的中断恢复依赖 04 章的两级 Checkpoint：回滚只发生在"本地盘 + PFS"都有副本的点。若 Checkpoint 太稀疏（每 10000 步），spike 后回滚等于丢掉 1 万步进度；太频繁又拖慢训练。工程平衡：周期保存兼顾恢复粒度与开销（1000 步级），异常现场单独存（最多 3 份）。

实验结果

（估算口径，非实测）

稳定性手段	可挽回进度	代价
仅周期 Checkpoint	丢失 ≤ 保存间隔	保存 IO 开销
+ spike 检测告警	人工干预及时	监控开销可忽略
+ 自动护栏（clip/LR 半衰）	多数 spike 自愈	极端发散时无效
+ 自动回滚	恢复最近稳定点	回滚丢 1000 步

工程启示

1. 双信号监控是底线：只盯 loss 会漏掉"梯度已爆炸、loss 还没跳"的窗口期；loss + grad_norm 双 EMA 是标准配置。
2. 护栏分级而不是一锤子：clip → 降 LR → 回滚，逐级温和干预，能自愈就不回滚——回滚是最后手段。
3. spike 现场要可复现：记录触发那几步的样本索引与权重快照，否则"训练白挂、原因成谜"。
4. 稳定性工程与存储、可观测性是一套：04 章存储、02 章 DCGM、本篇的 loss 监控要进同一监控体系，形成"发现-归因-恢复"闭环。

下一篇预告：预训练稳定了，推理侧的工程链也补得差不多了——下一篇回到 13 应用章，拆解"科学智能体落地：从工具调用到科学工作流自动化"。