A hybrid expert model training optimization method based on matrix routing and token allocation

Abstract

The application discloses a kind of hybrid expert model training optimization methods based on matrix routing and Token distribution, the method first in pre-training process, MoE model utilizes multiple expert submodels to process different input data, MoE model generates expert index by gate network, and generates dispatching index and combination index based on expert index.Second, based on the generation dispatching index and combination index, enter the dispatching token and combination token stage, token is dispatched to each expert for expert processing, the latter will be combined after expert processing token, generate combination token.Finally, after completing the pre-training of the above MoE model, fine-tuning is carried out, and training optimization is completed.The application ensures that the model fine-tuning precision remains unchanged, and the expert computing resources are fully utilized, effectively avoiding resource idleness.

Description

Technical Field

[0001] The present invention belongs to the technical field of neural networks, and in particular relates to a hybrid expert model training optimization method based on matrix routing and Token allocation. Background Art

[0002] In recent years, large models have demonstrated outstanding performance in tasks such as natural language processing (NLP), computer vision (CV), and recommender systems. Scaling model size has been recognized as a key technical approach to further improving model performance. Model parameter size has also rapidly increased from hundreds of billions to trillions, as seen in models like Llama 2 and GPT4. However, training and inference of these ultra-large models require extremely high computational and memory resources. For example, GPT-3, with 175 billion parameters, required 22 days of training on 1,000 A100 GPUs, costing over $1.2 million. Therefore, to reduce the computational cost of model training while scaling model size, the Mixture of Experts (MoE) technique has attracted widespread attention from AI practitioners. By activating only a subset of model parameters, MoE significantly reduces model computational requirements and has been successfully applied to models such as Mixtral and Grok-1.

[0003] The MoE model is a complex neural network structure that uses multiple specialized sub-models (experts) to process different input data. During the pre-training process of the MoE model, the input token is selected by the gating network to activate the indexes of the first k experts and realize the scheduling of tokens to experts. The scheduling stage generates a scheduling index, which is used to schedule the token to the corresponding expert, and each expert calculates it separately. Afterwards, a combination index is generated in the combination stage, and the output of the experts is aggregated based on these combination indexes to obtain the final result. However, when analyzing the pre-training and fine-tuning stages of MoE, it was found that it has two major challenges in terms of computational overhead and accuracy. First, in the pre-training stage, the generation of scheduling index and combination index leads to a lot of computational overhead. The expert index (DP, N) needs to be obtained through the addition, subtraction, multiplication and division operations of multiple high-dimensional matrices. The matrix dimensions are (DP, N, E*C), and the time complexity reaches O(N^3). Secondly, during the scheduling and combining operations, token scheduling and combining require complex matrix calculations. For example, in GShard, Switch Transformers, and Mixtral, the scheduling index requires a high-dimensional matrix multiplication: [DP,H,N]*[DP,N,E*C]=[DP,H,E*C], which significantly increases computational cost, reaching a time complexity of O(N^3). Finally, during the fine-tuning phase, token dropping is often used to alleviate the overfitting problem of MoE. However, small sample datasets contain a large number of zero vectors. Token allocation methods such as GShard, Switch Transformers, Tutel, and Mixtral are inappropriate during fine-tuning, resulting in a significant waste of computational and memory resources. This creates a conflict between accuracy and hardware resource requirements, especially under the top-k (k ≥ 2) selection strategy. Therefore, to address the computational efficiency and accuracy challenges of MoE models during pre-training and fine-tuning, optimizing the model's computational methods and resource allocation is particularly important. Summary of the Invention

[0004] This paper mainly addresses two key issues of the MoE model in the pre-training and fine-tuning stages: high computational overhead and precision loss. A hybrid expert model training optimization method based on matrix routing and token allocation is proposed to solve the above problems.

[0005] The present invention provides a hybrid expert model training optimization method based on matrix routing and token allocation. The analysis method is as follows:

[0006] Step 1: During pre-training, the MoE model uses multiple expert sub-models to process different input data. During this process, the MoE model first generates expert indices through a gating network (primarily composed of a softmax function), and then generates scheduling indices and combination indices based on the expert indices. During the scheduling index generation process, each token selects the corresponding expert based on the gating network, and the expert selection results for each token are summarized into a two-dimensional matrix M. This matrix M is implemented using the following formula:

