Distributed training optimization methods for models, electronic devices and storage media

Abstract

This invention relates to the field of artificial intelligence technology, providing a distributed training optimization method, electronic device, and storage medium for a model. The method includes: constructing computational and communication flows based on computational and communication operations in the model; and performing interleaved pipeline scheduling for at least two micro-batches of tasks during model training iterations. The interleaved pipeline scheduling includes: scheduling computational tasks from a first micro-batch to the computational flow for execution, while simultaneously scheduling communication tasks from a second micro-batch to the communication flow, or scheduling communication tasks from the first micro-batch to the communication flow for execution, while simultaneously scheduling computational tasks from the second micro-batch to the computational flow for execution, so that the tasks from the first and second micro-batches are processed in parallel. This invention achieves parallel processing of computational and communication tasks by performing interleaved pipeline scheduling on tasks from different micro-batches, improving hardware utilization and model training efficiency.

Description

Technical Field

[0001] This invention relates to the field of artificial intelligence technology, and in particular to a distributed training optimization method for models, an electronic device, and a storage medium. Background Technology

[0002] In recent years, large-scale language models based on the Transformer architecture have achieved significant success in natural language processing and other fields. To further improve model capacity and performance while controlling computational costs, the Mixture of Experts (MoE) model has been proposed. The MoE model dynamically distributes computational tasks to multiple parallel expert subnetworks through a router module, thereby effectively improving model performance with controllable computational costs. In multi-device distributed training scenarios, data needs to be exchanged globally (all-to-all) between different devices to ensure that data is correctly sent to the designated expert subnetwork for computation; this communication process is the main performance bottleneck.

[0003] Currently, the training and optimization of MoE models mainly rely on two mainstream approaches. One is the synchronous execution mode, which strictly executes serial communication and computation tasks. This mode results in idle computing units during communication and idle communication links during computation, leading to a serious waste of computing resources and low training efficiency. The other approach is to improve the load balancing among expert subnetworks through algorithm optimization (such as introducing auxiliary loss functions). However, this approach does not solve the problem of serial computation and communication, and the optimization effect is lagging, failing to immediately improve the efficiency of a single iteration, while also introducing additional computational overhead. Summary of the Invention

[0004] This invention provides a distributed training optimization method, electronic device, and storage medium for models, aiming to solve the problems of low hardware utilization and low training efficiency caused by the serial execution of computation and communication during the distributed training of large-scale models.

[0005] This invention provides a distributed training optimization method for models, comprising:

[0006] Based on the computational and communication operations in the model, computational streams and communication streams are constructed respectively. The computational stream is used to execute computational tasks, and the communication stream is used to execute communication tasks.

[0007] During the training iterations of the model, staggered pipeline scheduling is performed for tasks in at least two micro-batches.

[0008] The staggered pipeline scheduling includes: scheduling the computation tasks of the first micro-batch to the computation stream for execution, and simultaneously scheduling the communication tasks of the second micro-batch to the communication stream for execution, or scheduling the communication tasks of the first micro-batch to the communication stream for execution, and simultaneously scheduling the computation tasks of the second micro-batch to the computation stream for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, wherein the first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches.

[0009] According to the distributed training optimization method for a model provided by the present invention, the construction of computational flow and communication flow based on computational operations and communication operations in the model includes:

[0010] The computational operations in the model are reconstructed into at least one computational node, and the communication operations in the model are reconstructed into at least one communication node;

[0011] Based on the at least one computing node and the at least one communication node, the computing stream and the communication stream are constructed respectively. The computing stream is used to execute the computing tasks corresponding to the at least one computing node, and the communication stream is used to execute the communication tasks corresponding to the at least one communication node.

[0012] According to a distributed training optimization method for a model provided by the present invention, the step of reconstructing the computational operations in the model into at least one computational node includes:

[0013] The attention calculation, residual calculation, routing calculation, and rearrangement operation in the model are integrated into the first calculation node;

[0014] The fully connected layer computation and inverse rearrangement operation in the model are merged into a second computation node.

[0015] According to a distributed training optimization method for a model provided by the present invention, the step of reconstructing the communication operations in the model into at least one communication node includes:

[0016] The distribution operation in the model is reconstructed into a first communication node, and the merging operation in the model is reconstructed into a second communication node. Both the first and second communication nodes perform global exchange communication.

[0017] The distributed training optimization method for a model provided by the present invention further includes:

[0018] Based on the communication stream, obtain the asynchronous communication handle;

[0019] Before executing the target computation task on the computation stream, a wait operation is inserted. The wait operation is used to wait for the target communication task associated with the asynchronous communication handle to complete, and the input of the target computation task depends on the output of the target communication task.

[0020] According to a distributed training optimization method for a model provided by the present invention, when the target communication task is a distribution communication task in the forward stage, the target computation task is the computation task corresponding to the second computation node in the forward stage;

[0021] When the target communication task is a merged communication task in the reverse phase, the target computation task is the computation task corresponding to the second computation node in the reverse phase.

[0022] When the target communication task is a distribution communication task in the reverse phase, the target computing task is the computing task corresponding to the first computing node in the reverse phase.

[0023] The distributed training optimization method for a model provided by the present invention further includes:

[0024] Before executing the interleaved pipeline scheduling, a separate communication buffer is pre-allocated for the communication tasks associated with the communication flow;

[0025] If the target communication task associated with the asynchronous communication handle is not completed within the preset waiting time, the scheduling of one or more subsequent micro-batch tasks will be adjusted.

[0026] According to the distributed training optimization method for a model provided by the present invention, the parallel processing includes at least one of the following:

[0027] The reverse communication task of the first micro-batch and the forward computation task of the second micro-batch are processed in parallel.

[0028] The reverse computation task of the first micro-batch and the forward communication task of the second micro-batch are processed in parallel.

[0029] The present invention also provides a distributed training optimization device for models, comprising:

[0030] A construction unit is used to construct computational streams and communication streams based on computational and communication operations in the model, wherein the computational streams and the communication streams are executed asynchronously and in parallel.

[0031] An execution unit is configured to perform interleaved pipeline scheduling for at least two micro-batch tasks during the training iterations of the model.

[0032] The staggered pipeline scheduling includes: scheduling the computation tasks of the first micro-batch to the computation stream for execution, and simultaneously scheduling the communication tasks of the second micro-batch to the communication stream for execution, or scheduling the communication tasks of the first micro-batch to the communication stream for execution, and simultaneously scheduling the computation tasks of the second micro-batch to the computation stream for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, wherein the first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches.

