Turn Waste into Worth: Rectifying Top-$k$ Router of MoE
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2024
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866909114495401984 |
|---|---|
| author | Zeng, Zhiyuan Guo, Qipeng Fei, Zhaoye Yin, Zhangyue Zhou, Yunhua Li, Linyang Sun, Tianxiang Yan, Hang Lin, Dahua Qiu, Xipeng |
| author_facet | Zeng, Zhiyuan Guo, Qipeng Fei, Zhaoye Yin, Zhangyue Zhou, Yunhua Li, Linyang Sun, Tianxiang Yan, Hang Lin, Dahua Qiu, Xipeng |
| contents | 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%. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2402_12399 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Turn Waste into Worth: Rectifying Top-$k$ Router of MoE Zeng, Zhiyuan Guo, Qipeng Fei, Zhaoye Yin, Zhangyue Zhou, Yunhua Li, Linyang Sun, Tianxiang Yan, Hang Lin, Dahua Qiu, Xipeng Machine Learning Artificial Intelligence Computation and Language 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%. |
| title | Turn Waste into Worth: Rectifying Top-$k$ Router of MoE |
| topic | Machine Learning Artificial Intelligence Computation and Language |
| url | https://arxiv.org/abs/2402.12399 |