溢出或NaN问题

在进行工具定位前，请优先排除Checklist检查中的配置项问题和问题复现中的随机性问题。

案例1

某视觉模型从GPU迁移到NPU MindSpeed-LLM训练，从一开始就梯度溢出。

图1 梯度溢出打印的日志

从用户共享的训练截图中可以看到第0步梯度反向时逐层变大直至溢出。

定位方法：

使用dump采集工具采集第0步（溢出步）的mix级别数据，config.json配置如下：

{
    "task": "statistics",
    "dump_path": "/home/data_dump",
    "rank": [],
    "step": [0],
    "level": "mix",
    "enable_dataloader": false,
    "statistics": {
        "scope": [], 
        "list": [],
        "data_mode": ["all"],
        "summary_mode": "statistics"
    }
}

从训练截图中可看到每次self_attn反向之后梯度逐层变大，查看self_attn代码发现使用了FA算子，该算子历史上因使用规范引起的精度问题较多，优先查看dump中对应的反向数据，发现每次经过FA层反向后，norm值量级明显增大。

图2 dump采集的FA结果

快速验证，先在MindSpeed-LLM的训练配置中规避FA融合算子。
删除超参：-use_fused_attn

溢出消失，明确该问题为FA分支引入，但性能下降明显，需进一步明确FA精度原因。
通过查阅FA算子官网使用文档，分析FA算子在代码中的具体使用方式。
该问题为变长场景，原始输入为batch size=2，样例1的seqlen=3577，样例2的seqlen=1507，统一pad到3577长度，原始输入shape=[2, 3577, 32, 128]。

在进FA计算前，会将batch size和seq len做flatten，此时shape=[7154, 32, 128]，下一步去除其中的pad，因此Q和KV的输入长度变成了[5079, 32, 128]。

atten_mask字段要求：

atten_mask：Device侧的Tensor，可选参数，取值为1代表该位不参与计算（不生效），为0代表该位参与计算，数据类型支持BOOL、UINT8，数据格式支持ND格式，输入shape类型支持BNSS格式、B1SS格式、11SS格式、SS格式。varlen场景只支持SS格式，SS分别是maxSq和maxSkv。

按照官网说明，此时attention mask按照规则本该为[maxSq, maxSkv]，即[3577, 3577]，但实际客户代码中使用[query.shape[0], key.shape[0]，即[5079, 5079]，使用规范错误，导致算子底层执行计算时会按行读取，导致出现0、1的数值错位，最终导致梯度溢出。

解决方案：修正FA训练时传入的attention_mask。

结果：训练梯度溢出消失，loss正常收敛。

案例2

某多模态模型从GPU迁移到NPU后做微调，使用框架为FSDP，训练第2步loss出现NaN。

图3 NPU上运行结果

图4 GPU上运行结果

定位方法：

缩小规模。
该模型现网为128卡训练，做实验成本大，需首先缩小规模，减少层数后可在单机2卡稳定复现。
使用dump工具采集step1（最开始出现NaN的步数）的mix级别数据。
加入工具后，发现NaN问题消失。

去除工具并打开流同步进一步验证：
```
export ASCEND_LAUNCH_BLOCKING=1
```
开启流同步后问题也消失。

基于以上2个现象怀疑该FSDP模型训练存在内存踩踏问题。
缩小排查范围。
该模型由四个部分组成：vae，dit，denoiser和conditioner。

从训练完整模型改为只训练dit.transformer.layers，loss仍然有NaN，确认是transformer.layers的问题。
通过手动挂局部hook的方式打印梯度。
发现第1步loss的NaN不是第一现场，先出现NaN的是第0步post_attention_layernorm的反向梯度。

图5 post_attention_layernorm的反向梯度

与打开流同步的无NaN的梯度数据进行对比，除了input_layernorm和post_attention_layernorm层的weight和bias，其余的参数都能对上。

图6 有无NaN的梯度比对

对应的dump中的接口为Functional.layer_norm.10和Functional.layer_norm.11。
结合具体代码进行分析。
post_attention_layernorm对于图像和文本连续下发了两次。

图7 post_attention_layernorm代码

将其次数改为1次时候NaN消失，明确该问题出现在该算子重复调用时。

分析内存踩踏特征的方式是按异常数据是否存在规律性和连续性，所以需先采集对应数据再进行分析。
改用异步dump。
之前加入dump工具后NaN消失原因为观测到Tensor后，取统计量（min, max等）和落盘的操作会影响流上算子的执行，导致NaN不复现。

通过改异步dump方式，训练过程中工具不触发同步操作，改为在当前step训练结束后统一落盘，降低对算子执行顺序和流同步影响。

具体操作为：在config.json文件中加入async_dump: True的配置项。

重新采集Functional.layer_norm.10和Functional.layer_norm.11及其中间的torch.split.192反向数据，可在dump单个算子时复现NaN。
分析异步dump数据。
参考无loss NaN的dump.json文件，torch.split.192.backward的输入应为Functional.layer_norm.11的输出，而不开流同步时异步dump的torch.split.192.backward的输入与Functional.layer_norm.11的输出对不上，对比本该相等的2组数据：

图8 特征分析代码

发现刚好踩了size=2048（0-2047不等，2048-3071相等），满足内存踩踏特征。

图9 踩踏前后的数据差异
算子内存地址打印。
尝试通过修改torch_npu源码对算子的输入输出tensor对应的ptr地址和shape进行打印。

图10 ptr内存地址打印结果

从日志发现两个连续layernorm中，存在cast算子输出对concat算子输入的踩踏（两者地址一致）。

踩踏现场确认如下：

图11 踩踏发生的逻辑图

总结根因：缺失record的backend +多流并行的FSDP +连续下发的layernorm导致了内存踩踏。

解决方案：在torch_npu2.3的FSDP unshard流上添加record，确保流上的当前算子执行完成之前，tensor内存不会被下一个算子申请。

结果：loss NaN消失，正常收敛。

父主题： msprobe工具定位案例