[0033] The present invention also provides an electronic device, including a memory, a processor, and a computer program stored in the memory and running on the processor, wherein the processor executes the computer program to implement the distributed training optimization method of the model as described above.

[0034] The present invention also provides a non-transitory computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the distributed training optimization method for the model as described above.

[0035] The present invention also provides a computer program product, including a computer program that, when executed by a processor, implements the distributed training optimization method for the model as described above.

[0036] The distributed training optimization method, electronic device, and storage medium provided by this invention construct independent computation and communication streams based on computational and communication operations within the model, fundamentally creating conditions for parallel task processing. Furthermore, by performing interleaved pipeline scheduling on different micro-batch tasks, computational and communication tasks that would otherwise require sequential waiting can be executed concurrently. This effectively hides communication time within computation time, breaking the bottleneck of sequential execution of computation and communication. Since computational and communication tasks can be executed in parallel, the computing units of the computing device and the communication links responsible for data transmission can be busy simultaneously, greatly reducing device idle time, improving hardware utilization, and thus effectively improving model training efficiency. Moreover, compared to algorithm optimization schemes that rely on auxiliary loss functions, this invention is an optimization at the execution scheduling level. Its performance improvement is immediately apparent after the start of training iterations, without needing to wait for multiple iterations like algorithm optimization, solving the problem of insufficient real-time performance. More importantly, this invention improves efficiency by optimizing the task scheduling process without introducing any new computational load, such as calculating auxiliary loss, avoiding additional computational overhead. Attached Figure Description

[0037] To more clearly illustrate the technical solutions in this invention or related technologies, the accompanying drawings used in the description of the embodiments or related technologies will be briefly introduced below. Obviously, the accompanying drawings described below are some embodiments of this invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0038] Figure 1 This is a schematic diagram of the structure of the computing device provided by the present invention;

[0039] Figure 2 This is a flowchart illustrating the distributed training optimization method for models provided by this invention;

[0040] Figure 3 This is a schematic diagram of node reconstruction provided by the present invention;

[0041] Figure 4 This is a schematic diagram of the parallel processing of the first micro-batch of tasks and the second micro-batch of tasks provided by the present invention;

[0042] Figure 5 This is a schematic diagram of the structure of the distributed training and optimization device for models provided by the present invention;

[0043] Figure 6 This is a schematic diagram of the structure of the electronic device provided by the present invention. Detailed Implementation

[0044] To make the objectives, technical solutions, and advantages of this invention clearer, the technical solutions of this invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of this invention. All other embodiments obtained by those skilled in the art based on the embodiments of this invention without creative effort are within the scope of protection of this invention.

[0045] In recent years, large-scale language models based on the Transformer architecture have achieved significant success in natural language processing and other fields. To further improve model capacity and performance while controlling computational costs, the MoE model was proposed. The MoE model introduces multiple parallel expert subnetworks (hereinafter referred to as "experts"), and a routing module dynamically allocates input data (such as tokens) to some experts for processing, thereby significantly expanding the total number of model parameters without significantly increasing the computational cost per forward pass. It should be understood that in deep learning models, a token is the smallest unit of information processed by the model; its essence is the process of converting raw data (such as text, images, and audio) into discrete symbol sequences through word segmentation. For example, in text processing, each word, punctuation mark, or sub-word can be considered a token; in image processing, images are segmented into fixed-size pixel blocks, each block serving as a token; in audio processing, speech signals are segmented into short frames (such as 25ms audio segments), each frame serving as a token.

[0046] When training the MoE model in a distributed manner, data exchange between devices becomes crucial because the model is split and deployed across multiple computing devices. These computing devices can be GPUs (Graphics Processing Units), GPGPUs (General-purpose computing on Graphics Processing Units), TPUs (Tensor Processing Units), etc. Specifically, the data exchange process involves two critical all-to-all communication operations. First, after the routing module selects some experts for the token, an all-to-all communication operation is performed, sending the corresponding tokens to the devices where each expert is located for computation. After each expert completes its computation, another all-to-all communication operation is performed, sending the output results of each expert back to the original device for subsequent computation. It should be understood that all-to-all is a full-exchange communication mode between multiple devices, meaning that all senders broadcast information to all receivers, and all receivers also send back results to all senders.

[0047] Currently, the training and optimization of the MoE model mainly relies on the following two methods:

[0048] One approach is the synchronous execution mode. In this mode, the training process strictly follows a serial order. First, all devices perform all-to-all communication in parallel, distributing data to multiple experts. After all communication operations are completed, each device then performs computational tasks (such as matrix multiplication and accumulation). While this mode is simple to implement, it has significant drawbacks. Because computation and communication are executed entirely serially, when a device is performing all-to-all communication, its computing unit is idle and waiting; conversely, when computation is being performed, the communication bandwidth is not utilized. This serial dependency results in a large amount of idle time for computing resources such as GPUs during training, leading to low overall utilization and failing to fully leverage the multi-stream concurrent processing capabilities of modern hardware (such as GPUs).

[0049] Another approach is algorithmic optimization. This method attempts to improve load balancing at the algorithmic level, typically by introducing an auxiliary loss function (such as Balance Loss). This loss function aims to penalize uneven load distribution among experts, guiding the routing module to learn a more balanced token allocation strategy by continuously optimizing this loss term during training. However, this method also has limitations: First, it does not solve the fundamental problem of sequential computation and communication in a single iteration, and device idleness still exists. Second, the optimization effect requires multiple iterations to accumulate and become apparent, resulting in a delayed response to single or sudden load imbalances and an inability to immediately improve training efficiency. Finally, introducing the auxiliary loss function itself incurs additional computational overhead, typically increasing the total training time.

[0050] To address the aforementioned issues, this invention provides a distributed training optimization method for models. By effectively separating the computation flow and the communication flow, and executing interleaved pipeline scheduling across micro-batches, parallel processing of computation and communication tasks is achieved. This not only improves hardware utilization and model training efficiency but also instantly enhances the efficiency of a single iteration without additional computational overhead, thereby overcoming the aforementioned shortcomings.

[0051] Figure 1 This is a schematic diagram of the structure of the computing device provided by the present invention, as shown below. Figure 1 As shown, the execution entity of the distributed training optimization method for models provided by the present invention can be a computing device. The computing device 100 includes multiple processing modules 101 and memory 102. Here, the processing module 101 can be a streaming processor cluster (SPC). Each processing module 101 includes multiple computing units 103, and each computing unit 103 includes at least an on-chip cache 104 and a register 105.

