论坛报错活动第三十八期-使用MindSpore混合精度模式训练出现Loss NaN

经验分享 Tech Blogs 功能调试-Function Debugging
1

1 系统环境

硬件环境(Ascend/GPU/CPU): Ascend/GPU/CPU
MindSpore版本: mindspore=2.6.0
执行模式（PyNative/ Graph）: 不限
Python版本: Python=3.9
操作系统平台: Linux

2 报错信息

2.1 问题描述

使用MindSpore的混合精度训练一个模型时，遇到了Loss变为NaN的问题。在FP32模式下训练正常，但开启混合精度后几个epoch就会出现NaN。
AMP（自动混合精度）模式：训练初期正常，随后loss突变为NaN，梯度检查发现某些层出现极大值（1e+10级别）。

2.2 脚本信息

import mindspore as ms
from mindspore import nn, amp

# 模型定义
class LargeModel(nn.Cell):
    def __init__(self):
        super().__init__()
        self.backbone = create_backbone()  # 复杂骨干网络
        self.head = create_head()          # 分类头
       
    def construct(self, x):
        return self.head(self.backbone(x))

# 混合精度配置
model = LargeModel()
optimizer = nn.Adam(model.trainable_params(), learning_rate=0.001)
loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')

# 启用混合精度
model = amp.auto_mixed_precision(model, level="O2")  # O2级别

def train_step(data, label):
    def forward_fn(data, label):
        logits = model(data)
        loss = loss_fn(logits, label)
        return loss, logits
   
    grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)
    (loss, _), grads = grad_fn(data, label)
   
    # 这里出现NaN
    if ms.ops.isnan(loss).any():
        print("Loss is NaN!")
       
    optimizer(grads)
    return loss

3 根因分析

既然Float32训练正常,切换到混合精度后loss出现NaN.说明在计算过程中出现了溢出.

4 解决方案

两种解决方法:

1.找出对数值精度敏感的层,强制采用FP32运算

amp.auto_mixed_precision提供了参数 custom_fp，允许我们自定义哪些层使用FP16，哪些层强制使用FP32。

先将可能对数值精度敏感的关键模块强制设为FP32，观察问题是否消失。如果消失，再逐步缩小范围。

from mindspore.amp import custom_fp
model = amp.auto_mixed_precision(
    model,
    level="O2",
    custom_fp=custom_fp.CustomFP16(black_list=[nn.SoftmaxCrossEntropyWithLogits, CustomNorm])
)

2.全局梯度裁剪

clip_val = 1.0  # 初始阈值，根据监控调整
clip = nn.ClipByGlobalNorm(clip_val)

# 在训练步骤中，optimizer更新前应用
grads = clip(grads)
optimizer(grads)

如果训练收敛缓慢，且监控发现每一步梯度都被裁剪，说明阈值太小，应调大（如到3.0或5.0）

如果仍然偶尔出现NaN（在问题层固定为FP32后），可以微调阈值（如从1.0调到1.5）

目标是找到一个阈值，使得训练稳定，且梯度裁剪只是“偶尔”发生，用于应对突发的梯度波动.