Model training method, apparatus, system, and device

Abstract

The present disclosure relates to a model training method, device, system and equipment. In the execution of at least part of the forward process of the Nth batch of training samples, the dynamically generated calculation graph is divided into multiple sub-forward processes, and multiple sub-backward processes corresponding to the multiple sub-forward processes are obtained; after the forward process of the Nth batch of training samples is completed, and in the execution of at least part of the forward process of the Pth batch of training samples, at least one time period is used to execute one of the sub-forward processes in at least part of the forward process of the Pth batch of training samples and one of the multiple sub-backward processes of the Nth batch of training samples in parallel by using the calculation unit, and the sub-forward process and the sub-backward process executed in parallel, one of which belongs to the calculation operation, and the other belongs to the communication operation. Therefore, at least part of the communication cost overhead can be hidden in the model training process.

Description

Technical Field

[0001] This disclosure relates to the field of artificial intelligence, and in particular to a model training method, apparatus, system and device. Background Technology

[0002] Traditional AI infrastructure (AI) training frameworks typically employ parallelism methods for training models (such as hybrid expert models), including pipelined parallelism, tensor parallelism, and expert parallelism.

[0003] Pipeline parallelism divides the model into layers. During computation, the results from one layer are transmitted to the next layer via peer-to-peer (P2P) communication. Tensor parallelism and expert parallelism are similar strategies, both dividing the model within a layer and distributing it across different units (such as processes). They coordinate parallel computation between multiple units through interleaved communication scheduling, reducing the memory usage within individual units. While these parallelism methods can divide the model for large-scale parallel computation on a cluster, they all require communication between units as a cost.

[0004] Therefore, reducing communication overhead during model training is a problem that urgently needs to be solved. Summary of the Invention

[0005] One objective of this disclosure is to reduce communication overhead during model training.

[0006] According to a first aspect of this disclosure, a model training method is provided, comprising: during the execution of at least a partial forward process for a batch N of training samples, dividing a dynamically generated computation graph into multiple sub-forward processes and obtaining multiple sub-backward processes corresponding to the multiple sub-forward processes; after the execution of the forward process for the batch N of training samples is completed, and during the execution of the at least partial forward process for the batch P of training samples, using a computing unit to execute in parallel one sub-forward process of the at least partial forward process for the batch P of training samples and one sub-backward process of the multiple sub-backward processes for the batch N of training samples, wherein, during at least a time period, one of the sub-forward process and the sub-backward process executed in parallel is a computational operation and the other is a communication operation.

[0007] Optionally, the dynamically generated computation graph is divided into multiple sub-forward processes, including: identifying the first node in the computation graph that does not need to compute gradients based on the attribute information of each node in the computation graph; dividing the computation graph with the first node as the dividing position to obtain multiple sub-forward processes.

[0008] Optionally, dividing the computation graph with the first node as the dividing point includes: creating a new output tensor for each output tensor of the first node; and using the new output tensor as the input of the next sub-forward process.

[0009] Optionally, the method further includes: for the sub-forward process, creating a first pseudo-tensor of length zero and a second pseudo-tensor of length zero; using the first pseudo-tensor as the output of the sub-forward process and concatenating it with the output tensor of the sub-forward process; and using the second pseudo-tensor as the input of the next sub-forward process and concatenating it with the new output tensor.

[0010] Optionally, the model is a hybrid expert model, and the forward process of the hybrid expert model includes a routing phase, a distribution phase, a multilayer perceptron (MLP) phase, and an aggregation phase. The method further includes: in the forward process, not saving the forward calculation results of the MLP phase, and in the reverse process, recalculating the MLP parts other than the last layer in the MLP phase.

[0011] Optionally, the method further includes: not saving the forward calculation results of the distribution stage during the forward process, and recalculating the distribution stage during the reverse process.

[0012] According to a second aspect of this disclosure, a model training apparatus is also provided, comprising: a segmentation module, configured to segment a dynamically generated computation graph into multiple sub-forward processes during the execution of at least a portion of the forward process for the Nth batch of training samples, and to obtain multiple sub-backward processes corresponding to the multiple sub-forward processes; and a parallel execution module, configured to, after the execution of the forward process for the Nth batch of training samples is completed, and during the execution of the at least a portion of the forward process for the Pth batch of training samples, to execute in parallel for at least a time period one sub-forward process of the at least a portion of the forward process for the Pth batch of training samples and one sub-backward process of the multiple sub-backward processes for the Nth batch of training samples, using a computing unit, wherein one of the parallel-executed sub-forward process and the sub-backward process is a computational operation and the other is a communication operation.

[0013] According to a third aspect of this disclosure, a distributed training system for a model is also provided, comprising multiple computing units, wherein, during the execution of at least a portion of the forward process for the Nth batch of training samples, the computing units divide a dynamically generated computation graph into multiple sub-forward processes and obtain multiple sub-backward processes corresponding to the multiple sub-forward processes; after the execution of the forward process for the Nth batch of training samples is completed, and during the execution of the at least a portion of the forward process for the Pth batch of training samples, the computing units execute in parallel at least one sub-forward process of the at least a portion of the forward process for the Pth batch of training samples and one sub-backward process of the multiple sub-backward processes for the Nth batch of training samples for at least one time period, and the sub-forward process and the sub-backward process executed in parallel, wherein one of the sub-forward process and the sub-backward process is a computational operation and the other is a communication operation.

[0014] According to a fourth aspect of this disclosure, a computing device is provided, comprising: a processor; and a memory having executable code stored thereon, which, when executed by the processor, causes the processor to perform the method described in the first aspect above.

[0015] According to a fifth aspect of this disclosure, a computer program product is provided, including executable code that, when executed by a processor of an electronic device, causes the processor to perform the method described in the first aspect above.