[0007]

[0008] Among them, M ij is the element in the i-th row and j-th column of the matrix M, token i For the i-th token, expert j Indicated as the jth expert. To reduce the impact of unselected experts on the priority in subsequent processing, construct a matrix X of the same size as matrix M. Matrix M and matrix X (elements in the matrix are much larger than each element in M) are calculated as follows to make the non-selected positions have a lower priority:

[0009] M′=range(N)*M+(X+range(N))*(1-M)

[0010] Where N represents the length of X, and range represents a loop function that generates values ​​from 0 to N-1. In the generated matrix M′, non-zero vector values ​​are extracted through quick sorting and slicing. Then, all indices greater than or equal to X are reset to -1, indicating that the token is not processed by the expert. This completes the generation of the scheduling index, reducing the time complexity to O(Nlog(N)). The scheduling index is a three-dimensional matrix with dimensions (DP, E, C), where DP is the data parallel dimension, E is the number of experts, and C is the expert capacity.

[0011] For combined index generation, the algorithm calculates the cumulative sum Cumsum along the row dimension based on the matrix M using the parallel matrix processing characteristics of the GPU, while keeping the positions of the selected experts unchanged:

[0012] H=Cumsum(M,axis=1)×M

[0013] Add the sequence vector V of the corresponding expert to each column of the matrix H. The vector V is generated by the following formula:

[0014] V=Q×C

[0015] Q is a one-dimensional column vector, where each element is the number of rows minus 1. C is the expert capacity of the MoE, representing the number of tokens processed by each expert. After obtaining H, the combined index is obtained by summing each column and subtracting 1. This process reduces the time complexity of generating the combined index from O(N^3) to O(1), greatly improving computational efficiency. The scheduling index is a two-dimensional matrix with dimensions (DP, N).

[0016] Step 2: After generating the scheduling index and combination index, it is necessary to enter the scheduling token and combination token stage. The former will schedule the token to each expert for expert processing. After the expert processing, the token will be combined to generate a combination token. In the scheduling process and the combination process, it is necessary to search for the physical address of the token based on the scheduling index and combination index generated in step 1, so as to realize the scheduling and combination operations. This process usually requires the use of a high-dimensional matrix multiplication, and the time complexity reaches O(N^3). To this end, the processing of logical indexes is optimized by using the parallel computing capability of the GPU, which can significantly reduce the complexity. Specifically, the logical index includes two types: scheduling index and combination index. Based on the above information, each thread on the GPU runs independently, directly calculates the corresponding physical address from the logical index, and completes data migration and writing. Since there are two forms of logical indexes, each thread needs to independently calculate a unified target index based on the logical index. The present invention will first use the unified computing device architecture CUDA thread structure to obtain the token information that each thread block needs to process. First, an index write_index will be obtained, and its calculation formula is:

[0017] write_index=blockIdx.x*blockDim.x+threadIdx.x

[0018] +w×gridDim.x×blockDim.x

[0019] blockIdx.x, blockDim.x, threadIdx.x, and gridDim.x are the basic parameters of the CUDA thread structure; w is the loop variable, representing the token block being processed by the current thread. The multidimensional spatial coordinates of the data can be parsed based on the index write_index.

[0020]

[0021] k=writr_index mod output_dim2

[0022] Among them, output_dim0, output_dim1, and output_dim2 represent the logical index information of the input. For example, for the scheduling index (DP, E, C), each dimension represents an input. When a dimension does not exist, it is set to the default value of 1. In addition, i corresponds to the expert dimension, j corresponds to the token dimension, and k represents the specific feature dimension. The physical address will be calculated based on the three dimensions i, j, and k and the position of this token in the expert index (DP, N) (denoted as indices). The formula is as follows:

[0023] read_index=i*ouput_dim1*output_dim2+indices*output_dim2+k Through this formula, each thread independently completes the mapping from logical index to physical address, ensuring that data can be efficiently migrated, and obtains a unified target index read_index. After each thread calculates the physical address based on the target index, it will be used to write data at the target location:

[0024] output[write_index]=input[read_index]

[0025] Since each thread completes index calculation and data migration independently, the time complexity of the operation is reduced from O(N^3) to O(1), greatly reducing the time spent on physical address lookup.