[0052] Memory 102 can be High Bandwidth Memory (HBM) or other types of memory. On-chip cache 104 is a temporary memory with a smaller capacity than memory 102, but a faster data exchange speed. On-chip cache 104 can be a Gemm Main Buffer (GMB). Compared to on-chip cache 104, register 105 has a smaller capacity but a faster data exchange speed. Register 105 can specifically be a Thread Local Register (TLR).

[0053] It should be noted that, in addition to the above-described structures, the computing device 100 in this embodiment of the invention may also include other structures, and the invention does not specifically limit these.

[0054] based on Figure 1 The diagram illustrates the architecture of a computing device. This invention provides a distributed model training optimization method applicable to various scenarios, such as text processing, image processing, and audio processing. In different application scenarios, the model training data (i.e., the data corresponding to different micro-batches) has different physical meanings. For example, in a text processing scenario, the model training data can be text data related to text generation and text recognition. Similarly, in an image processing scenario, the model training data can be image data related to image preprocessing, image segmentation, and object detection. And in an audio processing scenario, the model training data can be audio data related to speech recognition and speech synthesis. The specific process of the distributed model training optimization method provided by this invention is described below.

[0055] Figure 2 This is a flowchart illustrating the distributed training and optimization method for models provided by this invention, as shown below. Figure 2 As shown, the method includes:

[0056] Step S10: Based on the computational and communication operations in the model, construct a computational stream and a communication stream respectively. The computational stream is used to execute computational tasks, and the communication stream is used to execute communication tasks.

[0057] It should be noted that the model described in the embodiments of this invention can be a large-scale hybrid expert (MoE) model, such as the Switch Transformer. Of course, the model in the embodiments of this invention is not limited to the MoE model, and its core idea can also be applied to other models that require large-scale distributed training and have significant computation and communication stages, such as traditional dense models.

[0058] Specifically, during model training, operations can be divided into two main categories: computational operations and communication operations. Computational operations refer to computations that primarily consume computing units (such as GPUs). In the MoE model, computational operations can include attention computation, residual computation, routing computation, fully connected layer computation (typically large-scale matrix multiplication), and rearrangement and inverse rearrangement operations for data alignment.

[0059] Communication operations refer to those operations that primarily consume bandwidth in the inter-device communication links, used to transfer data between different computing devices in a distributed training environment. In the distributed training of the MoE model, the most crucial communication operation is the All-to-All global exchange, which includes the distribution operation of distributing data to designated experts and the merging operation of collecting the expert's computation results.

[0060] Specifically, to achieve parallel execution of computation and communication, the operations during model training must first be decoupled. This means constructing independent computation and communication streams based on the computation and communication operations within the model. Constructing these streams involves creating two or more independent, parallel-executable instruction queues on a computing device (such as a GPU). A computation stream is a dedicated execution queue for scheduling and executing computational tasks, while a communication stream is a dedicated execution queue for scheduling and executing communication tasks. For example, a communication stream could be a communication queue created asynchronously with the computation stream using a specialized communication library.

[0061] Understandably, when computational operations in the model (such as attention computation and fully connected layer computation) are scheduled for execution, they constitute computational tasks. When communication operations in the model (such as distribution operations and merging operations) are scheduled for execution, they constitute communication tasks. By binding computational tasks and communication tasks to independent streams, decoupling between the two is achieved at both the hardware and software scheduling levels, creating conditions for subsequent parallel execution.

[0062] Step S20: In the training iteration of the model, staggered pipeline scheduling is performed for at least two micro-batch tasks.

[0063] The staggered pipeline scheduling includes: scheduling the computation tasks of the first micro-batch to the computation stream for execution, and simultaneously scheduling the communication tasks of the second micro-batch to the communication stream for execution, or scheduling the communication tasks of the first micro-batch to the communication stream for execution, and simultaneously scheduling the computation tasks of the second micro-batch to the computation stream for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, wherein the first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches.

[0064] Specifically, after the computation flow and communication flow are constructed, interleaved pipeline scheduling can be performed for at least two micro-batch tasks during the model training iterations to alternately schedule different micro-batch tasks to the computation flow and communication flow, thereby achieving parallel execution of computation tasks and communication tasks.

[0065] Here, the training iteration of the model refers to the complete process of the model completing one forward computation, backpropagation, and gradient update. In order to train efficiently in a distributed environment, a large training batch is usually split into multiple smaller micro-batches.

[0066] The aforementioned task involving at least two micro-batches refers to a situation where, in a single training iteration, two or more micro-batches of data are being processed on the pipeline, such as micro-batch A (or the first micro-batch) and micro-batch B (or the second micro-batch). Each micro-batch contains a series of computational and communication tasks.

[0067] Understandably, interleaved pipelined scheduling is an advanced scheduling strategy that differs from traditional serial execution. In traditional schemes, the next micro-batch B can only begin after all tasks (computation and communication) in a micro-batch A have been completed. Interleaved pipelined scheduling breaks down this barrier. It utilizes the aforementioned independent computation and communication flows to overlap the execution of tasks in different micro-batches. Its purpose is to achieve parallel processing of computation and communication tasks, thereby utilizing the idle computation time that would otherwise be used to wait for communication to complete, and vice versa, to maximize hardware utilization.

[0068] Specifically, when performing interleaved pipeline scheduling for tasks in any two of at least two micro-batches, the computation tasks of the first micro-batch can be scheduled to the computation stream for execution, while the communication tasks of the second micro-batch can be scheduled to the communication stream for execution. This means that while the computation stream is processing a computation task from the first micro-batch, the communication stream can simultaneously process a communication task from the second micro-batch. In this way, the computation of the first micro-batch and the communication of the second micro-batch are processed in parallel (i.e., overlapped execution).

[0069] When performing interleaved pipeline scheduling for tasks in any two of at least two micro-batches, communication tasks from the first micro-batch can be scheduled to the communication stream for execution, while computation tasks from the second micro-batch can be scheduled to the computation stream for execution. This means that while the computation stream is processing a computation task from the second micro-batch, the communication stream can simultaneously process a communication task from the first micro-batch, achieving parallel processing. This complements the above scenario, ensuring that overlap is maximized at different stages of the pipeline.

[0070] It should be noted that the first and second micro-batches mentioned above are relative concepts, not referring to two specific micro-batches, but rather to any two different micro-batches among at least two micro-batches. Typically, they can refer to any two micro-batches that are temporally adjacent or at different processing stages, and their referential relationship changes dynamically as the pipeline executes.