[0016] According to a sixth aspect of this disclosure, a non-transitory machine-readable storage medium is provided, on which executable code is stored, which, when executed by a processor of an electronic device, causes the processor to perform the method described in the first aspect above.

[0017] In the execution of at least a portion of the forward pass for the Nth batch of training samples, this disclosure divides the dynamically generated computation graph into multiple sub-forward passes and obtains multiple sub-backward passes corresponding to these sub-forward passes. After the forward pass for the Nth batch of training samples is completed, and while executing at least a portion of the forward pass for the Pth batch of training samples, a sub-forward pass from the at least portion of the forward pass for the Pth batch of training samples and a sub-backward pass from the multiple sub-backward passes for the Nth batch of training samples are executed in parallel by a computing unit for at least one time period. Furthermore, of the parallel-executed sub-forward and sub-backward passes, one is a computational operation and the other is a communication operation. Therefore, at least part of the communication cost can be masked during model training, thereby improving the efficiency of model training and increasing the utilization of computing units (e.g., clusters). Attached Figure Description

[0018] The above and other objects, features and advantages of this disclosure will become more apparent from the more detailed description of exemplary embodiments thereof taken in conjunction with the accompanying drawings, in which the same reference numerals generally represent the same components.

[0019] Figure 1 The diagram illustrates two mainstream structural designs for large language models.

[0020] Figure 2 A schematic flowchart of a model training method according to an embodiment of the present disclosure is shown.

[0021] Figure 3 The diagram illustrates the forward and backward processes during model training.

[0022] Figure 4A A schematic diagram of a computation graph is shown.

[0023] Figure 4B It shows the Figure 4A The diagram shows an effect after the computational graph has been segmented.

[0024] Figure 4C It shows the Figure 4B The diagram shown is a schematic of the optimized computation graph.

[0025] Figure 5 A flowchart illustrating the parallel execution of the forward and reverse processes is shown.

[0026] Figure 6 A schematic diagram of the recalculation part in a hybrid expert model is shown.

[0027] Figure 7 A schematic diagram of a model training apparatus according to an embodiment of the present disclosure is shown.

[0028] Figure 8 A schematic diagram of the structure of a computing device according to an embodiment of the present disclosure is shown. Detailed Implementation

[0029] Preferred embodiments of the present disclosure will now be described in more detail with reference to the accompanying drawings. While preferred embodiments of the present disclosure are shown in the drawings, it should be understood that the present disclosure may be implemented in various forms and should not be limited to the embodiments set forth herein. Rather, these embodiments are provided so that the present disclosure will be thorough and complete, and will fully convey the scope of the disclosure to those skilled in the art.

[0030] The following description uses a Mixture of Experts (MoE) model as an example. It should be understood that this disclosure can be used for communication overhead generated during the training of various model-splitting-based models, and is not limited to Mixture of Experts models.

[0031] This disclosure does not impose special requirements on the model splitting method (i.e., model parallelism technique). It can be either pipeline parallelism (PP) or tensor parallelism (TP). Pipeline parallelism refers to horizontally dividing the model into multiple sub-modules by layer or stage, with each computational unit responsible for running one of these sub-modules. Tensor parallelism refers to vertically splitting the parameter matrix of a single layer (by row or column), with each device storing a portion of the parameters and performing calculations with the same layer structure but different parameters. In tensor parallelism, the model structure remains unchanged, but each parameter is only 1 / tp_size (number of parallel devices), and each computational unit runs the entire process, requiring synchronization of results during execution.

[0032] Artificial intelligence is now widely used, with large language model services integrated into all aspects of life, such as daily searches and online shopping, which are accompanied by large model AI assistants.

[0033] Currently, the mainstream Large Language Model (LLM) architectures fall into two categories: dense models and hybrid expert models. For solving a complex task involving knowledge from multiple domains, the simplest approach is to bring together experts from those domains to collaboratively address the problem, then aggregate the results to obtain a high-quality answer. Hybrid expert models are designed based on this approach. The biggest structural difference between hybrid expert models and dense models is the use of a gating network and a multilayer perceptron (MLP) that incorporates multiple specialized sub-model structures.

[0034] Figure 1 The diagram illustrates two mainstream structural designs for large language models.

[0035] Figure 1 The left side shows the basic structure of a large language model, including a self-attention layer, a first normalization layer, a feedforward neural network (FFN), and a second normalization layer. Figure 1The right side shows schematic diagrams of feedforward neural networks using dense model structures and hybrid expert model structures, respectively. Figure 1 As shown on the right, the dense model includes only one FFN, while the hybrid expert model includes multiple FFNs, where each FFN can be regarded as an expert.

[0036] Gated networks use specific routing algorithms to assign tasks to specific experts, who then handle domain-specific problems within the multilayer perceptron. Therefore, a significant advantage of hybrid expert models over dense models is their ability to be pre-trained more efficiently with fewer resources, resulting in better model performance. However, the unique structure of hybrid expert models also makes them difficult to train effectively using traditional AI infrastructure training frameworks. Therefore, efficiently training hybrid expert models on large-scale clusters remains a hot topic in the industry.

[0037] Currently, one of the main challenges for traditional AI Infra training frameworks in efficiently training on large-scale clusters is how to fully mask the overhead of communication costs and improve the overall training efficiency of the cluster.

[0038] To enable hybrid expert models to leverage the computational performance of large-scale clusters, expert parallelism is employed, storing multiple expert models across different units within the cluster. Each time a model passes through a gating network, an allgather (summarizing distributed data across all nodes) or alltoall (each participating node sends data to and receives data from all other nodes) aggregation communication operation is required. This operation sends the current task to the unit where the expert model resides, as specified by the gating circuit. This process generates a large number of aggregation communication operations during training. This aggregation communication overhead increases with the scale of expert parallelism; for current mainstream large-scale hybrid expert model training, communication overhead accounts for 30%-50% of the overall training time. Therefore, aggregation communication often becomes the bottleneck for training hybrid expert models on large-scale clusters.