[0026] Step 3: After completing the pre-training of the above-mentioned MoE (mixture of experts) model, the model enters the fine-tuning stage to complete the training optimization. At this time, it usually faces two important challenges: the existence of zero vectors and the reasonable allocation of tokens. To address these problems, the present invention proposes an optimized token division and allocation method to improve resource utilization efficiency and ensure the fine-tuning accuracy of the model. In the fine-tuning stage, the proportion of zero vectors in small sample data may be high. Zero vectors are usually some inputs that lack valid information, and excessive processing of these zero vectors will lead to a waste of computing resources. In order to avoid this situation, the input tokens need to be analyzed and divided. This part of the data is mainly divided into two parts: a valid queue and an invalid queue. The valid queue contains all real tokens with non-zero importance values. The invalid queue contains zero vectors, which are used as fillers first, but will not directly participate in expert processing.

[0027] After completing the division of valid and invalid queues, tokens need to be efficiently allocated to MoE experts. The goal here is to prioritize valid tokens while maximizing resource utilization to avoid idle computing resources. Using the Top-k strategy (each token is assigned to at most k experts), tokens in the valid queue are sorted by importance and assigned to experts in order of importance. After valid tokens are allocated, some experts may still have idle computing slots. In order to ensure that all experts process the same number of tokens, tokens in the invalid queue need to be used to fill these vacancies. This method ensures that truly valid tokens are processed first, ensuring the accuracy of the model.

[0028] The beneficial effects of the present invention are:

[0029] The present invention addresses the problem of balancing computational efficiency and accuracy during pre-training and fine-tuning of sparse expert models (MoEs), and proposes several optimization methods, which bring significant technical advantages. First, by scheduling indexes and combining index generation algorithms, the time complexity of the original three-dimensional matrix multiplication is reduced from O(n^3) to O(Nlog(N)), thereby achieving efficient index generation. Secondly, the present invention utilizes a novel addressing method to further optimize the lookup complexity of the token physical address from O(N^3) to O(1), making full use of the parallel processing capabilities of the GPU and significantly improving data access speed. In addition, in the token allocation of the top-k strategy, the present invention introduces a new token allocation method, which changes the dynamic allocation method of tokens. While ensuring that valid tokens are processed, invalid tokens are discarded to ensure that expert computing resources are fully utilized while the model fine-tuning accuracy remains unchanged, effectively avoiding resource idleness. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] Figure 1 is the scheduling index generation algorithm graph;

[0031] Figure 2 It is a graph of the combined index generation algorithm;

[0032] Figure 3 It is an optimized design diagram for token distribution. DETAILED DESCRIPTION

[0033] The present invention will be further described below with reference to the accompanying drawings and specific implementation steps. A hybrid expert model training optimization method based on matrix routing and token allocation includes the following process:

[0034] Step 1: First, during the pre-training process, the token will be distributed to which expert in the MoE model through the gating network. After the expert processes the token, it will be recombined to generate the corresponding output content. This process first generates a scheduling index and a combination index to implement the token scheduling and combination process. Figure 1 In the process of generating the scheduling index, the present invention completes the generation of the scheduling index through the following four steps.

[0035] Step 1.1: Each token gets the expert selection of each token based on the gated network, that is Figure 1 As shown in step (a) of [1], after obtaining the corresponding expert index, a one-dimensional vector is formed, and these vectors are aggregated into a two-dimensional matrix M denoted as (DP, N), as shown in step (b). The following formula is shown:

[0036]

[0037] Step 1.2: After obtaining the matrix M, modify each non-zero element to obtain the number of valid elements before the column where the element is located, and assign it to the element, which is shown in step (c). Then, the present invention combines the matrix M with the range matrix and adds a maximum value X (X) to the unselected expert position.

[0038] N) to ensure that they have lower priority in subsequent steps, as shown in step (d), with the following formula:

[0039] M′=range(N)*M+(X+range(N))*(1-M)

[0040] In this case, range(N) generates a sequential matrix, where N represents the number of columns in M, and X is used to reduce the priority of non-selected positions. After generating the matrix M′, quick sort and slicing are performed. The slicing is primarily determined by the expert capacity. As shown in steps (e) and (f), non-zero vector values ​​are extracted.