[0071] For example, suppose that in a training iteration, micro-batch A, micro-batch B, and micro-batch C enter the pipeline sequentially. At a certain moment, micro-batch A completes its forward phase tasks and begins its reverse phase. Simultaneously, micro-batch B begins its forward phase tasks. In this scenario, micro-batch A can be considered the first micro-batch, and micro-batch B can be considered the second micro-batch. When micro-batch A performs its reverse phase communication tasks, micro-batch B performs its forward phase computation tasks; conversely, when micro-batch A performs its reverse phase computation tasks, micro-batch B performs its forward phase communication tasks. This achieves parallel processing.

[0072] As the pipeline progresses, micro-batch B completes its forward processing and begins executing its reverse processing tasks; simultaneously, a new micro-batch C enters the pipeline and begins executing its forward processing tasks. In this scenario, micro-batch B can be considered the first micro-batch, and micro-batch C can be considered the second micro-batch, thus achieving parallel computation and communication once again.

[0073] Therefore, the references to the first and second micro-batches are fluid. The scheduler can dynamically pair and parallelize any micro-batch performing a computational task with another micro-batch performing a communication task within a pipeline full of micro-batches, rather than being limited to a fixed pair of micro-batches. This mechanism ensures that as long as the pipeline is full (i.e., at least two micro-batches exist), the overlap between computation and communication can continue to occur, thereby minimizing hardware idle time throughout the training process and achieving stable and efficient pipelined operations.

[0074] The method provided in this invention constructs independent computation and communication streams based on computational and communication operations within the model, fundamentally creating conditions for parallel task processing. Furthermore, by performing interleaved pipeline scheduling on different micro-batch tasks, computational and communication tasks that would otherwise require sequential waiting can be executed concurrently. This effectively hides communication time within computation time, breaking the bottleneck of sequential execution of computation and communication. Since computational and communication tasks can be executed in parallel, the computing units of the computing device and the communication links responsible for data transmission can be busy simultaneously, greatly reducing device idle time, improving hardware utilization, and thus effectively improving model training efficiency. Moreover, compared to algorithm optimization schemes that rely on auxiliary loss functions, this invention is an optimization at the execution scheduling level. Its performance improvement is immediately apparent after the start of training iterations, without needing to wait for multiple iterations like algorithm optimization, thus solving the problem of insufficient real-time performance. More importantly, this invention improves efficiency by optimizing the task scheduling process without introducing any new computational load, such as calculating auxiliary loss, avoiding additional computational overhead.

[0075] Based on any of the above embodiments, step S10 specifically includes:

[0076] Step S11: Reconstruct the computational operations in the model into at least one computational node, and reconstruct the communication operations in the model into at least one communication node;

[0077] Step S12: Based on the at least one computing node and the at least one communication node, construct the computing stream and the communication stream respectively. The computing stream is used to execute the computing tasks corresponding to the at least one computing node, and the communication stream is used to execute the communication tasks corresponding to the at least one communication node.

[0078] It should be noted that the above embodiments propose the basic idea of ​​separating computation and communication tasks into different streams. However, in order to make this separation and subsequent pipeline scheduling more efficient, the embodiments of the present invention introduce a key step of node reconfiguration.

[0079] Node refactoring refers to the logical and implementation aggregation of multiple fine-grained operations in the original computation graph of the model, forming larger, more regular execution units, i.e., nodes. Its purpose is to reduce the startup overhead of kernel functions (such as kernels on GPUs) on computing devices, while improving locality through operator fusion and reducing unnecessary memory reads and writes. Furthermore, node refactoring can create coarser-grained, more clearly defined computation and communication blocks, thereby simplifying scheduling logic and making pipeline overlap more stable and efficient.

[0080] Specifically, the node reconstruction in this embodiment of the invention is divided into the reconstruction of computing nodes and the reconstruction of communication nodes. Here, a computing node is a logical unit formed by the fusion of one or more computing operations, and the computing tasks it executes mainly consume computing unit resources. A communication node is a logical unit composed of one or more communication operations, and the communication tasks it executes mainly occupy communication link resources.

[0081] Once the refactoring is complete, execution flows can be built based on these well-defined nodes. Specifically, the computation flow is used to execute computational tasks corresponding to at least one computation node, while the communication flow is used to execute communication tasks corresponding to at least one communication node. This means that the scheduling system submits the execution instructions of all computation nodes to the computation flow and the execution instructions of all communication nodes to the communication flow, thereby achieving separation between the two at the scheduling source.

[0082] Based on any of the above embodiments, in step S11, reconstructing the computational operations in the model into at least one computational node includes:

[0083] The attention calculation, residual calculation, routing calculation, and rearrangement operation in the model are integrated into the first calculation node;

[0084] The fully connected layer computation and inverse rearrangement operation in the model are merged into a second computation node.

[0085] Specifically, in order to maximize the granularity of computation tasks, when reconstructing computation nodes based on computational operations in the model, this embodiment of the invention integrates attention computation, residual computation, routing computation, and rearrangement operations in the Transformer layer of the model into a first computation node.

[0086] Here, attention computation is the core of the Transformer, used to calculate the correlation weights between different parts of the input sequence. Residual computation is a connection that adds the module's input and output, used to stabilize deep network training. Router computation is a unique operation of the MoE model, used to determine which expert to assign a token to based on its characteristics. The permute operation (MoE Permute) follows the routing decision, rearranging the order of tokens in memory according to the assignment results, ensuring that all tokens sent to the same expert are physically contiguous. This is a prerequisite for efficient execution of subsequent all-to-all communication.

[0087] Figure 3 This is a schematic diagram of node reconstruction provided by the present invention, as shown below. Figure 3As shown, the attention calculation, residual calculation, routing calculation, and reordering operation in the Transformer layer are merged into a single, larger computation node, namely the first computation node (denoted as Attention). The node reconstruction mechanism proposed in this embodiment of the invention is the first to realize the fusion of routing calculation and attention calculation in the MoE model, which can significantly reduce the GPU startup overhead and memory exchange of intermediate data, thereby reducing the overhead of stream synchronization.

[0088] Furthermore, such as Figure 3 As shown, the fully connected layer computation and the inverse repertoire operation in the Transformer layer of the model are merged into a second computation node. Here, the fully connected layer computation can be the MLP (Multi-Layer Perceptron) layer computation, which is the main computation of the expert network itself. It usually consists of two large-scale General Matrix Multiplications (GEMM) and a non-linear activation function, and is the most computationally intensive part of model training. The inverse repertoire operation (i.e., MoE Unpermute) is used to restore the expert-processed output to its original order in the sequence after the expert computation is completed and the data is returned through the communication node, so that subsequent layers can use and compute it correctly. These two operations are merged to form the second computation node (denoted as MLP), which represents the main expert computation load in the MoE model.

[0089] Based on any of the above embodiments, in step S11, reconstructing the communication operations in the model into at least one communication node includes:

