Megatron MoE GMM

背景与挑战

针对MoE单卡多专家计算，存在细碎的专家计算操作与通信，通过Grouped GeMM（Grouped General Matrix Multiplication）算子对多专家计算进行合并，提升MoE单卡多专家训练性能。

解决方案

通过调用GMM融合算子，对多个专家计算进行融合，达到加速效果。

使用场景

适用于以下典型场景：

• EP变小导致单卡专家数量增大以及DeepSeek MoE专家数量较多等场景。
• DeepSeek MoE finegrained expert单个专家较小、FFN规模不大以及TP变大导致单卡切分的计算变小。

使用方法

--moe-grouped-gemm

使用效果

Grouped GeMM算子通过减少内存访问、并行计算、降低计算复杂度等方面的优势，可以提高计算的效率和性能，达到提升模型训练性能的目的。

随着FFN规模提升，计算不再细碎，单专家计算效率提升，Grouped GeMM收益变小。当Grok模型的ffn_hidden_size为32768时，性能提升约5%；当ffn_hidden_size为16384时，性能提升约3%；当ffn_hidden_size为8192时，性能提升约6%。