[0041] Step 1.3: Reset all indices greater than or equal to X to -1, indicating that the token will not be processed by the expert, thereby completing the generation of the scheduling index. The matrix dimension of the scheduling index is (DP, E, C), where DP is the dimension of data parallelism, E is the number of experts, and C is the size of the expert capacity.

[0042] The overall implementation method is shown in Algorithm 1. The entire process utilizes the quick sort algorithm to reduce the time complexity of generating the scheduling index from O(N^3) to O(NlogN).

[0043]

[0044] Step 2: For composite index generation, the process is as follows Figure 2 shown.

[0045] Step 2.1: This step reuses the matrix M used in the scheduling index generation process, and the matrix dimensions are also (DP, N). Next, the invention replaces the target element based on the row value. The replacement rule is based on the row dimension. The cumulative sum of the matrix M is calculated and multiplied with the original matrix to ensure that the position of the expert remains unchanged, as shown in step (c). The formula is as follows:

[0046] H=Cumsum(M,axis=1)×M

[0047] After replacing the value of each element, it is necessary to add the corresponding expert sequence vector V to each column of the matrix H. The vector V is generated by the following formula:

[0048] V=Q×C

[0049] Q is a one-dimensional column vector, each element is the number of rows minus 1, and C is the expert capacity of MoE, which indicates the number of tokens that each expert can process.

[0050] Step 2.2: After obtaining H, the present invention sums each column and subtracts 1 to obtain the combined index, as shown in the following formula:

[0051] CombineIndex=ReduceSum(C,axis=0)-1

[0052] Algorithm 2 describes this process in detail. By leveraging the parallel matrix processing capabilities of the GPU accelerator, the time complexity is reduced from O(N^3) to O(1), effectively improving computational efficiency.

[0053]

[0054] Step 3: After the generation of the scheduling index and combination index is completed, it is necessary to enter the scheduling token and combination token stages. The scheduling token stage is responsible for assigning each token to the corresponding expert for processing, while the combination token stage integrates the output results of each expert. In these two stages, the core task is to efficiently find the physical address of the token based on the logical index (scheduling index and combination index) generated in steps 1 and 2 to complete the scheduling and combination of data. However, traditional methods use high-dimensional matrix multiplication to achieve physical address lookup, with a time complexity of up to O(N^3), which is inefficient for large-scale tasks.

[0055] To this end, the present invention utilizes the parallel computing capability of GPU and significantly optimizes the address calculation process by directly processing the logical index, reducing the complexity from O(N^3) to O(1). Figure 3 The specific method is as follows:

[0056] Step 3.1: In the CUDA architecture, each thread processes a Token block independently. Based on the GPU thread structure, the global index write_index of each thread can be calculated to determine the Token that the thread needs to process. The formula is:

[0057] write_index=blockIdx.x*blockDim.x+threadIdx.x

[0058] +w×gridDim.x×blockDim.x

[0059] Where blockIdx.x, blockDim.x, threadIdx.x, and gridDim.x are the basic parameters of the CUDA thread structure; w is the loop variable, representing the token block processed by the current thread. This formula allows each thread to clearly define its own task scope within the global computation.

[0060] Step 3.2: Based on write_index, the thread further parses the multidimensional space coordinates of the token in the logical index. Assuming the dimensions of the logical index are output_dim0, output_dim1, and output_dim2, the parsing formula is as follows:

[0061]

[0062] k=write_index mod output_dim2

[0063] Where i represents the expert dimension, j represents the token dimension, and k represents the feature dimension. For the scheduling index (DP, E, C) and combination index (DP, N), if some dimensions are missing, a default value of 1 can be set.

[0064] Step 3.3: Using the logical coordinates i, j, k and the position of the token in the expert index, the physical address is calculated as follows:

[0065] read_index=i*output_dim1*output_dim2+indices*output_dim2+k, where output_dim1 and output_dim2 represent the dimension size of the logical index.

[0066] Step 3.4: After each thread completes the calculation of the physical address, it directly extracts the content from the input data and writes it to the target address:

[0067] output[write_index]=input[read_index]

[0068] This method reduces the time complexity of physical address lookup and data migration from O(N^3) to O(1) by performing independent calculations on threads, significantly improving the efficiency of the scheduling and combination stages, and providing a key optimization method for the application of the MoE model in large-scale computing scenarios.

