Method for constructing routing reverse operator, computer device and readable storage medium
By dividing the forward routing operator modules and mapping the reverse submodule, the reverse routing operator is constructed, which solves the problem of low utilization of hardware computing resources and improves the training efficiency of the model.
Patent Information
- Application Number
- CN202510446318.0
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2025-04-10
- Publication Date
- 2025-08-12
- Estimated Expiration
- 2045-04-10
AI Technical Summary
When implementing routing reverse operators on artificial intelligence chips, the existing technology has the problem of low utilization of hardware computing resources, mainly due to frequent global memory access and redundant read and write operations.
By dividing the routing forward operator module, the algorithm formula of the reverse submodule is determined, and mapped to the sub-core of the artificial intelligence chip, the routing reverse operator is built to reduce data access to global memory.
It improves the utilization rate of hardware computing resources, reduces unnecessary intermediate results storage and reading, improves the computing efficiency of backpropagation, and thus improves the training efficiency of the model.
Smart Images

Figure CN119940405B_ABST
Abstract
Description
Technical Field
[0001] The present application relates to the field of artificial intelligence technology, and in particular to a method for constructing a routing reverse operator, a computer device, a computer-readable storage medium, and a computer program product. Background Art
[0002] In large models, the Mixture of Experts (MoE) pattern splits a complex model into multiple expert sub-models. These expert sub-models are relatively independent modules, each with its own specialized capabilities, capable of handling different types of tasks or data. MoE is an AI model based on the Transformer architecture. It replaces each feedforward network layer in the traditional Transformer model with an MoE layer. The gating network (Router) is a key component in the Mixture of Experts pattern. Its main function is to determine which expert sub-model to send the input token to for processing.
[0003] When implemented on AI chips (such as GPGPUs (General-Purpose Computing on Graphics Processing Units)), the gating network is called a Router operator. Its typical structure consists of single operators such as softmax and topk. While requiring minimal internal computation, it requires significant memory access, resulting in a memory-intensive architecture. In various deep learning frameworks, the Router functionality is typically implemented by connecting the required single operators.
[0004] The reverse operator of the router (also known as the routing reverse operator) is typically implemented through automatic reverse derivation within a deep learning framework. This involves performing a reverse gradient propagation for each single operator in the router according to the chain rule. This approach requires frequent data access to global memory, resulting in numerous redundant read and write operations and low utilization of hardware computing resources. Summary of the Invention
[0005] Based on this, it is necessary to provide a method for constructing a routing reverse operator, a computer device, a computer-readable storage medium, and a computer program product that can improve the utilization of hardware computing resources in the reverse propagation of the routing operator to address the above technical problems.
[0006] In a first aspect, the present application provides a method for constructing a route reverse operator, the method comprising:
[0007] Determine a reverse submodule included in a route reverse operator, wherein the reverse submodule is used to perform gradient calculation;
[0008] Calling each of the reverse submodules from the operator library, connecting each of the reverse submodules according to the dependency relationship, and mapping each of the reverse submodules to the corresponding sub-core in the artificial intelligence chip to construct a routing reverse operator;
[0009] The reverse submodule is constructed based on the derived algorithm formula after deducing the algorithm formula corresponding to the forward submodule, and the forward submodule is a submodule obtained by dividing the routing forward operator.
[0010] In one embodiment, the method further comprises:
[0011] Divide the routing forward operator into modules to obtain multiple forward sub-modules;
[0012] For any of the forward submodules, perform the following operations:
[0013] Determine the algorithm formula of the reverse submodule corresponding to the forward submodule according to the algorithm formula of the forward submodule;
[0014] Based on the algorithm formula of the reverse submodule, the operators included in the reverse submodule and the dependency relationships between the operators are determined, and the reverse submodule is constructed according to the operators and the dependency relationships between the operators.
[0015] In one embodiment, the input tensors, intermediate tensors and output tensors of the routing reverse operator are all stored in the global memory, and the internal tensors of the reverse submodule are stored in the designated memory, which is the memory corresponding to the sub-kernel that executes the reverse submodule.
[0016] In one embodiment, the reverse submodule included in the determining route reverse operator includes:
[0017] Obtain configuration parameters of the route reverse operator, and determine candidate reverse submodules whose configuration attributes in the configuration parameters are first attribute values;
[0018] The main path loss sub-template and each of the candidate reverse sub-modules are determined as the reverse sub-modules of the route reverse operator.
[0019] In one embodiment, before mapping each of the reverse submodules to a corresponding sub-core in the artificial intelligence chip to construct a route reverse operator, the method further includes:
[0020] Determining a position adjustment strategy for the transposition operator in the route reversal operator based on the computing resources of the artificial intelligence chip;
[0021] The position of the transposition operator in the routing reverse operator is adjusted based on the position adjustment strategy, and the designated dimension of the first target operator in the reverse submodule is re-designated according to the position adjustment strategy.
[0022] In one embodiment, determining a position adjustment strategy of a transposition operator in the route reversal operator based on computing resources of the artificial intelligence chip includes:
[0023] Determining a first hardware parallelism rate of the artificial intelligence chip based on the unspecified dimension before adjustment and the computing resources of the artificial intelligence chip;
[0024] After swapping the designated dimension and the non-designated dimension, determining a second hardware parallelism rate of the artificial intelligence chip based on the swapped non-designated dimension and computing resources of the artificial intelligence chip;
[0025] When the first hardware parallelism rate is less than the second hardware parallelism rate, a target position of the transposition operator in the route reversal operator is determined, and a position adjustment strategy is generated based on the target position of the transposition operator in the route reversal operator.
[0026] In one embodiment, determining the target position of the transpose operator in the route reversal operator includes:
[0027] Determining a second target operator from the route reverse operator, where the second target operator is an operator whose operation is dimension-independent;
[0028] Based on the position of the second target operator, a target position of the transpose operator in the route reversal operator is determined.
[0029] In one embodiment, the routing forward operator is divided into modules to obtain multiple forward sub-modules, including:
[0030] Get the routing forward operator graph corresponding to the routing forward operator;
[0031] Dividing the routing forward operator graph into submodules to obtain multiple operator graphs;
[0032] According to the multiple operator graphs, multiple forward submodules are obtained.
[0033] In a second aspect, the present application further provides a device for constructing a route reverse operator, the device comprising:
[0034] A first determining module, configured to determine a reverse submodule included in a route reverse operator, wherein the reverse submodule is configured to perform gradient calculation;
[0035] A first construction module is used to call each of the reverse sub-modules from the operator library, connect each of the reverse sub-modules according to the dependency relationship, and map each of the reverse sub-modules to the corresponding sub-core in the artificial intelligence chip to construct a routing reverse operator;
[0036] The reverse submodule is constructed based on the derived algorithm formula after deducing the algorithm formula corresponding to the forward submodule, and the forward submodule is a submodule obtained by dividing the routing forward operator.
[0037] In one embodiment, the apparatus further comprises:
[0038] A partitioning module is used to partition the routing forward operator into modules to obtain multiple forward sub-modules;
[0039] For any of the forward submodules, perform the following operations:
[0040] A second determining module is used to determine the algorithm formula of the reverse submodule corresponding to the forward submodule according to the algorithm formula of the forward submodule;
[0041] The second construction module is used to determine the operators included in the reverse submodule and the dependencies between the operators based on the algorithm formula of the reverse submodule, and construct the reverse submodule according to the operators and the dependencies between the operators.
[0042] In one embodiment, the input tensors, intermediate tensors and output tensors of the routing reverse operator are all stored in the global memory, and the internal tensors of the reverse submodule are stored in the designated memory, which is the memory corresponding to the sub-kernel that executes the reverse submodule.
[0043] In one embodiment, the first determining module is specifically configured to:
[0044] Obtain configuration parameters of the route reverse operator, and determine candidate reverse submodules whose configuration attributes in the configuration parameters are first attribute values;
[0045] The main path loss sub-template and each of the candidate reverse sub-modules are determined as the reverse sub-modules of the route reverse operator.
[0046] In one embodiment, the apparatus further comprises:
[0047] a fourth determining module, configured to determine a position adjustment strategy for the transposition operator in the route reversal operator based on the computing resources of the artificial intelligence chip;
[0048] An optimization module is used to adjust the position of the transposition operator in the routing reverse operator based on the position adjustment strategy, and re-specify the specified dimension of the first target operator in the reverse submodule according to the position adjustment strategy.
[0049] In one embodiment, the optimization module is specifically configured to:
[0050] Determining a first hardware parallelism rate of the artificial intelligence chip based on the unspecified dimension before adjustment and the computing resources of the artificial intelligence chip;
[0051] After swapping the designated dimension and the non-designated dimension, determining a second hardware parallelism rate of the artificial intelligence chip based on the swapped non-designated dimension and computing resources of the artificial intelligence chip;
[0052] When the first hardware parallelism rate is less than the second hardware parallelism rate, a target position of the transposition operator in the route reversal operator is determined, and a position adjustment strategy is generated based on the target position of the transposition operator in the route reversal operator.
[0053] In one embodiment, the optimization module is specifically configured to:
[0054] Determining a second target operator from the route reverse operator, where the second target operator is an operator whose operation is dimension-independent;
[0055] Based on the position of the second target operator, a target position of the transpose operator in the route reversal operator is determined.
[0056] In one embodiment, the partitioning module is specifically configured to:
[0057] Get the routing forward operator graph corresponding to the routing forward operator;
[0058] Dividing the routing forward operator graph into submodules to obtain multiple operator graphs;
[0059] According to the multiple operator graphs, multiple forward submodules are obtained.
[0060] In a third aspect, the present application further provides a computer device comprising a memory and a processor, wherein the memory stores a computer program, and the processor implements any of the above methods for constructing a routing reverse operator when executing the computer program.
[0061] In a fourth aspect, the present application further provides a computer-readable storage medium having a computer program stored thereon, wherein the computer program, when executed by a processor, implements any of the above methods for constructing a routing reverse operator.
[0062] In a fifth aspect, the present application also provides a computer program product, including a computer program, which implements any of the above methods for constructing a routing reverse operator when executed by a processor.
[0063] The above-mentioned construction method, computer device, computer-readable storage medium and computer program product of the routing reverse operator are obtained by constructing the routing reverse operator based on the reverse submodule in the operator library, and the reverse submodule is constructed based on the algorithm formula derived from the algorithm formula corresponding to the forward submodule. This specially derived reverse operator formula allows the reverse submodule to reduce certain transitional intermediate variables in the forward calculation through deduction. In this way, in the calculation of reverse propagation, there is no need to re-read and store this part of the data, reducing the access to this data in the global memory, that is, when performing reverse propagation, there is no need to perform gradient calculation on the entire operator graph of the routing operator, reducing unnecessary storage and reading of intermediate results, which can greatly improve the utilization rate of hardware computing resources, improve the calculation efficiency of reverse propagation, and thus improve the training efficiency of the model. BRIEF DESCRIPTION OF THE DRAWINGS
[0064] In order to more clearly illustrate the technical solutions in the embodiments of the present application or related technologies, the following briefly introduces the drawings required for use in the embodiments of the present application or related technical descriptions. Obviously, the drawings described below are only some embodiments of the present application. For ordinary technicians in this field, other related drawings can be obtained based on these drawings without paying any creative work.
[0065] Figure 1 is a forward propagation operator graph of a routing operator in a MoE model in one embodiment;
[0066] Figure 2 1 is a flow chart of a method for constructing a route reverse operator in one embodiment;
[0067] Figure 3 1 is a flow chart of a method for constructing a route reverse operator in one embodiment;
[0068] Figure 4 is a flow chart of step 302 in another embodiment;
[0069] Figure 5 An operator graph of a forward submodule obtained by dividing a routing operator in one embodiment;
[0070] Figure 6a is an operator graph of a route reverse operator in one embodiment;
[0071] Figure 6b is an operator graph of a main path submodule in one embodiment;
[0072] Figure 6c : is an operator graph of the z-loss submodule in one embodiment;
[0073] Figure 6d is an operator graph of a load balancing submodule in one embodiment;
[0074] Figure 7 202 is a flow chart of step 202 in one embodiment;
[0075] Figure 8 1 is a flow chart of a method for constructing a route reverse operator in another embodiment;
[0076] Figure 9 802 in one embodiment;
[0077] Figure 10 906 in one embodiment;
[0078] Figure 11a The operator graph after optimization of the main path submodule in one embodiment;
[0079] Figure 11b The figure is an operator graph after optimization of the z-loss submodule in one embodiment;
[0080] Figure 11c An operator graph after optimization of the load balancing submodule in one embodiment;
[0081] Figure 12 A schematic diagram of the structure of a GPGPU in one embodiment;
[0082] Figure 13 is a structural block diagram of a device for constructing a route reverse operator in one embodiment;
[0083] Figure 14 FIG. 1 is a diagram showing the internal structure of a computer device in one embodiment. DETAILED DESCRIPTION
[0084] In order to make the purpose, technical solutions and advantages of this application more clear, the following further describes this application in detail with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain this application and are not intended to limit this application.
[0085] In the Mixture of Experts (MoE) model, the routing operator is a key component, which is usually implemented with the help of lightweight neural networks (such as linear layers) or other structures. Different MoE models have some subtle differences in the specific design of the routing operator. Figure 1 As shown in Figure 2, the forward propagation process of the routing operator in a typical MoE model is shown. Figure 1In the forward operator shown, the input tensor (hereafter referred to as "Input") has a shape of [b, s / T, e] and is the input to the entire computation process. Here, b represents the batch size, which is the number of data samples processed in a single computation; s is the sequence length (seqlen); T is the total number of tokens; and e is the number of experts (expert num).
[0086] It should be noted that the routing operators and the shape expressions of different tensors in the embodiments of this application are merely examples of implementations of the embodiments of this application and are not to be construed as limitations on the shape expressions of routing operators and tensors. In practice, for routing operators of different structures, there is a corresponding relationship between the shape expressions of their tensors and the parameters of the routing operators, which is not specifically limited in the embodiments of this application.
[0087] The weight tensor (hereafter referred to as weight) is fed into the gating function along with the input. In deep learning, weight tensors are used to adjust the model's emphasis on different input features. Gating operates on the input and weight, outputting a logistic regression value, logits tensor 1 (hereafter referred to as logits1), with a shape of [b, s / T, e]. Logits are the raw prediction values before activation functions and are subsequently used for probability calculations, among other things.
[0088] The view function (hereafter referred to as view) performs a view transformation on logits1, changing the shape of the data and outputting the logistic regression value logits tensor 2 (hereafter referred to as logits2), with a shape of [b × s, e]. The seq_gather function is a sequentially parallel gather function, where gather is a basic function used to obtain data of a specified dimension (dim) and a specified index (index) from the original tensor. It processes logits2 to obtain the logistic regression value logits tensor 3 (hereafter referred to as logits3). This operation involves extracting information at specific locations or features from the sequence data.
[0089] The topk function is used to retrieve the k largest or smallest elements and their index positions in a tensor or array. For example, it performs a topk operation on logits3 to find the k most probable elements and their index positions in each sample, outputting a score of 1 (referred to as scores1 below) with a shape of [b × s, topk]. After the topk operation finds the k most probable elements in each sample, the argsort function sorts these elements and generates indices (referred to as indices below) with a shape of [b × s, topk]. Indices records the index position information of the relevant elements in the original data after the topk operation.
[0090] The normalized exponential function (hereinafter referred to as softmax) function applies the softmax activation function to the result of the topk operation, converting scores1 into probability values (hereinafter referred to as probs) with a shape of [b×s, topk]. The scatter function is a basic function that transfers data from one tensor to another based on the index. It transfers probs to another tensor based on the index obtained by indices, resulting in the topk_masked_gates tensor (hereinafter referred to as topk_masked_gates), which is a mask tensor used to mark the positions of the topk elements.
[0091] The topk function performs the topk operation on the input topk_masked_gates, where k = capacity, along dimension dim = 0. This operation finds the capacity (capacity is the expert capacity) elements with the highest probability along each dimension and outputs two results: a capacity_indices tensor (hereafter referred to as capacity_indices), which has a shape of [capacity, e] and records the indices of the capacity elements with the highest probability found in topk_masked_gates; and a capacity_probs tensor (hereafter referred to as capacity_probs), which has a shape of [capacity, e] and indicates the probability values of the corresponding capacity elements with the highest probability.
[0092] Perform a transpose operation on capacity_probs (hereinafter referred to as transpose) to swap the order of its dimensions to obtain the final probability tensor final_probs (hereinafter referred to as final_probs), with a shape of [e, capacity], which is the final output probability value. Perform a transpose operation on capacity_indices to obtain the final probability tensor final_indices (hereinafter referred to as final_indices), with a shape of [e, capacity], which is used to indicate the index value related to the final output probability value.
[0093] In the z_loss loss function, the smooth maximum function (hereinafter referred to as logsumexp) performs a logarithmic sum exponential operation on logits3 along the specified dimension (dim=1), and the output shape is [b×s,1]; the square function (hereinafter referred to as square) squares the result of logsumexp, and the output shape is [b×s,1]; the average function (hereinafter referred to as mean) calculates the average of the result of square, and after the coefficient transformation function (e×moe_aux_loss_coeff) transformation, the z_loss loss is obtained, with a shape of [1,1].
[0094] In aux_loss_load_balancing (load balancing loss), the softmax function applies the softmax activation function to logits3 again, with the data type specified as fp32. The output is score 2 (referred to as scores2 below), with a shape of [b × s, e]. The mean function calculates the mean of scores2 along dimension dim=0, resulting in a probs_mean_per_expert tensor (referred to as probs_mean_per_expert below), with a shape of [1, e]. The scatter function scatters the data according to indices, generating a masked topk_mask tensor (referred to as topk_mask below), with a shape of [b × s, e]. The sum function (referred to as sum below) sums topk_mask along dimension dim=0, resulting in a tokens_per_expert tensor (referred to as tokens_per_expert below), with a shape of [1, e]. This data is related to the number of tokens processed by each expert and is used for load balancing calculations. Finally, the probs_mean_per_expert and tokens_per_expert are multiplied and summed by the mean dot product + sum function (hereinafter referred to as dot mul+sum), and transformed by the coefficient transformation function (e×moe_aux_loss_coeff), and the load balancing loss (hereinafter referred to as aux_loss, also known as aux loss auxiliary loss function) is obtained, with a shape of [1, 1].
[0095] In other words, the routing operator involves multiple steps, including input processing, probability calculation, loss function calculation, and final result generation. These operations collectively contribute to the model's training and inference processes. The inputs, outputs, and intermediate data of these multiple steps are all stored in global memory. When implementing the routing operator's backpropagation process on a GPGPU using the chain rule in related technologies, gradient calculations are typically performed on the entire operator graph of the routing operator according to fixed rules. This requires frequent read and write operations on large amounts of data to global memory, potentially requiring multiple calculations and storage of the same intermediate results. This large number of redundant read and write operations results in low utilization of hardware computing resources.
[0096] The embodiment of the present application provides a method for constructing a routing reverse operator for gradient calculation during back propagation of a routing operator. The routing reverse operator constructed based on the embodiment of the present application does not need to perform gradient calculation on the entire operator graph of the routing operator during back propagation, thereby reducing unnecessary storage and reading of intermediate results, greatly improving the utilization rate of hardware computing resources, improving the computational efficiency of back propagation, and thereby improving the training efficiency of the model.
[0097] In one embodiment, Figure 2 As shown, a method for constructing a route reverse operator is provided. This embodiment uses the method applied to the host side as an example. It is understandable that the host side may include a CPU (Central Processing Unit). In this embodiment, the method includes the following steps 202 to 204, wherein:
[0098] Step 202: Determine a reverse submodule included in the route reverse operator, where the reverse submodule is used to perform gradient calculation.
[0099] In the embodiment of the present application, the forward routing operator is used for the forward propagation of the routing operator (ie, Router), and the reverse routing operator is used for the reverse propagation of the routing operator. The reverse routing operator includes at least one reverse submodule, which is an operator module used for gradient calculation in the reverse propagation. Figure 1 The embodiment of the present application is described using the routing forward operator shown as an example.
[0100] In an embodiment of the present application, the routing forward operator can be divided into modules in advance to obtain at least one forward sub-module, and reverse deduction can be performed based on the algorithm formula of each forward sub-module to obtain the algorithm formula of the reverse sub-module corresponding to each forward sub-module, and each reverse sub-module can be constructed based on the algorithm formula of each reverse sub-module and stored in the operator library.
[0101] For example, the routing forward operator can be analyzed to determine how the routing forward operator assigns weights to different expert networks based on the input data, and the decision-making role it plays in the forward propagation process of the entire model. For example, analyze what features (such as the dimension of the input vector, the numerical distribution, etc.) the routing forward operator uses to select experts and assign weights. Based on the analysis results, identify the parts of the routing forward operator with relatively independent functions, which usually correspond to different calculation steps or logical units. For example: referring to the aforementioned Figure 1 The example routing forward operator includes an input feature extraction sub-function, which is used to extract key features for routing decisions from the original input data, which may involve operations such as linear transformation and nonlinear activation; a weight calculation sub-function, which is used to calculate the assigned weight of each expert network based on the extracted features, which may use the softmax function, attention mechanism, etc.; an expert selection sub-function, which is used to select the expert network participating in subsequent calculations based on the calculated weights, which may use the topk function, etc.
[0102] Based on the identified sub-functions, the routing forward operator can be divided into at least one forward sub-module. It should be noted that when dividing, it is necessary to ensure that each forward sub-module has clear inputs and outputs and relatively independent functions. For example, Figure 1 The forward routing operator shown in the figure takes the raw input data as input and outputs the extracted feature vectors. The weight calculation module takes the extracted feature vectors as input and outputs the weights assigned to each expert network. The expert selection module takes the expert weights as input and outputs the index of the selected expert network.
[0103] Furthermore, for each forward submodule, reverse deduction can be performed based on its algorithm formula to derive the algorithm formula for each reverse submodule in backpropagation. For example, the algorithm formula for each forward submodule can be determined, and the variables and constants involved in the formula and the operational relationships between them can be clarified. For example, for the input feature extraction module, the linear transformation y = Wx + b (where x is the input vector, W is the weight matrix, b is the bias vector, and y is the output feature vector) is often used. The meaning and dimension of each parameter must be clearly defined.
[0104] After determining the algorithm formula of each forward submodule, we can use the chain rule to reversely deduce and obtain the reverse algorithm formula. For example: for the linear transformation y=Wx+b, assuming the loss function is L, according to the chain rule, we can deduce the reverse algorithm formula as .
[0105] After inferring the algorithm formula for each forward submodule, the corresponding reverse submodule can be constructed based on the algorithm formula and stored in the operator library. In this way, when implementing the routing reverse operator later, the reverse submodule included in the routing reverse operator can be determined based on the forward operator module included in the routing forward operator, and each reverse submodule can be called from the operator library to implement the routing reverse operator.
[0106] The following is a detailed introduction to the process of deriving the reverse submodule.
[0107] In an exemplary embodiment, referring to Figure 3 As shown, the process of deriving the reverse submodule may include the following steps 302 to 306, wherein:
[0108] Step 302: Divide the routing forward operator into modules to obtain multiple forward submodules;
[0109] For any forward submodule, do the following:
[0110] Step 304: Determine the algorithm formula of the reverse submodule corresponding to the forward submodule based on the algorithm formula of the forward submodule;
[0111] Step 306 : Based on the algorithm formula of the reverse submodule, determine the operators included in the reverse submodule and the dependencies between the operators, and construct the reverse submodule according to the dependencies between the operators.
[0112] A typical routing forward operator mainly includes the following core structures: softmax module, topk module, load balancing loss module and z-loss loss module, among which:
[0113] Softmax Module: The Router module calculates the scores of all experts for each input data point. To properly assign the input to the appropriate expert, these scores need to be normalized, a task accomplished by softmax. After softmax normalization, the input data is assigned to the expert with the highest score for subsequent calculations. This process helps assign appropriate weights to each expert based on the input, allowing the model to select the most suitable expert for processing based on the input features.
[0114] Topk module: In large-scale MoE models, requiring all experts to participate in the calculation would incur significant computational overhead. To effectively reduce this cost, the routing module uses Topk. Topk selects the k highest-scoring experts from all experts for activation, while unselected experts are excluded from subsequent calculations. This significantly reduces unnecessary computation while maintaining model performance, improving computational efficiency.
[0115] Load Balancing Loss Module: To ensure load balancing across experts and prevent some experts from being overused while others are idle, MoE models typically incorporate a load balancing term into their loss function. The load balancing loss calculates this load balancing loss and, by imposing constraints on the routing module, encourages it to distribute input data as evenly as possible across experts, thereby fully leveraging the capabilities of each expert and improving the overall generalization performance of the model.
[0116] Z-loss: To further assist the training process of the routing module and ensure the stability and diversity of expert selection, some MoE models introduce additional loss functions, including z-loss. Z-loss is a commonly used regularization loss used to assist in training the router, ensuring more stable and diverse expert selection. Its primary purpose is to alleviate the problem of irrational distribution of router outputs, avoiding over-reliance on a small number of experts or extreme allocation probabilities. Its core concept is to prevent extreme distributions in routing outputs by penalizing excessively large or small probabilities. Z-loss can constrain the routing module from various perspectives, making the model's expert selection more comprehensive and reasonable, avoiding over-reliance on specific experts, and thus enhancing the model's robustness and generalization capabilities.
[0117] These modules work together to form a complete routing forward operator. Based on this, the routing forward operator can be divided into modules to obtain corresponding multiple forward sub-modules.
[0118] In an exemplary embodiment, referring to Figure 4 As shown, in step 302, the routing forward operator is divided into modules to obtain multiple forward sub-modules, which may include the following steps 402 to 406, wherein:
[0119] Step 402: Obtain a routing forward operator graph corresponding to the routing forward operator;
[0120] Step 404: Divide the routing forward operator graph into submodules to obtain multiple operator graphs;
[0121] Step 406: Obtain multiple forward submodules based on the multiple operator graphs.
[0122] In the embodiment of the present application, the routing forward operator graph corresponding to the routing forward operator can be obtained, and the routing forward operator graph can be further divided into sub-modules based on the structure of the routing forward operator to obtain multiple operator graphs. For example, Figure 1 The routing forward operator graph shown in the figure is divided into submodules of the routing forward operator, and the division result can be shown as follows: Figure 5 As shown, Figure 5 Each dotted box in the figure corresponds to the operator graph of a submodule.
[0123] The routing forward operator graph is divided into multiple operator graphs, each of which corresponds to a forward submodule. Taking the typical Router operator structure used in the switch transformer model structure as an example, the forward submodules included in the routing forward operator graph can include the main path loss submodule, the load balancing submodule, and the z-loss loss submodule, which can be expressed as the following formula (1):
[0124] Formula (1)
[0125] in, Characterizes the total loss of the routing forward operator, Characterize the main path loss submodule, Characterize the load balancing loss submodule, Characterize the z-loss loss submodule, Characterizes the weight of the load balancing loss submodule, Characterizes the weight of the z-loss submodule. The main path loss submodule is the main part of the routing forward operator, mainly composed of single operators such as topk and softmax. Therefore, the main path loss submodule can also be further divided into more submodules, such as: softmax submodule, topk submodule, etc., which are not specifically limited in this embodiment of the application.
[0126] The load balancing loss submodule is typically an implementation of an auxiliary loss function, referred to as the aux loss function in this embodiment. The load balancing loss submodule is a key regularization term used to optimize the Router in the MoE model. The typical forward algorithm formula for the aux loss function is as follows:
[0127] Formula (2)
[0128] Among them, α is the scaling factor used to control the weight of load balancing loss in the total loss, N is the number of experts, Represents the probability that the Router assigns the current token to expert i, , The result of Router input x after softmax normalization calculation, , T is the total number of Tokens, x is the input of the Router, .
[0129] The typical forward calculation formula of the z-loss function is shown in formula (3).
[0130] Formula (3)
[0131] After obtaining the algorithm formula of each forward submodule, for any forward submodule, the algorithm formula of the corresponding reverse submodule can be obtained based on the algorithm formula of the forward submodule. Figure 5 The reverse process of each forward submodule of the routing forward operator shown is as follows:
[0132] The reverse algorithm formula of the main path loss submodule is Since the implementation of the main path varies in different routers, the reverse operator of the main path loss operator, with softmax and topk as the key parts, can be fused into multiple sub-paths, that is, fused into multiple reverse sub-modules. In the embodiment of this application, the reverse process of key operators such as softmax and topk is not elaborated in detail, and the reverse process can be performed based on the chain rule.
[0133] The reverse algorithm formula of the load balancing loss submodule is , for formula (2), according to the chain rule, we can get ,because The computational characteristics of It is not derivable, so we can get .in, , .
[0134] Therefore, the back propagation of the aux loss function is implemented as follows (IV):
[0135] Formula (4).
[0136] Among them, (meanreduce_bwd) represents the reverse implementation of meanreduce. Meanreduce refers to the function that performs aggregation operations on the input tensors. Here, it refers to the above formula for The operation of finding the mean. The reverse implementation of meanreduce requires evenly distributing the incoming gradient back to each element of the input. (softmax_bwd) refers to the reverse implementation of softmax. Please refer to the general softmax reverse implementation method. The specific derivation formula is not introduced in the examples of this application.
[0137] The algorithm formula of the z-loss loss function is reversed as follows: For ease of description, the calculation part of logsumexp in z-loss is recorded as LSE. LSE calculates the exponential sum of the input x and takes the logarithm, and records its square as SLSE. The formula is as follows: , .
[0138] Therefore, the forward formula of z-loss can be expressed as shown in formula (5):
[0139] Formula (5)
[0140] The derivation process of the reverse calculation formula corresponding to z-loss is as follows: ,in, , , .
[0141] Therefore, the back propagation of the z-loss function is implemented as shown in the following formula (6):
[0142] Formula (6)
[0143] In this way, by reverse-deriving the forward formulas of each forward submodule in the Router, we derive the reverse calculation formulas for each component of the Router. By analyzing the terms in the reverse calculation formulas, we can determine the operators included in the reverse submodules and their inputs and outputs. This means we can also determine the dependencies between the operators included in each reverse submodule. Based on these dependencies, we can then construct the reverse submodules.
[0144] For example, the obtained reverse calculation formula can be broken down into basic operation units, which are the operators contained in the reverse submodule. For example, operations such as multiplication, addition, and transposition in the formula can be regarded as different operators.
[0145] For each operator, its input and output can be determined based on the source and destination of the term in the reverse calculation formula. By analyzing the entire reverse calculation formula, all operators contained in each reverse submodule and their respective inputs and outputs can be determined.
[0146] Based on the inputs and outputs of operators, we can further analyze which operator outputs are the inputs of other operators, thereby determining the dependencies between operators. For example, if the input of operator B depends on the output of operator A, then operator B depends on operator A. When building the reverse submodule, operator A must be executed before operator B.
[0147] Graph structures (such as directed acyclic graphs) can be used to intuitively represent the dependencies between operators. Figure 5 The nodes in the diagram represent operators, and the directed edges represent the dependency between operators, from the output operator to the input operator. In this way, the logical relationship between the operators in the reverse submodule can be clearly demonstrated.
[0148] Furthermore, topologically sorting the dependency graph yields a linear operator execution order, thus forming the operator graph for the reverse submodule. This topological sorting ensures that when each operator is executed, all its dependent operators have already completed execution, thus satisfying inter-operator dependencies. By writing code based on the operator execution order derived from the topological sort, the reverse submodule can be implemented. In the code, each operator is called sequentially, passing the corresponding input data, ultimately implementing the functionality of the reverse submodule. Specific operator operations can be implemented using the operator libraries provided by deep learning frameworks (such as TensorFlow and PyTorch). These frameworks typically optimize and encapsulate common operators for direct call.
[0149] For example, refer to Figure 5The forward routing operator is based on the operator graph of the reverse routing operator constructed by each reverse submodule derived. Figures 6a to 6d As shown. Figure 6a The schematic diagram of the routing reverse operator shown in FIG shows that the routing reverse operator includes a main path submodule, a load balancing loss submodule, and a z-loss loss submodule, wherein d-logits-z_loss is the aggregated z_loss logits gradient tensor output by the z-loss loss submodule, with a shape of [b×s,e], and d-logtis-aux_loss is the aggregated aux-loss logits gradient tensor output by the load balancing loss submodule, with a shape of [b×s,e]. For example, the operator graph of the main path submodule refers to Figure 6b As shown in the figure, it mainly includes the reverse operator of the normalized exponential function (hereinafter referred to as softmax bwd, also known as the normalized exponential reverse function), the reverse operator of the topk function (hereinafter referred to as scatter), and other operators such as view function view, device transpose, and addition add. The operator diagram of the z-loss loss submodule is shown in the figure. Figure 6c As shown in the figure, it is the operator graph constructed based on the z-loss reverse function derived above. The operator graph of the load balancing loss submodule is shown in the figure. Figure 6d As shown, this is the operator graph constructed based on the aux loss reverse function derived above.
[0150] in, Figure 6b The corresponding main path submodule is composed of Figure 5 The main path loss submodule in the reverse derivation is obtained, Figure 6c The corresponding z-loss loss submodule is composed of Figure 5 The derivation result is derived from the z-loss loss submodule in formula (6). Figure 6d The corresponding load balancing loss submodule is composed of Figure 5 The load balancing loss submodule in is deduced to obtain the result, and the derivation result is described in Formula (IV). Figures 6a to 6d In the example, d-tensor (including d-Input, d-logits2, d-scores1, etc.) refers to the corresponding gradient tensor, for example, d-Input is the input gradient tensor.
[0151] exist Figure 6bIn the corresponding operator graph, since the topk operator is used in forward propagation, the gradient information needs to be passed back from the output of the topk operation to the input during backward propagation. Since the topk operation only selects some elements, the output gradient needs to be correctly distributed back to the corresponding position in the input tensor during backward propagation. The scatter operator can accomplish this task. The scatter operator can disperse the gradient information to the corresponding position in the input tensor based on the index information recorded during forward propagation. Therefore, the scatter operator is used for the topk position in the backward propagation process, that is, the scatter operator is the reverse implementation of the topk operator. The softmax bwd operator is the reverse function of softmax. When the scatter operator is used in forward propagation, the gradient information needs to be passed back from the output of the scatter operation to the input during backward propagation. Since the scatter operation scatters elements to different positions, the output gradient needs to be collected back to the corresponding position of the input tensor according to the index rule during forward propagation during backward propagation. The gather operator can collect elements at a specific position from the input tensor according to a given index. Therefore, the gather operator is used at the position corresponding to the scatter in the backward propagation process, that is, the gather operator is the reverse implementation of the scatter operator.
[0152] exist Figure 6c In the corresponding operator graph, based on formula (6), the reverse submodule corresponding to the operator graph includes three operators, one is the reverse implementation of softmax, which is represented as softmax in the operator graph, and the other is 2×LSE. i The corresponding representation is logsumexp in the operator graph, and the other is meanreduce_bwd (can also be expressed as reduce_bwd), which corresponds to mean_bwd in the operator graph, where mul is the fusion of the three operators.
[0153] exist Figure 6d In the corresponding operator graph, based on formula (IV), the reverse submodule corresponding to the operator graph includes three operators, one is the reverse implementation of softmax, which is represented as softmax_bwd in the operator graph, one is meanreduce_bwd, which corresponds to mean-bwd in the operator graph, and one is the implementation The operator corresponds to the dot mul sum in the operator graph.
[0154] In this way, based on the dependency relationship between tensors of operators, with the goal of reducing redundant memory transfer, each reverse sub-module is constructed, and each reverse sub-module can be stored in the operator library for subsequent calls.
[0155] Step 204: Call each reverse submodule from the operator library, connect each reverse submodule according to the dependency relationship, and then map each reverse submodule to the corresponding sub-core in the artificial intelligence chip to construct a routing reverse operator; wherein, the reverse submodule is constructed based on the algorithm formula derived from the algorithm formula corresponding to the forward submodule, and the forward submodule is a submodule obtained by dividing the routing forward operator.
[0156] In an embodiment of the present application, after determining each reverse submodule of the routing operator, each reverse submodule can be called from the operator library, and each reverse submodule can be connected in the order of the dependency relationship between the reverse submodules to construct a routing reverse operator. Further, each reverse submodule in each routing reverse operator is mapped to the corresponding sub-core in the artificial intelligence chip, and the reverse propagation of the routing operator can be performed based on the routing reverse operator. For example, if the output of a reverse submodule is the input of another reverse submodule, then they are connected according to this logic to finally obtain a complete routing reverse operator, and each reverse submodule is mapped to the sub-core of the artificial intelligence chip, for example, Figure 6b The multiple sub-modules divided by the main path loss sub-module are mapped to multiple sub-kernels (including sub-kernel function 1 (which can be expressed as subKernel-Main1), sub-kernel function 2 (which can be expressed as subKernel-Main2), sub-kernel function 3 (which can be expressed as subKernel-Main3), and sub-kernel function 4 (which can be expressed as subKernel-Main4)), the load balancing loss sub-module is mapped to the sub-kernel (including the sub-kernel function-auxloss, which can also be expressed as subKernel-auxloss), and the z-loss loss sub-module is mapped to the sub-kernel (including the sub-kernel function-auxloss, which can also be expressed as subKernel-zloss). In this way, the back propagation of the routing operator can be efficiently implemented on the artificial intelligence chip.
[0157] The method for constructing the routing reverse operator provided in the embodiment of the present application is adopted. Since the routing reverse operator is constructed based on the reverse submodule in the operator library, and the reverse submodule is constructed based on the algorithm formula obtained after deducing the algorithm formula corresponding to the forward submodule. This specially derived reverse operator formula enables the reverse submodule to reduce certain transitional intermediate variables in the forward calculation through deduction. In this way, in the calculation of reverse propagation, there is no need to re-read and store this part of the data, which reduces the access to this data in the global memory. That is, when performing reverse propagation, there is no need to perform gradient calculation on the entire operator graph of the routing operator, which reduces unnecessary storage and reading of intermediate results, can greatly improve the utilization rate of hardware computing resources, improve the calculation efficiency of reverse propagation, and thus improve the training efficiency of the model.
[0158] In an exemplary embodiment, the input tensors, intermediate tensors, and output tensors of the routing reverse operator are all stored in the global memory, and the internal tensors of the reverse submodule are stored in the designated memory, where the designated memory is the memory corresponding to the sub-kernel that executes the reverse submodule.
[0159] In the present application, refer to Figures 6b to 6d As shown, the tensor marked as 4 is the input tensor, which is located in the global memory, the tensor marked as 5 is the output tensor, which is also located in the global memory, the tensor marked as 6 is the intermediate tensor, which is the data exchanged between different subKernels and is located in the global memory, and the tensor marked as 7 is the internal tensor, which refers to the tensor calculated in the sub-kernel and is located in the designated memory corresponding to the sub-kernel. The designated memory is a memory structure closer to the core that is pre-divided for the sub-kernel, such as the shared cache (Shared Mem).
[0160] This way, after constructing a reverse routing operator for the reverse propagation of the routing operator, for any module within the reverse routing operator, the computational tensors within the subkernel executing that module can be stored in a memory structure closer to the core (such as a shared cache), thereby fully utilizing its high-speed memory access performance. For example, in subKernel-Main, when performing operations such as softmax reverse and topk reverse, storing intermediate results in the shared cache can reduce frequent accesses to global memory, thereby improving computational efficiency.
[0161] In the embodiment of the present application, after the reverse submodules are pre-derived and constructed, the corresponding route reverse operator can be constructed by freely and flexibly combining the reverse submodules during the construction of the route reverse operator. Figure 7 As shown, in step 202, determining the reverse submodule included in the route reverse operator may include the following steps 702 to 704, wherein:
[0162] Step 702: Obtain configuration parameters of the route reverse operator, and determine candidate reverse submodules whose configuration attributes in the configuration parameters are the first attribute value;
[0163] Step 704: determine the main path loss sub-template and each candidate reverse sub-module as the reverse sub-module of the route reverse operator.
[0164] In an embodiment of the present application, in actual application, a user can control whether to use a corresponding reverse submodule through the configuration parameters of the route reverse operator. For example, the configuration parameters of the route reverse operator can configure the configuration attributes of each route reverse operator. When the configuration attribute is a first attribute (e.g., True), the reverse submodule can be determined to be used and determined as a candidate reverse submodule. When the configuration attribute is a second attribute (e.g., False), the reverse submodule can be determined not to be used.
[0165] After determining each candidate reverse submodule, the main path loss submodule and each candidate reverse submodule can be connected in a corresponding order according to the dependency relationship, so that a route reverse operator can be constructed.
[0166] For example, referring to Table 1 below, different operator combinations corresponding to different configuration attributes of the reverse submodule are shown when the reverse submodule includes a load balancing loss submodule and a z-loss loss submodule. Among them, when it is desired to ensure that the load between multiple expert modules is relatively balanced, the load balancing loss module can be used. For example, in large-scale natural language processing tasks, different experts may be good at processing texts of different semantic types. If some experts always process texts with common semantics, while other experts are rarely used, a load balancing loss term can be introduced to optimize this situation, so that each expert can play his or her full role and improve the overall generalization ability and stability of the model. When it is necessary to adjust the probability distribution of the model output to make it more in line with the actual situation, the z-loss loss module can be used. For example, in applications such as medical diagnosis assistance and financial risk assessment, if the probability of the model output cannot accurately reflect the real uncertainty, serious consequences may occur. Z-loss can help adjust the model output so that the predicted probability distribution is more reasonable and improve the credibility and accuracy of the model prediction.
[0167] Table 1
[0168] serial number combination Load balancing loss module z-loss loss module 1 Main path loss submodule False False 2 Main path loss submodule + load balancing loss module True False 3 Main path loss submodule + z-loss loss module False True 4 Main path loss submodule + load balancing loss module + z-loss loss module True True
[0169] When deploying forward and reverse routing operators on AI chips, low computing resource utilization is often encountered. This is because the softmax and topk operators in the routing operators need to calculate data of a specified dimension (often the expert dimension (e) or the topk dimension (k)), while the AI chip's hardware resources parallelize data of other non-specified dimensions (such as batch size or sequence length).
[0170] In routing operators, the shape of input data is typically [batch_size, sequence_length, e], where batch_size and sequence_length are typically much larger than e or k. For example, batch_size might be 128, sequence_length might be 256, while e might be 16 and k might be 2.
[0171] The softmax and topk operations operate on the e dimension, meaning that each input data point (indexed by batch_size and sequence_length) requires independent calculation of the weights of e experts. The topk operation operates on the k dimension, meaning that each input data point requires selecting k experts from among the e experts. Consequently, if e = 16, only 16 data points can be processed in parallel on the e dimension. If k = 2, only 2 data points can be processed in parallel on the k dimension. However, for AI chips, the hardware capacity (hardware_capacity) typically ranges from tens to thousands, and this low-dimensional parallelization approach cannot fully utilize the hardware's computing resources.
[0172] For example, assuming the input data is of the shape [batch_size=128, sequence_length=256, e=16], and a softmax calculation is performed in the e dimension, each input data needs to calculate the weights of 16 experts. This results in a parallelization degree of only 16, while the hardware capacity of an AI chip may be 1024, meaning hardware resource utilization is only 16 / 1024. For a top-k calculation in the k dimension, each input data needs to select two experts from the 16 experts. This results in a parallelization degree of only 2, and hardware resource utilization is only 2 / 1024.
[0173] In the batch_size and sequence_length dimensions, the amount of data is large (128×256=32768), but these dimensions are not used for softmax and topk calculations, resulting in a waste of hardware resources.
[0174] Therefore, in order to further improve the utilization of hardware resources, in the embodiment of the present application, after the initial routing reverse operator is obtained by connection based on the determined reverse sub-module, the structure of each initial routing reverse operator can be optimized, the position of the transpose function in the initial routing reverse operator can be changed, and the e or topk dimension can be swapped with other dimensions to maximize the utilization of hardware resources.
[0175] The structural optimization process of the initial routing reverse operator is introduced below.
[0176] In an exemplary embodiment, referring to Figure 8 As shown, the method may further include the following steps 802 to 804, wherein:
[0177] Step 802: Determine a position adjustment strategy for the transposition operator in the routing reverse operator based on the computing resources of the artificial intelligence chip.
[0178] In an embodiment of the present application, the designated dimensions and non-designated dimensions of tensor data can be determined, and the hardware parallelism of the artificial intelligence chip can be determined based on the designated dimensions and non-designated dimensions and the computing resources of the artificial intelligence chip. The hardware parallelism reflects the degree to which the hardware can achieve parallel computing when the artificial intelligence chip processes related computing tasks. The higher the parallelism, the higher the computing efficiency may be in theory, and thus the position adjustment strategy for the transpose operator is determined based on the hardware parallelism.
[0179] In an exemplary embodiment, referring to Figure 9 As shown, in step 802, based on the computing resources of the artificial intelligence chip, determining the position adjustment strategy of the transposition operator in the route reverse operator may include the following steps 902 to 906, wherein:
[0180] Step 902: Determine a first hardware parallelism rate of the artificial intelligence chip based on the unadjusted non-specified dimension and the computing resources of the artificial intelligence chip.
[0181] Step 904: After swapping the designated dimension and the non-designated dimension, a second hardware parallelism ratio of the artificial intelligence chip is determined based on the swapped non-designated dimension and the computing resources of the artificial intelligence chip.
[0182] Step 906 : When the first hardware parallelism rate is less than the second hardware parallelism rate, determine the target position of the transpose operator in the route reversal operator, and generate a position adjustment strategy based on the target position of the transpose operator in the route reversal operator.
[0183] In an embodiment of the present application, a first hardware parallelism ratio can be calculated based on the current non-specified dimension and the computing resources of the artificial intelligence chip. This first hardware parallelism ratio represents the hardware parallelism ratio of the artificial intelligence chip before the structure is optimized. Subsequently, the non-specified dimension and the specified dimension can be swapped, that is, the specified dimension is used as the new non-specified dimension. Based on the new non-specified dimension and the computing resources of the artificial intelligence chip, a second hardware parallelism ratio is calculated. This second hardware parallelism ratio represents the hardware parallelism ratio of the artificial intelligence chip before and after the structure is optimized.
[0184] If the first hardware parallelism rate is less than the second hardware parallelism rate, it means that the hardware parallelism rate of the artificial intelligence chip is higher after the optimized structure is represented. In this case, the structure of the routing reverse operator can be optimized, and the position of the transpose operator in the routing reverse operator can be adjusted to achieve the exchange of specified dimensions and non-specified dimensions, thereby improving the hardware parallelism rate of the artificial intelligence chip.
[0185] In an exemplary embodiment, referring to Figure 10 As shown, in step 906, determining the target position of the transposition operator in the route reversal operator may include the following steps 1002 to 1004, wherein:
[0186] Step 1002: determining a second target operator from the reverse route operator, where the second target operator is an operator whose operation is independent of dimension;
[0187] Step 1004: Determine the target position of the transpose operator in the route reversal operator based on the position of the second target operator.
[0188] In an embodiment of the present application, a second target operator can be determined from the route reverse operator. The second target operator is an operator whose operation is independent of dimension. The computing characteristics of this type of operator do not depend on a specific dimension setting and are relatively stable when the dimension is adjusted. After determining the second target operator, a transposition operator can be set before the second target operator to swap the swapped specified dimension and non-specified dimension back and perform subsequent computing operations. That is, it can be determined that the target position of the transposition operator is before the second target operator, thereby generating a position adjustment strategy based on the target position.
[0189] Step 804 : Adjust the position of the transpose operator in the routing reverse operator based on the position adjustment strategy, and re-specify the designated dimension of the first target operator in the reverse submodule according to the position adjustment strategy.
[0190] In the embodiment of the present application, the position of the transposition operator in the route reverse operator is adjusted to the target position based on the position adjustment strategy, and the designated dimension of the first target operator in the reverse submodule can be re-designated to achieve the swapping of the original designated dimension and the non-designated dimension in the first target operator. For example: Figures 6a to 6d The route reverse operator shown is Figure 6b As shown, before the specified and non-specified dimensions are swapped, the transpose operator is placed before scatter, then scatter can be determined as the first target operator, and the sequence parallel gather reverse function (also expressed as seq_gatherbwd) is the second target operator. At this time, the input shape of scatter is [capacity, e], the specified dimension is e, and the non-specified dimension is capacity. After swapping the specified and non-specified dimensions, the resulting routing reverse operator is referenced Figures 11a to 11c As shown, Figure 11a is the adjusted main path submodule, Figure 11b is the adjusted z-loss loss submodule, Figure 11c This is the adjusted load balancing loss submodule.
[0191] In the back propagation path, the transpose operator will be located before seq_gather bwd. At this time, the input shape of scatter is [e, capacity]. Figure 11a In the scatter operator, you can re-set the specified dimension by using dim=1 to re-specify capacity as the specified dimension. Figure 11b In the output of d-logtis-aux_loss, the shape is adjusted to [e,b×s], and in Figure 11c In the mean-bwd operator of the load balancing loss submodule in , reassign e to the specified dimension by dim=1.
[0192] The construction method of the routing reverse operator provided in the embodiment of the present application is aimed at the Router reverse operator of the mixed expert mode in the large model training, and an efficient implementation method based on an artificial intelligence chip is designed to help improve the computational efficiency of the Router reverse operator. According to the basic structure of the Router forward operator, the embodiment of the present application derives the calculation formula of its reverse operator through the forward calculation formula, and splits the reverse operator into multiple freely combinable sub-modules. After the multiple sub-modules are connected in the corresponding order, the function of the Router reverse operator can be efficiently realized when running on the artificial intelligence chip, and the frequent memory access to the global memory within the Router operator is reduced as much as possible, and the calculation and memory access efficiency within the Router operator is optimized. In addition, the embodiment of the present application optimizes the structure of the existing Router reverse operator to further improve the parallel computing efficiency of hardware resources. Compared with the existing solution, the solution provided by the embodiment of the present application can greatly improve the computational efficiency of back propagation and improve the training effect of the model when running on the artificial intelligence chip.
[0193] It should be understood that, although the various steps in the flowcharts involved in the various embodiments described above are displayed in sequence according to the instructions of the arrows, these steps are not necessarily executed in sequence in the order indicated by the arrows. Unless otherwise specified herein, there is no strict order restriction on the execution of these steps, and these steps can be executed in other orders. Moreover, at least a portion of the steps in the flowcharts involved in the various embodiments described above can include multiple steps or multiple stages, and these steps or stages are not necessarily executed and completed at the same time, but can be executed at different times, and the execution order of these steps or stages is not necessarily to be carried out in sequence, but can be executed in turn or alternately with other steps or at least a portion of steps or stages in other steps.
[0194] In addition, it should be noted that the method for constructing a route reversal operator provided in the embodiment of the present application can be applied at least to fields such as speech processing, image processing, text processing, and video processing.
[0195] For example, in the field of speech processing, a MOE model can be trained to obtain a speech processing model; in the field of image processing, a MOE model can be trained to obtain an image processing model; in the field of text processing, a MOE model can be trained to obtain a text processing model; and in the field of video processing, a MOE model can be trained to obtain a video processing model. The method for constructing a routing reverse operator provided in the embodiments of the present application can be applied to the training process of speech processing models, text processing models, image processing models, and video processing models, accelerating the training efficiency of the models while improving the model accuracy, which can effectively improve the text processing accuracy, image processing accuracy, video processing accuracy, and speech processing accuracy.
[0196] In order to enable those skilled in the art to better understand the embodiments of the present application, the following will take the application of the MOE model to large-scale text classification tasks as an example to illustrate the embodiments of the present application. The application processes in the fields of speech processing, image processing, and video processing can refer to this example and will not be repeated in the embodiments of the present application.
[0197] In this example, during the data preprocessing phase, data collection, cleaning, and vectorization can be performed, where:
[0198] Data Collection: Collect large amounts of text data covering a wide range of categories, such as news articles, which may include news from different categories such as economics and culture. Text Cleaning: Remove noise from the text, such as special characters and stop words, to improve data quality. Text Vectorization: Convert the cleaned text into a computer-processable vector form. Common methods include the bag-of-words model, word frequency-inverse document frequency, or using pre-trained word vector models to map each word to a low-dimensional vector. For long texts, deep learning models can also be used to obtain contextual representations of the text.
[0199] The MOE model process includes an input layer, a gating network (a routing operator in the embodiment of the present application, for clarity of description, the routing operator will be used instead of the gating network in the following description), an expert network, and an output layer, wherein:
[0200] Input layer: The preprocessed text vector is used as the input of the MOE model. The shape may be [batch_size, sequence_length, embedding_dim], where batch_size is the number of samples processed at a time, sequence_length is the length of the text sequence, and embedding_dim is the vector dimension.
[0201] The routing operator is responsible for assigning input text to different experts. It receives the input vector and calculates the "fitness" (probability) of each expert for that input through linear transformations and activation functions (such as Softmax). For example, for a given input text, the gating network might calculate that Expert 1 has a fit probability of 0.2, Expert 2 has a fit probability of 0.3, and so on. To balance the expert load and prevent overuse of certain experts, the routing operator also incorporates load balancing and z-loss loss terms.
[0202] Expert Network: Multiple expert networks operate in parallel, each specializing in the classification of a specific type of text. Each expert receives a portion of the input data assigned by the gating network (weighted by the probability of adaptation), and performs feature extraction and classification on it. For example, Expert 1 may specialize in economic news, while Expert 2 specializes in cultural news. The expert network can be a simple multilayer perceptron (MLP) or a complex convolutional neural network or recurrent neural network.
[0203] Output layer: The outputs (classification results) of each expert network are weighted and summed according to the probabilities of the gating network to obtain the final classification result. For example, if expert 1 outputs a probability of 0.8 for the classification "economy" and expert 2 outputs a probability of 0.6 for the classification "culture," the gating network assigns a probability of 0.3 to expert 1 and a probability of 0.7 to expert 2. The final classification result is the weighted sum of the two expert results. After applying the Softmax function, the final probability distribution of each category is obtained.
[0204] During training, the cross-entropy loss function is used to measure the difference between the model's predictions and the true labels. Stochastic gradient descent is used to update the parameters of the routing operator and the expert network to minimize the loss function. During training, the backpropagation algorithm calculates gradients and adjusts model parameters to achieve more accurate classification of various text types.
[0205] During the reverse propagation process of the routing operator, in the embodiment of the present application, based on the functional modules adopted by the routing operator, the corresponding reverse submodule is called from the operator library to construct a routing reverse operator. The reverse submodule is an operator module for gradient calculation that is constructed based on the algorithm formula obtained after reverse deduction based on the algorithm formula of each functional module. The specific derivation process and the process of constructing the routing reverse operator based on the reverse submodule can refer to the relevant description of the aforementioned embodiment, which will not be repeated here in the embodiment of the present disclosure. Each reverse submodule of the routing reverse operator is mapped to multiple sub-kernels of the GPGPU for gradient calculation, thereby updating the parameters of the routing operator by gradient descent to minimize the loss function.
[0206] It should be noted that artificial intelligence chips have been widely used in deep learning model training. Artificial intelligence chips can include GPU (Graphics Processing Unit), GPGPU (General-Purpose Computing on Graphics Processing Units), TPU (Tensor Processing Unit), NPU (Neural Network Processing Unit), etc. Taking GPGPU as an example, Figure 12 A schematic structural diagram of a general-purpose graphics processing unit (GPGPU) is shown.
[0207] like Figure 12 As shown, a general-purpose graphics processor is actually an array of streaming processor clusters (SPC), including Figure 12 Stream processor clusters 1, ..., and M are shown, where M is a positive integer greater than 1. In a graphics processor, one stream processor cluster processes one computing task, or multiple stream processor clusters process one computing task. Multiple stream processor clusters share data through a global cache or global memory.
[0208] like Figure 12 As shown, taking stream processor cluster 1 as an example, a stream processor cluster includes multiple computing units (also called), such as Figure 12In the CU, there are CU 1, CU 2, ..., CU N, where N is a positive integer. Each CU is used to perform arithmetic and logical operations other than matrix multiplication, convolution, and other matrix operations, such as accumulation, reduction, conventional addition, subtraction, multiplication, and division. A CU includes multiple cores (also called computing cores or computing cores), each of which includes an arithmetic logic unit (ALU), a floating-point computing unit, etc. The computing core is used to perform specific computing tasks. In addition, the CU also includes registers (e.g. Figure 12 The register file in the computing unit and the shared cache are used to hierarchically store source data and destination data related to computing tasks. The shared cache in a computing unit is used to share data between the cores of the computing unit.
[0209] In parallel computing, computing tasks are generally executed by multiple threads. These threads are divided into multiple thread blocks before being executed in a general-purpose graphics processor (or parallel computing processor), and then distributed through the thread block distribution module ( Figure 12 (Not shown) Multiple thread blocks are distributed to various CUs. All threads in a thread block must be assigned to the same CU for execution. Thread blocks are also split into minimum execution warps (or simply warps). Each warp contains a fixed number (or less) of threads, for example, 32 threads. Multiple thread blocks can execute in the same CU or in different CUs.
[0210] In each computational unit, the warp scheduling / dispatching module ( Figure 12 (not shown) schedules and allocates thread warps so that the multiple compute cores of the compute unit can execute the warps. Depending on the number of compute cores in the compute unit, multiple warps in a thread block can execute simultaneously or in a time-sharing manner. Multiple threads in each warp execute the same instruction. Memory execution instructions are emitted to the shared cache within the compute unit or further emitted to the mid-level cache, global cache, or global memory for read and write operations.
[0211] In the embodiment of the present application, the parallel computing capability of GPGPU can be utilized to complete the training of the MOE model on GPGPU. This includes storing the input data, intermediate results, and output data related to the Router operator during the forward propagation process in the global memory of the GPGPU. These data need to be correctly loaded into the appropriate memory location before the reverse propagation begins for subsequent calculations. For example, the input feature tensor, the gating signal generated by the Router operator, and the output tensor after the Router processing, etc.
[0212] To better utilize the parallelism of GPGPUs, data is typically divided into multiple small blocks. Each block can be processed by one or more thread blocks. This allows gradient calculations to be performed simultaneously on different thread blocks, improving computational efficiency.
[0213] Based on the configuration of the Router operator, the reverse submodules included in the routing reverse operator can be determined, and each reverse submodule can be called from the operator library. The routing reverse operator corresponding to the Router operator is constructed according to the dependency connection, and each reverse submodule in the routing reverse operator is mapped to its corresponding subkernel. The subkernel is executed in each computing unit of the GPGPU to implement backpropagation calculation, so that each computing unit on the GPGPU is responsible for calculating the gradient of a portion of the data. For example, for operations such as matrix multiplication and activation functions involved in the Router operator, the corresponding reverse calculation kernel functions are implemented respectively. These kernel functions will calculate the gradient of the input of the current layer (Router operator layer) relative to the loss function based on the input gradient information (the gradient transmitted backward from the subsequent layer).
[0214] Because computations on GPGPUs are performed in parallel, different thread blocks may complete computations at different speeds. Synchronization is required to ensure that all thread blocks have completed their current computations before proceeding to the next step (such as gradient aggregation or data transfer).
[0215] After each computational unit calculates the local gradient, it needs to aggregate these local gradients to obtain the final gradient value. The aggregated gradient value will be used as the backpropagation result of the router operator and passed to the previous layer for further calculation.
[0216] Through the above steps, the back propagation of the Router operator is implemented on GPGPU, and the parallel computing capability of GPGPU is used to efficiently calculate the gradient, thereby improving the training efficiency of the MOE model.
[0217] Based on the same inventive concept, an embodiment of the present application further provides a device for constructing a routing reversal operator for implementing the aforementioned method for constructing a routing reversal operator. The solution provided by this device is similar to the solution described in the aforementioned method. Therefore, the specific limitations of the embodiments of the device for constructing one or more routing reversal operators provided below can be found in the limitations of the method for constructing a routing reversal operator described above and will not be repeated here.
[0218] In an exemplary embodiment, Figure 13As shown, a device 1300 for constructing a route reverse operator is provided, comprising: a first determining module 1302 and a first constructing module 1304, wherein:
[0219] A first determining module 1302 is configured to determine a reverse submodule included in a route reverse operator, wherein the reverse submodule is configured to perform gradient calculation;
[0220] The first construction module 1304 is configured to call each of the reverse submodules from the operator library, connect the reverse submodules according to the dependency relationship, and map each of the reverse submodules to a corresponding sub-core in the artificial intelligence chip to construct a route reverse operator;
[0221] The reverse submodule is constructed based on the derived algorithm formula after deducing the algorithm formula corresponding to the forward submodule, and the forward submodule is a submodule obtained by dividing the routing forward operator.
[0222] The construction device of the above-mentioned routing reverse operator is obtained by constructing the routing reverse operator based on the reverse submodule in the operator library, and the reverse submodule is constructed based on the algorithm formula derived from the algorithm formula corresponding to the forward submodule. This specially derived reverse operator formula enables the reverse submodule to reduce certain transitional intermediate variables in the forward calculation through deduction. In this way, in the calculation of reverse propagation, there is no need to re-read and store this part of the data, which reduces the access to this data in the global memory. That is, when performing reverse propagation, there is no need to perform gradient calculation on the entire operator graph of the routing operator, which reduces unnecessary storage and reading of intermediate results, can greatly improve the utilization rate of hardware computing resources, improve the calculation efficiency of reverse propagation, and thus improve the training efficiency of the model.
[0223] In one embodiment, the apparatus further comprises:
[0224] A partitioning module is used to partition the routing forward operator into modules to obtain multiple forward sub-modules;
[0225] For any of the forward submodules, perform the following operations:
[0226] A second determining module is used to determine the algorithm formula of the reverse submodule corresponding to the forward submodule according to the algorithm formula of the forward submodule;
[0227] The second construction module is used to determine the operators included in the reverse submodule and the dependencies between the operators based on the algorithm formula of the reverse submodule, and construct the reverse submodule according to the operators and the dependencies between the operators.
[0228] In one embodiment, the input tensors, intermediate tensors and output tensors of the routing reverse operator are all stored in the global memory, and the internal tensors of the reverse submodule are stored in the designated memory, which is the memory corresponding to the sub-kernel that executes the reverse submodule.
[0229] In one embodiment, the first determining module is specifically configured to:
[0230] Obtain configuration parameters of the route reverse operator, and determine candidate reverse submodules whose configuration attributes in the configuration parameters are first attribute values;
[0231] The main path loss sub-template and each of the candidate reverse sub-modules are determined as the reverse sub-modules of the route reverse operator.
[0232] In one embodiment, the apparatus further comprises:
[0233] a fourth determining module, configured to determine a position adjustment strategy for the transposition operator in the route reversal operator based on the computing resources of the artificial intelligence chip;
[0234] An optimization module is used to adjust the position of the transposition operator in the routing reverse operator based on the position adjustment strategy, and re-specify the specified dimension of the first target operator in the reverse submodule according to the position adjustment strategy.
[0235] In one embodiment, the optimization module is specifically configured to:
[0236] Determining a first hardware parallelism rate of the artificial intelligence chip based on the unspecified dimension before adjustment and the computing resources of the artificial intelligence chip;
[0237] After swapping the designated dimension and the non-designated dimension, determining a second hardware parallelism rate of the artificial intelligence chip based on the swapped non-designated dimension and computing resources of the artificial intelligence chip;
[0238] When the first hardware parallelism rate is less than the second hardware parallelism rate, a target position of the transposition operator in the route reversal operator is determined, and a position adjustment strategy is generated based on the target position of the transposition operator in the route reversal operator.
[0239] In one embodiment, the optimization module is specifically configured to:
[0240] Determining a second target operator from the route reverse operator, where the second target operator is an operator whose operation is dimension-independent;
[0241] Based on the position of the second target operator, a target position of the transpose operator in the route reversal operator is determined.
[0242] In one embodiment, the partitioning module is specifically configured to:
[0243] Get the routing forward operator graph corresponding to the routing forward operator;
[0244] Dividing the routing forward operator graph into submodules to obtain multiple operator graphs;
[0245] According to the multiple operator graphs, multiple forward submodules are obtained.
[0246] Each module in the aforementioned device for constructing a route reversal operator may be implemented in whole or in part through software, hardware, or a combination thereof. Each module may be embedded in or independent of a processor in a computer device in hardware form, or may be stored in a memory in a computer device in software form, so that the processor can call and execute the corresponding operations of each module.
[0247] In an exemplary embodiment, a computer device is provided. The computer device may be a terminal, and its internal structure diagram may be as shown in FIG. Figure 14 As shown. The computer device includes a processor, memory, an input / output interface, a communication interface, a display unit, and an input device. The processor, memory, and input / output interface are connected via a system bus, and the communication interface, display unit, and input device are connected to the system bus via the input / output interface. The processor of the computer device is used to provide computing and control capabilities. The memory of the computer device includes a non-volatile storage medium and internal memory. The non-volatile storage medium stores an operating system and a computer program. The internal memory provides an environment for the operation of the operating system and computer program in the non-volatile storage medium. The input / output interface of the computer device is used to exchange information between the processor and external devices. The communication interface of the computer device is used to communicate with external terminals via wired or wireless means, and the wireless means can be implemented via Wi-Fi, a mobile cellular network, near-field communication (NFC), or other technologies. When executed by the processor, the computer program implements a method for constructing a route reversal operator. The display unit of the computer device is used to form a visually visible image, and can be a display screen, a projection device, or a virtual reality imaging device. The display screen can be a liquid crystal display screen or an electronic ink display screen, and the input device of the computer device can be a touch layer covering the display screen, or a button, trackball or touchpad set on the computer device casing, or an external keyboard, touchpad or mouse.
[0248] Those skilled in the art will understand that Figure 14 The structure shown in the figure is only a block diagram of a part of the structure related to the present application scheme, and does not constitute a limitation on the computer device to which the present application scheme is applied. The specific computer device may include Figure 14 More or fewer components may be shown, or some components may be combined, or the components may be arranged differently.
[0249] In one embodiment, a computer device is further provided, including a memory and a processor. The memory stores a computer program, and the processor implements the steps in the above method embodiments when executing the computer program.
[0250] In one embodiment, a computer-readable storage medium is provided, on which a computer program is stored. When the computer program is executed by a processor, the steps in the above-mentioned method embodiments are implemented.
[0251] In one embodiment, a computer program product is provided, including a computer program, which implements the steps in the above method embodiments when executed by a processor.
[0252] It should be noted that the user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data used for analysis, stored data, displayed data, etc.) involved in this application are all information and data authorized by the user or fully authorized by all parties, and the collection, use and processing of relevant data must comply with relevant regulations.
[0253] Those skilled in the art will understand that all or part of the processes in the above-mentioned embodiments can be implemented by instructing the relevant hardware through a computer program. The computer program can be stored in a non-volatile computer-readable storage medium. When the computer program is executed, it can include the processes of the embodiments of the above-mentioned methods. In particular, any reference to memory, database, or other media used in the embodiments provided in this application can include at least one of non-volatile memory and volatile memory. Non-volatile memory can include read-only memory (ROM), magnetic tape, floppy disk, flash memory, optical memory, high-density embedded non-volatile memory, resistive random access memory (ReRAM), magnetic random access memory (MRAM), ferroelectric random access memory (FRAM), phase change memory (PCM), graphene memory, etc. Volatile memory can include random access memory (RAM) or external cache memory, etc. By way of illustration and not limitation, RAM can take various forms, such as static random access memory (SRAM) or dynamic random access memory (DRAM). The databases involved in the various embodiments provided herein may include at least one of a relational database and a non-relational database. Non-relational databases may include, but are not limited to, blockchain-based distributed databases. The processors involved in the various embodiments provided herein may be, but are not limited to, general-purpose processors, central processing units (CPUs), graphics processing units (GPUs), digital signal processors (DSPs), programmable logic devices (PLDs), quantum computing-based data processing logic devices, artificial intelligence (AI) processors, and the like.
[0254] The technical features of the above embodiments can be combined arbitrarily. In order to make the description concise, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this application.
[0255] The above-described embodiments merely represent several implementation methods of the present application. While the descriptions are relatively specific and detailed, they should not be construed as limiting the scope of the present application. It should be noted that a person of ordinary skill in the art may make various modifications and improvements without departing from the spirit of the present application, and these modifications and improvements fall within the scope of protection of the present application. Therefore, the scope of protection of the present application shall be determined by the appended claims.
Claims
1. A method for constructing a route reverse operator, characterized in that: The method comprises: Determine a reverse submodule included in a route reverse operator, wherein the reverse submodule is used to perform gradient calculation; Calling each of the reverse submodules from the operator library, connecting each of the reverse submodules according to the dependency relationship, and mapping each of the reverse submodules to the corresponding sub-core in the artificial intelligence chip to construct a routing reverse operator; Among them, the reverse submodule is constructed based on the reverse algorithm formula obtained by reverse deduction after reverse deduction of the forward algorithm formula corresponding to the forward submodule. The forward submodule is a submodule obtained after dividing the routing forward operator. The forward submodule includes a main path loss submodule, a load balancing loss submodule and a z-loss loss submodule. The reverse submodule includes a main path loss submodule. The forward submodule corresponds to the forward algorithm formula, and the reverse submodule corresponds to the reverse algorithm formula. The process of constructing the reverse submodule based on the reverse algorithm formula obtained by reverse deduction includes: For any reverse submodule, based on the reverse algorithm formula of the reverse submodule, the operators included in the reverse submodule and the dependencies between the operators are determined, and the reverse submodule is constructed according to the operators and the dependencies between the operators.
2. The method according to claim 1, characterized in that The method further comprises: Divide the routing forward operator into modules to obtain multiple forward sub-modules; For any of the forward submodules, perform the following operations: According to the forward algorithm formula of the forward submodule, a reverse algorithm formula of the reverse submodule corresponding to the forward submodule is determined.
3. The method according to claim 1 or 2, characterized in that The input tensors, intermediate tensors, and output tensors of the routing reverse operator are all stored in the global memory, and the internal tensors of the reverse submodule are stored in the designated memory, which is the memory corresponding to the sub-kernel that executes the reverse submodule.
4. The method according to claim 1 or 2, characterized in that The reverse submodule included in the route reverse operator includes: Obtain configuration parameters of the route reverse operator, and determine candidate reverse submodules whose configuration attributes in the configuration parameters are first attribute values; The main path loss sub-template and each of the candidate reverse sub-modules are determined as the reverse sub-modules of the route reverse operator.
5. The method according to claim 1 or 2, characterized in that Before mapping each of the reverse submodules to a corresponding sub-core in the artificial intelligence chip to construct a route reverse operator, the method further includes: Determining a position adjustment strategy for the transposition operator in the route reversal operator based on the computing resources of the artificial intelligence chip; The position of the transposition operator in the routing reverse operator is adjusted based on the position adjustment strategy, and the designated dimension of the first target operator in the reverse submodule is re-designated according to the position adjustment strategy.
6. The method according to claim 5, characterized in that The determining, based on the computing resources of the artificial intelligence chip, a position adjustment strategy of the transposition operator in the route reversal operator includes: Determining a first hardware parallelism rate of the artificial intelligence chip based on the unspecified dimension before adjustment and the computing resources of the artificial intelligence chip; After swapping the designated dimension and the non-designated dimension, determining a second hardware parallelism rate of the artificial intelligence chip based on the swapped non-designated dimension and computing resources of the artificial intelligence chip; When the first hardware parallelism rate is less than the second hardware parallelism rate, a target position of the transposition operator in the route reversal operator is determined, and a position adjustment strategy is generated based on the target position of the transposition operator in the route reversal operator.
7. The method according to claim 6, characterized in that Determining a target position of a transposition operator in the route reversal operator includes: Determining a second target operator from the route reverse operator, where the second target operator is an operator whose operation is dimension-independent; Based on the position of the second target operator, a target position of the transpose operator in the route reversal operator is determined.
8. The method according to claim 4, characterized in that The routing forward operator is divided into modules to obtain multiple forward submodules, including: Get the routing forward operator graph corresponding to the routing forward operator; Dividing the routing forward operator graph into submodules to obtain multiple operator graphs; According to the multiple operator graphs, multiple forward submodules are obtained.
9. A computer device comprising a memory and a processor, wherein the memory stores a computer program, wherein: When the processor executes the computer program, the steps of the method according to any one of claims 1 to 8 are implemented.
10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 8 are implemented.
11. A computer program product comprising a computer program, characterized in that When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 8 are implemented.
Citation Information
Patent Citations
Static generation method and device of reverse calculation graph, equipment and medium
CN117273115A