1-混合精度-梯度裁剪-EMA

1-混合精度、梯度裁剪与 EMA

1. 为什么训练稳定性在大规模训练中至关重要

大规模训练（深层网络、大 batch、百亿参数）带来了独特的数值稳定性挑战：

• 参数量大 → 梯度累积过程中数值上溢/下溢风险增加
• 层数深 → 梯度在反向传播中逐层放大/衰减（链式法则的乘法效应）
• 训练时间长 → 单次崩溃可能导致数十GPU/天的计算资源浪费
• 大模型断点重训成本高 → 必须从可靠的 checkpoint 恢复

本文档聚焦三个最核心的数值工程技巧：混合精度训练（AMP）、梯度裁剪、指数滑动平均（EMA），以及它们背后的数值问题检测与排查方法。

---

2. 混合精度训练

2.1 FP32、FP16、BF16 的区别

深度学习中的数值格式主要使用 IEEE 754 浮点数标准，核心构成是：1 位符号（sign）+ 位指数（exponent）+ 位尾数（mantissa）。

格式	符号位	指数位	尾数位	指数范围（以 2 为底）	精度（ULP）
FP32	1	8	23		约 7 位十进制
FP16	1	5	10		约 3.3 位十进制
BF16	1	8	7		约 2 位十进制

FP16 的问题：下溢（underflow）

FP16 的最小正规格化数约为 。在梯度反向传播中， 的连乘效应（ 为层数）使得深层网络的梯度可能远小于这个值，直接下溢为 0，导致参数几乎不更新。

FP16 的另一个问题：NaN 风险

FP16 的指数范围只有 ，当梯度乘以学习率（）超过 时，发生上溢（overflow），变为 Inf，再参与运算后变为 NaN。

BF16 的设计

BF16（Brain Float16，由 Google Brain 提出）保留了与 FP32 相同的 8 位指数，牺牲了 7 位尾数（相比 FP16 的 10 位）。这使得 BF16 的指数范围与 FP32 相同，从根本上消除了上溢风险；同时保留足够大的指数位来表达数量级，避免下溢。BF16 的精度较低，但在大多数深度学习训练中足够。

2.2 为什么混合精度能加速并省显存

混合精度训练的核心思想是：在能够保持数值稳定的前提下，用较低精度的格式存储和计算，以节省显存和算力。

显存节省来源

模型状态在混合精度训练中的存储格式：

状态类型	存储精度	说明
模型参数（权重）	FP16 或 BF16	前向/反向计算使用
优化器状态（一阶/二阶矩）	FP32	Adam 需要精确的一阶/二阶矩
梯度	FP16	与参数同精度存储
激活值	FP16	反向传播必需，显存占用大头
优化器状态副本	FP32	某些实现中为保证精度

以 Adam 优化器为例：标准 FP32 训练时，Adam 的二阶矩 需要 （FP32）的存储；但混合精度下 可以存为 FP16（实践中仍存 FP32 以保持精度）。整体而言，混合精度可将单卡显存占用减少约 30%~50%。

加速来源

现代 GPU（V100、A100、H100）的 Tensor Core 对 FP16/ BF16 矩阵运算有硬件加速，FP16 的矩阵乘法吞吐是 FP32 的 8~16 倍。混合精度利用 Tensor Core 加速，同时保持主权重和优化器状态为 FP32 以确保精度。

2.3 动态 Loss Scaling 的原理

数值下溢的来源

在 FP16 中，梯度 时会下溢为 0。神经网络训练中，loss 乘以系数（loss scaling factor）后，梯度被放大：

若 scaling factor 足够大（），则 ，使原本下溢的梯度落入 FP16 可表示的范围内。

动态 Loss Scaling 的算法

静态 Loss Scaling：手动设置一个固定 scaling factor（如 512 或 1024），简单但不够自适应。

动态 Loss Scaling（AMP 的默认做法）：

1. 初始阶段：使用较大的 loss scaling（如 ）
2. 检测：每一步检查是否存在 Inf/NaN（通过scaler.scale(loss).backward()）
3. 若发生下溢：scaler 减半（scaling factor ÷ 2），跳过当前 step
4. 若无异常连续 步：scaler 加倍（scaling factor × 2），允许使用更大的 scaling

# PyTorch AMP 的动态 Loss Scaling
scaler = torch.cuda.amp.GradScaler()
for data, target in dataloader:
    with torch.cuda.amp.autocast():
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()   # loss scaling + backward
    scaler.unscale_(optimizer)      # 取消 scaling 以便裁剪梯度
    scaler.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

