MoE model expert dynamic grouping training optimization method and system for text generation

Abstract

The invention discloses a MoE model expert dynamic grouping training optimization method and system aiming at text generation, and the method comprises the steps: firstly obtaining multi-field text data to generate corresponding vector representation, inputting the vector representation into a MoE module, and dynamically selecting an expert for calculation by a gating mechanism according to the input; secondly, constructing an expert grouping performance evaluation model, and performing expert grouping dynamic adjustment in the MoE model to obtain a global optimal expert grouping; and then on the basis of global optimal expert grouping, performing expert grouping level segmentation on AlltoAll communication and expert calculation in the MoE model, and completing MoE model parallel training of expert grouping. Finally, after parallel training, multi-field text input is independently calculated through an activated expert unit in the MoE module, output of all experts is weighted and aggregated according to gating distribution weights, and text representation is generated. According to the method, calculation fluctuation caused by uneven loads in the text generation process is relieved, and the training efficiency is remarkably improved.

Description

Technical Field

[0001] This invention relates to the field of distributed large model training, and in particular to a method and system for optimizing dynamic group training of experts for text generation MoE models, which is applicable to the distributed training acceleration of Mixture-of-Experts (MoE) models. Background Technology

[0002] With the rapid development of large-scale pre-trained models, the number of model parameters and structural complexity are constantly increasing, leading to a significant rise in computational and communication overhead during training. In multi-domain text generation tasks, data from different domains exhibit heterogeneity and long sequence characteristics, making it difficult for traditional dense models to balance computational efficiency and storage costs while maintaining performance.

[0003] Mixture of Experts (MoE) models, as a sparse neural network architecture, activate only a subset of experts through gating mechanisms, effectively reducing overall computational cost while maintaining model expressiveness. Compared to coarse-grained experts with larger parameter scales, fine-grained experts consist of more units with smaller parameters, enabling them to flexibly capture local patterns and subtle features, thus enhancing the model's ability to model heterogeneous data.

[0004] Compared to coarse-grained experts with a large parameter scale, fine-grained experts consist of more expert units with smaller parameters, enabling them to capture local patterns and subtle features more flexibly and enhance the model's ability to model heterogeneous data. However, fine-grained design also brings systemic challenges: first, it requires activating more experts, leading to a significant increase in all-to-all data volume; second, gating decisions cause load fluctuations, easily resulting in expert overload or idleness, reducing training efficiency and model performance.

[0005] To address the communication and load imbalance issues arising from fine-grained experts, existing research primarily focuses on optimization from both algorithmic and system perspectives. At the algorithmic level, some methods introduce auxiliary loss functions or improve routing strategies to balance expert activation; however, these methods often increase the complexity of the optimization objective, impacting model convergence efficiency. At the system level, research emphasizes improvements in communication and load scheduling. For example, DeepSpeed ​​employs a hierarchical full-interaction communication strategy, splitting cross-node communication into intra-node and inter-node levels to improve bandwidth utilization; FastMoE introduces a shadow expert mechanism combined with pipeline computation to achieve communication-computation overlap; and PipeMoE partitions inputs based on batch or sequence dimensions to improve resource utilization. However, these methods largely rely on static partitioning strategies, making it difficult to adapt to dynamic changes in expert load during training. The problems of load imbalance between pipelines and uneven load among experts within a pipeline due to input data partitioning still persist.

[0006] In summary, while existing research methods have reduced all-to-all communication overhead to some extent, issues such as single-granularity pipeline segmentation and uneven load distribution still exist in typical scenarios like fine-grained experts and long sequence inputs. Specifically, differences in feature distribution and length between texts from different domains can lead to some experts being overloaded while others are idle. The unevenness caused by segmentation along the batch or sequence dimension results in poor overall training acceleration. Summary of the Invention

