股票 
板块 
用户 
资讯 在大模型训练过程中，loss（损失）突变NaN是令开发者最头疼的问题之一。NaN意味着训练过程中的数值计算出现了异常（如溢出、除零等），导致梯度无法正常更新，训练被迫中断。
这类问题棘手的原因，主要体现在三个方面：
随机性：loss/ grad norm（梯度范数）可能在随机step出现NaN，重新运行训练不一定在相同 step 复现。
定位困难：loss/grad norm只是最终结果，它只能告诉你“训练出问题了”，但无法指出是哪个计算算子、哪张卡（Rank）率先产生了异常数据。
传播迅速：某个算子输出NaN后，异常数会沿着计算图和通信链路迅速扩散。等到loss/grad norm指标暴露问题时，源头往往已经被淹没在大量异常数据中。
本文以一次32卡MoE（混合专家）模型训练为例，介绍如何借助MindStudio精度调试工具，系统性地定位并复现这类随机出现的NaN问题。
MindStudio异常值定位方法
MindStudio精度调试工具可以检查算子输入输出中是否出现NaN，并将结果写到dump.json。定位分为整网数据采集、异常分析、算子输入输出数据采集和单算子用例复现四步，具体流程图如下图所示。
MoE模型训练随机NaN定位实践
案例背景：32卡（4xAtlas 800T A2）MoE模型训练在第8步日志中观测到loss为NaN；该问题复现存在随机性，重启训练后无法稳定在相同步数复现。
步骤一：整网数据采集
先覆盖全部32卡和全部step，持续检查API前向、反向的输入输出。这样既能回查step8之前的异常，也能保留跨卡分析需要的信息。
采集的接入方式也很简单，只需要在原有训练代码里增加少量调试逻辑：
{"task": "nan_check","dump_path": "./nan_check_dump","level": "L1"}import torchimport torch_npufrom msprobe.pytorch import PrecisionDebuggerdebugger = PrecisionDebugger(config_path="./nan_check_config.json")# 在现有训练循环中，覆盖完整的前向和反向计算。for batch in train_loader:optimizer.zero_grad()debugger.start(model)loss = train_step(batch)loss.backward()optimizer.step()debugger.stop()debugger.step()torch.npu.synchronize()
采集结果如下图所示：
步骤二：分析数据，缩小范围
训练日志在step8首次报告loss NaN，但异常可能更早就已产生。分析 dump数据，在采集数据中过滤NaN字段，step0～step5未发现NaN，首个NaN出现在step6。定位应从 step6 的异常计算开始，继续查它如何影响后续训练。
step6有32张卡的数据，逐卡查看还不够，还需要对齐同一次通信的输入输出。这里借助MindStudio Agent昇腾AI开发智能助手，加载 nan-overflow-detection Skill，分析通信链路，再沿输入已含 NaN 的 Rank 回查本地计算。
启动MindStudio Agent后，输入prompt：
请基于输入的训练dump数据，分析其中的NaN溢出，找出源卡和根因算子Agent会自动加载专家经验Skills，并定位源卡，分析结论如下：
在Rank3和Rank7两张卡上发现，池化模块里trueDiv反向环节的输入、输出数据都带上了NaN。它的输入梯度本身就已经异常，说明NaN并不是在这个计算步骤里生成的，真正的源头还在更前面。
顺着这条梯度链路继续往前追溯，索引选取反向输出的梯度却是正常的。由此我们锁定了异常区间：索引选取输出正常梯度之后、池化内trueDiv收到异常梯度之前，NaN就是在这一段计算中产生的。
步骤三：算子输入输出数据采集
借助调用堆栈，我们缩小范围到这两段计算逻辑来自VSA视频稀疏注意力模块。 它的前向计算顺序是：池化 → 注意力运算 → 索引选取；而梯度反向传播的路径刚好相反：先经过索引选取，再通过注意力运算，最后回到池化模块。结合前面锁定的异常区间，我们重点排查索引选取之后、池化内部计算之前的这段链路。
对比源码和采集数据后发现：模型里一部分自定义算子没有采集日志，中间梯度在哪一步变成NaN暂时看不到。我们可以参考官方文档，给自定义算子补上采集注册入口，拿到完整的输入输出记录，继续定位NaN最初产生的位置。
采集应覆盖池化、选块、Attention和结果整理的整个过程，以及对应的反向计算。标准PyTorch API由工具采集，自定义入口通过register_custom_api加入采集，代码示例如下：
import fastvideo_kernel.ops as vsa_ops# 注册池化、选块、Attention 和结果合并的自定义入口。for api in ("fused_block_mean", "fused_topk_mask","block_sparse_attn", "_combine_coarse_sparse"):debugger.register_custom_api(vsa_ops, api, api_prefix="VSA")# backward_ops 为训练代码实际加载的反向扩展模块。debugger.register_custom_api(backward_ops, "fusion_attention_grad", api_prefix="VSA")
backward_ops使用项目中已加载的自定义反向模块。注册后检查 dump 是否同时覆盖前向数据处理与反向调用，发现缺口再补。
重新运行，继续采集全部卡的数据。NaN 不一定在上次的 step 出现，需要保留通信链路和前后算子的记录，判断异常是在 VSA 内部产生，还是由上游传入。
{"task": "nan_check","dump_path": "./attention_nan_dump","level": "L1"}
补充采集数据之后，我们顺着执行顺序逐一核对每个计算环节的输入输出，在dump反向数据找到了NaN最早出现的位置：fusion_attention_grad输出的梯度值grad_v。这次计算接收的全部输入数据本身都是正常无异常的。由此锁定问题就出在这个计算环节，接下来，我们用本次采集到的数据单独复现它。
为了独立复现问题，我们把这次调用的相关数据导出存为case.pt 文件。文件里包含三类信息：原始调用的入参、命名参数，以及采集到的异常输出grad_v。前两类用来重新跑一遍计算，最后一份用来和计算结果做对比，验证异常。
步骤四：单算子用例复现
读取dump采集的pt文件里保存的入参信息，复用当时采集到的正常输入数据和原始调用配置，单独跑一遍 fusion_attention_grad计算，验证输出的grad_v是否会再次复现NaN异常。单算子用例用于确认问题根因。在输入无NaN、调用参数与训练时一致的前提下，单独运行后grad_v再次出现NaN，就确认本次NaN由fusion_attention_grad产生，问题定位到该反向算子。
小结
这类随机出现的NaN问题，不能只盯着报错日志，需要把训练告警和更早的计算痕迹串联起来分析。第8步loss报出异常，只是问题暴露的信号；而第6步首次出现NaN，才是我们溯源的起点。再结合多卡之间交互的数据，才能持续缩小范围，锁定需要重点排查的卡与模型代码位置。
借助MindStudio精度调试工具，我们可以留存每一步、每张卡的计算记录，顺着计算和多卡交互的依赖关系反向追查。遇到像VSA模块里没有采集覆盖的自定义计算逻辑，还可以灵活补充采集埋点；拿到异常现场的数据打包文件后，再单独复现可疑的计算单元。整个排查过程中，持续核对输入数据是否本身就有异常，以此分清NaN是在这里新生成的，还是从上游传递过来的。
相关文档
MindStudio Agent使用文档：https://gitcode.com/Ascend/msagent/blob/master/docs/zh/agent_guide/accuracy.md
MindStudio Probe使用示例：https://gitcode.com/Ascend/msprobe/blob/master/docs/zh/user_guide/dump/pytorch_nan_check_quickstart.md
自定义算子精度数据采集文档：
https://gitcode.com/Ascend/msprobe/blob/master/docs/zh/user_guide/dump/pytorch_data_dump_instruct.md#register_custom_api
扫描二维码一键直达msProbe开源社区
扫描二维码一键直达msAgent开源社区
扫描二维码一键直达MindStudio开源技术交流群
评论 (0)