2.4 混合精度的工程实现

PyTorch 的 autocast 与 GradScaler

# 训练循环（标准 AMP 写法）
scaler = torch.cuda.amp.GradScaler()
 
for epoch in range(num_epochs):
    for step, (x, y) in enumerate(dataloader):
        optimizer.zero_grad()
        with torch.cuda.amp.autocast(dtype=torch.float16):
            output = model(x)
            loss = criterion(output, y)
        scaler.scale(loss).backward()          # 1. 放大 loss
        scaler.unscale_(optimizer)              # 2. 取消 scaling（用于裁剪）
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)                 # 3. 更新（含梯度 unscaling）
        scaler.update()                         # 4. 更新 loss scaling factor

AMP 的工作流程分解

步骤	操作	说明
1	autocast() 进入 FP16 模式	前向传播中指定 ops 用 FP16 执行
2	scaler.scale(loss)	乘以 loss scaling factor，防止下溢
3	.backward()	反向传播，梯度以 FP16 存储
4	scaler.unscale_(optimizer)	将梯度恢复到正确尺度（用于裁剪/优化器）
5	clip_grad_norm_()	梯度裁剪
6	scaler.step(optimizer)	Adam 更新（用 unscaled 梯度）
7	scaler.update()	根据溢出情况更新 loss scaling factor

---

3. 梯度裁剪

3.1 梯度爆炸的数值根源

梯度爆炸（gradient explosion）的根本原因是反向传播中梯度的链式乘法。

设第 层输出 ，。反向传播中：

若每层权重的谱半径（最大奇异值），则每层反向传播时梯度被放大：

层后，梯度被放大为 倍。若某些层的 ，12 层后梯度将被放大 倍；24 层后 倍；100 层后指数爆炸不可控。

3.2 全局范数裁剪（Gradient Clipping by Global Norm）

全局范数裁剪将所有参数的梯度向量 视为一个整体，裁剪其全局范数：

裁剪公式：

即：如果全局范数超过 max_norm，将所有梯度按同一因子缩放，使全局范数恰好等于 max_norm（而非直接截断，截断会导致梯度信息损失）。

3.3 逐元素裁剪（Per-element Clipping）

另一种裁剪方式是对每个梯度分量独立裁剪：

这比全局范数裁剪更粗糙，可能将梯度方向扭曲（尤其当不同层的梯度方向不一致时），全局范数裁剪是标准做法。

3.4 为什么梯度裁剪能防止梯度爆炸

考虑一个简单场景：参数 的真实梯度为 ，但由于numerical instability，实际计算的梯度为 （，放大了 倍）。直接更新：

这相当于使用了 倍大的学习率，可能导致参数在 loss 曲面上 overshoot，甚至发散。

梯度裁剪将 限制在合理范围内（），使得即使梯度被放大 倍，更新步长仍受控制：

由于 被 归一化，等价于使用 而非 ，参数更新量被控制。

3.5 裁剪阈值如何选择

场景	推荐 max_norm	说明
通用视觉 CNN		经验值，1.0 是最常用的默认值
Transformer / BERT / GPT		Transformer 训练的标准配置（BERT 作者建议 ）
RNN / LSTM	或	RNN 中梯度爆炸风险更高，严格裁剪
强化学习（PPO）		PPO 中 policy gradient 梯度可能很大
大 batch 训练		大 batch 梯度更稳定，但也要防止罕见大梯度

判断方法：观察训练中 grad_norm 的分布。若大多数时候 grad_norm < 1.0，则 max_norm = 1.0 几乎不触发裁剪（正常）；若经常出现 grad_norm > 5.0 的峰值，说明存在梯度爆炸，需降低 max_norm 或配合学习率调整。

---

4. EMA（指数滑动平均）

4.1 定义与更新公式

EMA 在训练过程中对模型参数做滑动平均：

其中 是衰减系数（通常取 ）， 是第 步的当前参数。

展开递推：

这表明 是过去 步的加权平均，权重按 指数衰减。 意味着 EMA 窗口约覆盖过去 1000 步的参数。

4.2 为什么 EMA 常带来更平滑的验证表现

直觉上，EMA 参数 是一个”慢跟踪”版本——它不响应每个 step 的参数波动，而是以一定延迟平滑地跟随。这种平滑效应在训练中带来几个优势：