[0090] The distribution operation in the model is reconstructed into a first communication node, and the merging operation in the model is reconstructed into a second communication node. Both the first and second communication nodes perform global exchange communication.

[0091] Specifically, although the communication operations in the model are of a single type, refactoring them into independent communication nodes helps clarify their roles and dependencies in the pipeline. Specifically, the dispatch operation in the model can be refactored into the first communication node, and the merge operation into the second communication node.

[0092] Here, the dispatch operation refers to sending consecutive token data blocks from each device to the target device with the corresponding expert via global exchange communication (All-to-All) after the computation task corresponding to the first computing node has been completed. This operation is refactored into the first communication node (denoted as Dispatch).

[0093] The Combine operation refers to the process where, after each expert completes their calculations (i.e., after the calculation task corresponding to the second computing node is executed), this operation transmits the expert outputs distributed across different devices back to their respective token's original device via another global exchange communication (All-to-All). This operation is reconstructed into a second communication node (denoted as Combine).

[0094] The method provided in this invention not only separates computation and communication, but also aggregates the original fine-grained operations into coarse-grained, optimized computation and communication nodes through node reconstruction. This reconstruction lays a solid foundation for subsequent efficient interleaved pipeline scheduling. On the one hand, operator fusion reduces the overhead of computation itself; on the other hand, clear and regular node partitioning makes it easier for the scheduler to identify large blocks of tasks that can be parallelized, thereby achieving more sufficient and stable computation and communication overlap between different micro-batches.

[0095] Based on any of the above embodiments, the method further includes:

[0096] Based on the communication stream, obtain the asynchronous communication handle;

[0097] Before executing the target computation task on the computation stream, a wait operation is inserted. The wait operation is used to wait for the target communication task associated with the asynchronous communication handle to complete, and the input of the target computation task depends on the output of the target communication task.

[0098] It should be noted that the embodiments of the present invention further introduce a key technology for ensuring the correctness of data dependencies within a framework of parallel execution after the separation of computation and communication flows, namely, a dynamic synchronization control mechanism. In the foregoing embodiments, the present invention constructs independent computation and communication flows and reconstructs operations into regularized computation and communication nodes, thereby achieving parallel scheduling of tasks.

[0099] However, these parallel tasks are not completely independent; they have strict data dependencies. For example, the input to one computation task might happen to be the output of another communication task. Without control, a computation task might start executing before the data is ready, leading to computational errors. To address this, embodiments of the present invention propose a dynamic synchronization control mechanism designed to precisely manage these dependencies with minimal performance overhead.

[0100] Specifically, the core of the dynamic synchronization control mechanism lies in achieving precise synchronization across streams by utilizing asynchronous communication handles and inserting wait operations at key points. When a communication task (e.g., an All-to-All operation performed by a first or second communication node) is submitted to the communication stream, it can be started in asynchronous mode. In this mode, the function calling the communication operation will not block and wait for its completion, but will immediately return an asynchronous communication handle.

[0101] Here, an asynchronous communication handle is a lightweight object or identifier representing a communication task that is currently running in the background and has not yet been completed. It does not contain the communication result itself, but it can be used to query the task status or wait for its eventual completion. Specifically, to obtain an asynchronous communication handle, you can start the asynchronous operation and obtain the returned handle by passing a specific parameter (such as `async_op=True`) when calling a communication library function (such as `combine.backward`, i.e., the merge operation in the backpropagation phase). This handle is the asynchronous communication handle and can be denoted as `combine_bwd_handle`.

[0102] Understandably, the purpose of acquiring an asynchronous communication handle is to decouple the initiation and completion of tasks. After submitting the communication task to the communication stream, the initiator (such as the main scheduler thread) does not have to wait in place, but can immediately continue execution to schedule other unrelated tasks (e.g., submit another micro-batch of computation tasks to the computation stream), thereby maximizing parallelism.

[0103] After obtaining the asynchronous communication handle, a wait operation can be inserted before the computation stream executes the target computation task. The wait operation is a synchronous instruction (such as `combine_bwd_handle.wait()`). When the computation stream encounters this instruction, it pauses the execution of subsequent instructions until the event the instruction is waiting for occurs. Here, the event is the completion of the communication task represented by the associated asynchronous communication handle.

[0104] It should be understood that the target communication task and the target computation task are a pair of "producer-consumer" tasks with a direct data dependency. The target communication task executes on the communication stream, and its output is the input of the target computation task. The target computation task executes on the computation stream, and its input depends on the output of the target communication task. This means that the target computation task can only obtain all the necessary input data after the target communication task has completed; otherwise, it cannot execute correctly. For example, as... Figure 3 As shown, the computational task of the expert network (i.e., the computation of the fully connected layer) takes the token distributed to the current device as input, and therefore it must wait for the distribution communication task to complete.

[0105] Based on any of the above embodiments, when the target communication task is a distribution communication task in the forward phase, the target computing task is the computing task corresponding to the second computing node in the forward phase.

[0106] When the target communication task is a merged communication task in the reverse phase, the target computation task is the computation task corresponding to the second computation node in the reverse phase.

[0107] When the target communication task is a distribution communication task in the reverse phase, the target computing task is the computing task corresponding to the first computing node in the reverse phase.

[0108] It should be noted that, in order to achieve precise synchronization, this embodiment of the invention inserts waiting operations at specific stages of training. Model training mainly includes a forward phase and a backward phase. The forward phase refers to the process of data flowing from the input to the output to calculate the model's predicted values; the backward phase refers to the process of the loss gradient flowing from the output to the input to calculate and update the model parameters.

[0109] Specifically, when the target communication task is the forward distribution communication task (denoted as Dispatch.forward), the target computation task is the computation task corresponding to the second computation node in the forward stage (denoted as MLP.forward). Here, during forward computation, the first communication node (Dispatch) is responsible for distributing tokens to each expert. Only after these tokens are transmitted over the network and correctly received by the current device can the second computation node (MLP) responsible for executing expert computation begin its forward computation. Therefore, a wait operation on dispatch_fwd_handle (representing the asynchronous communication handle of the forward distribution communication task Dispatch.forward) must be inserted before executing MLP.forward on the computation stream.

[0110] When the target communication task is the merge communication task in the backward phase (denoted as Combine.backward), the target computation task is the computation task corresponding to the second computation node in the backward phase (denoted as MLP.backward). During backpropagation, the gradient computation order is the reverse of the forward phase. Gradients from subsequent layers first need to be transmitted back from the original position of the token to the expert's position through the second communication node (Combine). Only after the expert (corresponding to the MLP layer) receives these gradients can it compute its own parameter gradients and the gradients passed to the preceding modules. Therefore, before executing MLP.backward on the computation flow, it is necessary to wait for combine_bwd_handle (representing the asynchronous communication handle of the backward merge communication task Combine.backward) to complete.