[0007] This invention aims to address the problems of high all-to-all communication overhead and uneven load distribution between stages caused by pipeline stage partitioning in fine-grained expert architectures, thus blocking communication. It proposes a method and system for optimizing expert dynamic grouping training in MoE models for text generation. This method employs a parallel expert grouping approach at the MoE layer, interleaving communication and computation at the group-level granularity to achieve deep overlap. Simultaneously, it combines a real-time load-aware dynamic grouping strategy to adaptively optimize the load distribution of expert groups, overcoming the load unevenness limitations caused by existing methods that rely on device or sequence dimension segmentation, thereby improving overall training efficiency. This invention significantly improves the training efficiency and system scalability of large-scale fine-grained expert models in long-sequence task scenarios, providing an efficient and scalable solution for communication scheduling and load optimization in MoE systems. The method for optimizing expert dynamic grouping training in MoE models for text generation specifically includes the following steps:

[0008] Step 1: Acquire multi-domain text data. Each input data point undergoes preprocessing and feature encoding to generate a corresponding vector representation, which serves as the input to the MoE module. Within the MoE module, a gating mechanism dynamically selects a subset of experts for computation based on the input features, thereby achieving sparse activation.

[0009] Step 2: Construct an expert grouping performance evaluation model to guide subsequent expert grouping. This invention addresses the issues of expert load differences and computational communication bottlenecks encountered during the training of the MoE module by proposing an expert grouping performance evaluation model to guide expert grouping. Expert grouping involves grouping multiple experts deployed on the same GPU into groups, with experts within each group collaborating to complete tasks in the MoE module. The expert grouping scheme is denoted by (S, A), where S represents the number of expert groups within a single device, and A represents the expert-to-device allocation mapping. The performance evaluation model quantitatively evaluates candidate grouping schemes (S, A) by analyzing the impact of expert load distribution and grouping strategies on computation and communication. The model includes two key metrics:

[0010] (1) Calculation time T comp The computational load required for expert tasks is determined by the model's computational cost and the device's computing power (FLOPS). e·h·ffn_dim. Where N e The number of tokens processed by each expert, h is the hidden layer dimension of the model, and ffn_dim is the feedforward network expansion dimension.

[0011] (2) Communication time T comm Based on the load N of each expert e Given a group mapping A, calculate the number of tokens that each device in the group needs to send or receive, and take the maximum value as the communication volume N for that group. max The estimation is made by combining the transmission delay of the effective unit byte of the device (i.e., the reciprocal of the bandwidth) β, the fixed communication delay α, and the number of processes P.

[0012] The performance evaluation model, as shown in the following formula, is used to estimate the execution performance of the MoE module under any candidate grouping scheme. It consists of two parts: the completion time after communication and computation overlap, and the communication time when they cannot overlap.

[0013] T total =S·T stage +2·T comm (1)

[0014]

[0015] T stage Indicates the effective execution time within a phase.

[0016] Step 3: Based on the expert grouping performance evaluation model, dynamically adjust expert grouping within the MoE model based on real-time load. This method employs a phased search strategy: dynamically generating the optimal grouping scheme based on real-time expert load. The overall process consists of two stages: candidate grouping generation and grouping scheme search.

[0017] First, under the constraints of the total number of experts and the degree of expert parallelism, a candidate group set S is generated, requiring that each group number S is divisible by the number of experts carried by a single card. For each candidate group number S, this method uses a load-descending heuristic strategy to quickly generate an initial grouping scheme A: after sorting the experts in descending order of load, they are sequentially assigned to the group with the smallest current cumulative load, thus obtaining the initial maximum group load L. init This provides a pruning reference for subsequent searches.

[0018] Subsequently, under the condition of a fixed S, a state search and pruning mechanism is used to optimize the grouping scheme A. The model describes the expert-assigned states using a binary mask, progressively expanding the grouping state space. Each time, several unassigned experts are selected to generate new candidate states, and the current group load is recorded. Load balancing-oriented pruning is introduced during the search process: when the maximum group load of a candidate state exceeds the load limit L of the current optimal scheme... maxIf the branch expansion is terminated immediately, the computational complexity is significantly reduced.