[0069] Step 4: After pre-training the MoE (Mixture of Experts) model, the fine-tuning phase presents a major challenge: the presence of zero vectors leads to irrational token allocation. To address this, this paper proposes an optimized token allocation method based on the division of active and inactive queues, aiming to maximize resource utilization and ensure model fine-tuning accuracy.

[0070] Step 4.1: The data in the fine-tuning phase often contains a large number of zero vectors (invalid inputs), which lack effective information. If these zero vectors are directly used in the calculation, it will not only waste computing resources but also affect the performance of the model.

[0071] Therefore, the input data needs to be divided into the following steps:

[0072] Importance detection: The i-th input token is represented by T i , calculate its importance value I(T i ). The importance value can be calculated using the following formula:

[0073] I(T i )=|T i |2

[0074] where |T i |2 represents the L2 norm of the token. When I(T i )=0, indicating T i is the zero vector.

[0075] The queue is divided into valid queue and invalid queue. The valid queue means all I(T i )>0 (i.e., tokens with non-zero importance values). The invalid queue represents all I(T i By the above method, it is ensured that all real and valid tokens are identified and divided into the valid queue, while invalid tokens (zero vectors) enter the spare invalid queue.

[0076] Step 4.2: After completing the division of the valid queue and the invalid queue, tokens need to be assigned to the experts of the MoE model for processing. Based on the Top-k strategy (each token is assigned to at most k experts), first assign valid tokens first. Assume that the tokens in the valid queue are {T1, T2, ..., T m}, there are m valid tokens in total. According to the importance I(T i ) for {T1,T2,…,T m} to sort in descending order:

[0077] I(T1)≥I(T2)≥…≥I(T m )

[0078] After sorting is completed, for each T i , and assign them to k experts with idle computing resources in turn to ensure that these valid tokens can be processed first.

[0079] Step 4.3: After all real tokens have been allocated, the remaining empty positions in each expert are filled with zero vectors. This zero vector filling ensures that all experts can process a consistent number of tokens, achieving efficient GPU parallelization. Based on the token allocation, the present invention adjusts the token load according to the capabilities of each expert, ensuring that each expert receives the content they need to process. This systematic zero vector allocation ensures a balance between computational efficiency and model accuracy, and ensures optimal loading for each expert through balanced distribution.

[0080] The experimental results are shown in Table 1 and Table 2. Compared with Switch Transformers, this embodiment verifies the performance of the invention on various billion-level, tens of billions of levels, and hundreds of billions of levels models and three datasets: Wikitext, C4, and Lambada. It is concluded that the overall performance of the present invention is better than Switch Transformers, and when applied to Mixtral and GPT-MoE models, the present invention can achieve acceleration ratios of 1.39 times and 1.20 times, respectively. In terms of scalability, the present invention verifies scalability through different model parameters. Compared with Switch Transformers, as the model parameter size and the number of GPU accelerators increase, the present invention maintains stable acceleration performance. This shows that the present invention is more conducive to application to larger models.

[0081] Table 1: Speedup effect of Mixtral model in different scenarios

[0082]

[0083] Table 2: Speedup effect of the GPT-MoE model in different scenarios

[0084]

[0085]

Claims

1. A hybrid expert model training optimization method based on matrix routing and token allocation, characterized in that: The steps include: Step 1: During pre-training, the MoE model uses multiple expert sub-models to process different input data. The MoE model generates expert indexes through a gating network, and then generates scheduling indexes and combination indexes based on the expert indexes. The scheduling index generation process is as follows: each token selects the corresponding expert based on the gating network, and the expert selection results of each token are summarized into a two-dimensional matrix M. The matrix M is implemented by the following formula: Among them, M ij is the element in the i-th row and j-th column of the matrix M, token i For the i-th token, expert j Denotes the jth expert; Construct a matrix X of the same size as matrix M. The elements in matrix X are much larger than each element in M. Perform the following calculations on matrix M and matrix X: M′=range(N)*M+(X+range(N))*(1-M) Where N is the length of X, and range is a loop function that generates values ​​from 0 to N-1. In the generated matrix M′, non-zero vector values ​​are extracted through quick sorting and slicing. Then, all indices greater than or equal to X are reset to -1, indicating that the token is not processed by the expert. This completes the generation of the scheduling index, reducing the time complexity to O(Nlog(N)). The scheduling index is a three-dimensional matrix with dimensions (DP, E, C), where DP is the data parallel dimension, E is the number of experts, and C is the expert capacity. The process of generating the combined index is as follows: Based on the matrix M, the cumulative sum Cumsum is calculated along the row dimension using the parallel matrix processing characteristics of the GPU, while keeping the position of the selected experts unchanged: H=Cumsum(M,axis=1)×M Add the sequence vector V of the corresponding expert to each column of the matrix H. The vector V is generated by the following formula: V=Q×C Q is a one-dimensional column vector, each element is the number of rows - 1; After obtaining the matrix H, sum each column and subtract 1 to get the combined index; this process reduces the time complexity of generating the combined index from O(N^3) to O(1), where the scheduling index is a two-dimensional matrix with dimensions (DP, N); Step 2: Based on the generated scheduling index and combination index, enter the scheduling token and combination token stage, dispatch the token to each expert for expert processing, and after expert processing, combine the tokens to generate a combination token; Step 3: After completing the pre-training of the MoE model, perform fine-tuning to complete the training optimization.

2. The hybrid expert model training optimization method based on matrix routing and token allocation according to claim 1 is characterized in that: Step 2 also includes: in the token scheduling process and the token combination process, the physical address of the token is searched based on the scheduling index and the combination index generated in step 1. This process uses the parallel computing capability of the GPU to optimize the processing of logical indexes; the logical index includes two types: scheduling index and combination index. Each thread on the GPU runs independently, directly calculates the corresponding physical address from the logical index, and completes data migration and writing; the logical index exists in two forms, and each thread independently calculates a unified target index based on the logical index.

3. The hybrid expert model training optimization method based on matrix routing and token allocation according to claim 2 is characterized in that: Each thread independently calculates a unified target index based on the logical index. The specific process is as follows: Using the CUDA thread structure to obtain the token information processed by each thread block, first obtain an index write_index, which is calculated as follows: write_index=blockIdx.x*blockDim.x+threadIdx.x +w×gridDim.x×blockDim.x Among them: blockIdx.x, blockDim.x, threadIdx.x and gridDim.x are the basic parameters of the CUDA thread structure; w is the loop variable, which represents the Token block processed by the current thread; based on the index write_index, the multidimensional space coordinates of the data are parsed: k=write_indexmodoutput_dim2 Among them, output_dim0, output_dim1, and output_dim2 represent the logical index information of the input. When the dimension does not exist, it is set to the default value 1. In addition, i corresponds to the expert dimension, j corresponds to the token dimension, and k represents the specific feature dimension. The calculation of the physical address will be based on the three dimensions i, j, and k and the position indices of this token in the expert index (DP, N). The formula is as follows: read_index=i*ouput_dim1*output_dim2+indices*output_dim2+k Through this formula, each thread independently completes the mapping from logical index to physical address, and obtains a unified target index read_index. After each thread calculates the physical address based on the target index, it will be used to write data at the target location: output[write_index]=input[read_index] Each thread independently completes an index calculation and data migration, reducing the time complexity of the operation from O(N^3) to O(1).

4. The hybrid expert model training optimization method based on matrix routing and token allocation according to claim 3 is characterized in that: The specific implementation process of the fine-tuning is as follows: The input tokens are analyzed and divided into two parts: a valid queue and an invalid queue. The valid queue contains all real tokens with non-zero importance values; the invalid queue contains zero vectors, which are used as filler and are not directly involved in expert processing. After the division of valid and invalid queues is completed, tokens are allocated to MoE experts. Using the Top-k strategy, each token is allocated to at most k experts. The tokens in the valid queue are sorted by importance and allocated to experts in order of importance. After the valid tokens are allocated, if there are experts with idle computing positions, the vacancies are filled with tokens from the invalid queue.

5. The hybrid expert model training optimization method based on matrix routing and token allocation according to claim 4 is characterized in that: The division of the valid queue and the invalid queue is specifically as follows: The i-th input token is represented by T i , calculate its importance value I(T i ): I(T i )=|T i |2 where |T i |2 represents the L2 norm of the token; when I(T i )=0, indicating T i is the zero vector; The effective queue represents all I(T i )>0 token; Invalid queue means all I(T i )=0 token.

Images