Sparse Mixture of Experts (MoE) models are popular for training large language models due to their computational efficiency. However, the commonly used top-$k$ routing mechanism suffers from redundancy computation and memory costs due to the unbalanced routing. Some experts are overflow, where the exceeding tokens are dropped. While some experts are vacant, which are padded with zeros, negatively impacting model performance. To address the dropped tokens and padding, we propose the Rectify-Router, comprising the Intra-GPU Rectification and the Fill-in Rectification. The Intra-GPU Rectification handles dropped tokens, efficiently routing them to experts within the GPU where they are located to avoid inter-GPU communication. The Fill-in Rectification addresses padding by replacing padding tokens with the tokens that have high routing scores. Our experimental results demonstrate that the Intra-GPU Rectification and the Fill-in Rectification effectively handle dropped tokens and padding, respectively. Furthermore, the combination of them achieves superior performance, surpassing the accuracy of the vanilla top-1 router by 4.7%.

翻译：稀疏混合专家（MoE）模型因其计算效率而在训练大型语言模型中广受欢迎。然而，常用的top-$k$路由机制由于路由不平衡，存在冗余计算和内存开销的问题。部分专家过载，导致超出容量的令牌被丢弃；而部分专家空置，则需填充零值，对模型性能产生负面影响。为解决丢弃令牌和填充问题，我们提出了Rectify-Router，包括GPU内修正（Intra-GPU Rectification）和填充修正（Fill-in Rectification）。GPU内修正处理丢弃的令牌，将其高效路由到所在GPU内的专家，避免跨GPU通信。填充修正通过将填充令牌替换为路由得分高的令牌来解决填充问题。实验结果表明，GPU内修正和填充修正能分别有效处理丢弃令牌和填充问题。此外，两者结合可取得更优性能，相比原始top-1路由器，准确率提升4.7%。