去年年底我接了个挺典型的任务:一套基于Qwen系列架构的中文大模型微调代码,原本在8张GPU上已经跑得很稳,需要整体迁到8卡昇腾910B环境重新跑通全流程。当时看着标题一句话“昇腾大模型训练调试调优:模型训练全流程”,感觉就是把显卡换一下、跑起来的事。真动手后才发现,昇腾训练从环境准备、代码迁移、算子适配、分布式通信到性能调优和长稳保障,每一步都可能有独立的坑,而且这些坑之间有依赖关系,不按顺序排查会很痛苦。
这篇文章把我这趟迁移过程中最核心的调试和调优链路完整梳理一遍,把昇腾训练全流程拆成五个实操切入点:先搞清楚软件栈和硬件血缘、再改代码、再查算子问题、再做性能归因、最后收尾验证。内容面向正在做昇腾大模型训练,或者正准备从GPU平台迁移过来的团队,尤其是那些第一次接触NPU、不熟悉CANN/torch_npu这套生态的朋友。
1. 昇腾训练栈的硬件与软件配套:动手改代码前先把版本关系理清
1.1 一套昇腾训练任务到底涉及哪几层
很多人适应不了昇腾生态,是因为在GPU平台上写代码只需要面对PyTorch + CUDA这一层抽象,而在昇腾上训练,你实际面对的是一条更长的链路。最下面是昇腾的AI处理器(目前训练用得比较多的910系列),芯片上有AI Core计算单元、HBM显存,也有固定的内存和通信拓扑。芯片之上是驱动和固件,负责让系统识别到设备。再往上是CANN软件栈,这一层负责算子的调度、内存管理、图编译和Runtime。然后才是框架适配层,比如torch_npu就是把PyTorch的算子调度接到CANN上的插件。最后才是你真正写的训练脚本。
我刚接触时总以为“模型能跑起来,说明框架没问题”。但实际上脚本只是浮在最表面的那层,真正报出的错误可能来自CANN Runtime、可能是驱动不配套,甚至可能只是某条环境变量没被读到。所以遇到问题别急着改训练代码,先坐下来检查软硬件链路的状态,会省掉大量无用功。
1.2 版本不配套才是昇腾环境最大的“隐形杀手”
在GPU平台上,pip install torch总能装上一个基本能用的版本。昇腾不行,它要求驱动固件、CANN版本、Python版本、PyTorch版本、torch_npu版本之间保持严格配套。官方会给出配套表,但实际操作时很多问题不是“装不上”,而是“能装上,但跑到某个算子时崩溃或者结果异常”。这类问题最难查,因为它不会在一开始报错,通常训练到中途才出事,比如某个融合算子精度不对,或者卡在某个通信原语上。
我的习惯是进入任何昇腾环境后先跑一遍检查命令,确认版本信息后再开始折腾代码。
npu-smi info python -c "import torch; print(torch.__version__)" python -c "import torch_npu; print(torch_npu.__version__)"至少在报错时能拿出这三个版本号来判断是哪个组件的问题,而不是凭经验瞎猜。另外,如果团队里有多个镜像或多人共用一台训练机,一定要把这套版本关系固化到环境描述文件里,比如requirements.txt或者容器镜像,别让每个人各自装一份。
| 检查对象 | 建议命令 | 常见不匹配现象 |
|---|---|---|
| NPU设备状态 | npu-smi info | 设备掉线、显存查不到 |
| PyTorch版本 | python -c "import torch; print(torch.version)" | 算子行为异常、崩溃 |
| torch_npu插件 | python -c "import torch_npu; print(torch_npu.version)" | import后找不到NPU设备 |
| CANN/Runtime | 环境变量或自带的版本工具 | 初始化失败、算子加载失败 |
我那次迁移接手的机器上,CANN和torch_npu版本至少差了一个大版本。跑微小模型时看起来一切正常,一旦开完整训练,很快就出现某个attention相关的融合算子直接非法访问内存。换成配套版本后,同一个脚本什么问题都没有。先确认版本配套,再谈调试调优,这条规矩适用于昇腾全流程的每一个阶段。
2. GPU训练脚本迁移到NPU的第一轮代码改造清单
2.1 设备控制层面的改动:比想象中简单
如果只是单机多卡场景,模型代码从GPU迁移到昇腾,设备层面的改动其实不大。核心是把原来用cuda指代设备的地方整体替换成npu,并且让torch_npu在import时被加载。
import torch import torch_npu device = torch.device("npu" if torch.npu.is_available() else "cpu") model = model.to(device)昇腾芯片有自己的一套设备编号,环境变量叫ASCEND_RT_VISIBLE_DEVICES,用来控制当前进程能看到哪些卡。原来用CUDA_VISIBLE_DEVICES的启动脚本需要同步替换。这个变量在分布式任务里尤其重要,如果只改了代码没改启动脚本,很可能出现所有进程都抢同一张卡,或者压根找不到设备的情况。
2.2 分布式通信后端:NCCL换成HCCL
对于大模型训练,单卡几乎不可能承载完整流程,分布式通信是避不开的。GPU平台用NCCL做通信,昇腾平台上则使用HCCL。好在torch_npu做了不少兼容工作,PyTorch的distributed接口依然可用,但初始化通信后端时要把nccl改成hccl。
import torch.distributed as dist dist.init_process_group(backend="hccl", init_method="env://")这里有一个很容易被忽视的点:HCCL的底层通信和硬件拓扑强相关,首次初始化或跨机训练时容易卡很久,超时时间需要额外设置,比如通过HCCL_CONNECT_TIMEOUT环境变量调大连接超时时间。多机场景还需要确保节点间的网卡路由配置正确。我遇到过几次训练启动后一直卡在初始化阶段,最终定位到是容器内网卡信息缺失,HCCL找不到可用的通信口。排查思路是先确认各个进程都能看到预期的NPU设备,再确认跨节点网络连通性。
2.3 混合精度策略:不要沿用GPU上的apex习惯
在GPU上很多大模型训练会使用apex或者torch.amp来做混合精度。昇腾生态里,基于CUDA深度绑定apex通常不可用。如果是从老代码迁移,建议把混合精度统一改成PyTorch原生的torch.amp写法,并且在创建GradScaler时指定设备类型是npu。
scaler = torch.amp.GradScaler("npu")昇腾的910系列芯片对BF16的支持,在不同固件和算子实现上的表现有差异。不是所有算子都原生支持BF16,有些算子会回退到FP32,有些会走额外的转换路径。最稳妥的做法是选一个小验证集,分别在纯FP32训练和BF16混合精度训练下对比loss曲线,确认没有明显差异后再开全量训练。精度损失不会体现在每个算子上,有时候只会在Loss开始下降后的第几百步突然出现inf或者NaN。
2.4 数据加载也要适配:pin_memory不是白叫的
数据加载相关代码里,torch.Tensor.pin_memory()在GPU上可以把页锁定内存与GPU显存之间的拷贝加速,在昇腾上则没有完全对应的语义。如果迁移时没留意,轻则性能略降,重则在dataloader里报设备不支持的错。更合理的做法是把数据加载和增强放在CPU上完成,只把必须搬到NPU的batch数据执行.to("npu")。另外,某些数据预处理算子如果也跑在cuda上,需要明确改成由CPU或者NPU执行,避免残留对cuda底层指针的隐式依赖。
迁移第一轮的原则是:能跑通一个小规模用例,而不是一步到位追求性能。先把设备、通信、混合精度和数据入口这四处的显式依赖清干净,后面再逐步排查算子兼容问题。
3. 训练跑挂的现场排查:从loss不降到找到第一个NaN算子
3.1 先看权重初始化,不要动不动调学习率
迁移后在同一个模型上,最常遇到的不是跑不起来,而是“跑起来了但loss曲线不对”。我见过不少人在这一步开始疯狂调学习率、调warmup,其实帮助不大。正确的排查顺序应该是先确认模型结构和权重加载是完整的,再做一次“极小数据过拟合测试”,用一两条样本把一个batch跑到loss显著降低。这样能排除代码写错和权重没加载干净的问题。
如果极小过拟合正常,但全量数据上loss表现差,再看数据分布和随机种子。昇腾端上有些随机算子与GPU端的实现不同,同一个模型和同一个随机种子也可能跑出不一样的曲线。这种情况不是bug,只要数值范围合理即可。真正需要警惕的是loss完全不下降,或者一开始就NaN。
3.2 用梯度hook定位第一个出现NaN的位置
遇到NaN,我的第一反应不是去调precision,而是先找到第一个产生NaN的位置。神经网络的反向传播是有顺序的,越靠近输入侧的梯度如果变成NaN,通常说明某个上游算子已经算出了非法值。如果你只用torch.isnan(loss)去判断,只能知道结果坏了,不知道坏在哪。
可以在每个参数上挂一个梯度hook,在判断发生NaN的第一时间打印出参数名。
def register_nan_hook(model): for name, param in model.named_parameters(): param.register_hook( lambda grad, pname=name: print(f"NaN grad at {pname}") if torch.isnan(grad).any() else None )跑一次带hook的训练,通常能直接定位到第一个异常梯度。之后顺着这个参数往前找对应的算子,基本就是问题所在。CANN也提供了数据Dump能力,能在算子输入输出级别做检查,但对大多数模型问题,用梯度hook先缩小范围更高效。用这种方式我遇到过很多次真正的根因:某个在GPU上没问题的自定义算子,迁移到NPU后触发了数值溢出。
3.3 多卡训练卡住与HCCL通信问题
另一类高频问题是训练跑着跑着卡住,GPU平台可能只跟网络有关,昇腾上多了HCCL这个概念。多进程启动后,如果某个rank看不到其他rank,或者通信环建立不完整,就会出现“看起来像死锁”的情况。建议出现卡住首先看两点:一是各进程能否正常读到ASCEND_RT_VISIBLE_DEVICES,二是卡住的调用栈是否停在了wait或者allreduce等通信原语上。如果多个进程都停在同一通信操作,基本可以确认是通信初始化或拓扑发现问题。
实际操作时,我把正常情况下的超时参数调大一些,同时在代码里给通信原语加上超时或日志,方便后续观测。多卡任务的启动脚本也比单卡更容易出问题,shell里每个进程的环境变量必须一致且正确配置,rank分配要参考昇腾平台实际使用的启动方式。
3.4 看着像性能bug,实际是算子走偏了
还有一次印象很深的案例:8卡训练跑起来后,step time比理论上慢很多,卡与卡之间负载极不均匀。我以为是分布式通信或者数据加载的问题,优化一圈下来完全没有改善。后来用Profiling打点才发现,模型里有一些自定义的mask操作没被NPU原生算子覆盖,运行时做了大量transdata和数据搬移,等于在往瓶颈上不断加码。这类问题说明性能不佳时,不能只看通信层和数据层,算子下沉是否完整同样是训练全流程里必须排查的一环。
4. 性能调优真正陡峭的地方:从profile数据里找AI Core闲置的原因
4.1 先出Profiling再讨论调优方向
昇腾上的性能调优不是靠“觉得哪里慢就改哪里”,一定要先出可量化的profile数据。torch_npu.profiler跟PyTorch的profiler用法类似,可以统计到算子维度的时间消耗、AI Core利用率、通信耗时等。拿到数据后,我会重点看下面几个维度:step time是变大还是稳定、AI Core时间段占比多少、通信时间段占比多少、空闲时间(比如等待同步)占比多少。
有一种典型的坏情况是AI Core利用率看起来不低,但整个step time仍然很长,这往往说明指令发散严重、很多算子在互相等待或者存在大量低效算子调用。另一种相反的坏情况是AI Core利用率只有百分之二三十,但通信又不算多,这时问题通常出在算子之间的大量同步和同步点。这两种情况的优化方法完全不同,所以别跳步。
4.2 小算子过多和频繁同步是性能杀手
运行一个Transformer大模型,如果不开算子融合,每一层都会拆成大量很小的算子,比如add、layer_norm、residual、cast等。这些算子单个执行时间很短,但每次下发都要经历从CPU侧到NPU侧的调度开销,积累起来非常可观。在GPU上CUDA kernel的启动开销也不小,但昇腾的调度路径会让这类小算子问题更明显一些。
处理办法是尽量调用昇腾生态里已经融合好的大算子,例如FlashAttention的一体化实现,RMSNorm和残差连接的融合实现。很多情况下,把几个小算子替换成一个融合算子,整体性能就能提升不少。有时候一个transdata或者一个额外cast操作就会吃掉本应属于计算的AI Core时间片。
4.3 让通信和计算重叠加起来:DDP参数同步也能被藏掉
数据并行训练中,不同卡上的梯度需要在参数更新前做一次AllReduce。如果每次step结束就同步所有梯度,通信时间会直接曝露在训练路径上,拉长整个step time。业界通行做法是梯度分桶,让反向传播算出一部分梯度后就同步一部分,而不是等全部梯度算完再一次性同步。PyTorch自带的DDP其实已经在做类似的分桶通信,但迁移到昇腾后,需要确认下初始化的bucket size等参数是否适合当前模型和卡数。
另外一个很实用的思路是减少host与device之间的同步点。大模型调试时,如果为了打印中间loss频繁做.item()取数操作,就会自动触发同步,直接拖慢训练。我见过有人每个GPU的每个step都打印一堆指标,结果把整卡流水线打得支离破碎。建议把日志聚合到rank0,并控制打印频率,让训练主体逻辑里的同步点降到最少。把这些同步去掉后,通常能看到AI Core有效时间占比明显提升。
4.4 更大规模时的并行策略切换
如果只是单机8卡,数据并行往往已经够用。但模型参数一旦到几十B甚至上百B,单卡塞不下完整模型时,就必须考虑张量并行、流水线并行或者二者与数据并行的组合。昇腾多卡并行在通信拓扑上有自己的特点,比如同一个节点的卡间通信带宽远好于跨节点通信。所以并行切分时,尽量把通信量大的张量并行放在同一节点内,把跨节点留给流水线并行中更稀疏的通信。
从一个可用的调优顺序上讲,我建议先保证数据并行下单卡性能接近该卡的理论峰值,再去做张量并行,最后才尝试流水线并行。跳过前一步直接切3D并行的团队大多会遇到一个结果:模型能跑,但整体吞吐上不去,而且很难定位是并行策略问题还是算子本身没优化到位。
5. 训练全流程交给可靠阶段的最后一公里:精度对齐与长稳验证
5.1 把“看着像没坏”变成“量化可接受”的精度对齐
代码能跑通、loss在降,离真正可以交付还差一步:精度是否符合预期。尤其是在从GPU迁移到NPU的场景里,不能只用“loss降了”来判定模型没有问题。比较稳的做法是跑一段固定步数的小规模训练,保存GPU侧和NPU侧在同一验证集上的loss序列,观察曲线的间距和趋势是否一致。对梯度或中间激活做抽样比较,用余弦相似度来判断方向是否一致,通常比单点数值比较更能反映真实差异。
如果你有已保存的GPU侧checkpoint,可以直接在NPU上加载并继续训练一段,再和GPU侧对照eval指标。eval指标可以有微小差异,但不应出现数量级的抖动。如果差异过大,排除随机性后基本可以认定某些算子的实现精度低于可接受水平,需要回到算子替换或混合精度选项上做排查。
5.2 checkpoint搬运与增量训练的坑
跨硬件平台搬运checkpoint时,除了网络权重,优化器状态里可能记录了与设备相关的历史信息,比如exp_avg和exp_avg_sq,这些都只是数值,可以搬运,但需要注意状态字典的key是否完整。另一个常见的坑是模型代码里自定义的注册buffer或者随机状态也被保存进来,这些在NPU上加载时可能会产生不匹配。
增量训练的场景下,我通常会把加载后的model在少量样本上跑前几步,专门比较加载前后的loss曲线,确保checkpoint内容在NPU上生效且没有数据错位。续训不要求完全复现GPU上的每一步结果,但优化器动量状态如果因为key不对被重新初始化,就会出现“看起来续上了,实际上学习过程从零开始”的问题。
5.3 长稳压测:把训练挂上几天前值得做的事
跑通全流程之后,先别急着批量对多个实验排队。AI芯片在长稳运行中会有各类细碎问题,比如显存泄漏、通信卡死、单卡掉线、频繁重试导致训练时间大幅拉长。最直观的验证方式是把一个重要实验在目标配置下连跑3到5天,关注显存占用是否缓慢爬升、step时间的标准差是否在扩大、日志里是否出现偶发错误。
昇腾平台也有一系列环境变量用于调整通信超时和错误重试策略,在长稳阶段可以额外关注。我自己的习惯是把完整训练流程拆成启动、前几百步、中期和后期几个阶段,分别设置对应告警,再在训练看板上同时监控loss、吞吐和NPU状态。长稳通过后,整套方案才算真正具备交付能力,而不是只能在交互式调试时跑几分钟。
经过这个全流程的调优,现在再回头看昇腾大模型训练,“调试”和“调优”并不是两个割裂的阶段。调试解决的是能不能跑,调优解决的是跑得好不好,但两者共用同一条链路:尽快拿到有效的算子级数据和通信数据,让每一次改动都有可量化的反馈。我在实际项目里最大的体会是,昇腾调试的入口往往在版本配套和底层算子兼容,门槛比GPU生态高一些,但只要按环境、迁移、算子、通信、性能这条路径逐层排查,后面反而会越来越顺手。最后再分享一个小技巧:每次改动只动一个变量,记录下对应的step time和loss曲线,一段时间后回看这些记录,很多疑难问题都能快速归因到特定改动上,这也是我在昇腾平台调优时最依赖的工作方式。