Hybrid expert network training method, question-answering method, related equipment and program products
By integrating the second expert layer FC2 and the Unpermute operation layer in the MOE model into one layer, the problem of increased video memory usage during MOE model training is solved, and the video memory usage is reduced and the computing efficiency is improved.
Patent Information
- Application Number
- CN202510737431.4
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2025-06-04
- Publication Date
- 2025-09-05
- Estimated Expiration
- 2045-06-04
AI Technical Summary
The problem of increased video memory usage during MOE model training, especially when processing large batches of data, leads to a significant increase in video memory demand, resulting in a waste of computing resources and reduced training efficiency.
The second expert layer FC2 and the unpermute operation layer in the traditional MOE model are fused into one layer, defined as the first fusion layer. The input and probability Probs of the first fusion layer are only saved in the forward propagation process, and are used to calculate the gradient in the backward propagation process to reduce the memory usage.
It significantly reduces the amount of video memory used during training and improves computing efficiency, especially when processing large-scale data, where the video memory requirement is reduced by nearly 5 times.
Smart Images

Figure CN120258046B_ABST
Abstract
Description
Technical Field
[0001] The present application relates to the technical field of large model training, and more specifically, to a hybrid expert network training method, a question-answering method, related equipment and program products. Background Art
[0002] Large language models (LLMs) have experienced significant development in the past few years, with many adopting a mixture-of-expert (MoE) network architecture. The MoE architecture dynamically activates a subset of experts through expert routing, replacing dense FFN layers. This mechanism focuses computing resources on the most relevant parts, reducing unnecessary computational overhead. Compared to traditional dense models, MoE models offer more efficient training speeds.
[0003] Compared to traditional Dense models, the MoE model training process is more complex, which can easily lead to several issues. A common issue is increased memory usage. Specifically, during MOE model training, the activation values calculated at each layer during the forward propagation process must be stored in memory for gradient calculation during the backward propagation process. While the MOE model reduces computational complexity through sparse activations, each input token must be processed by multiple experts (e.g., top-k experts), resulting in an exponential increase in the storage requirements for intermediate activation values. For example, if the activation dimension after processing by a single expert is dff, each token in the MoE layer must store k×dff activation values (assuming k experts are selected), while the dense model only needs to store dff activation values. Therefore, the memory usage of activations generated during MOE model training increases linearly with batch size and input sequence length, leading to increased memory usage. Summary of the Invention
[0004] In view of the above problems, this application is proposed to provide a hybrid expert network training method, question-answering method, related equipment and program products to reduce the memory usage of the MOE model training process. The specific solution is as follows:
[0005] In a first aspect, a method for training a hybrid expert network is provided, wherein the hybrid expert network includes a permutation permute operation layer, a first expert layer FC1, an activation network, and a first fusion layer, wherein the first fusion layer is fused with a second expert layer FC2 and an unpermutation unpermute operation layer, and the method includes:
[0006] Get the subsequence of the current device input and the subsequences of other devices input to form a complete input sequence;
[0007] The complete input sequence is forward propagated through the Permute operation layer, the FC1, and the activation network in sequence to obtain the output of the activation network;
[0008] The output of the activation network is used as the input of the first fusion layer. In the forward propagation process of the first fusion layer, the processing of the FC2 and the Unpermute operation layers is sequentially performed, and the input of the first fusion layer and the probability Probs of the expert corresponding to each token in the complete input sequence are saved, where the Probs of the expert corresponding to each token in the complete input sequence is calculated by the routing module in the hybrid expert network;
[0009] In the back propagation process of the first fusion layer, based on the saved input of the first fusion layer and the Probs, the gradients of the parameters of the first fusion layer and the gradients of the Probs are respectively calculated;
[0010] The parameters of the hybrid expert network are updated according to the gradient of the parameters of each layer in the back propagation process.
[0011] In one possible design, in another implementation of the first aspect of the embodiments of the present application, the Permute operation layer and the FC1 constitute a second fusion layer;
[0012] In the forward propagation process of the second fusion layer, a subsequence of the current device input is saved, subsequences of other device inputs are obtained, and the subsequences are combined with the subsequences of the current device input to form the complete input sequence. The Permute operation layer and FC1 are sequentially performed on the complete input sequence, and the processing results are used as the input of the activation network;
[0013] In the back propagation process of the second fusion layer, a subsequence of input from other devices is obtained, and the gradient of the parameters of the second fusion layer is calculated based on the saved subsequence of input from the current device and the subsequence of input from other devices.
[0014] In one possible design, in another implementation of the first aspect of the embodiments of the present application, the process of calculating the gradient of the second fusion layer parameters based on the saved subsequence of the current device input and the subsequence of the other device input includes:
[0015] Calculating the gradient of the subsequence input by the current device;
[0016] Calculating a gradient of a weight parameter based on the subsequence input by the other device and the subsequence input by the current device;
[0017] The process of calculating the gradient of the subsequence input by the current device is performed in parallel with the process of obtaining the subsequences input by other devices.
[0018] In one possible design, in another implementation of the first aspect of the embodiments of the present application, the process of respectively calculating the gradient of the first fusion layer parameters and the gradient of the Probs based on the saved input of the first fusion layer and the Probs includes:
[0019] Get the gradient of the output calculated by other devices, and merge it with the gradient of the output calculated by the current device to obtain the gradient of the complete output grad_output;
[0020] Performing a Permute operation on the Probs to obtain rearranged Probs, and performing a dot product between the rearranged Probs and the input of the first fusion layer to obtain a scaled input;
[0021] Based on the grad_output, calculate the gradient of the scaled input, the gradient of the weight parameter, and the gradient of the rearranged Probs;
[0022] Calculating the gradient of the input of the first fusion layer according to the gradient of the scaled input;
[0023] The gradient of the Probs before rearrangement is calculated according to the gradient of the Probs after rearrangement.
[0024] In one possible design, in another implementation of the first aspect of the embodiments of the present application, the process of obtaining the gradient of the output calculated by other devices is performed in parallel with the process of obtaining the scaled input.
[0025] In one possible design, in another implementation of the first aspect of the embodiments of the present application, in the back propagation process of the second fusion layer, the process of obtaining a subsequence input by another device includes:
[0026] In the back propagation process of the second fusion layer, the all-gather communication method is used to obtain the subsequences input by other devices.
[0027] In one possible design, in another implementation of the first aspect of the embodiments of the present application, the process of obtaining the gradient of the output calculated by other devices includes:
[0028] Use the all gather communication method to obtain the gradient of the output calculated by other devices.
[0029] In a second aspect, a question-answering method is provided, comprising:
[0030] Get question information;
[0031] Sending the question information to the configured hybrid expert model to obtain answer information output by the hybrid expert model;
[0032] The large hybrid expert model is a model trained using the hybrid expert network training method described in any one of the first aspects above.
[0033] In a third aspect, a hybrid expert network training device is provided, wherein the hybrid expert network includes a permutation permute operation layer, a first expert layer FC1, an activation network, and a first fusion layer, wherein the first fusion layer is fused with a second expert layer FC2 and an unpermutation unpermute operation layer, and the device includes:
[0034] A computing unit is configured to obtain a subsequence of the current device input and a subsequence of the input of other devices to form a complete input sequence; forward propagates the complete input sequence sequentially through the Permute operation layer, the FC1, and the activation network to obtain the output of the activation network; uses the output of the activation network as the input of the first fusion layer, and sequentially performs processing of the FC2 and the Unpermute operation layer in the forward propagation process of the first fusion layer, and saves the input of the first fusion layer and the probability Probs of the expert corresponding to each token token in the complete input sequence, wherein the Probs of the expert corresponding to each token in the complete input sequence is calculated by the routing module in the hybrid expert network; in the back propagation process of the first fusion layer, based on the saved input of the first fusion layer and the Probs, respectively calculates the gradient of the first fusion layer parameters and the gradient of the Probs;
[0035] The parameter updating unit is used to update the parameters of the hybrid expert network according to the gradient of the parameters of each layer in the back propagation process.
[0036] In a fourth aspect, an electronic device is provided, comprising: a memory and a processor;
[0037] The memory is used to store programs;
[0038] The processor is used to execute the program to implement the various steps of the hybrid expert network training method described in any one of the first aspects of this application, or to implement the various steps of the question-answering method described in the second aspect.
[0039] In the fifth aspect, a readable storage medium is provided, on which a computer program is stored. When the computer program is executed by a processor, the various steps of the hybrid expert network training method described in any one of the first aspects of this application are implemented, or the various steps of the question-answering method described in the second aspect are implemented.
[0040] In the sixth aspect, a computer program product is provided, comprising a computer program, which, when executed by a processor, implements the various steps of the hybrid expert network training method described in any one of the aforementioned first aspects of this application, or implements the various steps of the question-answering method described in the aforementioned second aspect.
[0041] A traditional hybrid expert network consists of a permute layer, a first expert layer FC1, an activation network, a second expert layer FC2, and an unpermute layer. To speed up computation, the permute layer in the forward propagation process rearranges the inputs to align the hidden states of tokens to be processed by the same expert. This is then processed by FC1 and the activation network. The unpermute layer weights the results of processing by different experts for the same token, summing them according to the corresponding probabilities Probs. The unpermute operation is then performed to obtain the output of the current device. To calculate the gradients of the FC2 parameters and Probs during the backward propagation process, the FC2 input (i.e., the activation network output) must be saved during the forward propagation process for FC2 gradient calculation. Furthermore, the Probs and FC2 outputs must be saved during the forward propagation process for Probs gradient calculation during the backward propagation process. Based on the above introduction to the technical solution of the present application, it can be seen that the present application fuses the second expert layer FC2 and the unpermute operation layer in the traditional hybrid expert network into one layer (the first fusion layer). In this way, only the input and Probs of the first fusion layer need to be saved during the forward propagation process. During the backward propagation process, the input of the first fusion layer can be used to calculate the gradient of the parameters of the first fusion layer, as well as the gradient of Probs, without the need to additionally save the output of FC2 during the forward propagation process, which greatly reduces the amount of video memory occupied by the activations that need to be saved. BRIEF DESCRIPTION OF THE DRAWINGS
[0042] Various other advantages and benefits will become apparent to those skilled in the art upon reading the detailed description of the preferred embodiment below. The accompanying drawings are for illustration purposes only and are not to be considered as limiting the present application. The same reference symbols are used throughout the drawings to represent the same components. In the drawings:
[0043] Figure 1The schematic diagram of the overall processing framework of the MLP module in the existing MoE model is shown as an example;
[0044] Figure 2 Example Figure 1 Schematic diagram of the calculation process corresponding to the framework;
[0045] Figure 3 A schematic diagram of an implementation system architecture of the hybrid expert network training method provided in an embodiment of the present application;
[0046] Figure 4 A flow chart of a hybrid expert network training method provided in an embodiment of the present application;
[0047] Figure 5 A schematic diagram illustrating a calculation process of the hybrid expert network training method of the present application is provided;
[0048] Figure 6 A schematic diagram of the structure of a hybrid expert network training device provided in an embodiment of the present application;
[0049] Figure 7 A schematic diagram of the structure of an electronic device provided in an embodiment of the present application. DETAILED DESCRIPTION
[0050] The following will be combined with the drawings in the embodiments of this application to clearly and completely describe the technical solutions in the embodiments of this application. Obviously, the embodiments described are only part of the embodiments of this application, not all of the embodiments. Based on the embodiments in this application, all other embodiments obtained by ordinary technicians in this field without making creative efforts are within the scope of protection of this application.
[0051] Compared to traditional dense model training, the MOE model is more complex and prone to problems such as load imbalance, increased communication complexity, and increased memory usage. This application focuses on solving the problem of increased memory usage during MOE model training.
[0052] First, the memory usage problem of the existing MOE model training process is analyzed.
[0053] In some implementations, MoE model training generally adopts an expert parallel solution, that is, different experts are distributed on different devices (computing nodes), and each device only stores and calculates its own experts. For example, a training strategy with strict load balancing is used, that is, each expert group calculates all tokens. In the existing solution, the overall process of the MLP module in the MoE model is as follows: Figure 1 shown.
[0054] In tensor and sequence parallel mode, each segment input subsequence ( Figure 1The input subsequence 1 (Input1) and input subsequence 2 (Input2) in the example are passed through the gated Gate network to obtain the corresponding probability distribution logits output, and the logits of each subsequence are fused ( Figure 1 The complete probability distribution Total Logits is obtained by combining the data in the example of All Gather communication operation. Then, the expert ID corresponding to each token is obtained through the routing network Routing. Figure 1 Indices in), the probability that each token corresponds to an expert (corresponding Figure 1 Probs in), and the number of tokens that each expert needs to process (corresponding to Figure 1 At the same time, we need to AllGather the hidden states of all tokens in all input subsequences Input to obtain the complete input sequence Total Input.
[0055] In order to speed up the calculation process, the hidden states can be rearranged according to the indices by rearranging the permuted operation layer to arrange the hidden states of the tokens that need to be processed by the same expert together to obtain the permuted result. Then the calculation process of the first expert layer, activation network, and second expert layer is carried out. The first and second expert layers can use a fully connected layer network, for example Figure 2 FC1 Gmm and FC2 Gmm in the figure. FC represents a fully connected layer, and Gmm stands for Grouped Matrix Multiplication (GMM). By using grouped matrix multiplication in the fully connected layer, the first and second expert layers can efficiently handle multi-expert parallel computations, improving computational efficiency. The Geglu activation function can be used for network activation.
[0056] The output of the second expert layer serves as the input to the unpermuted layer. The unpermuted layer weights the results of different expert processing for the same token (weights are probabilities) and sums them. The weighted sum is then subjected to an unpermuted operation, referencing the position mapping relationship row_id_map. The processed results are then distributed to the corresponding processing devices using ReduceScatter. The row_id_map records the position mapping relationship between the hidden state of each token before and after permutation. This mapping relationship allows the unpermuted operation to restore the data to its pre-permuted dimensions.
[0057] In order to more intuitively show the changes in the corresponding data in the above calculation process, this application Figure 2 Will Figure 1 The computational flow shown in the figure is presented, with the data shape information annotated. Here, S represents the length of the complete input sequence, B represents the batch size (Batch_Size), H represents the dimension of the token's hidden state vector, H_e represents the output dimension of the experts in the first expert layer, T represents the tensor parallelization scale (TP Size), and TopK represents the number of selected experts.
[0058] Combine Figure 2 As shown in the figure, the subsequence of the current device input, namely the local input, is obtained, and its data shape is [S / T, B, H]. The complete input sequence Total Input is further obtained through the All Gather communication operation, and its data shape is [S, B, H]. The Permuted result is obtained after processing in the Permute operation layer, and its data shape is After the first expert layer FC1 Gmm is processed, the output FC1 is obtained, and its data shape is After further activation network processing, taking the activation network using the Geglu activation function as an example, this activation function can reduce the dimension of the hidden layer state vector by half, and the data shape of the Geglu output result is The output of Geglu is further sent to the second expert layer FC2 Gmm for processing. The second expert layer selected here can restore the dimension of the hidden layer state vector to size H. The data shape of the FC2 output result is The FC2 output is further processed in the Unpermute layer, yielding the Unpermuted result, whose data shape is [S, B, H]. The Reduce Scatter is then distributed to other devices, yielding the current device's output, known as the Local Output, whose data shape is [S / T, B, H].
[0059] Figure 2 The following table shows the activation data shape generated by each layer during the forward propagation of the MOE model. In order to facilitate gradient calculation during the backward propagation process, the main activation memory information that needs to be saved is shown in Table 1:
[0060] Table 1
[0061]
[0062] It should be noted that the above Table 1 only illustrates some of the main activations that need to be saved, not all activations. In addition to the main activations shown in Table 1, Figure 2 Indices, row_id_map, Probs, etc. also need to be saved. However, the size of these parameters is much smaller than the main activation shown in Table 1, which can be clearly seen from the shape of the parameters.
[0063] Taking sequence length S=8192, TopK=8, H=6144, and H_e=2048 as an example, the amount of active memory that needs to be saved as shown in Table 1 is 1920MB (assuming BF16 type storage). It can be seen that this part of the active memory is still very large.
[0064] In order to reduce the amount of video memory occupied during the MOE model training phase, this application provides an improved hybrid expert network training solution.
[0065] This application provides a hybrid expert network training method that can be applied to Figure 3 The system architecture shown in FIG. 1 may include a terminal 100 and a server 200. The server 200 may include one or more servers ( Figure 3 (This section includes a server as an example).
[0066] The terminal 100 or the server 200 can be used alone to execute the hybrid expert network training method provided in the embodiment of the present application. In addition, the terminal 100 and the server 200 can also be used in conjunction to execute the hybrid expert network training method provided in the embodiment of the present application.
[0067] This embodiment provides a hybrid expert network training method that improves the hybrid expert network structure by fusing the second expert layer FC2 and the unpermute operation layer in the traditional MOE model into a single layer, defined as the first fusion layer. This allows the calculation of the Prob value gradient from the unpermute operation to the input before the first fusion layer. This means that the Prob value gradient can be calculated using the input of the first fusion layer, eliminating the need to save the FC2 output and significantly reducing the active video memory usage.
[0068] In a possible implementation, the improved hybrid expert network includes a permutation permute operation layer, a first expert layer FC1, an activation network and a first fusion layer, wherein the first fusion layer is integrated with a second expert layer FC2 and an anti-permutation unpermute operation layer. Figure 4 , a hybrid expert network training method provided in an embodiment of the present application is introduced, and the method may specifically include the following steps:
[0069] Step S100: Obtain a subsequence input by the current device and a subsequence input by other devices to form a complete input sequence.
[0070] Step S110: The complete input sequence is forward propagated through the Permute operation layer, FC1, and the activation network in sequence to obtain the output of the activation network.
[0071] Specifically, the Permute layer rearranges the complete input sequence, grouping the hidden states of tokens processed by the same expert together to speed up subsequent computations. Afterwards, the first expert layer, FC1, and the activation network process produce the output of the activation network.
[0072] Among them, the activation network can adopt various types of activation functions, such as Geglu activation function, relu activation function, etc. The Geglu activation function is used as an example for explanation in the subsequent embodiments of this application.
[0073] Step S120: Use the output of the activated network as the input of the first fusion layer. In the forward propagation process of the first fusion layer, perform the processing of the FC2 and Unpermute operation layers in sequence, and save the input of the first fusion layer and the probability Probs of the expert corresponding to each token token in the complete input sequence.
[0074] The Probs of each token corresponding to the expert in the complete input sequence is calculated by the routing module in the hybrid expert network. This process can be referred to Figure 1 The relevant introduction of the corresponding embodiments will not be repeated here.
[0075] In this embodiment, the second expert layer FC2 and the Unpermute operation layer are fused to form the first fusion layer, with the output of the activation network serving as the input to the first fusion layer. During the forward propagation process, the first fusion layer sequentially performs the FC2 and Unpermute operations on the input to obtain the output of the current device. Furthermore, to facilitate gradient calculation during the backward propagation process, the first fusion layer also saves the input of the first fusion layer (i.e., the output of the activation network) and the aforementioned Probs during the forward propagation process.
[0076] Step S130: In the back propagation process of the first fusion layer, based on the saved input and Probs of the first fusion layer, the gradients of the parameters of the first fusion layer and the gradient of Probs are calculated respectively.
[0077] Specifically, since the present application fuses the FC2 and Unpermute operations into one layer, the gradient of Probs can be calculated based on the input of the first fusion layer, without saving the output of FC2 for calculating the gradient of Probs as in the prior art, thus saving video memory usage.
[0078] Step S140: Update the parameters of the hybrid expert network according to the gradient of the parameters of each layer in the back propagation process.
[0079] During the backpropagation process, the gradients of the parameters of each layer are calculated from the output layer forward, and are used to update the parameters of the hybrid expert network. Since this embodiment combines the traditional FC2 and Unpermute operation layers into one layer, the calculation strategy for the backpropagation process of the first fused layer is described in step S130 above. The parameter gradients of the other layers can be calculated according to the existing backpropagation calculation strategy, or other improved algorithms can be used for calculation, which is not limited in this embodiment.
[0080] Compared with the traditional MOE model training process, the traditional MOE model training process requires saving the input of FC2 (i.e., the output of the activation network) during the forward propagation process for the gradient calculation of FC2 during the backward propagation process. In addition, the output of Probs and FC2 must be saved during the forward propagation process for the gradient calculation of Probs during the backward propagation process. The method provided in this embodiment fuses the second expert layer FC2 and the Unpermute operation layer in the hybrid expert network into one layer (the first fusion layer). In this way, only the input of the first fusion layer and Probs need to be saved during the forward propagation process. During the backward propagation process, the input of the first fusion layer can be used to calculate the gradient of the parameters of the first fusion layer and the gradient of Probs, without the need to save the output of FC2 during the forward propagation process. This greatly reduces the amount of video memory occupied by the activations that need to be saved.
[0081] In some possible implementations, in order to further reduce the amount of video memory occupied during the training process, the structure of the hybrid expert network can be further improved in this embodiment. Specifically:
[0082] The Permute operation layer and FC1 can be fused into one layer, defined as the second fusion layer.
[0083] In the forward propagation process of the second fusion layer, the subsequence of the current device input is saved, and the subsequences of the other device inputs are obtained. Together with the subsequence of the current device input, the complete input sequence is formed. The Permute operation layer and FC1 are performed on the complete input sequence in sequence, and the processing results are used as the input of the activation network.
[0084] In the back propagation process of the second fusion layer, a subsequence of input from other devices is obtained, and the gradient of the parameters of the second fusion layer is calculated based on the saved subsequence of input from the current device and the obtained subsequence of input from other devices.
[0085] Using the method of this embodiment, in the second fusion layer, the forward propagation process only needs to save the subsequence of the current device input, and its data shape is [S / T, B, H]. When calculating the gradient through backpropagation, it first communicates with other devices to obtain the subsequence of the other device input, and combines it with the local subsequence to form a complete input sequence, and then the gradient of the second fusion layer parameters can be calculated. Compared with the traditional MOE model that uses a structure with a Permute operation layer and FC1 separated, the traditional MOE model training process needs to save the activation data shape for the FC1 layer. By comparison, it can be seen that after the Permute operation layer and FC1 are fused into one layer (the second fusion layer), the memory usage of the activations required to be saved in the forward propagation process is further reduced.
[0086] In one possible implementation, in the back propagation process of the second fusion layer, the communication process of obtaining the subsequence of input from other devices can be executed in parallel with the gradient calculation process. That is, the communication process and the gradient calculation process can be overlapped, thereby improving the operating efficiency.
[0087] Specifically, the process of calculating the gradient of the second fusion layer parameters based on the saved subsequence of the current device input and the obtained subsequence of the other device input may include:
[0088] First, the gradient of the subsequence of the current device input is calculated.
[0089] After obtaining the subsequences of the inputs from other devices, a complete input sequence is obtained based on the subsequences of the inputs from other devices and the subsequence of the input from the current device, and the gradient of the weight parameter is calculated based on the complete input sequence.
[0090] The process of calculating the gradient of the subsequence input by the current device and the communication process of obtaining the subsequences input by other devices can be executed in parallel, thereby improving operation efficiency.
[0091] In one possible implementation, the backpropagation process of the second fusion layer, which communicates to obtain subsequences of inputs from other devices, can use an all-gather communication method. Other communication methods can also be used to obtain subsequences of inputs from other devices. All-gather is a collective communication operation designed to gather data from each process (or node) across all processes. This means that at the end of the all-gather operation, each process has the data of all other processes.
[0092] In some embodiments of the present application, the aforementioned step S130, in the back propagation process of the first fusion layer, is further described, in which the gradients of the first fusion layer parameters and the gradients of Probs are respectively calculated based on the saved input and Probs of the first fusion layer. The specific implementation process may include the following sub-steps:
[0093] S1. In the back propagation process of the first fusion layer, the gradient of the output calculated by other devices is obtained and fused with the gradient of the output calculated by the current device to obtain the gradient of the complete output grad_output.
[0094] This process can use the all gather communication method to obtain the gradient of the output calculated by other devices. Of course, other communication methods can also be used to obtain the gradient of the output calculated by other devices.
[0095] S2. Perform a Permute operation on Probs to obtain rearranged Probs, and perform a dot product between the rearranged Probs and the input of the first fusion layer to obtain a scaled input.
[0096] Specifically, since the input of the first fusion layer is processed by the previous permute operation layer and is in a permuted state, this step also requires the permutation operation of Probs to obtain the permuted Probs.
[0097] Define the input of the first fusion layer as x, and the rearranged Probs as Permuted_Probs:
[0098] Permuted_Probs=Permute(Probs,row_id_map).
[0099] The definition of row_id_map is as described above and will not be repeated here.
[0100] Furthermore, the rearranged Probs is multiplied by the input x of the first fusion layer to obtain the scaled input scaled_x:
[0101] .
[0102] S3. Based on the gradient grad_output of the complete output obtained in step S1, calculate the gradient of the scaled input scaled_x, the gradient of the weight parameter w, and the gradient of the rearranged Probs.
[0103] Since scaled_x is in the Permuted state, grad_output must also be in the same Permuted state. Therefore, the Permute operation is performed on grad_output to obtain the rearranged grad_output:
[0104] Permuted_grad=Permute(grad_output,indices).
[0105] The definition of indices is as described above and will not be repeated here.
[0106] Furthermore, the gradient of scaled_x is calculated. Since there are multiple experts, group matrix multiplication Gmm can be used. The gradient of scaled_x is expressed as scaled_dx, and the calculation formula is:
[0107] scaled_dx=Gmm(Permuted_grad,weight.T).
[0108] Among them, weight represents the expert's parameter, and weight.T means transposing the weight.
[0109] Calculate the gradient dw of weight w:
[0110] dw=Gmm(scaled_x.T,Permuted_grad).
[0111] Among them, scaled_x.T indicates the transpose operation of scaled_x.
[0112] Calculate the gradient of Permuted_Probs:
[0113] .
[0114] S4. Calculate the gradient dx of the input x of the first fusion layer according to the gradient scaled_dx of the scaled input, and calculate the gradient dProbs of Probs before permutation according to the gradient Permuted_dProbs of Probs after permutation.
[0115] Specifically:
[0116] ;
[0117] dProbs=Permute_bwd(Permuted_dProbs,row_id_map).
[0118] The Permute_bwd operation is the reverse of the Permute operation. Since Permuted_Probs is obtained by performing the Permute operation on Probs, the gradient of Probs in this step can be converted to performing the Permute_bwd operation on Permuted_dProbs.
[0119] In one possible implementation, the process of obtaining the gradient of the output calculated by other devices in step S1 and the process of obtaining the scaled input in step S2 can be performed in parallel. That is, the process of the current device communicating with other devices to obtain the output gradient and the process of the current device calculating the scaled input scaled_x can be overlapped, thereby improving operating efficiency.
[0120] Reference Figure 5 , which illustrates the calculation process after adopting the improved hybrid expert network training method of this application, and annotates the data shape information.
[0121] Get the subsequence of the current device input, that is, the local input, as the input of the second fusion layer.
[0122] In the second fusion layer (Permute & FC1):
[0123] Forward Propagation: Step 1. Save the local input for gradient calculation in the backward propagation process (i.e., save local input for backward). The saved local input subsequence has a data shape of [S / T, B, H]. Step 2. Use the All Gather communication method to obtain subsequences of input from other devices and combine them with the subsequence of the current device's input to form the complete input sequence. Steps 3-4. Perform the Permute operation and the FC1 operation on the complete input sequence.
[0124] Backward Propagation Process: Step 1. Use the All Gather communication method to obtain subsequences of input from other devices and combine them with the saved subsequence of the current device's input to form the complete input sequence, also known as All Gather Input. Step 2. Calculate the gradient of the local input subsequence x, also known as Compute grad of x. Step 3. Distribute the gradient of the local input subsequence x to other devices, also known as Reduce Scatter grad of x. Step 4. Calculate the gradient of the weight parameter w, also known as Compute grad of w.
[0125] The forward propagation process obtains the output of the second fusion layer, whose data shape is The output of the second fusion layer is used as the input of the activation network. The activation network can use a variety of activation functions. Figure 5 The Geglu activation function is used as an example to illustrate.
[0126] After the Geglu activation network is processed, the output is obtained, and its data shape is The output of the Geglu activation network is used as the input of the first fusion layer.
[0127] In the first fusion layer (FC2 & Unpermute):
[0128] Forward propagation process: Steps 1-2. After FC2 and Unpermute operations, the local output is obtained. Step 3. The input and Probs of the first fusion layer are saved to facilitate the gradient calculation in the back-propagation process. Figure 5 The saved input data shape of the first fusion layer is further illustrated in Step 4: Distribute the output of the current device to other devices, which is called Reduce Scatter Output.
[0129] Backword process: Step 1. Use the All gather communication method to obtain the gradient of the output calculated by other devices, and fuse it with the gradient of the output calculated by the current device to obtain the gradient of the complete output grad_output, also known as All gather grad_output. Step 2. Perform the Permute operation on Probs. Step 3. Perform the dot multiplication of the rearranged Probs and the input x of the first fusion layer to obtain the scaled input (corresponding to Figure 5 middle ). Step 4. Based on grad_output, compute the gradient of the scaled input scaled_x, the gradient of the weight parameter w, and the gradient of the permuted Probs, i.e., Compute grad of scaled_x,w,Permuted_Probs. Step 5. Based on the gradient of the scaled input, compute the gradient of the input x of the first fusion layer. Based on the gradient of the permuted Probs, compute the gradient of the pre-permuted Probs, i.e., Compute grad of x,Probs.
[0130] based on Figure 5 The improved hybrid expert model training method shown in Table 2 requires saving the main activation memory information:
[0131] Table 2
[0132]
[0133] It should be noted that the above Table 2 only illustrates some of the main activations that need to be saved, not all activations. In addition to the main activations shown in Table 2, Figure 5 Indices, row_id_map, Probs, etc. also need to be saved. However, the size of these parameters is much smaller than the main activation shown in Table 2, which can be clearly seen from the shape of the parameters.
[0134] To compare the amount of main activation data required to be saved with the traditional hybrid expert network training method (i.e., comparing Tables 1 and 2), this embodiment also takes sequence length S=8192, TopK=8, H=6144, and H_e=2048 as an example. The amount of activation memory required to be saved as shown in Table 2 is 396MB (assuming BF16 type storage). Compared to the 1920MB of activation memory required to be saved as shown in Table 1, the amount of activation memory used by the improved method of this application is reduced by nearly 5 times. Moreover, the larger the H, T, and TopK, the higher the benefit of using the improved method of this application.
[0135] Some embodiments of the present application further provide a question-answering method, which can be implemented based on a large hybrid expert model. The large hybrid expert model can be trained using the hybrid expert network training method of the aforementioned embodiment. The question-answering method includes the following steps:
[0136] First, obtain the question information.
[0137] The question information is further sent to the hybrid expert model to obtain the answer information output by the hybrid expert model.
[0138] The hybrid expert network training device provided in an embodiment of the present application is described below. The hybrid expert network training device described below and the hybrid expert network training method described above can be referenced to each other.
[0139] See also Figure 6 , Figure 6 This is a structural diagram of a hybrid expert network training device disclosed in an embodiment of the present application. In this embodiment, the hybrid expert network includes a permutation permute operation layer, a first expert layer FC1, an activation network, and a first fusion layer. The first fusion layer is integrated with a second expert layer FC2 and an unpermutation unpermute operation layer. Figure 6 As shown, the hybrid expert network training device may include:
[0140] The computing unit 11 is used to obtain a subsequence of the current device input and a subsequence of the input of other devices to form a complete input sequence; forward propagates the complete input sequence sequentially through the Permute operation layer, the FC1, and the activation network to obtain the output of the activation network; uses the output of the activation network as the input of the first fusion layer, and in the forward propagation process of the first fusion layer, sequentially performs the processing of the FC2 and the Unpermute operation layer, and saves the input of the first fusion layer and the probability Probs of the expert corresponding to each token token in the complete input sequence, wherein the Probs of the expert corresponding to each token in the complete input sequence is calculated by the routing module in the hybrid expert network; in the back propagation process of the first fusion layer, based on the saved input of the first fusion layer and the Probs, respectively calculates the gradient of the first fusion layer parameters and the gradient of the Probs;
[0141] The parameter updating unit 12 is used to update the parameters of the hybrid expert network according to the gradient of the parameters of each layer in the back propagation process.
[0142] In one possible implementation, the Permute operation layer and the FC1 constitute a second fusion layer; in the forward propagation process of the second fusion layer, the computing unit saves a subsequence of the current device input, obtains a subsequence of the other device input, and composes the subsequence and the subsequence of the current device input into the complete input sequence, performs the processing of the Permute operation layer and the FC1 on the complete input sequence in sequence, and uses the processing result as the input of the activation network; in the backward propagation process of the second fusion layer, obtains a subsequence of the other device input, and calculates the gradient of the second fusion layer parameters based on the saved subsequence of the current device input and the subsequence of the other device input.
[0143] In one possible implementation, the process of calculating the gradient of the second fusion layer parameter by the calculation unit based on the saved subsequence of the current device input and the subsequence of the other device input includes:
[0144] Calculating the gradient of the subsequence input by the current device;
[0145] Calculating a gradient of a weight parameter based on the subsequence input by the other device and the subsequence input by the current device;
[0146] The process of calculating the gradient of the subsequence input by the current device is performed in parallel with the process of obtaining the subsequences input by other devices.
[0147] In one possible implementation, the calculation unit calculates the gradient of the first fusion layer parameters and the gradient of the Probs based on the saved input of the first fusion layer and the Probs, respectively, including:
[0148] Get the gradient of the output calculated by other devices, and merge it with the gradient of the output calculated by the current device to obtain the gradient of the complete output grad_output;
[0149] Performing a Permute operation on the Probs to obtain rearranged Probs, and performing a dot product between the rearranged Probs and the input of the first fusion layer to obtain a scaled input;
[0150] Based on the grad_output, calculate the gradient of the scaled input, the gradient of the weight parameter, and the gradient of the rearranged Probs;
[0151] Calculating the gradient of the input of the first fusion layer according to the gradient of the scaled input;
[0152] The gradient of the Probs before rearrangement is calculated according to the gradient of the Probs after rearrangement.
[0153] In one possible implementation, the process of the computing unit obtaining the gradient of the output calculated by other devices is performed in parallel with the process of obtaining the scaled input.
[0154] In one possible implementation, the process of the computing unit acquiring the subsequences input by other devices during the back propagation process of the second fusion layer includes:
[0155] In the back propagation process of the second fusion layer, the all-gather communication method is used to obtain the subsequences input by other devices.
[0156] In one possible implementation, the process of the computing unit obtaining the gradient of the output calculated by other devices includes:
[0157] Use the all gather communication method to obtain the gradient of the output calculated by other devices.
[0158] An electronic device is also provided in an embodiment of the present application. Figure 7 , which shows a schematic diagram of the structure of an electronic device suitable for implementing the embodiment of the present application. The electronic device in the embodiment of the present application may include but is not limited to terminals such as mobile phones, tablet computers, computers, etc. Figure 7 The electronic device shown is merely an example and should not limit the functions and scope of use of the embodiments of the present application.
[0159] like Figure 7 As shown, the electronic device may include a processing device (such as a central processing unit, a graphics processing unit, etc.) 601, which can perform various appropriate actions and processes according to the program stored in the read-only memory (ROM) 602 or the program loaded from the storage device 608 to the random access memory (RAM) 603 to implement the hybrid expert network training method of the aforementioned embodiment of the present application, or implement the question-answering method of the aforementioned embodiment of the present application. When the electronic device is powered on, the RAM 603 also stores various programs and data required for the operation of the electronic device. The processing device 601, ROM 602 and RAM 603 are connected to each other via a bus 604. The input / output (I / O) interface 605 is also connected to the bus 604.
[0160] Typically, the following devices may be connected to the I / O interface 605: an input device 606 including, for example, a touch screen, a touchpad, a keyboard, a mouse, a camera, a microphone, an accelerometer, a gyroscope, etc.; an output device 607 including, for example, a liquid crystal display (LCD), a speaker, a vibrator, etc.; a storage device 608 including, for example, a memory card, a hard disk, etc.; and a communication device 609. The communication device 609 may allow the electronic device to communicate with other devices wirelessly or by wire to exchange data. Figure 7 The electronic device is shown with various devices, but it should be understood that it is not required to implement or possess all of the devices shown. More or fewer devices may be implemented or possessed instead.
[0161] An embodiment of the present application also provides a computer program product including computer-readable instructions. When the computer-readable instructions are executed on an electronic device, the electronic device implements any one of the hybrid expert network training methods or question-answering methods provided in the embodiments of the present application.
[0162] A computer-readable storage medium is also provided in an embodiment of the present application. The storage medium carries one or more computer programs. When the one or more computer programs are executed by an electronic device, the electronic device can implement any hybrid expert network training method or question-answering method provided in the embodiment of the present application.
[0163] It should also be noted that the device embodiments described above are merely illustrative, wherein the units described as separate components may or may not be physically separate, and the components displayed as units may or may not be physical units, that is, they may be located in one place, or they may be distributed across multiple network units. Some or all of the modules may be selected according to actual needs to achieve the purpose of the present embodiment. In addition, in the drawings of the device embodiments provided in this application, the connection relationship between the modules indicates that there is a communication connection between them, which can be specifically implemented as one or more communication buses or signal lines.
[0164] Through the description of the above embodiments, those skilled in the art can clearly understand that the present application can be implemented by means of software plus necessary general hardware, and of course can also be implemented by special hardware including application-specific integrated circuits, special CPUs, special memories, special components, etc. In general, all functions performed by computer programs can be easily implemented with corresponding hardware, and the specific hardware structures used to implement the same function can also be diverse, such as analog circuits, digital circuits or special circuits, etc. However, for the present application, software program implementation is a better implementation method in most cases. Based on this understanding, the technical solution of the present application is essentially or the part that contributes to the prior art can be embodied in the form of a software product, which is stored in a readable storage medium, such as a computer's floppy disk, USB flash drive, mobile hard disk, ROM, RAM, magnetic disk or optical disk, etc., and includes a number of instructions to enable a computer device (which can be a personal computer, training equipment, or network equipment, etc.) to execute the methods described in each embodiment of the present application.
[0165] In the above embodiments, all or part of the embodiments may be implemented by software, hardware, firmware, or any combination thereof. When implemented by software, all or part of the embodiments may be implemented in the form of a computer program product.
[0166] The computer program product includes one or more computer instructions. When the computer program instructions are loaded and executed on a computer, the process or function described in the embodiment of the present application is generated in whole or in part. The computer can be a general-purpose computer, a special-purpose computer, a computer network, or other programmable devices. The computer instructions can be stored in a computer-readable storage medium, or transmitted from one computer-readable storage medium to another computer-readable storage medium. For example, the computer instructions can be transmitted from a website, a computer, a training device or a data center by wired (such as coaxial cable, optical fiber, digital subscriber line (DSL)) or wireless (such as infrared, wireless, microwave, etc.) mode to another website, computer, training device or data center. The computer-readable storage medium can be any available medium that a computer can store or a data storage device such as a training device, a data center that includes one or more available media integrations. The available medium can be a magnetic medium, (for example, a floppy disk, a hard disk, a tape), an optical medium (for example, a DVD), or a semiconductor medium (for example, a solid-state drive (SSD)).
[0167] The various embodiments in this specification are described in a progressive manner, and each embodiment focuses on the differences from other embodiments. The various embodiments can be combined as needed, and the same or similar parts can be referenced to each other.
Claims
1. A hybrid expert network training method, characterized in that: The hybrid expert network includes a permutation permute operation layer, a first expert layer FC1, an activation network, and a first fusion layer, wherein the first fusion layer is fused with a second expert layer FC2 and an unpermutation unpermute operation layer, and the method includes: Get the subsequence of the current device input and the subsequences of other devices input to form a complete input sequence; The complete input sequence is forward propagated through the Permute operation layer, the FC1, and the activation network in sequence to obtain the output of the activation network; The output of the activation network is used as the input of the first fusion layer. In the forward propagation process of the first fusion layer, the processing of the FC2 and the Unpermute operation layers is sequentially performed, and the input of the first fusion layer and the probability Probs of the expert corresponding to each token in the complete input sequence are saved to the video memory, where the Probs of the expert corresponding to each token in the complete input sequence is calculated by the routing module in the hybrid expert network. In the back propagation process of the first fusion layer, based on the input of the first fusion layer and the Probs stored in the video memory, respectively calculate the gradient of the first fusion layer parameters and the gradient of the Probs; Update the parameters of the hybrid expert network according to the gradient of the parameters of each layer during the back propagation process; Wherein, the Permute operation layer and the FC1 constitute the second fusion layer; In the forward propagation process of the second fusion layer, a subsequence of the current device input is saved, subsequences of other device inputs are obtained, and the subsequences are combined with the subsequences of the current device input to form the complete input sequence. The Permute operation layer and FC1 are sequentially performed on the complete input sequence, and the processing results are used as the input of the activation network; In the back propagation process of the second fusion layer, communication with other devices is performed to obtain subsequences of inputs from other devices, and the gradients of the parameters of the second fusion layer are calculated based on the saved subsequences of inputs from the current device and the subsequences of inputs from the other devices.
2. The method according to claim 1, characterized in that The process of calculating the gradient of the second fusion layer parameters based on the saved subsequence of the current device input and the subsequence of the other device input includes: Calculating the gradient of the subsequence input by the current device; Calculating a gradient of a weight parameter based on the subsequence input by the other device and the subsequence input by the current device; The process of calculating the gradient of the subsequence input by the current device is performed in parallel with the process of obtaining the subsequences input by other devices.
3. The method according to claim 1, characterized in that The process of respectively calculating the gradient of the first fusion layer parameters and the gradient of the Probs based on the saved input of the first fusion layer and the Probs includes: Get the gradient of the output calculated by other devices, and merge it with the gradient of the output calculated by the current device to obtain the gradient of the complete output grad_output; Performing a Permute operation on the Probs to obtain rearranged Probs, and performing a dot product between the rearranged Probs and the input of the first fusion layer to obtain a scaled input; Based on the grad_output, calculate the gradient of the scaled input, the gradient of the weight parameter, and the gradient of the rearranged Probs; Calculating the gradient of the input of the first fusion layer according to the gradient of the scaled input; The gradient of the Probs before rearrangement is calculated according to the gradient of the Probs after rearrangement.
4. The method according to claim 3, characterized in that The process of obtaining the gradient of the output calculated by other devices is performed in parallel with the process of obtaining the scaled input.
5. The method according to claim 1, wherein In the back propagation process of the second fusion layer, the process of obtaining subsequences input from other devices includes: In the back propagation process of the second fusion layer, the all-gather communication method is used to obtain the subsequences input by other devices.
6. The method according to claim 3, characterized in that The process of obtaining the gradient of the output calculated by other devices includes: Use the all gather communication method to obtain the gradient of the output calculated by other devices.
7. A question-answering method, characterized in that: include: Get question information; Sending the question information to the configured hybrid expert model to obtain answer information output by the hybrid expert model; The large hybrid expert model is a model trained by the hybrid expert network training method according to any one of claims 1 to 6.
8. A hybrid expert network training device, characterized in that: The hybrid expert network includes a permute operation layer, a first expert layer FC1, an activation network, and a first fusion layer. The first fusion layer is integrated with a second expert layer FC2 and an unpermute operation layer. The device includes: A computing unit is configured to obtain a subsequence of the current device input and a subsequence of the input of other devices to form a complete input sequence; forward propagates the complete input sequence sequentially through the Permute operation layer, the FC1, and the activation network to obtain the output of the activation network; uses the output of the activation network as the input of the first fusion layer, and sequentially performs processing of the FC2 and the Unpermute operation layer in the forward propagation process of the first fusion layer, and saves the input of the first fusion layer and the probability Probs of the expert corresponding to each token token in the complete input sequence to the video memory, wherein the Probs of the expert corresponding to each token in the complete input sequence is calculated by the routing module in the hybrid expert network; in the back propagation process of the first fusion layer, based on the input of the first fusion layer and the Probs stored in the video memory, respectively calculates the gradient of the first fusion layer parameters and the gradient of the Probs; A parameter updating unit, configured to update the parameters of the hybrid expert network according to the gradient of the parameters of each layer in the back propagation process; Wherein, the Permute operation layer and the FC1 constitute the second fusion layer; During the forward propagation process of the second fusion layer, the computing unit saves a subsequence of the current device input, obtains a subsequence of the input of other devices, and forms the complete input sequence with the subsequence of the current device input, sequentially performs the processing of the Permute operation layer and the FC1 on the complete input sequence, and uses the processing result as the input of the activation network; during the backward propagation process of the second fusion layer, the computing unit communicates with other devices to obtain a subsequence of the input of other devices, and calculates the gradient of the parameters of the second fusion layer based on the saved subsequence of the current device input and the subsequence of the input of the other devices.
9. An electronic device, characterized in that: include: memory and processor; The memory is used to store programs; The processor is used to execute the program to implement the various steps of the hybrid expert network training method as described in any one of claims 1 to 6, or to implement the various steps of the question-answering method as described in claim 7.
10. A readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, it implements the various steps of the hybrid expert network training method as described in any one of claims 1 to 6, or implements the various steps of the question-answering method as described in claim 7.
11. A computer program product comprising a computer program, characterized in that When the computer program is executed by a processor, it implements the various steps of the hybrid expert network training method according to any one of claims 1 to 6, or implements the various steps of the question-answering method according to claim 7.
Citation Information
Patent Citations
Checkpoint selection method and device based on DNN model and storage medium
CN114692829A
Neural network model training method and device, electronic equipment and storage medium
CN115688917A