[0019] After obtaining the optimal solution A for all candidate group numbers S, the performance evaluation model constructed in step 1 is invoked to calculate the execution time T for each combination (S, A). total T was ultimately selected. total The smallest combination is the globally optimal expert grouping scheme (S) * A * This method achieves an optimal balance between computation and communication overhead by ensuring load balancing while also considering communication and computation efficiency.

[0020] Step 4: Parallel Training Method for the MoE Model Based on Expert Grouping. During training iterations, the All-to-All communication and expert computation in the MoE model are segmented at the expert group level. The optimized expert grouping scheme from Step 2 (S...) is adopted. * A * Splitting the data into groups avoids the load imbalance problem caused by traditional uniform partitioning and reduces resource waste due to waiting. Furthermore, partitioning based on expert groups reduces the amount of communication required for a single All-to-All communication session. Since there are no data dependencies between groups, communication and computation can overlap, further improving the overall execution efficiency of the MoE module.

[0021] Sub-step 4.1: All-to-All Communication Decomposition Based on Expert Grouping. The purpose of All-to-All communication is to send each token to its target expert for computation. To decompose communication and computation at the group level, the input tokens need to be re-partitioned according to the expert's group. First, the input tokens are assigned to the experts according to the routing (i.e., the expert grouping scheme (S)). * A * The tokens are segmented into expert-level subsets; then, the tokens corresponding to experts within the same group G are merged to obtain group-level subsequences, each subsequence constituting an independent micro-batch B. Therefore, AlltoAll communication only needs to be completed between experts within the group, transmitting only the tokens of that group's micro-batch, significantly reducing the amount of communication per transaction.

[0022] Sub-step 4.2: Parallel Method Between Expert Subgroups. Each micro-batch contains two types of operations: intra-group All-to-All communication and intra-group expert computation. Based on the group-driven partitioning method, there is no data dependency between different subgroups. Each micro-batch can be executed asynchronously as an independent scheduling unit. Therefore, the various micro-batches B can be executed interleaved on the time axis, thus forming a pipeline structure in which communication and computation overlap, achieving overlap between computation and communication, and improving the execution efficiency of the MoE module.

[0023] To characterize this interleaved execution of communication and computation, the symbol "||" is used to denote parallel execution between operations. Therefore, the execution relationship of the k-th micro-batch can be represented as:

[0024] Dispatch(B k )||Compute(B k-1 Combine(B) k-2 ),2≤k≤S-1 (3)

[0025] Among them, Dispatch(B k ) represents the Dispatch AlltoA11 communication stage of the k-th micro-batch, Compute(B k-1 ) represents the expert computation stage of the (k-1)th micro-batch, Combine(B,) k-2 ) represents the (k-2)th micro-batch Combine AlltoA1l communication stage.

[0026] Step 5: After the multi-domain text input is independently calculated by each activated expert unit in the MoE module, the outputs of each expert are weighted and aggregated according to the gating weights to generate the final text representation.

[0027] In another aspect, the present invention provides an expert dynamic grouping training optimization system for MoE models of text generation, comprising the following modules:

[0028] The multi-domain text acquisition and input module is used to acquire multi-domain text data and generate corresponding vector representations to input into the MoE module. In the MoE module, the gating mechanism dynamically selects experts for calculation based on the input features to achieve sparse activation.

[0029] The performance evaluation module constructs an expert group performance evaluation model based on computation time and communication time.

[0030] The expert grouping module is used to dynamically adjust the expert grouping based on the real-time load in the MoE model according to the expert grouping performance evaluation model, so as to obtain the globally optimal expert grouping scheme.

[0031] The parallel training module, based on the globally optimal expert grouping scheme, performs expert group-level segmentation of All-to-All communication and expert computation in the MoE model, and completes the parallel training of the expert-grouped MoE model.

[0032] The text generation output module is used to independently compute the outputs of the experts activated in the MoE module after parallel training of multi-domain text inputs. The outputs of each expert are weighted and aggregated according to the gating weights to generate the final text representation.

[0033] Beneficial effects of this invention:

