1. 昇腾AI训练中grad_norm异常现象解析第一次在昇腾平台上跑大模型训练任务时看到日志里突然蹦出grad_normNAN的报错我整个人都是懵的。grad_norm这个看似简单的指标实际上是模型训练健康的晴雨表。简单来说它就是所有参数梯度向量的二范数相当于给整个模型的学习方向做了个综合体检。当这个值变成NAN时就像体检报告上突然出现异常待查四个大字意味着反向传播过程中出现了数值不稳定。常见症状往往伴随着loss曲线突然飙升或直接变成NAN模型参数更新完全失控。我遇到过最棘手的情况是在128卡分布式训练时某个迭代步突然出现grad_norm异常导致整个训练任务中断。这时候首先要区分是硬件问题还是软件问题——就像医生要先判断病人是外伤还是内伤。通过简单的梯度检查代码可以快速确认是单卡故障还是普遍现象# 快速检查梯度异常的工具函数 def check_grad_nan(model): for name, param in model.named_parameters(): if param.grad is not None and torch.isnan(param.grad).any(): print(f发现异常梯度参数: {name}) return True return False在昇腾NPU环境下硬件问题通常表现为特定计算卡上的持续异常而软件问题往往与特定数据触发的算子行为相关。有个很实用的经验如果grad_norm异常是偶发的大概率是数据问题如果是持续出现的就要重点检查硬件状态。2. 分布式环境下的问题定位技巧在Megatron-LM这样的分布式框架里排查grad_norm异常就像在迷宫里找一只会隐身的兔子。首先要理解框架的并行策略——TPTensor Parallelism、PPPipeline Parallelism、DPData Parallelism三个维度的并行会让问题定位变得复杂。我常用的方法是修改MegatronOptimizer的clip_grad_norm方法插入分布式rank信息打印def clip_grad_norm(self, clip_grad, check_for_nan_in_grad): params self.get_parameters() for param in params: if torch.isnan(param.grad).any(): from megatron.core import parallel_state tp_rank parallel_state.get_tensor_model_parallel_rank() pp_rank parallel_state.get_pipeline_model_parallel_rank() dp_rank parallel_state.get_data_parallel_rank() print(f异常卡位置: TP{tp_rank}, PP{pp_rank}, DP{dp_rank}) break grads_for_norm self.get_main_grads_for_grad_norm() return clip_grad_norm_fp32(params, grads_for_norm, clip_grad, check_for_nan_in_grad, model_parallel_groupself.get_model_parallel_group())这个技巧帮我节省了大量排查时间。有一次在256卡训练时发现只有TP3的所有卡都报NAN其他TP组的卡正常这就把问题范围缩小到了特定张量并行组。进一步检查发现是某块昇腾910B的HBM显存出现了硬件故障更换故障卡后问题立即消失。对于更隐蔽的软件问题需要建立参数名追踪机制。原始Megatron-LM的优化器参数没有名称信息我改写了get_param_groups函数增加了参数名持久化功能# 在megatron/optimizer/init.py中修改get_param_groups param_id_name_map {} for group in param_groups: for name in group[names]: param_id_name_map[len(param_id_name_map)] name # 保存到对应rank的JSON文件 with open(fparam_map_tp{tp_rank}pp{pp_rank}dp{dp_rank}.json, w) as f: json.dump(param_id_name_map, f)3. 算子级问题诊断实战当确定不是硬件问题后真正的挑战才开始。特定数据触发的算子bug就像训练过程中的幽灵时隐时现。我总结出一套钩子函数诊断法通过注册前向和反向钩子来捕捉异常数据。以常见的Linear层为例我们可以这样设置诊断钩子def debug_hook(module, input, output): # 前向传播检查 if any(torch.isnan(t).any() for t in input if isinstance(t, torch.Tensor)): print(f前向输入含NAN: {module.__class__.__name__}) torch.save(input, nan_input.pt) if torch.isnan(output).any(): print(f前向输出含NAN: {module.__class__.__name__}) torch.save({input:input, output:output}, nan_io.pt) # 注册钩子示例 for name, module in model.named_modules(): if isinstance(module, nn.Linear): module.register_forward_hook(debug_hook)反向传播的检查更为关键因为grad_norm异常往往源于反向计算。这里需要特别注意昇腾自定义算子的行为差异def bwd_debug_hook(module, grad_input, grad_output): # 反向传播检查 nan_in_grad any(torch.isnan(g).any() for g in grad_input if isinstance(g, torch.Tensor)) nan_out_grad any(torch.isnan(g).any() for g in grad_output if isinstance(g, torch.Tensor)) if nan_in_grad or nan_out_grad: print(f检测到反向传播异常: {module.__class__.__name__}) torch.save({ grad_input: grad_input, grad_output: grad_output, module_state: module.state_dict() }, nan_grads.pt) # 注册反向钩子 problem_module.register_backward_hook(bwd_debug_hook)在实际项目中我曾用这个方法发现昇腾平台上一个自定义GELU算子的边界条件问题——当输入值超过某个阈值时反向传播会产生NAN。通过保存的异常数据我们很快复现并修复了这个问题。4. 系统化的诊断流程设计经过多次实战我总结出一套标准化的诊断流程初级筛查运行梯度检查脚本确认异常范围# 分布式环境下的梯度检查命令 python -m torch.distributed.launch --nproc_per_node8 grad_check.py硬件诊断通过昇腾工具检查硬件状态# 使用昇腾诊断工具 npu-smi info -t board -i 0 -c 0数据追踪在训练脚本中植入以下诊断代码# 在训练循环中加入诊断点 if torch.isnan(grad_norm): print(fStep {step}出现grad_norm异常) torch.save({ inputs: batch, model_state: model.state_dict(), optim_state: optimizer.state_dict() }, fcrash_step_{step}.pt) break算子隔离通过逐步注释模型组件定位问题算子最小复现用保存的异常数据构建测试用例对于复杂的分布式训练我还设计了一个异常传播分析工具可以可视化NAN值在计算图中的传播路径def trace_nan_propagation(model, input): with torch.autograd.set_detect_anomaly(True): output model(input) try: loss output.sum() loss.backward() except RuntimeError as e: print(f异常追踪: {str(e)}) # 这里可以添加更详细的堆栈分析这套方法在多个大模型训练项目中成功定位了包括昇腾算子精度问题、数据加载器线程竞争、混合精度训练不稳定等各种疑难杂症。特别是在千亿参数模型的训练中系统化的诊断流程可以节省数天的调试时间。