NPU卡训练时报NaN，如何定位数据溢出根源？

Question

在NPU卡上进行深度学习训练时，模型输出出现NaN，常由数值溢出引发。如何判断是数据预处理、梯度爆炸还是NPU定点计算精度限制导致的溢出？特别是在使用低精度（如FP16或BF16）时，激活值或梯度异常增长是否与归一化缺失、学习率过高或权重初始化不当有关？如何通过插入调试算子、监控张量范围及利用NPU提供的溢出检测工具快速定位根本原因？

Answer 1

在NPU卡上深度学习训练中NaN问题的系统性诊断与根因定位

1. 问题背景与现象描述

在基于NPU（神经网络处理单元）进行深度学习模型训练时，模型输出出现NaN（Not a Number）是常见的稳定性问题。该现象通常由数值溢出引发，尤其是在使用低精度浮点格式（如FP16或BF16）时更为显著。溢出可能源于多个层面：数据预处理不当、梯度爆炸、权重初始化不合理、归一化缺失、学习率设置过高，或NPU定点计算过程中的精度截断。

2. 溢出来源的层级划分

为系统排查，可将潜在原因划分为以下三类：

• 数据预处理层：输入数据未归一化、存在极端离群值、标签错误等。
• 模型训练动态层：梯度爆炸、激活值激增、学习率过大、权重初始化不当。
• NPU硬件计算层：低精度算术运算导致的舍入误差累积、定点数范围溢出、NPU固件对异常值处理机制不足。

3. 判断路径：从表象到根因的递进分析

采用“由浅入深”的排查策略，逐步缩小问题范围。

1. 确认是否在训练初期即出现NaN（指示初始化或学习率问题）。
2. 检查输入数据分布，是否存在inf或极大值。
3. 监控前向传播中各层激活值的最大绝对值。
4. 记录反向传播中各层梯度的L2范数。
5. 对比FP32与FP16/BF16下的训练行为差异。
6. 启用NPU提供的溢出检测工具（如华为Ascend的AICORE异常捕获）。
7. 插入调试算子（如DebugPrint或CheckNumerics）定位具体算子。
8. 验证归一化层（BatchNorm/LayerNorm）是否生效。
9. 评估优化器是否启用梯度裁剪（Gradient Clipping）。
10. 检查权重初始化策略（如Xavier、He初始化）是否适配当前激活函数。

4. 数据预处理相关性分析

问题类型	典型表现	检测方法	解决方案
未归一化输入	第一层激活值>1e5	打印input.max()	添加Z-score归一化
标签异常	Loss初始即为inf	检查label.min/max	清洗标签数据
离群样本	单步训练后NaN	逐样本训练测试	数据增强+异常检测
数据类型不匹配	NPU转换时报错	dtype检查	显式转换为float32

5. 梯度爆炸与模型动态分析

当使用低精度训练时，梯度若未受控，极易在累加过程中超出FP16表示范围（~65504）。常见诱因包括：

• 学习率过高（如>1e-3用于Adam）
• 缺失LayerNorm或BatchNorm
• ReLU类激活导致激活值无界增长
• 循环结构（RNN/LSTM）中长期依赖积累

可通过以下代码片段监控梯度：

def log_gradients(named_params):
    for name, param in named_params:
        if param.grad is not None:
            grad_norm = param.grad.data.norm(2).item()
            if grad_norm > 1e4:
                print(f"[Warning] Large gradient in {name}: {grad_norm}")

6. NPU定点计算与精度限制影响

NPU常采用定制化数据通路，支持INT8/FP16/BF16等格式。其内部可能使用块浮点（Block Floating Point）或定点量化，导致动态范围受限。例如：

• FP16指数位仅5bit，易发生上溢（>65504）或下溢（<6e-5）
• 某些NPU算子融合可能导致中间结果截断
• 缺乏IEEE 754标准兼容的NaN传播机制

7. 调试工具链与根因定位流程图

graph TD A[训练出现NaN] --> B{是否首步即发生?} B -- 是 --> C[检查数据输入与标签] B -- 否 --> D[插入CheckNumerics算子] C --> E[打印input.max/min] D --> F[定位首个输出NaN的算子] F --> G{是否为MatMul/Conv?} G -- 是 --> H[检查权重初始化] G -- 否 --> I[检查Activation函数] H --> J[改用He/Xavier初始化] I --> K[替换ReLU为LeakyReLU] J --> L[启用梯度裁剪] K --> L L --> M[切换至FP32验证] M --> N{是否仍出现NaN?} N -- 否 --> O[确认为精度问题] N -- 是 --> P[检查NPU驱动与固件版本]

8. 实践建议与防御性编程

为提升训练鲁棒性，建议实施以下措施：

• 始终在输入端进行标准化：(x - mean) / std
• 使用梯度裁剪：torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
• 启用混合精度训练中的损失缩放（Loss Scaling）
• 定期保存中间张量用于离线分析
• 利用NPU厂商提供的Profiling工具（如Ascend Profiler）监控Tensor范围
• 在关键算子后插入tf.debugging.check_numerics或PyTorch等效实现
• 避免使用sigmoid或tanh在深层网络开头
• 对Embedding输出做l2_normalize
• 设置学习率warmup阶段以平滑初始梯度
• 记录每层输出的mean和std用于趋势分析