[0034] To address the issues of single pipeline segmentation granularity and uneven load distribution in scenarios with fine-grained experts and long sequence inputs, this invention proposes a multi-domain text generation method and system for the MoE model based on dynamic expert grouping. This method constructs a performance evaluation model based on expert grouping to quantitatively evaluate candidate grouping schemes. During training, a load-aware expert grouping search algorithm is employed, combined with real-time load data, to dynamically determine the optimal number of groups and expert allocation scheme. This adapts to the dynamic changes in token distribution within the MoE model, maintaining a balanced computational load across subgroups. Micro-batches are divided according to expert groups, enabling parallel processing and computation-communication overlap between groups. This effectively alleviates computational fluctuations caused by uneven load distribution and significantly improves training efficiency. Attached Figure Description

[0035] Figure 1 This is a flowchart illustrating the implementation process of the present invention;

[0036] Figure 2 This is a schematic diagram of parallel training according to the present invention. Detailed Implementation

[0037] The present invention will be further described below with reference to the accompanying drawings and specific implementation steps:

[0038] An expert dynamic grouping training optimization method for MoE models of text generation, such as Figure 1 and Figure 2 The process, as shown in the diagram, includes the following steps:

[0039] Step 1: Acquire multi-domain text data, including various types of text input such as news, Q&A, social media, and technical documents. Each input data piece undergoes preprocessing and feature encoding to generate a corresponding vector representation, which is then used as input to the MoE module. In the MoE module, a gating mechanism dynamically selects a subset of experts for computation based on the input features, thereby achieving sparse activation.

[0040] Step 2: Construct an expert grouping performance evaluation model to guide expert grouping. In this step, considering the computational and communication characteristics of the MoE model under a fine-grained expert structure, a performance evaluation model for expert grouping is proposed to quantitatively evaluate the execution efficiency of different expert grouping schemes. The performance evaluation model models both communication overhead and expert computational overhead, comprehensively reflecting the impact of grouping strategies on the overall execution time. The inputs to this evaluation model are the expert grouping scheme (S, A) and the load N of each expert. e .

[0041] (1) Calculation time T compThe expert computation mainly consists of two consecutive matrix multiplications. The first multiplication projects the token received by each expert from the hidden dimension h to the feedforward network's extended dimension ffn_dim, and the second multiplication projects it back from the extended dimension to the hidden dimension. Therefore, the total computational cost for a single expert is approximately 2N. e ·h·ffn_dim. Where N e This represents the number of tokens received by the expert. Assuming the device's computing power is FLOPS, the overall computation time can be expressed as:

[0042]

[0043] This model reveals the decisive impact of expert grouping strategies on computational latency: when the distribution is uneven, an excessive load on one expert will lead to T... comp It was significantly magnified.

[0044] (2) Communication time T comm In the AlltoAll communication process, suppose there are P parallel processes, d type N represents the number of bytes occupied by a single numeric element. max If N is the maximum amount of data to be sent in the expert group, then the amount of data that each card needs to send in one communication is N. max ·h·d type Bytes. According to the distributed communication delay model, the communication overhead can be approximated as:

[0045]

[0046] Where α represents the message initiation delay, and β represents the transmission delay per unit byte (i.e., the reciprocal of bandwidth). This model shows that when the load is unbalanced, it will lead to N max It is too large, thus becoming a bottleneck in the overall communication.

[0047] After expert groups are formed, each group corresponds to a pipeline stage, and communication and computation operations between groups overlap through scheduling. The total execution time characteristics of the system evolve with the number of groups S: when S=1, the system degenerates into serial execution, and the total time is a simple sum of communication and computation; when S=2, preliminary overlap can be formed; when S>2, an all-to-all communication and computation overlap stage appears, involving dispatch and combine. Therefore, the effective execution time within a stage can be expressed as T. stage :

[0048]

[0049] Based on the above communication and computation modeling, the total execution time of the grouping scheme can be obtained as follows, where T commThis represents the communication overhead that cannot be masked by computation within any stage and must be completed sequentially:

[0050] T total =S·T stage +2·T comm (4)

