Turn Waste into Worth: Rectifying Top- k Router of M o E (2024)

Abstract: 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 empty, 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%.

## introduction sparse mixture of experts (moe) is gaining popularity as a model architecture for training large language models @xcite @xcite @xcite owing to its computational efficiency. in a sparse moe model, each token is assigned to one or more experts based on a routing mechanism. the top-k router is currently the most widely used routing mechanism, where tokens are directed to the experts with the top-k scores. however, top-k router is unbalanced, where the number of tokens routed to different gpus is not the same. in order to achieve a balanced workload across gpus, the top-k routing imposes a maximum limit on the number of tokens that each expert can process, resulting in any tokens that exceed this limit being dropped and empty experts being filled with zeros, which negatively affects overall model performance @xcite . previous studies have attempted to address the balance issue in routing by introducing auxiliary loss mechanisms @xcite @xcite . but there are drawbacks to the way, the performance degradation due to dropped tokens is still significant @xcite . even some methods have made improvements to propose absolutely balanced routers, but they have been found to underperform the original top-k routing methodology @xcite . rather than focusing on improving the balance of the top-k router, we propose an alternative approach called the rectify-router, which rectifies top-k router by post-processing the dropped tokens and padding from the top-k router. we propose two rectify-routers: the intra-gpu rectification and the fill-in rectification. the intra-gpu rectification is designed to handle the dropped tokens, while the fill-in rectification specifically addresses the padding issue. post-processing the dropped tokens with another router may bring expensive communication cost. therefore, we propose the intra-gpu rectification which routes the dropped tokens to the experts within the gpu where they are located, eliminating the need for inter-gpu communication. our empirical experiments have demonstrated that the intra-gpu rectification effectively handles the post-processing of dropped tokens and is more efficient than the commonly used routers, in terms of communication. to address the padding issue, we present the fill-in rectification, which replace padding tokens with the tokens that have high routing scores. fillin rectification first identifies the optimal expert for each token based on the routing scores and subsequently selects the tokens with the highest routing score to replace the padding for each expert. by employing fill-in rectification, tokens with the higher routing scores receive more computational allocation. the intra-gpu rectification and fill-in rectification are orthogonal approaches that can be seamlessly combined. our experiments have demonstrated their effectiveness in handling dropped tokens and padding. furthermore, combing the intra-gpu rectification and fill-in rectification yield improved performance compared to using them individually. ## related works the routing of moe can be classified into two categories: balanced and unbalanced. the balanced routing assigns the same number of tokens to each expert, while the unbalanced routing does not make sure that the number of tokens received by each expert is the same. unbalanced routing top-k routing was the most commonly used unbalanced routing proposed by @xcite , which greedily assigns tokens to experts, according to the token-expert assignment scores. numerous moe models have adopted top-k routing, including switch transformer @xcite , glam @xcite , st-moe @xcite , flan-moe @xcite , and nllb @xcite , to name just a few. it is worth noting that many unbalanced routing methods are variations or derivatives of top-k routing. for example, switch transformer @xcite argues in favor of using top-1 routing instead of top-2 routing for improved efficiency. st-moe @xcite and limoe @xcite propose auxiliary loss functions to enhance the stability of moe during training. additionally, scomoe @xcite and gating-dropout @xcite improve the efficiency of top-k routing by designing hierarchical routing systems based on the hierarchical structure of the communication topology. the routing method proposed in this paper is also a variation of top-k routing. however, unlike the aforementioned approaches, our objective is to address the issues of dropped tokens and padding that arise from unbalanced routing. switch transformer @xcite tackles the problem of dropped tokens by increasing the capacity of experts, allowing each expert to handle more tokens. while this approach reduces the number of dropped tokens, it introduces additional overhead in terms of both speed and memory. on the other hand, megablocks @xcite addresses the challenges of padding and dropped tokens by gathering all experts onto the same gpu and employing model parallelism rather than expert parallelism. however, the model parallelism is shown to be more expensive than the expert parallelism by tutel @xcite . ## preliminary in this section, we will introduce expert parallelism, top-k routing, and two prevalent challenges that emerge while employing top-k routing: padding and dropped tokens. expert parallelism and top-k routing in expert parallelism, experts are distributed across gpus uniformly. if there are n experts and k gpus, each gpu contains k/n experts. the process of transmitting tokens to the respective experts entails inter-gpu communication. top-k routing greedily assigns tokens to experts according to the routing score: since each token undergoes processing by multiple experts, the outputs of these experts for the same token are consolidated through linear combination. the combining weights are determined by the normalized routing scores, as defined in eq. ( 1 ): the top-k routing approach exhibits an inherent imbalance, wherein the distribution of tokens among different experts is not uniform. however, the current distributed framework exclusively supports balanced computation across gpus. consequently, there exists a limitation on the maximum number of tokens that each expert can receive, which is referred to as the capacity. the capacity is determined by the capacity factor, which is typically set to k for top-k routing @xcite . mathematically, the capacity can be expressed as: @xmath0. ## method in this paper, we introduce a novel approach to address both the dropped tokens and padding as-sociated with top-k routing by utilizing rectify-routers. specifically, we propose two rectify-routers: the intra-gpu rectification and the fill-in rectification, which are visualized in figure 2 . the intra-gpu rectification is designed to efficiently post-process the dropped tokens, while the fill-in rectification is dedicated to addressing the padding problem. ## conclusion in this paper, we present the rectify-router, a method to tackle dropped tokens and padding in moe models. by introducing the intra-gpu rectification and the fill-in rectification, we effectively handle the issues of dropped tokens and padding, respectively. experimental results demonstrate the individual effectiveness of both techniques and the synergistic performance improvement when they are combined. furthermore, our methods prove to be effective in diverse settings, including varying numbers of experts, different expert capacities, and even without the load-balance loss.