[0039] To reduce the overhead of collection communication, we first need to solve an implementation problem.

[0040] The following explanation uses PyTorch as an example. It should be understood that the technical solution disclosed herein can also be applied to other similar deep learning frameworks. Traditional AI infrastructure frameworks are generally based on PyTorch and involve secondary development. Within the framework, the forward and backward propagation of each component needs to be implemented, and the order between components needs to be adjusted and handed over to the underlying PyTorch. PyTorch then constructs a dynamic computation graph to ensure the normal training of the model and the completion of forward and backward computations.

[0041] To ensure computational accuracy, PyTorch generates a computation graph that only generates the corresponding backward computation process after each process has completed its forward computation according to a predetermined flow, without allowing for additional control operations. Therefore, controlling the execution timing of the current training phase during model training within the existing framework is extremely difficult.

[0042] To address this, this disclosure proposes an interrupt switching strategy. "Interruption" in the interrupt switching strategy refers to dividing the dynamically generated computation graph during a single execution of the forward process by the computation unit into multiple sub-forward processes and corresponding multiple sub-reverse processes. "Switching" in the interrupt switching strategy refers to the computation unit executing a sub-forward process from a later training batch and a sub-reverse process from an earlier training batch in parallel. Furthermore, in the parallel execution of the sub-forward and sub-reverse processes, one involves computation, and the other involves communication. This can mask the communication overhead during model training, thereby improving the efficiency of model training and increasing the utilization of computation units (such as clusters).

[0043] Figure 2 A schematic flowchart of a model training method according to an embodiment of the present disclosure is shown.

[0044] Figure 2 The method shown can be executed by any computing unit in the cluster used for distributed training of the model. The computing unit in the cluster can also be called a computing node, worker node, or cluster node. A computing unit can be a physical or virtual independent computing unit. For example, a computing unit can be a GPU (Graphics Processing Unit). Another example is a process. It should be understood that due to the large size of the training samples, the entire training sample is usually divided into multiple batches during training. Training samples can be one or more combinations of text data, image data, audio data, or video data. Before being used for training, these data can be appropriately preprocessed to transform them into multiple tokens suitable for training. Preprocessing operations include, but are not limited to, data cleaning, standardization, or structuring.

[0045] See Figure 2 In step S210, during the execution of at least part of the forward process for the Nth (N is a positive integer) batch of training samples, the dynamically generated computation graph is divided into multiple sub-forward processes, and multiple sub-backward processes corresponding to the multiple sub-forward processes are obtained.

[0046] This implementation first employs a method of segmenting the dynamically generated computation graph in PyTorch to generate multiple sub-forward processes. These sub-forward processes can include multiple computation stages (e.g., attention computation stages) and communication stages (e.g., dispatch communication stages). In a forward process, after completing a sub-forward process, the reverse input of that sub-forward process and the computation graph can be saved.

[0047] Figure 3 The diagram illustrates the forward propagation and backward propagation processes during model training. It should be understood that forward propagation is the process of calculating the model's predicted values. Given input data, a series of calculations are performed through the layers and activation functions in the model to obtain the model's output. Backpropagation is the process of calculating the gradient of the loss function with respect to the model parameters. Specifically, it involves calculating the gradient of the loss function with respect to the parameters of each layer and then passing this gradient back to the optimizer for updates. Through forward propagation, the model's predicted values ​​can be calculated; through backpropagation, the gradient of the loss function with respect to the model parameters can be calculated and updated using the optimizer. By repeatedly performing these two processes, the model's performance can be progressively optimized, thereby achieving the training of a deep learning model.

[0048] Figure 3 The middle section shows the forward pass during model training. Figure 3 The right side shows the automatically generated backward computation stage by PyTorch. To better control the timing of computation and communication stages during model training, the original PyTorch computation graph can be directly split. Since the computation graph is destroyed, no corresponding backward computation stage will be generated. Therefore, the forward computation stage can be saved. For example, as shown... Figure 3 As shown on the left, after a certain sub-forward process (i.e., communication or computation at a certain stage in the forward process) is completed, the corresponding sub-reverse process can be saved in the form of a stack. The reverse processes are then executed sequentially according to their execution timing after the entire forward process is completed.

[0049] During model training, PyTorch uses its Autograd (Automatic Differentiation Engine) to automatically generate the corresponding Autograd computation graph for the forward process. This allows for automatic differentiation during the backward process, ensuring proper model training. PyTorch's Autograd computation graph uses a hybrid execution technique to record and generate the forward computation flowchart, which in turn generates the backward computation graph. However, for use cases where the forward and backward processes execute in parallel and computation and communication overlap, it's necessary to divide a single forward process into multiple sub-forward processes and obtain corresponding sub-backward processes. PyTorch itself doesn't support precise splitting of forward and backward processes, and manual splitting is inconvenient, prone to graph splitting errors and difficult to debug, and results in increased GPU memory usage due to multiple storage operations for backward results.

[0050] In view of this, this disclosure proposes a method for automatically segmenting computational graphs. An example of the automatic segmentation principle is as follows.

[0051] In PyTorch, the computation graph is a directed acyclic graph (DAG) consisting of nodes and edges. Nodes in the computation graph represent tensors or functions, and edges represent dependencies between tensors and functions. Each tensor object contains a pointer `grad_fn` to the node that generated its function; `grad_fn` represents which function operation generated the tensor. In PyTorch, `requires_grad` is a very important attribute used to identify whether a tensor needs to retain gradient information during computation, i.e., whether gradients need to be computed. When training deep learning models using PyTorch, it is usually necessary to set `requires_grad = True` for trainable parameters (i.e., trainable tensors) in the model so that gradients can be computed during backpropagation. `requires_grad` is a attribute specific to tensor nodes; function nodes do not need to compute gradients and do not have the `requires_grad` attribute.