[0051] Step 3: Based on the expert grouping performance evaluation model, dynamically adjust the expert groups in the MoE model according to real-time load. To address the issue of unbalanced expert computational load caused by dynamic changes in token distribution during training, this step first collects the real-time computational load values ​​of all experts at the beginning of the iteration, and denots the computational load of the e-th expert in the current iteration as N. e This forms the load vector N = [N1, N2, ..., N]. E [E], where E represents the total number of experts. Based on this, a phased search strategy is adopted to optimize expert grouping: the first phase generates a set of candidate group numbers based on the total number of experts and parallelism constraints, and enumerates each candidate value S one by one; the second phase searches for the optimal expert allocation scheme A given the number of groups S. Finally, the time performance evaluation model T established in step 1 is called... total The execution time of each candidate grouping scheme (S, A) is calculated, and the grouping configuration with the shortest execution time is selected as the final result. The search for the optimal allocation scheme A mainly includes the following two key steps:

[0052] Sub-step 3.1: Generate an initial grouping scheme based on descending order allocation. Given the number of candidate groups S (capacity of each group g = E / S), define the cumulative load of the k-th group as... Among them G k Let k represent the set of experts contained in the k-th group. Sort the load in descending order, prioritize processing experts with higher loads, and assign them to the group k with the lowest current cumulative load. * In the middle, update G k* The experts included in the group and the load L of the group k* ,Right now

[0053]

[0054] Once all experts have been assigned, an initial grouping scheme is obtained. The load of this scheme will be used as the pruning threshold L in subsequent search processes. init When the maximum grouping load of a candidate grouping scheme exceeds the threshold, the expansion of that branch is terminated early, thereby significantly reducing the search space and computational complexity.

[0055] Sub-step 3.2: State-space modeling and optimization solution. With the number of candidate groups S, the expert assignment problem is modeled as a state-space search. The expert assignment states are encoded as E-bit binary masks, where each state is represented by the mask ∈ {0,1}. E This means that in the mask, the i-th bit being 1 indicates that expert i has been assigned to a group, and 0 indicates that it has not been assigned. The initial state mask = 0 indicates that all experts are unassigned, and the final state mask = 2. E -1 indicates that all experts have been assigned. In each state, dp[mask][k] represents the maximum group load when the expert assignment state is masked and k groups have been formed.

[0056]

[0057] The optimal grouping scheme A is solved through state transitions. For the current state (mask, k), enumerate all unassigned expert subsets R and select a subset G of size g per group as a new grouping candidate, generating a new state mask'. The new maximum group load is the maximum of dp[mask][k] and the expert load of the newly added subset G:

[0058]

[0059] During state transition, load pruning is introduced if the maximum group load L' max The threshold L has been exceeded. init If the state is positive, then the state and its subsequent extended branches are discarded without further searching; otherwise, take the currently calculated L'. max Update the state with the smaller value from the previously recorded maximum group load in this state:

[0060] dp[mask'][k+1]=min(dp[mask'][k+1],L' max (8)

[0061] By repeatedly performing this state transition operation, the remaining experts are gradually assigned to various groups until S complete groups are formed. This method can systematically explore all possible grouping schemes, while effectively eliminating obviously imbalanced candidates through pruning, thus ensuring search efficiency.

[0062] Step 4: Parallel Training Method for MoE Model Based on Expert Grouping. During training iterations, the MoE model is trained using the obtained expert grouping scheme (S...). * A *The method divides the All-to-All communication and expert computation in the MoE model into groups, with each group only processing experts and their corresponding tokens, thereby reducing the amount of All-to-All communication per transaction. At the same time, by taking advantage of the fact that there is no data dependency between groups, each group's micro-batch is treated as an independent scheduling unit. Through staggered scheduling of communication and computation, a pipeline parallel mode is formed, thereby improving the overall execution efficiency.