1. 减少参数振荡：训练末期的 loss 曲线趋于平稳，但参数仍在极小值附近振荡（尤其大学习率时）。EMA 平滑了这种振荡，使验证集上的预测更稳定。
2. 对抗短时过拟合：若训练在某个 step 出现异常波动（如梯度爆炸导致参数跳到错误位置），EMA 不会立即跟随——因为 的权重系数意味着只有极小一部分异常波动传导到 EMA 参数。
3. ensemble 效应：EMA 相当于对过去 1000 步的参数做了”软集成”，类似于在参数空间上的滑动平均集成（parameter averaging），这在理论上可证明能收敛到更宽的极小值。

4.3 EMA 的数学解释：平滑化与置信度

从贝叶斯角度，EMA 可解释为对参数后验分布的蒙特卡洛近似。当前参数 是对真实最优参数的无偏估计，但其方差较大。EMA 通过加权平均，将多个高方差样本组合成一个低方差的估计：

时，EMA 的方差约为单个参数估计的 1000 倍平滑——即方差缩减为原来的 。

4.4 EMA 在工程中的使用

class EMA:
    def __init__(self, model, decay=0.9999):
        self.shadow = {}
        self.decay = decay
        # 初始化 shadow 参数为当前模型参数的副本
        self.register(model)
 
    def register(self, model):
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()
 
    def update(self, model):
        for name, param in model.named_parameters():
            if param.requires_grad:
                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
                self.shadow[name] = new_average.clone()
 
    def apply_shadow(self, model):
        # 用 EMA 参数覆盖当前模型参数
        for name, param in model.named_parameters():
            if param.requires_grad:
                param.data = self.shadow[name].clone()

使用流程：

• 训练时：每 step 调用 ema.update(model) 更新 EMA shadow
• 推理时：先 ema.apply_shadow(model)，再做推理（或者直接加载 EMA 参数做推理）
• 注意：EMA 参数不参与训练，训练仍是正常的 Adam/SGD；EMA 只是提供了一个更好的参数快照

---

5. 残差网络、Transformer、RNN 中的数值不稳定来源

5.1 残差网络（ResNet）

ResNet 的数值不稳定来源主要是残差路径信号方差在深层累积：

若残差分支 输出方差 ，则 的方差 的方差，逐层放大。解决方案：

• 在残差分支中使用 BatchNorm 控制输出方差（BatchNorm 使 的输出方差接近 1）
• 使用 Zero-init shortcut（将残差分支输出权重初始化为 0，使初始时残差路径等效于恒等映射）

5.2 Transformer

Transformer 的数值不稳定来源是注意力分数的尺度问题：

若 过小， 的数值尺度可能很大，导致 softmax 的输入指数爆炸；若过大，则 softmax 输入趋近于 0，梯度消失。

训练不稳定的其他来源：

• 残差累积：Transformer 有 个残差路径，每个残差路径的输出方差在层间叠加。若 SubLayer 输出方差 ，信号逐层放大
• LayerNorm 位置：Post-norm 在残差后做归一化，在训练初期可能因为 的信号叠加导致第一层输出方差过大，需要 warmup 配合稳定训练
• Pre-norm 的优势：Pre-norm 在 SubLayer 前做归一化，使得每一层的输入方差被控制，从根本上缓解深层信号放大问题

5.3 RNN / LSTM

RNN 的数值不稳定主要表现为梯度消失和梯度爆炸，根源在于时间步展开后的链式乘法：

若循环权重矩阵 的谱半径 ，则反向传播时梯度指数增长（爆炸）；若 ，则指数衰减（消失）。

LSTM 通过门控机制（input、forget、output gate）来控制信息的保留和丢弃，将梯度传播路径分解为多个可控的门控操作，从而缓解梯度消失/爆炸。但 LSTM 的门控参数仍然可能因为初始化不当导致问题——例如 forget gate 偏置过小导致信息过快衰减。

---

6. 梯度裁剪与 AMP 的协同使用

6.1 为什么需要先 unscale 再 clip

GradScaler.unscale_() 的作用是将缩放后的梯度恢复到正确的数值尺度，以便进行裁剪。流程是：

1. scaler.scale(loss).backward() → 梯度以 FP16 存储，但数值被 loss scaling factor 放大
2. scaler.unscale_(optimizer) → 将梯度恢复到接近 FP32 的数值尺度（但仍是 FP16 存储）
3. clip_grad_norm_() → 在正确的数值尺度下做全局范数裁剪
4. scaler.step(optimizer) → scaler 会自动根据 unscale 后的梯度决定是否跳过该 step

若不执行 unscale 而直接 clip，则 clip 的阈值（如 max_norm=1.0）对应的是缩放后的梯度尺度，与设计意图不符。