[0111] When the target communication task is the dispatch communication task in the reverse phase (denoted as Dispatch.backward), the target computation task is the computation task corresponding to the first computation node in the reverse phase (denoted as Attention.backward). After the expert (corresponding to MLP) completes the reverse computation, the resulting gradient needs to be transmitted back to the original position of the token through the reverse operation of the first communication node (Dispatch). Only after the first computation node (Attention) receives these returned gradients can it continue its reverse computation. Therefore, before executing Attention.backward on the computation flow, it is necessary to wait for dispatch_bwd_handle (representing the asynchronous communication handle of the reverse dispatch communication task Dispatch.backward) to complete.

[0112] The method provided in this invention establishes an efficient and accurate dynamic synchronization mechanism. It uses asynchronous communication handles and delayed waiting to perform necessary synchronization waiting only at the last moment before data is actually needed, rather than blocking the entire process at the beginning of the communication task. This "use-as-you-go" strategy solves the data dependency problem in parallel execution while maximizing the overlap window between computation and communication. This is a key technical guarantee for achieving an efficient interleaved pipeline, thereby maximizing training performance while ensuring computational correctness.

[0113] Based on any of the above embodiments, the method further includes:

[0114] Before executing the interleaved pipeline scheduling, a separate communication buffer is pre-allocated for the communication tasks associated with the communication flow;

[0115] If the target communication task associated with the asynchronous communication handle is not completed within the preset waiting time, the scheduling of one or more subsequent micro-batch tasks will be adjusted.

[0116] It should be noted that, to further improve the stability and robustness of the distributed training process, this embodiment of the invention also provides a fault tolerance mechanism. In complex distributed computing environments, communication delays or failures are common problems that can cause the entire training process to stall or crash. This embodiment of the invention effectively addresses such abnormal situations by introducing timeout detection and buffer pre-allocation mechanisms.

[0117] Specifically, in this embodiment of the invention, before performing interleaved pipeline scheduling, an independent communication buffer is pre-allocated for communication tasks. Specifically, before training begins or during the initialization phase of each iteration, the system pre-allocates a dedicated buffer in the memory of the computing device (such as a GPU) for each expert's communication operations (such as All-to-All Dispatch and Combine). The advantages of this approach are: first, it avoids the runtime overhead and latency caused by dynamically allocating memory during training; second, by providing independent memory space for each communication task, it eliminates execution conflicts or blocking that may result from memory resource contention, ensuring the smooth parallel execution of the computation and communication flows.

[0118] Furthermore, this embodiment of the invention introduces a timeout detection mechanism in dynamic synchronization control. To ensure the correctness of data dependencies, the computation flow synchronizes the communication flow by waiting before executing tasks that depend on the communication results. However, indefinite waiting can cause the system to stalemate when communication is abnormal. Therefore, the waiting operation in this embodiment of the invention is configured with timeout monitoring for a preset duration.

[0119] Specifically, the system sets a reasonable timeout threshold for each waiting operation. During training, a timer is started when the computation stream performs a waiting operation. If the target communication task associated with the asynchronous communication handle it is waiting for has not been completed within the preset time (e.g., due to network congestion or node failure causing an All-to-All communication timeout), the system will no longer wait but will trigger preset fault tolerance processing.

[0120] The core of this fault-tolerant processing lies in adjusting the scheduling of subsequent tasks. For example, the system can mark a micro-batch experiencing a communication failure as failed and discard it from the current training pipeline, while continuing to schedule other normal micro-batches to avoid the entire training process being blocked by a single point of failure. In other implementations, the system can also attempt to reschedule failed communication tasks to the communication stream or report the anomaly to a higher-level training management module, which will then decide whether to reassign expert routes or adjust training parameters. In this way, the embodiments of the present invention not only achieve efficient computational communication overlap but also ensure the robustness of the training process in the face of underlying hardware or network fluctuations, significantly improving the availability and reliability of large-scale model training.

[0121] Based on any of the above embodiments, the parallel processing includes at least one of the following:

[0122] The reverse communication task of the first micro-batch and the forward computation task of the second micro-batch are processed in parallel.

[0123] The reverse computation task of the first micro-batch and the forward communication task of the second micro-batch are processed in parallel.

[0124] Specifically, Figure 4 This is a schematic diagram illustrating the parallel processing of the first micro-batch and the second micro-batch of tasks provided by the present invention, as shown below. Figure 4 As shown, this illustrates the situation where two consecutive micro-batches (such as the first and second micro-batches) perform interleaved pipeline scheduling on the computation and communication flows after the training iterations have entered a stable state. In this stable state, the pipeline is "filled," meaning that while the second micro-batch is executing its forward propagation task, the first micro-batch (whose forward propagation has been completed) is executing its backward propagation task. This embodiment of the invention utilizes this time window to achieve the overlap of computation and communication tasks.

[0125] In one embodiment, such as Figure 4 As shown, when processing the tasks of the first micro-batch and the second micro-batch in parallel, the reverse communication task of the first micro-batch and the forward computation task of the second micro-batch can be processed in parallel. This parallel mode utilizes the downtime of computation in one micro-batch to execute the reverse communication of the other micro-batch.

[0126] For details, please refer to Figure 4 The first line shows the overlapping portion. When the second micro-batch is executing its forward computation task on the computation stream, i.e., the computation task corresponding to the first computation node (Attention.forward), the hardware's computation unit is occupied. Simultaneously, the second micro-batch is executing its reverse communication task, i.e., the reverse operation of the second communication node (Combine.backward) is scheduled to be executed on a separate communication stream. Since the computation stream and communication stream can work in parallel, the computation process of Attention.forward and the communication process of Combine.backward achieve temporal overlap.

[0127] Similarly, refer to Figure 4 The third line overlaps. While the second micro-batch is executing its other forward computation task on the computation stream, namely the computation corresponding to the second computation node (MLP.forward), another reverse communication task of the first micro-batch, namely the reverse operation of the first communication node (Dispatch.backward), is being executed on the communication stream. Both also achieve parallel processing.

[0128] In another embodiment, such as Figure 4As shown, when processing the tasks of the first micro-batch and the second micro-batch in parallel, the reverse computation task of the first micro-batch and the forward communication task of the second micro-batch can be processed in parallel. This parallel mode complements the above situation, ensuring that hardware resources are fully utilized at each stage of the pipeline.