[0063] Sub-step 4.1: All-to-All Communication Decomposition Based on Expert Grouping. All-to-All communication involves sending each input token to its target expert. To achieve communication decomposition, the tokens need to be re-divided before communication. Let the total number of input tokens be N, the total number of experts be E, and the routing function be r(i), where r(i) represents the target expert number corresponding to the i-th token. The input tokens are sorted according to the routing results and divided into expert-level subsets T = {T1, T2, ..., T...}. E}, t i Let be the i-th element of set T, where each expert e corresponds to a set of tokens. Where T... e It is the set of all tokens for the target expert e; L e For T e Number of tokens in the set:

[0064] T e ={t i |r(i)=e},L e =|T e | (9)

[0065] Furthermore, according to mapping relationship A * Determine the group to which each expert belongs, and define the group set {G1, G2, ..., G...} s}, where each group G s Includes several experts ε s Those belonging to the same group G s The group-level subsequence T is obtained by merging the token sets corresponding to all experts. Gs (i.e., group-level micro-batch) and Token quantity L s :

[0066]

[0067] Sub-step 4.2: Parallel Method Among Expert Subgroups. Each micro-batch contains independent communication and computation tasks, and its execution flow includes three stages: the intra-group Dispatch All-to-All communication stage, where the intra-group Token data is sent to the target expert's device according to the routing results; the expert computation stage, where the target expert performs feedforward computation (such as FFN) on the received Token; and the intra-group Combine All-to-All stage, where the expert computation results are sent back to the original device of the Token. Since there is no data dependency between different expert subgroups, the execution stages of each micro-batch can be staggered in time to form a pipeline. Specifically, in the steady-state stage of the pipeline, micro-batch B... k Compute Stage and Micro Batch B k+1 The DispatchAlltoAll Stage can be executed in parallel; meanwhile, B k The Combine AlltoAll Stage can also be combined with B k+1 The Compute Stages are parallelized, thus forming a three-stage overlapping pipeline on the timeline.

[0068] While this method significantly improves parallel efficiency, it inevitably suffers from bubbles: during pipeline startup, computation waits for the first group's communication to complete; during pipeline termination, communication cannot overlap with new computation tasks. To address this issue, this method introduces computation tasks from shared experts during idle periods to fill pipeline gaps, alleviate tail latency, and thus reduce resource idle time.

[0069] Step 5: After parallel training, the multi-domain text inputs are independently computed by the activated expert units in the MoE model. The outputs of each expert are weighted and aggregated according to the gating weights to generate the final text representation.

[0070] In another aspect, the present invention also provides an expert dynamic grouping training optimization system for MoE models of text generation, comprising the following modules:

[0071] The multi-domain text acquisition and input module is used to acquire multi-domain text data and generate corresponding vector representations to input into the MoE module. In the MoE module, the gating mechanism dynamically selects experts for calculation based on the input features to achieve sparse activation.

[0072] The performance evaluation module constructs an expert group performance evaluation model based on computation time and communication time. It consists of two parts: the completion time after the communication and computation overlap, and the communication time that cannot overlap.

[0073] The expert grouping module is used to dynamically adjust expert groups based on real-time load in the MoE model according to the expert grouping performance evaluation model. It obtains the globally optimal expert grouping scheme through candidate grouping generation and grouping scheme search.

[0074] The parallel training module, based on the globally optimal expert grouping scheme, performs expert group-level segmentation of All-to-All communication and expert computation in the MoE model, and completes the parallel training of the expert-grouped MoE model.

[0075] The text generation output module is used to independently compute the outputs of the experts activated in the MoE module after parallel training of multi-domain text inputs. The outputs of each expert are weighted and aggregated according to the gating weights to generate the final text representation.

[0076] experiment:

[0077] This experiment was conducted on an NVIDIA H800 GPU cluster (8 cards, 80GB VRAM per card) and a GPU-accelerator K100 GPU cluster (8 cards, 80GB VRAM per card), using NVLink high-speed interconnect to ensure low-latency communication between nodes. The distributed training framework was implemented based on Megatron-LM, with customized modifications to its MoE (Mixture of Experts) module code. Core modifications included adding configuration parameters such as split_num and group_size related to expert grouping pipeline overlap and dynamic expert grouping optimization, thereby improving the parallel overlap ratio of communication and computation and reducing the synchronization waiting overhead caused by uneven load among experts.