[0052] Therefore, based on the attribute information of each node in the computation graph, the first node in the computation graph that does not require gradient calculation can be identified. The computation graph can then be segmented using this first node as the splitting point to obtain multiple sub-forward processes. The first node can be, for example, a Function node that does not have the `requires_grad` attribute. In this disclosure, the operation represented by a Function node can include computational operations and communication operations. For example, nodes corresponding to collective communication operations such as dispatch and combine in the training process of a hybrid expert model can all be considered Function nodes. Segmenting the computation graph using the first node as the splitting point can mean splitting the graph using the output of the first node (i.e., the Tensor nodes connected to the first node) as the breakpoint, cutting off the edges between the output of the first node and subsequent nodes. Through this segmentation, multiple sub-forward processes can be obtained, each corresponding to an execution stage. The Function involved in this execution stage can be either a computational operation or a communication operation.

[0053] In PyTorch, each Function node contains two methods: `forward()` and `backward()`, used for calculating the loss during forward propagation and the gradient during backward propagation in the computation graph. During forward propagation, Autograd dynamically creates corresponding nodes by overloading various Tensor operations (such as addition, multiplication, and matrix multiplication), connecting them to form a Directed Acyclic Graph (DAG). At this point, the structure of the computation graph depends on the specific computation path. When `loss.backward()` is called, the Autograd engine traverses the computation graph backward from the endpoint to the starting point, executing the `backward()` function of each node in turn to calculate the gradient. The gradient is propagated using the chain rule, enabling automatic differentiation. Therefore, dividing the computation graph into multiple sub-forward processes is equivalent to automatically obtaining multiple sub-backward processes corresponding to these sub-forward processes.

[0054] The output tensor of each sub-forward process is used as the input tensor of the next sub-forward process. To maintain the integrity of the sub-forward processes, when splitting the computation graph, a new output tensor can be created for each output tensor of the first node. Each output tensor of the first node is the output of the sub-forward process corresponding to the first node, and the newly created output tensor is the input of the next sub-forward process.

[0055] Below is an example code for splitting a computational graph.

[0056] args_a_old:[tensor_1_old,int_2,tensor_nograd_3]

[0057] args_b_old:{'a':tensor_4_old,'b':float_5,'c':tensor_6_old}

[0058] args_a_new,args_b_new=break_graph(args_a_old,args_b_old)

[0059] args_a_new:[tensor_1_new,int_2,tensor_nograd_3]

[0060] args_b_new:{'a':tensor_4_new,'b':float_5,'c':tensor_6_new}

[0061] The code above can be described as follows: The `args_a_old` and `args_b_old` are disconnected because the preceding part is computation, while the subsequent part involves communication between these two `args`. `args_a` and `args_b` are output tensors, but their forms may be complex, such as lists or dictionaries. The objects within these lists or dictionaries need to be extracted. If a tensor has `requires_grad=True` (gradient computation is required), a new tensor needs to be created using a detach operation. This new tensor has no connection to the previous computation graph. Then, the new tensor is replaced with its original position within `args`, ensuring that the new `args_a_new` is completely disconnected from the previous graph.

[0062] Figure 4A A schematic diagram of a computation graph is shown. Figure 4A Nodes a and b in the code are tensors that correspond to args_a_old and args_b_old, respectively, which need to be cut off in the code above.

[0063] Figure 4B It shows the Figure 4A The diagram shows an effect after the computational graph has been segmented.

[0064] See Figure 4BNode a1 is equivalent to node a, i.e., args_a_old in the code above. Node b1 is equivalent to node b, i.e., args_b_old. Node a2 is the new tensor after performing a detach operation on node a, i.e., args_a_new in the code above. Node b2 is the new tensor after performing a detach operation on node b, i.e., args_b_new in the code above. The detach operation returns a new tensor that shares the same memory space as the original tensor, but is not tracked in the computation graph; that is, it does not participate in backpropagation. It should be noted that... Figure 4B This is merely a schematic diagram illustrating the logically expected effect.

[0065] To reduce storage overhead, this disclosure proposes creating a first pseudo-tensor and a second pseudo-tensor of zero length for each sub-forward process. The first pseudo-tensor is used as the output of the sub-forward process and connected to its output tensor; the second pseudo-tensor is used as the input of the next sub-forward process and connected to its new output tensor. Since there are no edges between the first and second pseudo-tensors, the computation graph can be broken, effectively segmenting it. Furthermore, the zero-length first pseudo-tensor can serve as the starting point for calling the corresponding sub-reverse process without storing the actual tensor, thus further reducing storage (graphical memory) consumption.

[0066] Figure 4C It shows the Figure 4B The diagram shown is a schematic of the optimized computation graph.

[0067] See Figure 4C fake_1 corresponds to the first pseudo-tensor mentioned above, and fake_2 corresponds to the second pseudo-tensor mentioned above. fake_1 serves as... Figure 4C The upper half represents the starting point of the sub-reverse process corresponding to the sub-forward process. fake_2 serves as... Figure 4C The lower half represents the endpoint of the sub-reverse process corresponding to the sub-forward process. It should be understood that... Figure 4C The fake_1 at the end of the first graph needs to be stored for reverse processing. The fake_2 in the second graph doesn't need to be stored; only the gradients (grad) of a2 and b2 during the reverse processing from a2 and b2 to fake_2 need to be stored. Then, when executing fake_1, the gradients of a2 and b2 need to be obtained from a2 and b2, and sent back along the edges from fake_1 to a1 and b1. This allows the process to start. Figure 4CThe upper half represents the reverse process corresponding to the sub-forward process. That is, the two tensor nodes a1 and b1 are connected to fake_1, allowing fake_1 to pass the gradients obtained from a2 and b2 to a1 and b1 based on this connection during the corresponding reverse process. Figure 4C In the lower half of the sub-forward process, fake_2 is connected to the tensor nodes a2 and b2 respectively. This allows a2 and b2 to pass their gradients to fake_2 based on this connection during the corresponding backward process. fake_2 then stores the gradients of a2 and b2 for use in the backward process from fake_1 to a1 and b1. There is no connection between fake_1 and fake_2, so the gradients stored by fake_2 can be passed to fake_1 manually or by another executor.