6.2 NaN 和 Inf 的来源分析

NaN / Inf 在训练中出现的典型来源：

来源	原因	表现
除零	归一化分母 中	激活/梯度出现 NaN
上溢（overflow）	FP16 指数范围 ， 超出范围	Inf → NaN
下溢（underflow）	梯度太小（）被表示为 0	训练停滞（而非 NaN）
Log(0)	BCE / Cross-Entropy 中的 ，	NaN
Softmax 溢出	logits 太大（ → ）	Inf → NaN
梯度爆炸	某 step 梯度范数急剧增大，后续计算变为 Inf	loss 或 grad norm 突变

---

7. 工程实践：监控与排查

7.1 训练日志中应该监控的指标

每个训练 step 应记录以下指标（建议每 10~100 step 记录一次）：

指标	正常范围	异常信号
Loss	单调不增（或轻微振荡下降）	突然 NaN / Inf；突然跳升
Grad Norm（全局梯度范数）	（依任务而异）	（梯度爆炸）；（梯度消失）
LR（当前学习率）	符合 schedule	恒为 0（scaler 持续 halving）
Overflow Count（scaler 跳 step 次数）	0 或极少数	连续多步 overflow → loss scaling 配置错误
Activation Stats（各层激活均值/方差）	均值接近 0（除 ReLU 正区间），方差稳定	某层方差突增 10 倍 → 信号放大
Weight Norm（权重范数）	缓慢增长或恒定	快速单调增长 → 学习率过大

7.2 如何定位某一步开始变坏

步骤 1：二分搜索定位故障点

若训练在 step 时出现 NaN，但之前 步正常，可用二分法定位：

# 伪代码：checkpoint 逐步加载测试
# 从第 T/2 步加载，检查是否出现异常
# 若正常：从 T*3/4 步加载测试
# 若异常：从 T/4 步加载测试
# 递归直到定位到具体 step

步骤 2：逐层激活值分析

若定位到某 step ，在这一步插入逐层激活统计：

for name, param in model.named_parameters():
    print(f"{name}: mean={param.mean():.6f}, std={param.std():.6f}, "
          f"max={param.abs().max():.6f}")

检查哪一层的激活/梯度在 时出现异常峰值。

步骤 3：梯度成分分析

有时 NaN 不是来自激活，而是来自某些参数的梯度爆炸（如某个 Embedding 的梯度异常大）。通过记录 grad_norm 的分解：

grad_norms = {}
for name, param in model.named_parameters():
    if param.grad is not None:
        grad_norms[name] = param.grad.norm().item()
# 排序，找出norm最大的参数
sorted_norms = sorted(grad_norms.items(), key=lambda x: x[1], reverse=True)
print(sorted_norms[:5])

7.3 Checkpoint 与恢复训练

何时保存 Checkpoint

场景	保存频率	保存内容
正常训练	每 个 epoch（）	模型参数 + 优化器状态 + EMA +Scaler
大模型（> 1B）	每 1K~10K steps（防止崩溃损失超过 1 hour）	同上 + random state
验证集评估	每次保存 checkpoint 时同步评估	仅模型参数 + EMA

Checkpoint 应包含的状态

checkpoint = {
    'epoch': epoch,
    'step': step,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'ema_shadow': ema.shadow if use_ema else None,
    'scaler_state_dict': scaler.state_dict() if use_amp else None,
    'random_state': torch.get_rng_state(),  # 随机种子
    'torch.cuda_rng_state': torch.cuda.get_rng_state_all(),  # GPU 随机状态
}

断点重训的关键注意事项

1. Optimizer state 必须同时恢复：Adam 的 momentum 和 variance 在恢复后必须与断点前完全一致，否则参数更新轨迹会分叉，导致收敛到不同的局部最优
2. Scaler state 必须恢复：AMP 的 loss scaling 因子必须恢复，否则会导致异常的 loss scaling 行为
3. BatchNorm 统计量恢复：若从 checkpoint 恢复时 running_mean/var 与训练状态不一致，会导致推理结果突变
4. 随机种子必须一致：若训练中使用了随机增强（如 data augmentation 的随机裁剪），恢复时若随机种子不一致，相同的输入数据会经过不同的增强，导致结果不同

---

8. 训练稳定性排错清单

问题分类与处理

A. Loss 突然变成 NaN / Inf