[0078] The training dataset used was Wikitext-103, which contains approximately 103 million high-quality English Wikipedia articles covering multiple fields such as news, technology, and history. It exhibits long-sequence characteristics and is a commonly used benchmark for long-sequence language modeling tasks. This dataset effectively examines the performance of the MoE model in handling large-scale tokens, adapting to differences in input features across multiple domains, and implementing expert load balancing strategies. The model structure used was based on the DeepSeekV3-MoE design, with adjustments made to some parameter sizes and expert configurations according to experimental requirements. Table 1 lists the main hyperparameter settings of the DeepSeekV3 model used in this experiment.

[0079] Table 1

[0080]

[0081]

[0082] To verify the performance improvement effect of the expert-dynamic grouping-based pipelined parallel training method for the MoE model during training, this experiment conducted comparative tests. The experiment focused on evaluating the impact of different expert group numbers (split_num) on the parallel efficiency of MoE submodules and the overall training time. Under the same training hyperparameters and the same checkpoint, Megatron-LM and the optimized version of this invention were run respectively, and the training time of MoE submodules and the single-step training time (StepTime) were statistically analyzed. In the single-machine environment of H800 server, the experimental results are shown in Tables 2 and 3. The single-step training time was improved by 14%, and in long sequence scenarios, it can be improved by 22%.

[0083] Table 2

[0084]

[0085] In a standalone environment of a domestically produced DCU K100 server, with the current split_num=2 setting, the speed of the MoE portion is increased by 7% to 10%.

[0086] Table 3

[0087]

[0088] The optimization scheme of this invention is based on load-aware expert dynamic grouping and parallel strategy between groups, which significantly reduces communication latency and effectively alleviates the problems of single pipeline segmentation granularity and load imbalance and instability. When dealing with multi-domain text input and ultra-long sequence scenarios, the optimization advantages are more obvious. It can take into account the differences in features of different domains and the load fluctuation of long sequences, ensure the stable acceleration of the MoE model training process, and significantly improve the overall training efficiency.

Claims

1. A method for dynamic grouping training and optimization of expert models for text generation MoE models, characterized in that, Includes the following steps: Step 1: Acquire multi-domain text data, generate corresponding vector representations, and input them into the MoE model. The gating mechanism dynamically selects experts for calculation based on the input to achieve sparse activation. Step 2: Construct an expert group performance evaluation model; Step 3: Based on the expert grouping performance evaluation model, dynamically adjust the expert grouping in the MoE model based on real-time load to obtain the globally optimal expert grouping scheme; Step 4: Based on the globally optimal expert grouping scheme, perform expert group-level segmentation of All-to-All communication and expert computation in the MoE model to complete the parallel training of the expert-grouped MoE model; Step 5: After parallel training, the multi-domain text inputs are independently computed by the activated expert units in the MoE model. The outputs of each expert are weighted and aggregated according to the gating weights to generate the final text representation.

2. The expert dynamic grouping training optimization method for MoE model of text generation according to claim 1, characterized in that, Step 2 is implemented as follows: Expert grouping involves grouping multiple experts deployed on the same GPU into groups, with experts within each group collaborating to complete tasks in the MoE module. The expert grouping scheme is denoted by (S, A), where S represents the number of expert groups within a single device, and A represents the expert-to-device allocation mapping. The performance evaluation model quantifies the candidate grouping scheme (S, A) by analyzing the impact of expert load distribution and grouping strategies on computation and communication. The expert grouping performance evaluation model includes two key metrics: (1) Calculation time T comp The computational load required for expert tasks is determined by the model's computational load and the device's computing power (FLOPS). e ·h·ffn_dim; where N e The number of tokens processed by each expert, h is the hidden layer dimension of the model, and ffn_dim is the feedforward network expansion dimension; (2) Communication time T comm Based on the load N of each expert e Given a group mapping A, calculate the number of tokens that each device in the group needs to send or receive, and take the maximum value as the communication volume N for that group. max The estimation is made by combining the effective unit byte transmission delay β, the fixed communication delay α, and the number of processes P of the device; The expert grouping performance evaluation model is as follows, used to estimate the execution performance of the MoE module under any candidate grouping scheme; it consists of two parts: the completion time after communication and computation overlap, and the communication time that cannot overlap. T total =S·T stage +2·T comm T stage Indicates the effective execution time within a phase.