[0068] In step S220, after the forward process of the Nth batch of training samples is completed, and during the execution of at least part of the forward process of the Pth batch of training samples (P can be a positive integer), a sub-forward process of at least part of the forward process of the Pth batch of training samples and a sub-reverse process of multiple sub-reverse processes of the Nth batch of training samples are executed in parallel using a computing unit for at least one time period. Furthermore, one of the sub-forward process and the sub-reverse process executed in parallel is a computational operation and the other is a communication operation.

[0069] The "at least partial forward process" refers to the forward process configured for a computational unit to execute. During model training after configuration, the forward process and corresponding backward process executed by the same computational unit are determined. Furthermore, the same computational unit may execute only a portion of its configured forward process or the complete forward process, depending on the model parallelism employed. For example, in a pipelined parallel scenario, a single computational unit executes only a portion of the forward process, while in a tensor parallel scenario, a single computational unit may execute the entire forward process.

[0070] At least a portion of the forward pass of the Nth batch of training samples is executed first, followed by at least a portion of the forward pass of the Pth batch of training samples. There may be one, two, or more batches of training samples between the Nth and Pth batches. For example, P = N + 1.

[0071] The reverse process is essentially an automatic differentiation process based on the chain rule, meaning that different sub-reverse processes within the same training batch are executed in a specific order. Therefore, although sub-forward and sub-reverse processes from different batches are executed in parallel within the same computational unit, sub-reverse processes within the same batch are executed strictly in sequence, ensuring that the computational logic remains error-free.

[0072] For example, the at least part of the forward process handled by the computing unit can be divided into alternating computation and communication phases. Correspondingly, the reverse process corresponding to the partial forward process can also be divided into alternating communication and computation phases. The computing unit can execute one computation phase and one communication phase in parallel at each step, wherein the computation phase originates from the at least part of the forward process for the Pth batch of training samples, and the communication phase originates from the reverse process corresponding to the at least part of the forward process for the Nth batch of training samples. Although the execution times of the parallel computation phase and communication phase in each step may not be completely consistent, overall, complete overlap of computation and communication during training can be achieved. Furthermore, to address the issue that the execution times of the parallel computation phase and communication phase in each step may not be completely consistent, appropriate means (e.g., adjusting hyperparameters) can be employed to ensure that the execution times of the parallel computation phase and communication phase in each step are as consistent as possible.

[0073] It should be understood that the computing unit can also perform the computation graph segmentation operation mentioned above during the execution of at least part of the forward process of the Pth batch of training samples, so that the reverse process corresponding to the at least part of the forward process of the Pth batch of training samples can be executed in parallel with the at least part of the forward process of subsequent batches of training samples.

[0074] Figure 5 A flowchart illustrating the parallel execution of the forward and reverse processes is shown.

[0075] like Figure 5 As shown, the hybrid expert model can be broken down into five parts: the self-attention stage, the preprocessing stage before the MLP, the dispatch communication stage, the MLP computation stage, and the combine communication stage. Each of these stages has a forward process and a backward process. This can be decomposed into two models (i.e.,...) Figure 5 After using Model 1 and Model 2, both models will run simultaneously. It's important to understand that simultaneous running of the two models refers to running them simultaneously on different training batches (i.e., different data), not on the same training batch. For example, in the first time period, the forward process represented by Model 1 can be executed for data a, and the backward process represented by Model 2 can be executed for data b; in the subsequent second time period, the backward process represented by Model 2 can be executed for data a, and the forward process represented by Model 1 can be executed for data c; in the subsequent third time period, the forward process represented by Model 1 can be executed for data d, and the backward process represented by Model 2 can be executed for data c.

[0076] Taking the execution of the forward process represented by Model 1 for data a and the simultaneous execution of the backward process represented by Model 2 for data b as an example, the parallel operation of Model 1 and Model 2 can be divided into four steps. Step 1: When Model 1 (i.e., the forward process) runs, it first enters the attention and preprocessing computation stage before MLP, while Model 2 (i.e., the backward process) first enters the backward communication stage of combine. Step 2: After Model 1 completes the computation, it enters the forward communication stage of dispatch, while Model 2, after completing the communication, enters the backward computation stage of MLP. Step 3: After Model 1 completes the communication, it enters the forward computation stage of MLP, while Model 2 enters the backward dispatch communication stage. Step 4: After Model 1 completes the computation, it enters the forward communication stage of combine, while Model 2, after completing the communication, enters the preprocessing and attention computation stage of the backward MLP.

[0077] In each step of the above implementation process, the two models alternate between a computation phase and a communication phase, achieving complete overlap between computation and communication, thus improving the efficiency of utilizing the entire cluster's computational and communication resources. Although the computation time and communication time may not be completely consistent in each overlapping computation and communication phase, the complete overlap of computation and communication time can be achieved as much as possible by adjusting hyperparameters (number of expert parallelisms) and the number of resources allocated to communication and computation, thereby reducing the overall training time required and improving training efficiency and the utilization rate of computing units.

[0078] Thus, combined Figures 1 to 5 An example is provided on how to effectively mask communication overhead to improve the efficiency of model training.

