MoE Token重排性能优化

问题分析

在DeepSpeed MoE（Mixture of Experts，混合专家）框架中，Token重排操作原通过两个BatchMatmul实现，其时间复杂度为O(s²)。然而，由于矩阵的稀疏特性，原方法在重排计算中产生了不必要的计算负担，存在潜在的优化空间。

解决方案

通过采用等价且更加高效的PyTorch API: index_select，重排操作的时间复杂度得以降低至O(s)，从而显著提升了模型训练的整体性能。

解决思路

• 重排过程：在 top1gating 或 top2gating 函数中，计算出每个专家选择的 Token 索引 expert_select_token_idx，其形状为[E*C]。在 MoE 前向传播过程中，依据此索引，通过 index_select API 实现 Token 的重排。
• 反重排过程：top1gating 或 top2gating 函数还需计算每个 Token 在各专家输出中的索引位置 token_rearrange_ec_idx，形状为[S]。在 MoE 前向传播过程中，Token 经过专家处理后，通过 index_select API 从[E*C,M]形状的专家输出中恢复出 Token 的输出[S,M]，最后乘以 Token 对应专家的选择权重，以得到 MoE 层的最终输出。

使用场景

当进入MoE层的实际序列长度达到或超过8K时，此优化策略将展现出显著的效果。

使用方法

要启用该优化算法，仅需在配置中设置以下参数：

--enable-token-rearrange-opt

同时，确保已正确设置了MoE相关的参数，例如：

--expert-model-parallel-size    N     # 设置专家模型并行度大小，缺省值为1  
--num-experts   N     # 设置专家数量，缺省值为None

使用效果

预期性能提升幅度大约在2%至3%之间，具体收益取决于模型的具体配置和数据集特性。