--- title: 大模型训练中首步 Loss 为 NaN 的根因分析与解决方案-官方技术文章-昇腾社区 description: 在基于 PyTorch 框架、CANN 框架及 Atlas 800T A2 硬件平台的大模型训练任务中，采用 8 卡并行训练方案，模型结构包含 MoE（Mixture of Experts）模块。训练过程中，模型在首步（step 1）前向传播正常，但反向传播阶段梯度出现 NaNÿ keywords: 大模型训练中首步,Loss,NaN,的根因分析与解决方案,官方技术文章,昇腾社区,背景概述,问题现象描述 url: https://www.hiascend.com/developer/techArticles/20260713-3 section: (其他) --- # 大模型训练中首步 Loss 为 NaN 的根因分析与解决方案-官方技术文章-昇腾社区 URL: https://www.hiascend.com/developer/techArticles/20260713-3 描述: 在基于 PyTorch 框架、CANN 框架及 Atlas 800T A2 硬件平台的大模型训练任务中，采用 8 卡并行训练方案，模型结构包含 MoE（Mixture of Experts）模块。训练过程中，模型在首步（step 1）前向传播正常，但反向传播阶段梯度出现 NaNÿ 关键词: 大模型训练中首步,Loss,NaN,的根因分析与解决方案,官方技术文章,昇腾社区,背景概述,问题现象描述 官方技术文章 [了解详情](https://www.hiascend.com/zh/developer/techArticles) 大模型训练中首步 Loss 为 NaN 的根因分析与解决方案 大模型训练中首步 Loss 为 NaN 的根因分析与解决方案 模型训练 发表于: 2026/07/13 8 0 ## 背景概述 在基于 PyTorch、CANN训练任务中，采用 8 卡并行训练方案，模型结构包含 MoE模块。训练过程中，模型在首步前向传播正常，但反向传播阶段梯度出现 NaN，导致 Loss 发散，训练中断。该问题具有典型“首步 NaN”特征，且仅在 NPU 上出现，GPU 端训练稳定，具备明确的对比基准。为快速定位根因，采用数据 Dump 与跨设备比对分析策略，最终定位至自定义算子调用中的数据类型不匹配问题。 ## 问题现象描述 ● 训练阶段：首步（step 1）前向计算正常，Loss 为有效数值。 ● 反向传播阶段：梯度计算中出现 NaN，Loss 为 NaN。 ● 异常位置：反向传播过程中，`Module._fsdp_wrapped_module.module.gpt.transformer.model.layers.2._fsdp_wrapped_module.mlp.moe.M12DevDualMoeBlock`模块输出梯度为 NaN。 ● 对比现象：GPU 端训练全程稳定，Loss 与 NPU 保持一致，说明问题非模型结构或数据本身导致，而是硬件/算子执行差异所致。 ## 故障排查 ### 1. 数据采集与初步分析 采用`msprobe`工具采集 NPU 与 GPU 的混合级（mix-level）Dump 数据，包含模型所有模块的输入输出、PyTorch API 调用的 Tensor 信息。 通过搜索`nan`关键字，发现部分`squeeze`算子的输入输出为 NaN，但进一步排查其调用栈，确认输入 Tensor 由`torch.empty`初始化，排除初始化问题。 ### 2 关键异常点定位：to算子输入为空 在调用栈中发现`to`算子输入为 shape 为 0 的 Tensor，该现象在 GPU 端同样存在，说明非 NPU 独有问题，排除数据处理逻辑为根本原因。 ### 3 采用标杆数据比对法定位异常模块 由于 GPU 端训练稳定，可作为golden path，采用`msprobe compare`命令进行 NPU 与 GPU 数据比对： msprobe compare -tp -gp 筛选所有模块级数据，发现： ● 前向传播阶段，NPU 与 GPU 数据基本对齐； ● 反向传播阶段，`M12DevDualMoeBlock`模块输出梯度在 NPU 上出现 NaN，而 GPU 输出正常。 进一步查看该模块的输入输出 Dump 数据，确认： ● 输入梯度为有效数值，无 NaN； ● 输出梯度为 NaN，表明问题发生在该模块内部计算逻辑。 ### 3.4 深入模块代码与算子分析 查看`M12DevDualMoeBlock`模块内部实现，重点关注调用的`fused_npu_moe`模块。该模块作为`hidden_states`的梯度来源之一，其输出存在 NaN。 通过分析调用栈与算子行为，发现： ● `gate`层输出的`hidden_states`梯度正常； ● 但`fused_npu_moe`模块中，`unpermute`算子输出为 NaN； ● 该`unpermute`算子依赖`probs`输入，且其底层实现要求输入数据类型为`bf16`； ● 实际传入的`probs`数据类型为`fp32`，超出算子支持范围。 关键发现：`permute`算子（或其底层实现）对输入数据类型有严格限制（仅支持`bf16`），但调用方未进行类型转换，导致算子行为异常，输出 NaN。 ## 问题根因 `permute`算子调用时，输入`probs`数据类型为`fp32`，而该算子底层实现仅支持`bf16`类型输入，导致计算异常，输出梯度为 NaN。 尽管`torch_npu`接口未对非法类型输入进行拦截，且该算子未在官方文档中明确列出，但其行为符合底层算子（如`mindspeed`自定义算子）的类型约束。因此，问题本质为算子调用时数据类型不匹配。 ## 解决方案 ### 修复措施 在调用`fused_npu_moe`模块前，对`probs`输入显式执行类型转换： probs = probs.to(torch.bfloat16) ### 验证结果 ● 修复后，首步反向传播梯度恢复正常，Loss 为有效数值； ● 多步训练过程中 Loss 与 GPU 端保持一致，未再出现 NaN； ● 模型训练流程稳定，可正常推进。 ## 总结 1. 在存在稳定标杆（如 GPU）的场景下，推荐使用 NPU 与 GPU 数据比对法，可有效规避 NaN 干扰，快速定位异常模块。 2. 自定义算子调用需严格遵循类型约束，尤其在 NPU 平台，底层算子对数据类型敏感，建议在调用前显式进行`cast`操作。 3. 建议在开发阶段增加类型校验逻辑，对关键算子输入进行类型检查，避免运行时因类型不匹配导致不可预测行为。 边框设置 无框线 边距 宽度 1磅 颜色 自由布局设置 整体布局 子模块 评论 修订记录 对正文进行的文本增删、样式修改都将标记为修订 自定义多级列表 列表设置 - 1 - 2 - 3 - 4 - 5 - 6 - 7 - 8 - 9 - 1. - a. - i. - 1. - a. - i. - 1. - a. - i. 前缀 无 序号 1 2 3 ... 后缀 . 编号格式 列表显示 继承层级 不继承 位置 对齐方式 默认 单元格边距 边距 默认 左边距 cm 右边距 cm 上边距 cm 下边距 cm 点赞 0 本页内容 背景概述 [了解详情](https://www.hiascend.com/#2T3eFTOb5v4YJBPTTtktsl) 问题现象描述 [了解详情](https://www.hiascend.com/#56vk32DBol63ZGa6aXP0AB) 故障排查 [了解详情](https://www.hiascend.com/#3OcwMYym5LkU1NXL4mOCnr) 问题根因 [了解详情](https://www.hiascend.com/#5IajLC1GIw7JVLYKcdXknM) 解决方案 [了解详情](https://www.hiascend.com/#23LUSSf3lCtn1MmhkiKqUp) 总结 [了解详情](https://www.hiascend.com/#24BGWUDGOZk8tUFcOWc44M)