[0129] For details, please refer to Figure 4 The second line overlaps. While the second micro-batch is executing its forward communication task, i.e., the forward operation (Dispatch.forward) of the first communication node, this task runs on the communication stream and consumes communication bandwidth. Simultaneously, the first micro-batch is executing its reverse computation task on the computation stream, i.e., the reverse operation (MLP.backward) of the second computation node. The communication process of Dispatch.forward and the computation process of MLP.backward overlap in time.

[0130] Similarly, refer to Figure 4 The fourth line overlaps. While the second micro-batch is executing its forward communication task, i.e., the forward operation (Combine.forward) of the second communication node, the first micro-batch is executing its backward computation task on the computation flow, i.e., the backward operation (Attention.backward) of the first computation node, and the two are parallelized again.

[0131] The method provided in this invention cleverly overlaps the forward phase task of one micro-batch with the reverse phase task of another micro-batch. Whether it's the overlap of computational and communication tasks, or the overlap of communication and computational tasks, it significantly reduces the idle waiting time of computing devices such as GPUs. This strategy of hiding the execution time of one task within the execution time of another is the fundamental reason why this invention can overcome the bottlenecks of existing technologies and significantly improve hardware utilization and end-to-end training efficiency.

[0132] The distributed training optimization device for models provided by the present invention is described below. The distributed training optimization device described below can be referred to in correspondence with the distributed training optimization method described above.

[0133] Based on any of the above embodiments Figure 5 This is a schematic diagram of the structure of the distributed training and optimization device for models provided by the present invention, as shown below. Figure 5 As shown, the device includes:

[0134] The construction unit 510 is used to construct a computational stream and a communication stream based on the computational and communication operations in the model, respectively. The computational stream is used to execute computational tasks, and the communication stream is used to execute communication tasks.

[0135] Execution unit 520 is used to perform interleaved pipeline scheduling for at least two micro-batch tasks during the training iteration of the model.

[0136] The staggered pipeline scheduling includes: scheduling the computation tasks of the first micro-batch to the computation stream for execution, and simultaneously scheduling the communication tasks of the second micro-batch to the communication stream for execution, or scheduling the communication tasks of the first micro-batch to the communication stream for execution, and simultaneously scheduling the computation tasks of the second micro-batch to the computation stream for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, wherein the first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches.

[0137] The apparatus provided in this invention constructs independent computational and communication streams based on computational and communication operations within the model, fundamentally creating conditions for parallel task processing. Furthermore, by performing interleaved pipeline scheduling on different micro-batch tasks, computational and communication tasks that would otherwise require sequential waiting can be executed concurrently. This effectively hides communication time within computation time, breaking the bottleneck of sequential execution of computation and communication. Since computational and communication tasks can be executed in parallel, the computing units of the computing device and the communication links responsible for data transmission can be busy simultaneously, greatly reducing device idle time, improving hardware utilization, and thus effectively improving model training efficiency. Moreover, compared to algorithm optimization schemes that rely on auxiliary loss functions, this invention is an optimization at the execution scheduling level. Its performance improvement is immediately apparent after the start of training iterations, without needing to wait for multiple iterations like algorithm optimization, solving the problem of insufficient real-time performance. More importantly, this invention improves efficiency by optimizing the task scheduling process without introducing any new computational load, such as calculating auxiliary loss, avoiding additional computational overhead.

[0138] Based on any of the above embodiments, the construction unit 510 includes:

[0139] The node reconstruction subunit is used to reconstruct the computational operations in the model into at least one computational node, and the communication operations in the model into at least one communication node;

[0140] An execution flow construction subunit is used to construct the computing flow and the communication flow based on the at least one computing node and the at least one communication node, respectively. The computing flow is used to execute the computing tasks corresponding to the at least one computing node, and the communication flow is used to execute the communication tasks corresponding to the at least one communication node.

[0141] Based on any of the above embodiments, the node reconstruction subunit is specifically used for:

[0142] The attention calculation, residual calculation, routing calculation, and rearrangement operation in the model are integrated into the first calculation node;

[0143] The fully connected layer computation and inverse rearrangement operation in the model are merged into a second computation node.

[0144] Based on any of the above embodiments, the node reconstruction subunit is specifically used for:

[0145] The distribution operation in the model is reconstructed into a first communication node, and the merging operation in the model is reconstructed into a second communication node. Both the first and second communication nodes perform global exchange communication.

[0146] Based on any of the above embodiments, the device further includes a synchronization unit, which is used for:

[0147] Based on the communication stream, obtain the asynchronous communication handle;

[0148] Before executing the target computation task on the computation stream, a wait operation is inserted. The wait operation is used to wait for the target communication task associated with the asynchronous communication handle to complete, and the input of the target computation task depends on the output of the target communication task.

[0149] Based on any of the above embodiments, when the target communication task is a distribution communication task in the forward phase, the target computing task is the computing task corresponding to the second computing node in the forward phase.

[0150] When the target communication task is a merged communication task in the reverse phase, the target computation task is the computation task corresponding to the second computation node in the reverse phase.

[0151] When the target communication task is a distribution communication task in the reverse phase, the target computing task is the computing task corresponding to the first computing node in the reverse phase.

[0152] Based on any of the above embodiments, the device further includes:

[0153] A buffer allocation unit is used to pre-allocate independent communication buffers for communication tasks associated with the communication flow before executing the interleaved pipeline scheduling;

[0154] The scheduling adjustment unit is used to adjust the scheduling of one or more subsequent micro-batch tasks if the target communication task associated with the asynchronous communication handle has not been completed within the preset duration of the waiting operation.

[0155] Based on any of the above embodiments, the parallel processing includes at least one of the following:

[0156] The reverse communication task of the first micro-batch and the forward computation task of the second micro-batch are processed in parallel.

[0157] The reverse computation task of the first micro-batch and the forward communication task of the second micro-batch are processed in parallel.

[0158] Figure 6 An example is a schematic diagram of the physical structure of an electronic device, such as... Figure 6 As shown, the electronic device may include: a processor 610, a communication interface 620, a memory 630, and a communication bus 640, wherein the processor 610, the communication interface 620, and the memory 630 communicate with each other through the communication bus 640. The processor 610 can call logical instructions in the memory 630 to execute a distributed training optimization method for the model. This method includes: constructing a computational stream and a communication stream based on computational operations and communication operations in the model, respectively. The computational stream is used to execute computational tasks, and the communication stream is used to execute communication tasks. In the training iteration of the model, interleaved pipeline scheduling is performed for at least two micro-batches of tasks. The interleaved pipeline scheduling includes: scheduling the computational tasks of the first micro-batch to the computational stream for execution, and simultaneously scheduling the communication tasks of the second micro-batch to the communication stream for execution, or scheduling the communication tasks of the first micro-batch to the communication stream for execution, and simultaneously scheduling the computational tasks of the second micro-batch to the computational stream for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel. The first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches.