检查项	处理
Grad Norm 正常但 Loss NaN	检查 loss 计算本身（如 in BCE）；检查是否在某些位置用了 torch.log(sigmoid(logit)) 而非 F.binary_cross_entropy_with_logits
Grad Norm 爆炸前兆（>）	启用梯度裁剪 max_norm=1.0；降低学习率
Loss 在某个 layer 后 NaN	该层的激活或梯度数值溢出；检查该层权重初始化；确认无除零
FP16 训练出现 NaN	开启 AMP bf16（dtype=torch.bfloat16）；或使用 GradScaler 的动态 scaling
检查输入数据是否有 NaN	在数据加载后加 assert not torch.isnan(x).any()

B. Loss 完全不下降（训练停滞）

检查项	处理
Grad Norm 极小（）	梯度消失；检查激活函数（是否饱和）；检查初始化；检查是否有 detach() 错误截断了梯度
Grad Norm 正常但 Loss 不降	学习率过低；模型结构问题（如某层无梯度）；损失函数配置错误
梯度被 detach() 截断	检查是否有手动 .detach() 调用；检查是否有 inplace 操作阻断梯度
学习率被 GradScaler 持续 halving	Loss scaling factor 持续下降（get_scale() < 1）→ 切换到 BF16 或调高初始 loss scaling

C. Loss 振荡不收敛

检查项	处理
学习率过大	降低学习率 5~10 倍
Batch size 太小	增大 batch size（或梯度累积步数）
梯度裁剪过于频繁（grad norm 经常 ）	降低 max_norm；或配合降低学习率
学习率 schedule 过于激进	改用 cosine decay 或更慢的 step decay

D. 训练到后期突然发散

检查项	处理
学习率衰减后发散	Cosine decay 的末期学习率过低（）；保持末期学习率在 量级
EMA 震荡	可能对当前任务过大；尝试 或
BatchNorm 统计量漂移	在新数据集上 fine-tune 时重新计算 BN 统计量

---

9. 默认工程配置建议

9.1 通用深度学习训练

# 混合精度 + 梯度裁剪 + EMA 的标准配置
scaler = torch.cuda.amp.GradScaler()
ema = EMA(model, decay=0.9999)
 
for epoch in range(num_epochs):
    for step, (x, y) in enumerate(dataloader):
        optimizer.zero_grad()
        with torch.cuda.amp.autocast(dtype=torch.float16):
            output = model(x)
            loss = criterion(output, y)
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)                      # 1. unscale
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 2. 裁剪
        scaler.step(optimizer)                          # 3. 更新
        scaler.update()
        ema.update(model)                              # 4. EMA 更新
 
# 验证时使用 EMA 参数
ema.apply_shadow(model)
val_loss = evaluate(model, val_loader)

9.2 Transformer 大模型训练（GPT / LLaMA）

# 使用 BF16（若硬件支持 A100 / H100）
with torch.cuda.amp.autocast(dtype=torch.bfloat16):
    output = model(x)
    loss = criterion(output, y)
 
# 梯度裁剪：Transformer 推荐 max_norm=1.0
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
 
# EMA decay：0.9999 ~ 0.99999（参数量越大，EMA decay 越接近 1）
ema = EMA(model, decay=0.9999)
 
# 学习率：warmup + cosine decay，末期不低于 base_lr / 10
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=2000,
    num_training_steps=total_steps
)

9.3 CV 检测/分割任务

# 目标检测场景：batch size 可能较小，使用 GroupNorm 替代 BatchNorm
model = build_detection_model(use_group_norm=True)
 
scaler = torch.cuda.amp.GradScaler()
# 检测任务建议 max_norm=10.0（比 Transformer 更宽松）
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0)
 
# SGD + Momentum 配置
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

---

10. 不同硬件上的精度选择

硬件	推荐精度	说明
V100	FP16（Tensor Core）+ GradScaler	V100 的 BF16 支持不完整，优先 FP16
A100	BF16（原生 Tensor Core 支持）	A100 对 BF16 有完整硬件支持，无上溢风险
H100	BF16 + FP8（可选）	H100 新增 FP8 Tensor Core，大幅加速
RTX 3090 / 4090	FP16	无 BF16 原生支持，需要软件模拟（不推荐）
CPU 训练	FP32	Tensor Core 不可用，只能用 AVX-512（无加速）

A100 + BF16 的优势：BF16 的指数范围与 FP32 相同，从根本上消除了 FP16 中梯度上溢导致 NaN 的风险，同时保留了比 FP32 更低的显存占用和更高的吞吐。因此，有 A100/H100 的场景，强烈推荐使用 BF16 而非 FP16。