This API supports load balancing and expert pruning. The mapped expert table and mask can be transferred to the mixture of experts (MoE) layer for data distribution and processing.
Load balancing: To address load imbalance, this operator can map the top K logical expert IDs to the rank IDs for each token. The computation method is as follows:
The following describes how to compute the rank to which the ith token is sent.
[object Object]Expert pruning: The top K experts to which tokens are sent can be pruned based on the threshold. The computation method is as follows:
[object Object]with shape[object Object]is broadcasted to become[object Object]with shape[object Object], where experts corresponding to[object Object]in the[object Object]dimension will be directly pruned.[object Object]elements in[object Object]will also be pruned if they meet conditions.[object Object]
Each operator has calls. First, [object Object] is called to obtain the workspace size required for computation and the executor that contains the operator computation process. Then, [object Object] is called to perform computation.
Parameters
[object Object]Returns
[object Object]: status code. For details, see .The first-phase API implements input parameter verification. The following errors may be thrown.
[object Object]
Deterministic computation:
[object Object]defaults to a deterministic implementation.
API mapping and calling sequence: This API must be used together either with
[object Object]and[object Object]/[object Object], with a fixed calling sequence ([object Object]→[object Object]→[object Object]/[object Object]);Or with
[object Object]and[object Object]/[object Object], with a fixed calling sequence ([object Object]→[object Object]→[object Object]/[object Object]). For details, see .Parameter consistency requirements: The values of the
[object Object]and[object Object]parameters used during the calling must be consistent for all ranks and at different network layers, and must be consistent with the corresponding parameters of[object Object]and[object Object]/[object Object].Hardware-related definitions: [object Object]Atlas A3 training products/Atlas A3 inference products[object Object]: In this scenario, a single rank contains dual dies. Therefore, the "rank" in the parameter description indicates a single die.
Restriction on the shape format:
[object Object]: Number of tokens output by the rank, which must be in the range (0, 512].[object Object]: Number of selected top experts, which must be in the ranges (0, 16] and (0,[object Object]].[object Object]: Number of MoE experts, which must be in the range (0, 1024].[object Object]: Number of columns in the mapping table[object Object]. The value range is [2,[object Object]+ 1]. The first column indicates the number of deployed instances of logical experts (value > 0), and the following[object Object]– 1 columns indicate the corresponding rank IDs.- Restriction on the total number of instances: The total number of MoE expert instances deployed on all ranks cannot exceed 1024. That is,
[object Object]([object Object]indicates the number of instances deployed on a single rank). - Consistency of the number of instances per rank: The number of expert instances deployed on each rank must be the same.
The following uses the [object Object]Atlas A3 training series products/Atlas A3 inference series products[object Object] and Ascend 950PR/Ascend 950DTas examples to describe how to invoke the MoeUpdateExpert, MoeDistributeDispatchV2, and MoeDistributeCombineAddRmsNorm operators.
The following example is for reference only. For details, see .