[0079] Another challenge for traditional AI Infra training frameworks in efficiently training on large-scale clusters is how to balance the use of GPU memory and computational cost to improve the efficiency of cluster utilization.

[0080] With the rise of large-scale models, their size has become increasingly massive. Although GPUs are continuously releasing new models with ever-growing memory, the increase in GPU memory cannot keep pace with the growth rate of large model parameters. Excessively large model parameters necessitate the use of more GPUs for parallel computation, leading to additional communication overhead during training. The traditional approach is to use a recalculation strategy. During the forward pass of large model training, the calculated activation values ​​are saved and then used to calculate the gradients of the weights during the backward pass, resulting in significant memory fragmentation. Recalculation, on the other hand, does not save these activation values ​​but recalculates them during the backward pass. While this reduces computational efficiency, it saves memory and reduces communication overhead. Therefore, balancing memory usage and computational cost is one of the major challenges for AI Infra training frameworks on large-scale clusters.

[0081] For hybrid expert model scenarios, this disclosure models and analyzes the memory usage and computational cost among the components in the hybrid expert model, and proposes a more reasonable recomputation method to balance memory usage and computational cost.

[0082] Throughout the training process, the results of the forward computation are used in the gradient update operation on the weight matrix during the backward computation. Therefore, a large amount of forward computation results need to be saved before and after the computation, resulting in significant memory overhead. A common approach to solve this problem is to recalculate the model, without saving the forward computation results, and instead recalculate the forward computation during the backward computation to produce the required output. However, this approach leads to significant computational redundancy.

[0083] Figure 6 A schematic diagram of the recalculation part in a hybrid expert model is shown.

[0084] See Figure 6 The forward process of the hybrid expert model includes the routing phase, the distribution phase, the MLP phase (i.e., the dispatch phase), and the aggregation phase (i.e., the combine phase).

[0085] For the MLP layer of the hybrid expert model, the first step is to select which specific expert to be assigned to based on the gating network. Then, the data is sent to the corresponding process for computation through the dispatch phase of the aggregate communication. Each expert will go through three steps: matrix multiplication fc1 (first fully connected layer), activation function calculation, and matrix multiplication fc2 (second fully connected layer). Finally, the results are summarized through the combine phase of the aggregate communication.

[0086] Recompiling the entire model directly is equivalent to re-executing the entire forward pass, resulting in significant computational and communication costs. If the MLP is recomputed directly, it would lead to matrix multiplication by fc1 in the first fully connected layer, requiring the activation function to be preserved. Furthermore, two dispatch and combine phases are needed during the backward pass.

