Megatron MoE AllGather Dispatcher性能优化

背景与挑战

• gather及scatter算子替换

  在Megatron MoE中的Allgather分支，存在使用gather及scatter的操作。gather及scatter功能为沿dim轴根据索引逐元素进行取值或赋值操作，此操作会有大量的随机地址，对性能造成巨大影响。

  在Megatron MoE中对gather及scatter的调用主要通过以下方式：

  self.global_local_map = global_local_map.view(-1, 1).expand(-1, hidden_states.shape[-1])
  local_hidden_states = torch.gather(global_hidden_states, 0, self.global_local_map)

  通过expand操作，扩展index的维度，再通过扩展后index对hidden_states进行逐元素取值或赋值。

• 异步通信

  在Allgather Dispatcher分支中，permute函数开头分别对hidden_states、max_ind及max_prob三个数据进行allgather通信，这些操作会串行执行，但各计算任务之间并非串行依赖关系。

解决方案

• gather及scatter算子替换

  由于index是通过expand进行扩展的，因此它的每一行内容都是一致的，而用户没有必要使用gather及scatter进行逐元素的操作，可通过index算子以及indexput算子进行逐行操作，对gather及scatter进行等价替换。

• 异步通信

  通过对通信任务进行重新排序，并使用async=True参数进行异步下发，达到计算和通信并行的目的。

使用场景

本优化策略适用于部署了Mcore MoE（Mixture of Experts）架构的深度学习模型，并且已经在系统配置中启用了--moe-token-dispatcher-type allgather参数。

使用方法

设置如下参数即开启AllGather Dispatcher并行优化。

--moe-permutation-async-comm

使用效果

根据实际测试数据显示，类DeepSeekV2十亿参数级别的MoE模型，采用上述优化措施后，端到端训练过程中的性能提升了约10%。这意味着不仅加快了模型收敛速度，同时也降低了达到相同精度水平所需的计算资源消耗。