作者：昇腾实战派
知识地图：https://blog.csdn.net/Lumos_Lovegood/article/details/161455142

背景概述

在大规模混合专家模型（MoE）的训练过程中，计算效率与显存贮瓶颈是经常遇到的挑战。以 Qwen3.5 MoE 为代表的模型结构，在训练时尤其需要高性能的算子实现和显存优化策略。我们在实际训练过程中，发现原有的 PyTorch 原生小算子可能会影响性能。为此，我们系统性地开展了一系列训练优化实验，重点引入 Triton 算子、优化显存异步卸载机制，并验证不同并行训练策略下的精度一致性，最终显著提升训练速度并保障收敛性。

具体软硬件信息如下：

组件	版本
vLLM	0.18.0
vLLM-Ascend	0.18.0
torch	2.9.0 + cpu
torch-npu	2.9.0
transformers	5.3.0
veRL	main
CANN	8.5.0
设备	Atlas 800T A2

一、GDN Triton 算子接入优化

原训练流程中，线性注意力层所使用的 GDN实现是基于 PyTorch 原生小算子实现的，成为训练速度的瓶颈之一。通过将其替换为mindspeed上Triton版本的GDN算子可以极大程度的提升训练速度。

mindspeed算子链接： https://gitcode.com/Ascend/MindSpeed/tree/master/mindspeed/lite/ops/triton

具体实施步骤如下：

1. 集成 Triton 算子实现
   将 MindSpeed 仓库中的 Triton 版 GDN 算子（路径：mindspeed/lite/ops/triton/）集成至训练代码库中。需注意区分不同版本的实现，调整相应头文件引用，避免 ImportError。
   

2. 替换模型中的算子调用
   在transformers的 src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py 中，将原有的 PyTorch 实现替换为 Triton 版本的 GDN 算子：
   源代码
   

   替换后
   

3. 修复梯度异常问题
   在模型中的 verl/models/transformers/chunk_gated_delta_rule.py 中发现梯度出现 NaN，经分析定位到特定计算语句，调整后恢复正常：
   

4. 类型对齐调整
   在多处出现张量 dtype 不一致导致的运行时错误，通过显式统一数据类型解决。
   

二、训练性能提升对比

在 Triton-ascend 3.2.0 环境下进行性能对比测试，使用 Triton 算子后训练速度提升显著：

• 训练耗时：单次更新耗时明显减低
• 收敛性保持一致：Reward 曲线与原有实现趋势一致，精度无损失

性能对比如下图所示：

图中红色为Mindspeed，绿色为torch原生小算子

三、Async offload

通过 Profiling 分析发现，激活卸载（Activation Offload）过程中存在同步等待，导致空闲时间。我们引入异步激活卸载机制，并结合梯度检查点（Gradient Checkpointing）策略，显著降低峰值显存占用，同时尽可能减少对训练速度的影响。该策略在类似结构的 MoE 模型中已验证有效，目前正在进一步测试中。

四、其他融合算子优化

除了 GDN 算子之外，我们还接入了多种昇腾融合算子以进一步提升性能，包括：

• RMS Norm
• RoPE（Rotary Position Embedding）
• MoE 路由算法

这些优化已集成至主分支，可通过替换 modeling 文件中的对应模块直接启用。

五、训练精度验证与策略对比

在扩展到 20K 训练规模后，我们开展了多组精度对比实验，主要结果如下：

5.1 FSDP 与 Megatron 后端对比

实验中发现，FSDP 训练策略在 NPU 和 GPU 环境下均出现约 35 步后训练不稳定的现象，且 Reward 与准确率（Accuracy）指标出现不一致：Reward 上升但 Accuracy 显著下降。通过对比实验推测，可能是 FSDP 与 Megatron 在分布式策略实现上的差异所致。

红⾊为GPU-Megatron
绿⾊为NPU-FSDP
⻩⾊为GPU-FSDP

5.2 优化策略尝试

针对上述问题，我们尝试了两种优化方案：

1. MoE Router 精度提升至 FP32
   将 MoE 路由计算强制为 FP32 类型，缓解梯度异常，Reward 有一定提升。
   可参考 https://github.com/verl-project/verl/pull/5249、https://github.com/mikequan0425/verl/commit/b569fbbdcf9d38557406f0adbfc88d0324ee38b6
   

   

2. 算法替换：GRPO → GSPO
   切换策略优化算法，实验显示初期 Reward 和 Accuracy 有所改善，但仍在 35 步左右仍会出现问题。

5.3 切换至 Veomni Engine

最终我们将训练后端从 Legacy Worker 迁移至 Veomni Engine。新环境下：

• Reward 曲线正常上升，收敛趋势良好；
• 30 步时 Accuracy 超过Megatron 基线；

结果对比如下：

因FSDP后端实现与veomni相近，得知开源社区上在 Engine worker 上对齐了 Veomni 和 Megatron 的训练精度，且先前实验均基于legacy worker拉起，计划如下：

1. 打通Veomni Engine进行训练
2. FSDP backend 从legacy worker切换至engine worker，由于main分支已经计划删除legacy worker，这时今后无法规避的任务，顺带在这次尝试打通

从实验结果可以发现：
Veomni Engine reward上升正常，精度符合预期

结论

通过系统性的算子优化、训练策略调整与精度验证，我们成功将 Triton 算子应用于 Qwen3.5 MoE 训练中，在保持收敛性的前提下显著提升训练效率。Veomni Engine 的引入进一步解决了分布式训练中的稳定性问题，为后续大规模扩展提供了可靠基础。