[0159] Furthermore, the logical instructions in the aforementioned memory 630 can be implemented as software functional units and, when sold or used as independent products, can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the part that contributes to related technologies, or a portion of the technical solution, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods described in the various embodiments of the present invention. The aforementioned storage medium includes various media capable of storing program code, such as USB flash drives, portable hard drives, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0160] On the other hand, the present invention also provides a computer program product, which includes a computer program that can be stored on a non-transitory computer-readable storage medium. When the computer program is executed by a processor, the computer can execute the model distributed training optimization method provided by the above methods. The method includes: constructing a computational stream and a communication stream based on computational operations and communication operations in the model, wherein the computational stream is used to perform computational tasks and the communication stream is used to perform communication tasks; and performing interleaved pipeline scheduling for at least two micro-batches of tasks during the training iteration of the model. The interleaved pipeline scheduling includes: scheduling the computational tasks of a first micro-batch to the computational stream for execution, and simultaneously scheduling the communication tasks of a second micro-batch to the communication stream for execution, or scheduling the communication tasks of the first micro-batch to the communication stream for execution, and simultaneously scheduling the computational tasks of the second micro-batch to the computational stream for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, wherein the first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches.

[0161] In another aspect, the present invention also provides a non-transitory computer-readable storage medium storing a computer program thereon, which, when executed by a processor, implements a model distributed training optimization method provided by the above methods. This method includes: constructing a computational stream and a communication stream based on computational and communication operations in the model, wherein the computational stream is used to perform computational tasks and the communication stream is used to perform communication tasks; and, during the training iteration of the model, performing interleaved pipeline scheduling for at least two micro-batches of tasks; wherein the interleaved pipeline scheduling includes: scheduling the computational tasks of a first micro-batch to the computational stream for execution, and simultaneously scheduling the communication tasks of a second micro-batch to the communication stream for execution, or scheduling the communication tasks of the first micro-batch to the communication stream for execution, and simultaneously scheduling the computational tasks of the second micro-batch to the computational stream for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, wherein the first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches.

[0162] The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate. The components shown as units may or may not be physical units; that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of this embodiment according to actual needs. Those skilled in the art can understand and implement this without any creative effort.

[0163] Through the above description of the embodiments, those skilled in the art can clearly understand that each embodiment can be implemented by means of software plus necessary general-purpose hardware platforms, and of course, it can also be implemented by hardware. Based on this understanding, the above technical solutions, in essence or the parts that contribute to the related technology, can be embodied in the form of software products. This computer software product can be stored in a computer-readable storage medium, such as ROM / RAM, magnetic disk, optical disk, etc., and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute the methods described in the various embodiments or some parts of the embodiments.

[0164] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, and not to limit them; although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some of the technical features; and these modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A distributed training optimization method for a model, characterized in that, include: Based on the computational and communication operations in the model, computational streams and communication streams are constructed respectively. The computational stream is used to execute computational tasks, and the communication stream is used to execute communication tasks. During the training iterations of the model, staggered pipeline scheduling is performed for tasks in at least two micro-batches. The staggered pipeline scheduling includes: scheduling the computation tasks of the first micro-batch to the computation stream for execution, and simultaneously scheduling the communication tasks of the second micro-batch to the communication stream for execution, or scheduling the communication tasks of the first micro-batch to the communication stream for execution, and simultaneously scheduling the computation tasks of the second micro-batch to the computation stream for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, wherein the first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches; The method further includes: Based on the communication stream, obtain the asynchronous communication handle; When executing a target computation task on the computation stream, a wait operation is inserted before the target computation task. The wait operation is used to wait for the target communication task associated with the asynchronous communication handle to complete. The input of the target computation task depends on the output of the target communication task.

2. The distributed training and optimization method for models according to claim 1, characterized in that, The computational and communication operations based on the model are used to construct computational and communication flows, respectively, including: The computational operations in the model are reconstructed into at least one computational node, and the communication operations in the model are reconstructed into at least one communication node; Based on the at least one computing node and the at least one communication node, the computing stream and the communication stream are constructed respectively. The computing stream is used to execute the computing tasks corresponding to the at least one computing node, and the communication stream is used to execute the communication tasks corresponding to the at least one communication node.

3. The distributed training optimization method for models according to claim 2, characterized in that, The step of reconstructing the computational operations in the model into at least one computational node includes: The attention calculation, residual calculation, routing calculation, and rearrangement operation in the model are integrated into the first calculation node; The fully connected layer computation and inverse rearrangement operation in the model are merged into a second computation node.

4. The distributed training and optimization method for models according to claim 2, characterized in that, The step of reconstructing the communication operations in the model into at least one communication node includes: The distribution operation in the model is reconstructed into a first communication node, and the merging operation in the model is reconstructed into a second communication node. Both the first and second communication nodes perform global exchange communication.

5. The distributed training optimization method for models according to claim 3, characterized in that, When the target communication task is a distribution communication task in the forward phase, the target computing task is the computing task corresponding to the second computing node in the forward phase; When the target communication task is a merged communication task in the reverse phase, the target computation task is the computation task corresponding to the second computation node in the reverse phase. When the target communication task is a distribution communication task in the reverse phase, the target computing task is the computing task corresponding to the first computing node in the reverse phase.

6. The distributed training optimization method for models according to claim 1, characterized in that, Also includes: Before executing the interleaved pipeline scheduling, a separate communication buffer is pre-allocated for the communication tasks associated with the communication flow; If the target communication task associated with the asynchronous communication handle is not completed within the preset waiting time, the scheduling of one or more subsequent micro-batch tasks will be adjusted.

7. The distributed training and optimization method for models according to any one of claims 1 to 6, characterized in that, The parallel processing includes at least one of the following: The reverse communication task of the first micro-batch and the forward computation task of the second micro-batch are processed in parallel. The reverse computation task of the first micro-batch and the forward communication task of the second micro-batch are processed in parallel.

8. An electronic device comprising a memory, a processor, and a computer program stored in the memory and running on the processor, characterized in that, When the processor executes the computer program, it implements the distributed training optimization method for the model as described in any one of claims 1 to 7.

9. A non-transitory computer-readable storage medium having a computer program stored thereon, characterized in that, When the computer program is executed by the processor, it implements the distributed training optimization method for the model as described in any one of claims 1 to 7.

Images