[0087] Considering the reverse process, the last layer in the MLP (corresponding to Figure 5 The output of the second fully connected layer in the MLP is useless for the ensemble communication combine phase. Therefore, for the last layer in the MLP, only the input tensor of that layer is needed for the gradient calculation. Thus, during the forward pass, the forward computation results of the MLP phase can be omitted, and during the backward pass, the MLP parts other than the last layer in the MLP phase are recomputed.

[0088] In other words, although the computation results of the MLP are not saved, when recompiling the MLP, it is considered that the computation result of the last layer in the MLP is meaningless. Therefore, only the part of the MLP excluding the last fully connected layer is recomputed. This reduces the recompiling overhead.

[0089] Generally, in the forward computation process, the dispatch phase is more time-consuming than the matrix multiplication fc1 and fc2 calculations. In addition, since the communication generates more GPU memory overhead, the results after the dispatch phase can be saved on each process, and the matrix multiplication fc1 and activation function can be recalculated to calculate the input for the gradient calculation of matrix multiplication fc2. The dispatch phase can be omitted in the backward process to reduce one communication operation. Although it will slightly increase the GPU memory consumption, it can reduce the amount of computation that needs to be recalculated and improve the overall computational efficiency of the training framework.

[0090] In some exemplary embodiments, the forward computation results of the distribution phase may not be saved during the forward process, and the distribution phase may be recomputed during the reverse process. That is, the distribution phase may also be recomputed.

[0091] It should be understood that whether to recompile the distribution stage can be determined by analyzing the relationship between the recompiling overhead of the distribution stage and the memory overhead of saving the calculation results of the distribution stage (i.e., the distribution results) based on the actual situation. When it is determined that the recompiling overhead of the distribution stage is less than or significantly less than the memory overhead, the distribution stage should be recomputed.

[0092] For example, the decision to recalculate for the distribution stage can be determined based on the number of experts selected in the MLP stage. For instance, if the number of experts selected in the MLP stage is greater than or equal to a first threshold, it indicates a large number of experts are selected, and each expert needs to undergo dispatch communication, resulting in significant memory overhead. In this case, recalculation for the distribution stage is warranted. Conversely, if the number of experts selected in the MLP stage is less than the first threshold, it indicates a smaller number of experts are selected, resulting in less memory overhead, and recalculation for the distribution stage is unnecessary. For example, the first threshold can be set to 2.

[0093] The model training method disclosed herein can also be implemented as a model training device. Figure 7 A schematic diagram of a model training apparatus according to an embodiment of the present disclosure is shown. The functional units of the model training apparatus can be implemented by hardware, software, or a combination of hardware and software that implement the principles of the present disclosure. Those skilled in the art will understand that... Figure 7 The described functional units can be combined or divided into sub-units to achieve the principles of the invention described above. Therefore, the description herein supports any possible combination, division, or further limitation of the functional units described herein. A brief description of the functional units that the reading device may have and the operations that each functional unit can perform is given below. For details involved, please refer to the relevant descriptions above; they will not be repeated here.

[0094] See Figure 7 The model training device 700 may include a segmentation module 710 and a parallel execution module 720.

[0095] The segmentation module 710 is used to segment the dynamically generated computation graph into multiple sub-forward processes during the execution of at least a portion of the forward process for the Nth batch of training samples, and to obtain multiple sub-backward processes corresponding to the multiple sub-forward processes.

[0096] The parallel execution module 720 is used to execute, after the forward process of the Nth batch of training samples is completed, and during the execution of at least part of the forward process of the Pth batch of training samples, a sub-forward process of at least part of the forward process of the Pth batch of training samples and a sub-reverse process of the multiple sub-reverse processes of the Nth batch of training samples are executed in parallel using a computing unit for at least one time period. The sub-forward process and the sub-reverse process executed in parallel are respectively a computational operation and a communication operation.

[0097] In some exemplary embodiments, the segmentation module 710 can identify the first node in the computation graph that does not require gradient calculation based on the attribute information of each node in the computation graph, and segment the computation graph with the first node as the segmentation position to obtain multiple sub-forward processes.

[0098] In some further exemplary embodiments, the segmentation module 710 creates a new output tensor for each output tensor of the first node, where each output tensor of the first node is the output of the sub-forward process corresponding to the first node; and uses the new output tensor as the input of the next sub-forward process.

[0099] In some further exemplary embodiments, the segmentation module 710 may use the first pseudo-tensor as the output of the sub-forward process and connect it to the output tensor of the sub-forward process; and use the second pseudo-tensor as the input of the next sub-forward process.

[0100] In some exemplary embodiments, the model is a hybrid expert model. The forward process of the hybrid expert model includes a routing phase, a distribution phase, an MLP phase, and an aggregation phase. During the forward process, the model training device 700 does not save the forward calculation results of the MLP phase, and during the reverse process, it recalculates the MLP parts other than the last layer in the MLP phase.

[0101] In some further exemplary embodiments, during the forward process, the model training device 700 does not save the forward computation results of the distribution phase, and during the reverse process, it recompiles the distribution phase.

[0102] This disclosure also proposes a distributed training system for a model, comprising multiple computing units. During the execution of at least a portion of the forward pass for the Nth batch of training samples, each computing unit divides a dynamically generated computation graph into multiple sub-forward passes and obtains multiple sub-backward passes corresponding to the sub-forward passes. After the forward pass for the Nth batch of training samples is completed, and during the execution of the at least partial forward pass for the Pth batch of training samples, the computing unit executes, in parallel at least one sub-forward pass in the at least partial forward pass of the Pth batch of training samples and one sub-backward pass in the multiple sub-backward passes of the Nth batch of training samples for at least one time period. Furthermore, of the sub-forward pass and the sub-backward pass executed in parallel, one is a computational operation and the other is a communication operation. Details regarding the operations that the computing unit can perform can be found in the relevant description above.

[0103] Figure 8 A schematic diagram of a computing device according to an embodiment of the present disclosure is shown, which can be used to implement the above-described model training method.

[0104] See Figure 8 The computing device 800 includes a memory 810 and a processor 820.

[0105] Processor 820 may be a multi-core processor or may contain multiple processors. In some embodiments, processor 820 may include a general-purpose main processor and one or more special-purpose coprocessors, such as a graphics processing unit (GPU), a digital signal processor (DSP), etc. In some embodiments, processor 820 may be implemented using custom circuitry, such as an application-specific integrated circuit (ASIC) or a field-programmable gate array (FPGA).

[0106] Memory 810 may include various types of storage units, such as system memory, read-only memory (ROM), and permanent storage devices. ROM may store static data or instructions required by the processor 820 or other modules of the computer. Permanent storage devices may be read-write storage devices. Permanent storage devices may be non-volatile storage devices that retain stored instructions and data even when the computer is powered off. In some embodiments, permanent storage devices use mass storage devices (e.g., magnetic or optical disks, flash memory) as permanent storage devices. In other embodiments, permanent storage devices may be removable storage devices (e.g., floppy disks, optical drives). System memory may be a read-write storage device or a volatile read-write storage device, such as dynamic random access memory. System memory may store some or all of the instructions and data required by the processor during operation. Furthermore, memory 810 may include any combination of computer-readable storage media, including various types of semiconductor memory chips (DRAM, SRAM, SDRAM, flash memory, programmable read-only memory), and disks and / or optical disks may also be used. In some implementations, memory 810 may include a removable storage device that is readable and / or writable, such as a laser disc (CD), a read-only digital versatile optical disc (e.g., DVD-ROM, dual-layer DVD-ROM), a read-only Blu-ray disc, an ultra-high density optical disc, a flash memory card (e.g., SD card, mini SD card, Micro-SD card, etc.), a magnetic floppy disk, etc. Computer-readable storage media do not contain carrier waves or transient electronic signals transmitted wirelessly or via wired connections.

[0107] The memory 810 stores executable code, which, when processed by the processor 820, enables the processor 820 to execute the model training method described above.

[0108] The model training methods, apparatus, systems and devices according to this disclosure have been described in detail above with reference to the accompanying drawings.

[0109] Furthermore, the method according to the present invention can also be implemented as a computer program or computer program product, which includes computer program code instructions for performing the steps defined in the above-described method of the present invention.

[0110] Alternatively, the present invention can also be implemented as a non-transitory machine-readable storage medium (or computer-readable storage medium, or machine-readable storage medium) storing executable code (or computer program, or computer instruction code) thereon, which, when executed by a processor of an electronic device (or computing device, server, etc.), causes the processor to perform the various steps of the method described above according to the present invention.

[0111] Those skilled in the art will also understand that the various exemplary logic blocks, modules, circuits, and algorithm steps described in connection with the disclosure herein can be implemented as electronic hardware, computer software, or a combination of both.

[0112] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of systems and methods according to various embodiments of the present invention. In this regard, each block in a flowchart or block diagram may represent a module, segment, or portion of code containing one or more executable instructions for implementing a specified logical function. It should also be noted that in some alternative implementations, the functions marked in the blocks may occur in a different order than those marked in the drawings. For example, two consecutive blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in the block diagrams and / or flowcharts, and combinations of blocks in the block diagrams and / or flowcharts, can be implemented using a dedicated hardware-based system that performs the specified function or operation, or using a combination of dedicated hardware and computer instructions.

[0113] The various embodiments of the present invention have been described above. These descriptions are exemplary and not exhaustive, nor are they limited to the disclosed embodiments. Many modifications and variations will be apparent to those skilled in the art without departing from the scope and spirit of the described embodiments. The terminology used herein is chosen to best explain the principles, practical application, or improvement of the technology in the market, or to enable others skilled in the art to understand the embodiments disclosed herein.

Claims

1. A model training method, comprising: In performing at least part of the forward process for the Nth batch of training samples, the dynamically generated computation graph is divided into multiple sub-forward processes, and multiple sub-backward processes corresponding to the multiple sub-forward processes are obtained. The nodes in the computation graph represent tensors or functions, the edges in the computation graph represent the dependencies between tensors and functions, and the operations corresponding to the function nodes are computation operations or communication operations. The multiple sub-forward processes include a computation phase and a communication phase. After the forward process of the Nth batch of training samples is completed, and during the execution of at least a portion of the forward process of the Pth batch of training samples, a sub-forward process of the at least portion of the forward process of the Pth batch of training samples and a sub-reverse process of the plurality of sub-reverse processes of the Nth batch of training samples are executed in parallel using a computing unit for at least one time period. Furthermore, of the sub-forward process and the sub-reverse process executed in parallel, one is a computational operation and the other is a communication operation. The process of dividing the dynamically generated computation graph into multiple sub-forward processes includes: identifying the first node in the computation graph that does not require gradient calculation based on the attribute information of each node in the computation graph; dividing the computation graph with the first node as the dividing position to obtain the multiple sub-forward processes.

2. The method according to claim 1, wherein the computation graph is segmented with the first node as the segmentation position, comprising: Create a new output tensor for each output tensor of the first node; The new output tensor is used as the input to the next sub-forward process.

3. The method according to claim 2, further comprising: For the aforementioned sub-forward process, a first pseudo-tensor of length zero and a second pseudo-tensor of length zero are created; The first pseudo tensor is used as the output of the sub-forward process and connected to the output tensor of the sub-forward process. The second pseudo-tensor is used as the input to the next sub-forward process and connected to the new output tensor.

4. The method according to claim 1, wherein, The model is a hybrid expert model, and its forward process includes a routing phase, a distribution phase, a multilayer perceptron (MLP) phase, and an aggregation phase. The method further includes: During the forward pass, the forward computation results of the MLP stage are not saved, and during the reverse pass, the MLP parts other than the last layer in the MLP stage are recomputed.

5. The method according to claim 4, wherein, The method further includes: During the forward process, the forward calculation results of the distribution stage are not saved, and during the reverse process, the distribution stage is recalculated.

6. A model training device, comprising: The segmentation module is used to segment the dynamically generated computation graph into multiple sub-forward processes during at least a part of the forward process for the Nth batch of training samples, and to obtain multiple sub-backward processes corresponding to the multiple sub-forward processes. The nodes in the computation graph represent tensors or functions, the edges in the computation graph represent the dependencies between tensors and functions, and the operation corresponding to the node representing the function is a computation operation or a communication operation. The multiple sub-forward processes include a computation phase and a communication phase. A parallel execution module is configured to, after the forward process of the Nth batch of training samples is completed, and during the execution of at least a portion of the forward process of the Pth batch of training samples, execute, within at least a time period, a sub-forward process of the at least portion of the forward process of the Pth batch of training samples and a sub-reverse process of the plurality of sub-reverse processes of the Nth batch of training samples in parallel using a computing unit. Furthermore, of the sub-forward process and the sub-reverse process executed in parallel, one is a computational operation and the other is a communication operation. The segmentation module identifies the first node in the computation graph that does not require gradient calculation based on the attribute information of each node in the computation graph, and segments the computation graph with the first node as the segmentation position to obtain the multiple sub-forward processes.

7. A distributed training system for a model, comprising multiple computing units, wherein, In performing at least part of the forward process for the Nth batch of training samples, the computing unit divides the dynamically generated computation graph into multiple sub-forward processes and obtains multiple sub-backward processes corresponding to the multiple sub-forward processes. The nodes in the computation graph represent tensors or functions, the edges in the computation graph represent the dependencies between tensors and functions, and the operation corresponding to the node representing the function is a computation operation or a communication operation. The multiple sub-forward processes include a computation phase and a communication phase. After the computational unit completes the forward process of the Nth batch of training samples, and during the execution of at least a portion of the forward process of the Pth batch of training samples, it executes in parallel for at least a time period a sub-forward process of the at least portion of the forward process of the Pth batch of training samples and a sub-reverse process of the plurality of sub-reverse processes of the Nth batch of training samples. Furthermore, of the sub-forward process and the sub-reverse process executed in parallel, one is a computational operation and the other is a communication operation. The computation unit identifies the first node in the computation graph that does not require gradient calculation based on the attribute information of each node in the computation graph, and divides the computation graph with the first node as the dividing position to obtain the multiple sub-forward processes.

8. A computing device, comprising: processor; as well as A memory having executable code stored thereon, which, when executed by the processor, causes the processor to perform the method as described in any one of claims 1 to 5.

9. A computer program product comprising executable code that, when executed by a processor of an electronic device, causes the processor to perform the method as claimed in any one of claims 1 to 5.

10. A non-transitory machine-readable storage medium having executable code stored thereon, which, when executed by a processor of an electronic device, causes the processor to perform the method as claimed in any one of claims 1 to 5.

Images