变废为宝：修正 MoE 的 Top-$k$ 路由
机器学习
2024-02-22 v2 人工智能
计算与语言
摘要
稀疏混合专家 (MoE) 模型因其计算效率而在训练大型语言模型中广受欢迎。然而，常用的 top- 路由机制由于路由不平衡而遭受冗余计算和内存成本的困扰。一些专家溢出，导致超出的 token 被丢弃；而另一些专家空闲，被迫用零填充，这对模型性能产生负面影响。为了解决 token 丢弃和填充问题，我们提出了 Rectify-Router，包含 GPU 内修正 (Intra-GPU Rectification) 和填充修正 (Fill-in Rectification)。GPU 内修正处理被丢弃的 token，将其高效路由到其所在 GPU 内的专家，以避免 GPU 间通信。填充修正通过将填充 token 替换为具有高路由分数的 token 来解决填充问题。实验结果表明，GPU 内修正和填充修正分别有效地处理了被丢弃的 token 和填充问题。此外，两者的结合实现了卓越的性能，其准确率超过了 vanilla top-1 路由器 4.7%。
引用
@article{arxiv.2402.12399,
title = {Turn Waste into Worth: Rectifying Top-$k$ Router of MoE},
author = {Zhiyuan Zeng and Qipeng Guo and Zhaoye Fei and Zhangyue Yin and Yunhua Zhou and Linyang Li and Tianxiang Sun and Hang Yan and Dahua Lin and Xipeng Qiu},
journal= {arXiv preprint arXiv:2402.12399},
year = {2024}
}