Prime Intellect 正式发布了 Prime Mega MoE，这是一个专为混合专家（MoE）模型训练设计的融合、重叠 GPU 内核。该内核建立在 DeepSeek 开源的 DeepGEMM 库中的 Mega MoE 基础之上，通过补充缺失的反向传播逻辑，使该融合内核能够直接应用于模型训练，而不仅限于推理阶段。
- 官方介绍：https://www.primeintellect.ai/blog/prime-mega-moe
- GitHub：https://github.com/PrimeIntellect-ai/prime-rl
核心痛点：MoE 训练中的通信瓶颈
在标准的 MoE 模型训练设置中，每个 token 会被路由到位于不同 GPU 上的少量专家。这一过程要求 GPU 在每一层执行以下独立操作：发送 token、运行专家矩阵乘法、接收结果。
这种分离的操作模式会导致 GPU 在等待传输层完成时处于空闲状态。前向传播和反向传播中反复出现的等待时间，构成了 MoE 模型训练的主要通信瓶颈，直接增加了整体的 GPU 小时消耗。
技术原理：计算与通信的深度融合
DeepSeek 的 Mega MoE 内核最初通过一次内核调用取代了上述操作链，融合了专家并行分发、两个专家线性层、SwiGLU 激活函数以及结果合并。它还利用仅 GPU 暂存缓冲区跳过了中间张量的生成，并在数据传输期间对已到达的 token 组进行计算。
Prime Intellect 在此基础上编写了缺失的反向传播内核，并同样采用了重叠（Overlapping）执行策略。该内核支持多种数据类型，包括 BF16、Blackwell 架构专属的 MXFP8 以及 Hopper 架构专属的块状 FP8。此外，prime-rl 配置还提供了一个 fp8_bwd 模式，允许采用 BF16 进行前向传播，同时使用 FP8 进行反向传播以进一步优化显存与速度。
性能表现：训练步骤显著提速
根据 Prime Intellect 在 8 节点 H200 集群（32 至 128 块 GPU）、FP8 精度及 16,384 token 序列长度下的自报测试数据，Prime Mega MoE 展现了明确的加速效果：
- 完整训练步骤：相较于 prime-rl 基线（使用 torch all-to-all 加 DeepGEMM FP8 专家），完整训练步骤时间缩短了 32% 至 34%。相较于 TorchTitan，提速幅度达到 48% 至 55%。
- 吞吐量提升：在 Qwen3-30B-A3B 模型上，吞吐量从 5,337 token/s/GPU 提升至 7,886 token/s/GPU；在 Qwen3-235B-A22B 模型上，从 1,094 提升至 1,697 token/s/GPU。
- 单 MoE 层延迟：前向传播的加速效果最为显著，在 Qwen3-30B-A3B 上从 12.0 毫秒降至 3.9 毫秒（提升 3.11 倍）。反向传播同样获得优化，延迟从 15.4 毫秒降至 8.2 毫秒（提升 1.87 倍）。
工程部署与生态兼容
Prime Mega MoE 作为分发后端集成在 prime-rl 训练器中。开发者可以通过修改配置文件轻松启用该功能：
- 在配置文件的
[trainer.model.moe.dispatch]部分，设置type = "mega_moe"。 - 根据硬件条件，从
bf16、mxfp8、fp8或fp8_bwd中选择所需的精度模式。
需要注意的是，训练内核随 prime-kernels 提供的预构建 DeepGEMM wheel 一同发布。标准的 DeepGEMM 构建版本不包含这些训练所需的扩展组件。
与其他融合 MoE 方案相比，Prime Mega MoE 在提供完整训练（反向传播）支持的同时，保持了计算与通信的深度融合，且采用宽松的 Apache 2.0 许可，为训练大规模 MoE 模型的团队提供了切实的成本与时间优化方案。
© 版权声明
文章版权归作者所有，未经允许请勿转载。
