第 14 天凌晨两点,终端还在刷 HCCL AllReduce failed。
这套 VL-Mamba 多模态模型已经在 PyTorch 上训练了半年。迁移前,我们的设想很直接:把 torch.cuda 换成 torch_npu,先让训练跑起来,再慢慢调性能。标准的 ResNet 和 Transformer 层确实能过,自定义的动态位置编码一接入,错误就开始成片出现。
20 天里,API 替换只占很短一段,时间主要耗在算子重写、数值对齐和多卡通信。
API 替换只够走到自定义算子门口
第一天,torch_npu 很快跑通了常规网络层。到了 Custom Ops,兼容层的边界立刻显出来。torch.compile 在这套 NPU 环境里表现不稳定,终端给出的也不再是容易追踪的 Python Traceback,而是 AICore Error、Vector CPU Core mismatch。
在这套任务里,torch_npu 解决了接口接入;进入 CANN 的算子执行路径后,复杂 Custom Ops 仍要单独处理。
第 3 天,我们停止继续堆 API 兼容补丁,转向 MindSpore,复杂逻辑直接改写为 Ascend C 算子。这个决定增加了前期工作量,却让错误回到可定位的算子边界:输入、输出、数据类型、内存访问和核函数行为,都可以逐项核对。
公开的 MindSpore 自定义算子开发实战——从 CUDA 到 Ascend C 的迁移与优化 也沿着 AKG 与 Ascend C 展开:案例处理的是 LayerNorm,从 CUDA 迁到 Ascend C,再用双缓冲和向量化优化。我们的对象是多模态位置编码,可借鉴的是同一套方法——先把算子写清楚,再看性能。
单卡能收敛,八卡仍会翻车
多卡训练最先暴露的是 FlashAttention。在这套 910C 集群上,默认实现一旦涉及跨 Die 交互,耗时就明显上升。原来在 A100 与 CUDA 环境里没有排到前面的 Attention 通信,迁移后成了主要瓶颈之一。
第 10 天又出现了更隐蔽的问题:单卡 Loss 正常,扩到 8 卡后,验证集 Loss 开始震荡。排查一周,线索最终落到 TopK 的舍入行为。在这次环境中,CUDA 默认采用向最近偶数舍入,Ascend 某些模式存在细微差异。误差单次很小,累积到训练过程里就不能忽略。
处理方式有两条:在 MindSpore 中明确指定算子的数值行为;走自定义算子时,在 Ascend C 里手动处理舍入逻辑。
这类问题最好在扩容前拦住。先核对单算子的输入输出,再看整图,最后扩到多卡。单卡 Loss 正常,只能说明最基础的执行路径没有出错。
35% 到 78%,差在算子路径
我们保留了三组项目数据:
| 方案 | 平均 MFU | 现场表现 |
|---|---|---|
| PyTorch + CUDA A100 | 约 78% | 显存占用稳定,调试效率高 |
PyTorch + torch_npu |
约 35% | 显存碎片明显,频繁触发 GC;CPU 预处理和算子编译等待较多 |
| MindSpore + 自定义 Ascend C 算子 | 约 78% | 吞吐量达到初迁组的 2.3 倍;主要瓶颈转到 HCCL |
这些数字只对应这套模型、这批算子和当时的集群配置,不能当作通用芯片对比。A100 基线和 910C 终迁组都测到约 78% MFU,也不代表两套硬件在所有任务上等价。
终迁组的提升主要来自手写 Ascend C 算子。多模态模型里的稀疏矩阵运算改用向量化加载和双缓冲后,内存访问减少,计算侧的等待随之下降。到这个阶段,最显眼的瓶颈已经转到 HCCL 通信。
在这套任务里,芯片参数之外,算子是否进入高效执行路径、通信是否拖住计算,同样影响最终结果。
三个动作,越早做越省时间
先用 torch_npu 摸清兼容边界。 标准层和基础训练流程可以先跑起来,自定义算子则要尽早列清单。核心逻辑越早进入 MindSpore 和 Ascend C,后面隔着兼容层查错的时间越少。
把数值回归放在扩容前。 torch.nn.Linear 和 mindspore.nn.Dense 要按两套实现来验证,尤其是 AMP 场景下的 Loss Scaling、TopK 等数值敏感环节。接口名称接近,不代表输出能直接对齐。
HCCL 问题按拓扑查。 配置多机多卡训练时,使用 msrun 启动,并逐项核对网卡绑定和 IB/RoCE 配置。计算单元出现空转,先看通信等待时间,再判断是不是算子本身过慢。
从一个算子开始,比整模型硬搬更稳
对这次 VL-Mamba,迁移顺序已经很清楚:标准层先跑通,关键自定义算子单独迁到 Ascend C;单卡完成数值回归,再上 8 卡排查 HCCL;最后用 MFU 和通信等待时间判断还有多少调优空间。
第 3 天放弃纯 API 替换,是整个项目的转折点。
要开始类似迁移,可以先挑一个最影响训练的自定义算子,做 CUDA 与 Ascend C 的输入输出对齐,再检查双缓冲、向量化和通信开销。范围小,结果容易验证,出问题也容易回退。
参考资料
注:MFU、吞吐与报错来自这次项目记录,公开资料用于说明算子迁移方法。