In MoE architectures, the number of input tokens distributed across experts can vary significantly, resulting in imbalanced AlltoAll communication and uneven expert workload distribution. NPUs hosting hot experts suffer from insufficient compute and communication resources, while those with cold experts are prone to underutilization, leading to performance degradation. The load balancing feature is designed to reduce NPU resource imbalance and improve model inference performance.
MindIE supports two load balancing modes: static load balancing in redundancy mode and forcible load balancing.
Static load balancing in redundancy mode: Redundant experts are deployed to share the load of hot experts, thereby facilitating effective load balancing.
Forcible load balancing: Mock the outputs of the top k operator by replacing the original top k outputs with fake tensors that ensure absolute load balancing among experts. This mode only provides a theoretical upper limit for load balancing. It changes the actual routing of model experts and cannot be used in official services.
- The Atlas 800I A2/A3 inference server supports this feature.
- The DeepSeek R1/V3 and Qwen-MoE models support this feature.
- This feature is applicable exclusively during All2All collective communication in MoE architectures (
[object Object]is set to[object Object]in the model configuration file). In the prefill-decode disaggregation scenario, prefill and decode instances usually use different collective communication modes. Therefore, parameters for configuring load balancing must be set separately. - Forcible load balancing only represents the theoretical upper limit of load balancing and cannot be used in official services.
- Static load balancing in redundancy mode is implemented by deploying redundant experts on NPUs hosting routing experts. Deploying one additional redundant expert on each NPU requires 2.4GB extra graphics memory.
Procedure for static load balancing in redundancy mode: Collect expert hotspot information, generate redundant expert deployment tables, and set load balancing parameters.
- Forcible load balancing does not require collecting expert hotspot information or generating redundant expert deployment tables. It can be enabled by setting the load balancing parameters alone.
This step is to obtain expert hotspot distribution in the actual service data or dataset.
Users can export expert hotspot data in actual service scenarios as a .csv file by setting two environment variables
[object Object]and[object Object]without enabling load balancing. The hotspot information of the prefill and decode phases is saved separately to generate a redundant expert deployment table for each phase.The procedure is as follows:
Add the following environment variables to the
[object Object]and[object Object]fields in[object Object]:[object Object][object Object]
Run model inference services and generate files containing hotspot information.
[object Object]
After the hotspot information is generated, manually gather the expert hotspot information on all servers into a single folder. Alternatively, set the file export paths on all servers to shared disk paths.
After the hotspot information is collected, each NPU generates a .csv file that contains a matrix ([object Object] × the number of experts per NPU). Each number in the matrix represents the number of tokens processed by experts in that layer. The matrix is appended to the collection file at an interval of eight tokens.
Based on the collected expert hotspot information, use the [object Object] component of the tool to generate a redundant expert deployment table.
The following describes how to install the
[object Object]component.[object Object]The installation is successful if the following information is displayed:
[object Object]Refer to to generate a redundant expert deployment table using the
[object Object]component. The typical 8-server 64-device configuration is as follows:[object Object]msit provides two load balancing algorithms: compute-communication load balancing (C2LB) and speculative-moe interface algorithms. Currently, the optimal result is obtained by using the speculative-moe level 2 mixed algorithm (al 5).
[object Object]
Load balancing parameters can be configured by modifying the [object Object] file in the [object Object] installation directory. Modify the [object Object], [object Object], [object Object], [object Object], [object Object], and [object Object] parameters in the [object Object] field. By default, load balancing is disabled. The typical configuration is as follows:
The parameters are described as follows.