3. The expert dynamic grouping training optimization method for MoE model of text generation according to claim 2, characterized in that, The specific implementation process of step 3 is as follows: The expert grouping is dynamically adjusted based on real-time load, and a phased search strategy is adopted to dynamically generate the optimal grouping scheme based on real-time expert load. The overall process is divided into two stages: candidate grouping generation and grouping scheme search. First, under the constraints of the total number of experts and the degree of expert parallelism, a candidate group set S is generated, requiring that each group S is divisible by the number of experts carried by a single card. For each candidate group S, an initial grouping scheme A is quickly generated using a load descending heuristic strategy: experts are sorted in descending order of load and then sequentially assigned to the group with the smallest current cumulative load, resulting in the initial maximum group load L. init ; Subsequently, under the condition of fixed S, the grouping scheme A is optimized by state search and pruning mechanism; The model describes the expert assignment states using binary masks, progressively expanding the group state space. Each time, several unassigned experts are selected to generate new candidate states, and the current group load is recorded. Load balancing guided pruning is introduced during the search process: when the maximum group load of a candidate state exceeds the load limit L of the current optimal solution... max If this happens, immediately terminate the branch expansion; After obtaining the optimal solution A for all candidate group numbers S, the constructed performance evaluation model is invoked to calculate the execution time T for each combination (S, A). total T was ultimately selected. total The smallest combination is the globally optimal expert grouping scheme (S) * A * ).

4. The expert dynamic grouping training optimization method for MoE model of text generation according to claim 3, characterized in that, The specific implementation process of parallel training of the MoE model by the expert group is as follows: All-to-All Communication Decomposition Based on Expert Grouping: The purpose of All-to-All communication is to send each token to its target expert to perform computation, and to re-divide the input tokens according to the expert's group; first, the input tokens are processed using an expert grouping scheme (S... * A * The tokens are split into expert-level subsets; then the tokens corresponding to experts in the same group G are merged to obtain group-level subsequences, and each subsequence constitutes an independent micro-batch B; AlltoAll communication is only completed between experts within the group and only the tokens of the micro-batch of that group are transmitted. Parallelism between expert subgroups: Each micro-batch contains two types of operations: intra-group All-to-All communication and intra-group expert computation; based on the group-driven partitioning method, there is no data dependency between different subgroups, and each micro-batch can be executed asynchronously as an independent scheduling unit. Each micro-batch B is executed in an interleaved manner on the time axis, forming a pipeline structure in which communication and computation overlap, thus achieving the overlap of computation and communication.

5. A MoE model expert dynamic grouping training optimization system for text generation, used to implement the MoE model expert dynamic grouping training optimization method according to any one of claims 1 to 4, characterized in that, Includes the following modules: The multi-domain text acquisition and input module is used to acquire multi-domain text data and generate corresponding vector representations to input into the MoE module. In the MoE module, the gating mechanism dynamically selects experts for calculation based on the input features to achieve sparse activation. The performance evaluation module constructs an expert group performance evaluation model based on computation time and communication time. The expert grouping module is used to dynamically adjust the expert grouping based on the real-time load in the MoE model according to the expert grouping performance evaluation model, so as to obtain the globally optimal expert grouping scheme. The parallel training module, based on the globally optimal expert grouping scheme, performs expert group-level segmentation of AlltoAll communication and expert computation in the MoE model, and completes parallel training of the expert-grouped MoE model. The text generation output module is used to independently compute the outputs of the experts activated in the MoE module after parallel training of multi-domain text inputs. The outputs of each expert are weighted and aggregated according to the gating weights to generate the final text representation.

Images