Model distributed training optimization method, electronic equipment and storage medium

Abstract

The invention relates to the technical field of artificial intelligence, and provides a model distributed training optimization method, electronic equipment and a storage medium, and the method comprises the steps: respectively constructing a calculation flow and a communication flow based on the calculation operation and the communication operation in a model; in the training iteration of the model, executing staggered pipeline scheduling for at least two micro-batches of tasks; wherein the staggered pipeline scheduling comprises the following steps: scheduling a first micro-batch of calculation tasks to a calculation flow for execution, and scheduling a second micro-batch of communication tasks to a communication flow for execution, or scheduling a first micro-batch of communication tasks to the communication flow for execution, and scheduling the second micro-batch of calculation tasks to the calculation flow for execution, the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel. According to the method, the parallel processing of the calculation task and the communication task is realized by executing the staggered assembly line scheduling on the tasks of different micro-batches, and the hardware utilization rate and the model training efficiency are improved.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of artificial intelligence, and in particular to a model distributed training optimization method, an electronic device and a storage medium. BACKGROUND

[0002] In recent years, large-scale language models based on the Transformer architecture have achieved remarkable success in natural language processing and other fields. In order to further improve the model capacity and performance while controlling the computing cost, a Mixture of Experts (MoE) model is proposed. The MoE model dynamically allocates computing tasks to multiple parallel expert sub-networks through a Router module, thereby effectively improving the model performance under controllable computing cost. In the distributed training scenario of multiple devices, data needs to be exchanged between different devices through All-to-All communication to ensure that the data can be correctly sent to the designated expert sub-network for calculation, and this communication process is the main performance bottleneck.

[0003] Currently, the training optimization of the MoE model mainly relies on two mainstream schemes. One is the synchronous execution mode, which strictly performs serial communication and computing tasks. This mode causes the computing unit to be idle when communicating and the communication link to be idle when computing, resulting in serious waste of computing resources and low training efficiency. The other scheme is to improve the load balancing between expert sub-networks through algorithm optimization (such as introducing auxiliary loss functions), but this scheme does not solve the serial problem of computing and communication, and the optimization effect has a lag, which cannot instantly improve the efficiency of a single iteration, and also introduces additional computing overhead. SUMMARY

[0004] The present application provides a model distributed training optimization method, an electronic device and a storage medium, aiming to solve the defects of low hardware utilization and low training efficiency caused by the serial execution of computing and communication in the distributed training of large-scale models.

[0005] The present application provides a model distributed training optimization method, comprising: Based on the computing operations and communication operations in the model, a computing flow and a communication flow are constructed respectively, the computing flow is used to execute computing tasks, and the communication flow is used to execute communication tasks; In the training iteration of the model, for at least two micro-batch tasks, an interleaved pipeline scheduling is performed; The interleaved pipeline scheduling includes: scheduling a first micro-batch of computing tasks to the computing flow execution while scheduling a second micro-batch of communication tasks to the communication flow execution, or scheduling the first micro-batch of communication tasks to the communication flow execution while scheduling the second micro-batch of computing tasks to the computing flow 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 being any two different micro-batches of the at least two micro-batches.

[0006] According to the model distributed training optimization method provided by the application, the computing operation and the communication operation in the model are used to construct a computing flow and a communication flow, respectively, and the method comprises the following steps: The computing operation in the model is reconstructed into at least one computing node, and the communication operation in the model is reconstructed into at least one communication node. Based on the at least one computing node and the at least one communication node, the computing flow and the communication flow are constructed respectively, the computing flow is used to execute the computing task corresponding to the at least one computing node, and the communication flow is used to execute the communication task corresponding to the at least one communication node.

[0007] According to the model distributed training optimization method provided by the application, the computing operation in the model is reconstructed into at least one computing node, and the method comprises the following steps: The attention computing, the residual computing, the routing computing and the rearrangement operation in the model are fused into a first computing node. The fully connected layer computing and the inverse rearrangement operation in the model are fused into a second computing node.

[0008] According to the model distributed training optimization method provided by the application, the communication operation in the model is reconstructed into at least one communication node, and the method comprises the following steps: 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, the first communication node and the second communication node both performing global exchange communication.

[0009] According to the model distributed training optimization method provided by the application, the method further comprises the following steps: Based on the communication flow, an asynchronous communication handle is acquired. Before executing a target computing task on the computing flow, a waiting operation is inserted, the waiting operation being used to wait for the execution of a target communication task associated with the asynchronous communication handle to be completed, and the input of the target computing task depending on the output of the target communication task.

[0010] According to the model distributed training optimization method provided in the application, when the target communication task is a forward-stage distribution communication task, the target computing task is a computing task corresponding to the second computing node in the forward stage; When the target communication task is a reverse-stage merging communication task, the target computing task is a computing task corresponding to the second computing node in the reverse stage; When the target communication task is a reverse-stage distribution communication task, the target computing task is a computing task corresponding to the first computing node in the reverse stage.

[0011] According to the model distributed training optimization method provided in the application, the method further comprises: Before the interleaved pipeline scheduling is performed, an independent communication buffer is pre-allocated for a communication task associated with the communication flow; If the target communication task associated with the asynchronous communication handle is not executed completely within the preset time length of the waiting operation, the scheduling of one or more subsequent micro-batches of tasks is adjusted.

[0012] According to the model distributed training optimization method provided in the application, the parallel processing comprises at least one of the following: The reverse communication task of the first micro-batch is processed in parallel with the forward computing task of the second micro-batch; The reverse computing task of the first micro-batch is processed in parallel with the forward communication task of the second micro-batch.

[0013] The application further provides a model distributed training optimization device, comprising: A construction unit is configured to construct a computing flow and a communication flow based on computing operations and communication operations in a model, and the computing flow and the communication flow are executed asynchronously and in parallel; An execution unit is configured to perform interleaved pipeline scheduling for tasks of at least two micro-batches in a training iteration of the model; The interleaved pipeline scheduling comprises scheduling a computing task of a first micro-batch to the computing flow for execution, while scheduling a communication task of a second micro-batch to the communication flow for execution, or scheduling the communication task of the first micro-batch to the communication flow for execution, while scheduling a computing task of the second micro-batch to the computing flow 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 being any two different micro-batches in the at least two micro-batches.

[0014] The application further provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and running on the processor, wherein the processor implements the model distributed training optimization method according to any one of the above when executing the computer program.

[0015] The application further provides a non-transitory computer-readable storage medium, which stores a computer program, wherein the computer program is executed by a processor to implement the model distributed training optimization method according to any one of the above.

[0016] The application further provides a computer program product, comprising a computer program, wherein the computer program is executed by a processor to implement the model distributed training optimization method according to any one of the above.

[0017] The model distributed training optimization method, the electronic device and the storage medium provided by the application fundamentally create conditions for parallel processing of tasks by constructing independent computing streams and communication streams based on computing operations and communication operations in the model. On this basis, the originally serially waiting computing tasks and communication tasks can be overlapped by performing staggered pipeline scheduling on different micro-batch tasks, which effectively hides the communication time in the computing time, thereby breaking the serial execution bottleneck of computing and communication. Since the computing tasks and the communication tasks can be executed in parallel, the computing units of the computing device and the communication links responsible for data transmission can be in a busy state at the same time, greatly reducing the idle time of the device and improving the hardware utilization, thereby effectively improving the model training efficiency. In addition, compared with the algorithm optimization scheme relying on auxiliary loss functions, the application is an optimization at the execution scheduling level, and the performance improvement effect can be immediately realized after the start of training iteration, without waiting for multiple iterations as in algorithm optimization, solving the problem of insufficient real-time performance. More importantly, the application improves efficiency by optimizing the task scheduling process without introducing any new computational load, such as computing auxiliary loss, avoiding additional computational overhead. BRIEF DESCRIPTION OF DRAWINGS

[0018] In order to more clearly illustrate the technical solutions in the application or the related art, the following will briefly introduce the drawings needed to be used in the embodiments or the related art description. Obviously, the drawings in the following description are some embodiments of the application, and for those skilled in the art, other drawings can also be obtained without creative labor.

[0019] Figure 1 is a structural schematic diagram of a computing device provided by the application; Figure 2 is a flow schematic diagram of the model distributed training optimization method provided by the application; Figure 3is a schematic diagram of node reconstruction provided by the present application; Figure 4 is a schematic diagram of parallel processing of tasks of a first micro-batch and tasks of a second micro-batch provided by the present application; Figure 5 is a structural schematic diagram of a model distributed training optimization device provided by the present application; Figure 6 is a structural schematic diagram of an electronic device provided by the present application. DETAILED DESCRIPTION

[0020] In order to make the objects, technical solutions and advantages of the present application clearer, the technical solutions in the present application will be described clearly and completely below with reference to the drawings in the present application. Obviously, the described embodiments are part of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative labor fall within the scope of protection of the present application.

[0021] In recent years, large-scale language models based on the Transformer architecture have achieved remarkable success in natural language processing and other fields. In order to further improve the model capacity and performance while controlling the computing cost, the MoE model is proposed. The MoE model introduces multiple parallel expert sub-networks (hereinafter referred to as "experts") and dynamically allocates input data (such as tokens) to some experts for processing by a routing module, thereby significantly expanding the total parameter quantity of the model without significantly increasing the single forward computation quantity. It should be understood that in a deep learning model, a token refers to the smallest unit of information processed by the model, which is essentially the process of converting raw data (such as text, images, audio) into a discrete symbol sequence by tokenization. For example, in the text processing scenario, each word, punctuation, or subword can be regarded as a token; in the image processing scenario, the image is divided into fixed-size pixel blocks, and each block is regarded as a token; in the audio processing scenario, the speech signal is cut into short-time frames (such as 25ms audio segments), and each frame is regarded as a token.

[0022] When training a MoE model in a distributed manner, data exchange between devices becomes crucial since the model is split and deployed on multiple computing devices. Here, the 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 between devices involves two key All-to-All communications. First, after the routing module selects the partial experts for a token, an All-to-All communication operation is performed to send the corresponding token to the device where each expert is located for computation. After each expert completes the computation, another All-to-All communication operation is performed to transmit the output 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, i.e., a communication mode in which all senders broadcast information to all receivers, and all receivers also feed back the results to all senders.

[0023] Currently, the training optimization of MoE models mainly relies on the following two schemes: One 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 to distribute data to multiple experts; after all communication operations are completed, each device then uniformly performs the computation task (such as matrix multiplication and accumulation). Although this mode is simple to implement, it has significant defects. Since the computation and communication are completely serially executed, when the device performs All-to-All communication, its computing unit is in an idle waiting state; conversely, when performing computation, the communication bandwidth is not utilized. This serial dependence results in a large amount of idle period of GPU and other computing resources during the training process, with low overall utilization, which cannot fully utilize the multi-stream concurrent processing capability possessed by modern hardware (such as GPU).

[0024] Another is the algorithm optimization mode. This scheme attempts to improve the load balancing problem from the algorithm level, and the typical practice is to introduce a auxiliary loss function (such as Balance Loss). The loss function aims to punish the unbalanced allocation of experts, and by constantly optimizing this loss term during training, it guides the routing module to learn a more balanced token allocation strategy. However, this method also has limitations: first, it does not solve the fundamental problem of serial computation and communication in a single iteration, and device idleness still exists. Second, the optimization effect needs to be accumulated through multiple iterations to appear, and the response to single or burst load imbalance is lagging, and cannot immediately improve training efficiency. Finally, the introduction of auxiliary loss function itself will bring additional computational overhead, usually increasing the total training time.

[0025] To solve the above problems, the present application provides a model distributed training optimization method, which effectively separates the computation flow and the communication flow, and performs interleaved pipeline scheduling across micro-batches, realizing parallel processing of computation tasks and communication tasks, which not only improves hardware utilization and model training efficiency, but also immediately improves single iteration efficiency without additional computational overhead, thereby overcoming the above defects.

[0026] Figure 1 The present application provides a structure diagram of a computing device, as shown in Figure 1 The execution subject of the model distributed training optimization method provided by the present application can be a computing device, which includes a plurality of processing modules 101 and a memory 102. Here, the processing module 101 can be a streaming processor cluster (SPC). Each processing module 101 includes a plurality of computing units 103, and each computing unit 103 includes at least an on-chip cache 104 and a register 105.

[0027] The memory 102 can be a high bandwidth memory (HBM), or other types of memory. The on-chip cache 104 is a temporary memory, and its capacity is smaller than that of the memory 102, but the data exchange speed is faster than that of the memory 102. The on-chip cache 104 can be a general matrix main buffer (GMB). Compared with the on-chip cache 104, the capacity of the register 105 is smaller than that of the on-chip cache 104, but the data exchange speed is faster than that of the on-chip cache 104. The register 105 can be a thread local register (TLR).

[0028] It should be noted that the computing device 100 in the embodiment of the present application can include other structures in addition to the above structure, and the present application does not make specific limitations thereto.

[0029] Based on Figure 1 Based on the architecture diagram of the computing device shown in FIG. 1, the embodiment of the present application provides a model distributed training optimization method, which can be applied to various scenarios such as text processing, image processing, audio processing, etc. In different application scenarios, the model training data (i.e. data corresponding to different micro-batches) has different physical meanings. For example, in the text processing scenario, the model training data can be text data related to text generation, text recognition, etc. For another example, in the image processing scenario, the model training data can be image data related to image preprocessing, image segmentation, target detection, etc. For another example, in the audio processing scenario, the model training data can be audio data related to speech recognition, speech synthesis, etc. The specific process of the model distributed training optimization method provided by the present application will be introduced below.

[0030] Figure 2 is a flowchart of the model distributed training optimization method provided by the present application, as shown in FIG. 2, the method comprises the following steps: Figure 2 Step S10, based on the computing operation and the communication operation in the model, a computing flow and a communication flow are respectively constructed, the computing flow is used for executing a computing task, and the communication flow is used for executing a communication task.

[0031] It should be noted that the model in the embodiment of the present application can be a large-scale hybrid expert (MoE) model, such as Switch Transformer, etc. Of course, the model in the embodiment of the present application is not limited to the MoE model, and the core idea thereof can also be applied to other models that need to be trained in a large-scale distributed manner and have significant computing and communication stages, such as traditional dense (Dense) models.

[0032] Specifically, in the model training process, the operation can be divided into two categories, namely computing operation and communication operation. Among them, the computing operation refers to the operation mainly consuming the computing unit of the computing device (such as GPU). In the MoE model, the computing operation can include attention calculation, residual calculation, routing calculation, full connection layer calculation (usually large-scale matrix multiplication calculation), and rearrangement and inverse rearrangement operation for data alignment, etc.

[0033] ​The communication operation refers to an operation mainly occupying the bandwidth of the communication link between devices, and is used to transmit data between different computing devices in a distributed training environment. In the distributed training of the MoE model, the most core communication operation is the all-to-all communication, which includes a distribution operation of distributing data to a specified expert and a merging operation of collecting the computing results of the experts back.

[0034] Specifically, in order to realize the parallel execution of computing and communication, it is first necessary to decouple the operations in the model training process, that is, to construct independent computing streams and communication streams according to the computing operations and communication operations in the model. Here, constructing the computing stream and the communication stream refers to creating two or more independent, parallel-executable instruction queues on a computing device (such as a GPU). Among them, the computing stream is an execution queue specially used for scheduling and executing computing tasks. The communication stream is an execution queue specially used for scheduling and executing communication tasks. For example, the communication stream can be a communication queue created by using a special communication library and asynchronous with the computing stream.

[0035] It can be understood that when the computing operations (such as attention calculation, fully connected layer calculation, etc.) in the model are scheduled and executed, computing tasks are constituted. When the communication operations (such as distribution operation, merging operation, etc.) in the model are scheduled and executed, communication tasks are constituted. By binding the computing tasks and the communication tasks to independent streams respectively, the decoupling of the two is realized from the hardware and software scheduling level, creating conditions for subsequent parallel execution.

[0036] Step S20, in the training iteration of the model, for the tasks of at least two micro-batches, an interleaved pipeline scheduling is performed; Among them, the interleaved pipeline scheduling includes: scheduling the computing tasks of a first micro-batch to the computing stream for execution, while 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, while scheduling the computing tasks of the second micro-batch to the computing 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 being any two different micro-batches in the at least two micro-batches.

[0037] Specifically, after the computing stream and the communication stream are constructed, in the training iteration of the model, for the tasks of at least two micro-batches, an interleaved pipeline scheduling is performed to alternately schedule the tasks of different micro-batches to the computing stream and the communication stream, so as to realize the parallel execution of the computing tasks and the communication tasks.

[0038] Here, a model training iteration refers to the complete process of the model completing forward computation, backpropagation, and gradient update. To achieve efficient training in a distributed environment, a large training batch is often split into multiple smaller micro-batches.

[0039] The aforementioned at least two micro-batches of tasks means that in one training iteration, two or more micro-batches of data are being processed in the pipeline, for example, micro-batch A (or the first micro-batch) and micro-batch B (or the second micro-batch). Each micro-batch consists of a series of computational and communication tasks.

[0040] It's understandable that staggered pipeline scheduling is an advanced scheduling strategy that differs from traditional serial execution. In traditional schemes, all tasks (computation and communication) of a micro-batch A must complete before tasks of the next micro-batch B can begin. Staggered pipeline scheduling breaks this barrier. It leverages the independent computation and communication flows established above to overlap tasks from different micro-batches. Its purpose is precisely to achieve parallel processing of computation and communication tasks, thereby utilizing idle computation time that would otherwise be spent waiting for communication to complete, and vice versa, to maximize hardware utilization.

[0041] Specifically, when performing interleaved pipeline scheduling for tasks in any two micro-batches 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 of the first micro-batch, the communication stream can simultaneously process a communication task of 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).

[0042] When performing interleaved pipeline scheduling for tasks from any two micro-batches in at least two micro-batches, the communication tasks from the first micro-batch can be scheduled to the communication stream for execution, while the 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 maximum overlap between different stages of the pipeline.

[0043] It should be noted that the terms "first micro-batch" and "second micro-batch" are relative terms and do not refer to specific micro-batches. Instead, they refer to any two different micro-batches within at least two micro-batches. Generally, they can refer to any two micro-batches that are temporally adjacent or at different processing stages, and their reference relationship changes dynamically as the pipeline executes.

[0044] For example, it is assumed that in one training iteration, micro-batch A, micro-batch B and micro-batch C enter the pipeline in turn. At a certain moment, micro-batch A finishes the task of the forward stage and starts to enter the reverse stage, while micro-batch B starts to perform the task of the forward stage. In this scenario, micro-batch A can be regarded as the first micro-batch, micro-batch B can be regarded as the second micro-batch, and when micro-batch A performs the communication task of the reverse stage, micro-batch B performs the calculation task of the forward stage, and when micro-batch A performs the calculation task of the reverse stage, micro-batch B performs the communication task of the forward stage, thereby realizing parallel processing.

[0045] As the pipeline advances, after micro-batch B finishes the processing of the forward stage, it starts to perform the task of the reverse stage; at the same time, a new micro-batch C enters the pipeline and starts to perform the task of the forward stage. In this scenario, micro-batch B can be regarded as the first micro-batch, and micro-batch C can be regarded as the second micro-batch, thereby realizing the parallel of calculation and communication again.

[0046] Therefore, the reference of the first micro-batch and the second micro-batch is dynamic, and the scheduler can dynamically pair and parallelize any micro-batch performing a calculation task with another micro-batch performing a communication task in the pipeline filled with micro-batches, without being limited to fixed two micro-batches. This mechanism ensures that the overlap of calculation and communication can continue to occur as long as the pipeline is filled (i.e., there are at least two micro-batches), thereby minimizing the hardware idle time during the entire training process and realizing stable and efficient pipeline operation.

[0047] The method provided by the embodiment of the application fundamentally creates conditions for parallel processing of tasks by constructing independent calculation streams and communication streams based on calculation operations and communication operations in the model. On this basis, the calculation task and the communication task that originally need to be serially waited can be overlapped by performing staggered pipeline scheduling on the tasks of different micro-batches. This way effectively hides the communication time in the calculation time, thereby breaking the serial execution bottleneck of calculation and communication. Since the calculation task and the communication task can be executed in parallel, the calculation unit of the calculation device and the communication link responsible for data transmission can be in a busy state at the same time, greatly reducing the idle time of the device and improving the hardware utilization rate, thereby effectively improving the model training efficiency. In addition, compared with the algorithm optimization scheme relying on auxiliary loss functions, the present application is an optimization at the execution scheduling level, and the performance improvement effect can be immediately realized after the start of the training iteration, without waiting for multiple iterations as in algorithm optimization, thereby solving the problem of insufficient real-time performance. More importantly, the present application improves the efficiency by optimizing the task scheduling process and does not introduce any new calculation amount, such as calculating auxiliary loss, thereby avoiding additional calculation overhead.

[0048] Based on any of the above embodiments, step S10 specifically includes: Step S11, reconstructing the computing operations in the model into at least one computing node and reconstructing the communication operations in the model into at least one communication node; Step S12, based on the at least one computing node and the at least one communication node, respectively constructing the computing flow and the communication flow, the computing flow being used to execute the computing task corresponding to the at least one computing node, and the communication flow being used to execute the communication task corresponding to the at least one communication node.

[0049] It should be noted that the above embodiments propose the basic idea of separating computing and communication tasks into different flows. However, in order to make this separation and subsequent pipeline scheduling more efficient, the embodiments of the present application introduce the key step of node reconstruction.

[0050] Node reconstruction refers to logically and implementationally aggregating multiple fine-grained operations in the original computation graph of the model to form larger and more regular execution units, i.e., nodes. The purpose is to reduce the startup overhead of kernel functions on computing devices (such as kernels on GPUs), and to improve locality through operator fusion to reduce unnecessary memory read and write. In addition, through node reconstruction, computing blocks and communication blocks with coarser granularity and clearer boundaries can be formed, thereby simplifying the scheduling logic and making pipeline overlap more stable and efficient.

[0051] Specifically, the node reconstruction in the embodiments of the present application is divided into reconstruction of computing nodes and reconstruction of communication nodes. Here, a computing node is a logical unit fused by one or more computing operations, and the computing task it executes mainly consumes computing unit resources. A communication node is a logical unit composed of one or more communication operations, and the communication task it executes mainly occupies communication link resources.

[0052] After reconstruction is completed, execution flows can be constructed based on these clearly defined nodes. Specifically, the computing flow is used to execute the computing task corresponding to the at least one computing node, and the communication flow is used to execute the communication task corresponding to the at least one communication node, which means that the scheduling system will submit the execution instructions of all computing nodes to the computing flow, and submit the execution instructions of all communication nodes to the communication flow, thereby realizing the separation of the two at the scheduling source.

[0053] Based on any of the above embodiments, in step S11, the reconstructing the computing operations in the model into at least one computing node includes: fusing the attention computation, residual computation, routing computation and rearrangement operation in the model into a first computing node; fusing the fully connected layer computation and inverse rearrangement operation in the model into a second computing node.

[0054] Specifically, in reconstructing the computing nodes according to the computing operations in the model, in order to maximize the granularity of the computing task, the embodiments of the present application fuse the attention computing, residual computing, routing computing and rearrangement operation in the Transformer layer of the model into a first computing node.

[0055] Here, the attention computing is the core of the Transformer, which is used to calculate the correlation weight between different parts of the input sequence. The residual computing is a connection that adds the module input and output, which is used to stabilize the deep network training. The routing computing is an operation specific to the MoE model, which is used to determine which expert to assign according to the features of the input token. The rearrangement (i.e. MoE Permute) operation is to rearrange the order of the tokens in the memory according to the assignment result after the routing decision, so that all tokens sent to the same expert are physically continuous, which is the premise of efficient execution of the subsequent All-to-All communication.

[0056] Figure 3 is a schematic diagram of node reconstruction provided by the present application, as Figure 3 shown, the attention computing, residual computing, routing computing and rearrangement operation in the Transformer layer are fused into a single, larger computing node, i.e. a first computing node (denoted as Attention). The node reconstruction mechanism proposed by the embodiments of the present application first realizes the fusion of routing computing and attention computing in the MoE model, which can significantly reduce the startup overhead of the GPU and the memory exchange of the intermediate data, thereby reducing the overhead of stream synchronization.

[0057] Further, as Figure 3 shown, the fully connected layer computing and inverse rearrangement operation in the Transformer layer of the model are fused into a second computing node. Here, the fully connected layer computing can be MLP (Multi-Layer Perceptron) layer computing, which is the main computing of the expert network itself, usually consisting of two large-scale general matrix multiplication (GEMM) and a nonlinear activation function, and is the largest part of the model training. The inverse rearrangement operation (i.e. MoE Unpermute) is used to restore the output processed by the expert to its original order in the sequence after the expert computing is completed and the data is transmitted back through the communication node, so that the subsequent layer can be correctly used and calculated. The two operations are fused to form a second computing node (denoted as MLP), which represents the main expert computing load in the MoE model.

[0058] Based on any of the above embodiments, in step S11, the communication operation in the model is reconstructed into at least one communication node, including: The distribution operation in the model is reconstructed into a first communication node, and the combination operation in the model is reconstructed into a second communication node, and the first communication node and the second communication node both perform global exchange communication.

[0059] Specifically, although the communication operation in the model is of a single type, reconstructing it into an independent communication node helps to clarify its role and dependency relationship in the pipeline. Specifically, the distribution operation in the model can be reconstructed into a first communication node, and the combination operation in the model can be reconstructed into a second communication node.

[0060] Here, the distribution (Dispatch) operation refers to sending the continuous token data block on each device to the target device with the corresponding expert through global exchange communication (All-to-All) after the execution of the first computing node corresponding computing task. This operation is reconstructed into a first communication node (denoted as Dispatch).

[0061] The combination (Combine) operation refers to transmitting the expert output result distributed on different devices back to the original device of each token through another global exchange communication (All-to-All) after the completion of the calculation of each expert (i.e., after the execution of the second computing node corresponding computing task). This operation is reconstructed into a second communication node (denoted as Combine).

[0062] The method provided by the embodiment of the application not only separates the calculation and the communication, but also aggregates the original fine-grained operation into coarse-grained, optimized computing nodes and communication nodes through the node reconstruction manner. This reconstruction lays a solid foundation for subsequent efficient interleaved pipeline scheduling. On the one hand, the operator fusion reduces the overhead of the calculation itself; on the other hand, the clear and neat node division makes it easier for the scheduler to identify the large block tasks that can be parallel, so as to realize more sufficient and stable calculation and communication overlap between different micro-batches.

[0063] Based on any of the above embodiments, the method further includes: Based on the communication flow, an asynchronous communication handle is acquired; Before executing the target computing task on the computing flow, a waiting operation is inserted, the waiting operation is used for waiting for the execution of the target communication task associated with the asynchronous communication handle to be completed, and the input of the target computing task depends on the output of the target communication task.

[0064] It should be noted that the embodiment of the present application further introduces a key technology for ensuring the correctness of data dependency in the framework of parallel execution after the separation of the computing flow and the communication flow, that is, a dynamic synchronization control mechanism. In the foregoing embodiment, the present application constructs independent computing flow and communication flow, and reconstructs the operation into regular computing nodes and communication nodes, thereby realizing parallel scheduling of tasks.

[0065] However, these parallel tasks are not completely independent, and they have strict data dependency. For example, the input of a computing task may be exactly the output of another communication task. Without control, the computing task may start to execute when the data is not ready, resulting in calculation errors. In this regard, the embodiment of the present application proposes a dynamic synchronization control mechanism, which aims to accurately manage such dependencies with minimal performance cost.

[0066] Specifically, the core of the dynamic synchronization control mechanism is to realize accurate synchronization across flows by using an asynchronous communication handle and inserting a waiting operation at a key point. When a communication task (for example, an All-to-All operation performed by the first communication node or the second communication node) is submitted to the communication flow, it can be started in an asynchronous mode. In this mode, the function that calls the communication operation will not block and wait for its completion, but will return an asynchronous communication handle immediately.

[0067] Here, the asynchronous communication handle is a lightweight object or identifier that represents a communication task that is being executed in the background and has not yet completed. It does not contain the communication result itself, but can be used to query the task status or wait for its final completion. Specifically, to obtain the asynchronous communication handle, a specific parameter (such as async_op=True) can be passed when calling the communication library function (such as combine.backward, which is the merge operation in the backward propagation phase), and the asynchronous operation can be started and the returned handle is obtained, which is the asynchronous communication handle and can be denoted as combine_bwd_handle.

[0068] It can be understood that the purpose of obtaining the asynchronous communication handle is to decouple the initiation and completion of the task. After submitting the communication task to the communication flow, the initiator (such as the main scheduling thread) can not have to wait in place, but can immediately continue to execute and schedule other unrelated tasks (for example, submit another micro-batch computing task to the computing flow), thereby maximizing parallelism.

[0069] After obtaining the asynchronous communication handle, a wait operation can be inserted before the compute stream executes the target compute task. The wait operation is a synchronization instruction (e.g., combine_bwd_handle.wait ()), which, when the compute stream encounters the instruction, suspends the execution of subsequent instructions until the event that the instruction is waiting for occurs. In this case, the event is the completion of the execution of the communication task represented by the asynchronous communication handle associated with the instruction.

[0070] It should be understood that the target communication task and the target compute task are a pair of "producer-consumer" tasks with a direct data dependency relationship. The target communication task is executed on the communication stream, and its output is the input of the target compute task. The target compute task is executed on the compute stream, and its input depends on the output of the target communication task, which means that the target compute task must obtain all the input data it needs after the target communication task is completed, otherwise it cannot be executed correctly. For example, as shown in FIG. 1, the compute task of the expert network (i.e., the full connection layer computation) has its input as the token distributed to the current device, so it must wait for the completion of the distribution communication task. Figure 3

[0071] Based on any of the above embodiments, when the target communication task is a distribution communication task in the forward phase, the target compute task is the compute task corresponding to the second compute node in the forward phase; When the target communication task is a combination communication task in the backward phase, the target compute task is the compute task corresponding to the second compute node in the backward phase; When the target communication task is a distribution communication task in the backward phase, the target compute task is the compute task corresponding to the first compute node in the backward phase.

[0072] It should be noted that in order to achieve precise synchronization, the embodiments of the present application insert a wait operation at a specific stage of training. Model training mainly includes a forward (forward) phase and a backward (backward) phase. Among them, the forward phase refers to the process of data flowing from the input end to the output end, and the model predicting the value; the backward phase refers to the process of loss gradient flowing from the output end to the input end, and the model parameters are calculated and updated.

[0073] ​In particular, when the target communication task is the forward stage dispatch 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, in the forward computation, the first communication node (Dispatch) is responsible for dispatching tokens to each expert. Only when these tokens are transmitted through the network and correctly received by the current device, can the second computation node (MLP) responsible for performing expert computation start its forward computation. Therefore, a wait operation on dispatch_fwd_handle (denotes the asynchronous communication handle of the forward dispatch communication task Dispatch.forward) must be performed before MLP.forward on the computation flow.

[0074] When the target communication task is the backward stage combine communication task (denoted as Combine.backward), the target computation task is the computation task corresponding to the second computation node in the backward stage (denoted as MLP.backward). In the backward propagation, the computation order of the gradient is opposite to that in the forward stage. The gradient from the subsequent layer first needs to be transmitted back to the location of the expert from the original location of the token through the second communication node (Combine). After the expert (corresponding to the MLP layer) receives these gradients, it can compute its own parameter gradient and the gradient transmitted to the previous module. Therefore, before MLP.backward on the computation flow is executed, combine_bwd_handle (denotes the asynchronous communication handle of the backward combine communication task Combine.backward) must be waited for to complete.

[0075] When the target communication task is the backward stage dispatch communication task (denoted as Dispatch.backward), the target computation task is the computation task corresponding to the first computation node in the backward stage (denoted as Attention.backward). After the expert (corresponding to the MLP) completes the backward computation, the obtained gradient needs to be transmitted back to the original location of the token through the backward operation of the first communication node (Dispatch). After the first communication node (Attention) receives these transmitted back gradients, it can continue its backward computation. Therefore, before Attention.backward on the computation flow is executed, dispatch_bwd_handle (denotes the asynchronous communication handle of the backward dispatch communication task Dispatch.backward) must be waited for to complete.

[0076] The method provided by the embodiment of the present application establishes an efficient and accurate dynamic synchronization mechanism. It performs necessary synchronization waiting only at the last moment before the data is really needed through the asynchronous communication handle and the delayed waiting mode, instead of blocking the whole process at the beginning of the communication task. This use-and-wait strategy solves the data dependency problem in parallel execution, while maximizing the overlap window of computation and communication, which is a key technical guarantee for realizing efficient interleaved pipeline, thereby maximizing the training performance on the premise of ensuring the correctness of the computation.

[0077] Based on any of the above embodiments, the method further comprises: allocating an independent communication buffer for the communication task associated with the communication flow before performing the interleaved pipeline scheduling; if the target communication task associated with the asynchronous communication handle is not executed within the preset time length of the waiting operation, adjusting the scheduling of the task of one or more subsequent micro-batches.

[0078] It should be noted that, in order to further improve the stability and robustness of the distributed training process, the embodiment of the present application further provides a fault-tolerant mechanism. In a complex distributed computing environment, communication delay or failure is a common problem, which may cause the entire training process to stall or crash. The embodiment of the present application can effectively deal with such abnormal situations by introducing a timeout detection and buffer pre-allocation mechanism.

[0079] Specifically, the embodiment of the present application pre-allocates an independent communication buffer for the communication task before performing the interleaved pipeline scheduling. Specifically, before the training starts or in the initialization phase of each iteration, the system will pre-allocate a dedicated buffer in the memory of the computing device (such as GPU) for each expert's communication operation (such as All-to-All Dispatch and Combine). The advantages of this are: first, it avoids the runtime overhead and delay caused by dynamic allocation of memory during training; second, by providing an independent memory space for each communication task, it eliminates the execution conflicts or blockages that may be caused by memory resource competition, providing a guarantee for the smooth parallel execution of the computation flow and the communication flow.

[0080] In addition, the embodiment of the present application also introduces a timeout detection mechanism in dynamic synchronization control. In order to ensure the correctness of data dependency, the computation flow will synchronize the communication flow through a waiting operation before executing a task that depends on the communication result. However, indefinite waiting will cause the system to be in a deadlock when the communication is abnormal. Therefore, the waiting operation in the embodiment of the present application is configured with a preset time length of timeout monitoring.

[0081] Specifically, the system sets a reasonable timeout threshold for each waiting operation. During the training process, when the compute stream executes a waiting operation, the timer is started. If the target communication task associated with the waiting asynchronous communication handle has not been completed (e.g., an All-to-All communication times out due to network congestion or node failure) within the preset time period, the system no longer continues to wait, but triggers the preset fault-tolerant processing.

[0082] The core of the fault-tolerant processing is to adjust the scheduling of subsequent tasks. For example, the system can mark the current micro-batch that has a communication failure as failed, and discard it from the current training pipeline, while continuing to schedule subsequent other normal micro-batches to avoid blocking the entire training process by a single failure point. In some other implementations, the system can also try to reschedule the failed communication task to the communication stream, or report the exception to the higher-level training management module, and let it decide whether to re-perform expert routing allocation or adjust the training parameters. In this way, the embodiments of the present application not only realize efficient computation and 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.

[0083] Based on any of the above embodiments, the parallel processing includes at least one of: parallel processing the backward communication task of the first micro-batch and the forward computation task of the second micro-batch; parallel processing the backward computation task of the first micro-batch and the forward communication task of the second micro-batch.

[0084] Specifically, Figure 4 is a schematic diagram of parallel processing of tasks of a first micro-batch and tasks of a second micro-batch provided by the present application, as Figure 4 shown, which shows the case of staggered pipeline scheduling of two consecutive micro-batches (such as the first micro-batch and the second micro-batch) on the compute stream and the communication stream after the training iteration enters a stable state. In this stable state, the pipeline is "filled", that is, when 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. The embodiments of the present application make use of this time window to realize the overlap of computation tasks and communication tasks.

[0085] In an embodiment, as Figure 4 shown, when parallel processing the tasks of the first micro-batch and the tasks of the second micro-batch, the backward communication task of the first micro-batch and the forward computation task of the second micro-batch can be parallel processed. This parallel mode takes advantage of the gap of one micro-batch performing computation to execute the backward communication of another micro-batch.

[0086] For details, please refer toFigure 4 The first row of overlapping parts. When the second micro-batch is executing its forward computation task on the computation stream, i.e. Attention.forward corresponding to the first computation node, the hardware's computation unit is occupied. Meanwhile, the second micro-batch is executing its backward communication task, i.e. Combine.backward corresponding to the second communication node, which is scheduled to be executed on the independent communication stream. Since the computation stream and the communication stream can work in parallel, the computation process of Attention.forward and the communication process of Combine.backward achieve time overlap.

[0087] Similarly, referring to the fourth row of overlapping parts of Figure 4 When the second micro-batch is executing its forward communication task, i.e. Combine.forward corresponding to the second communication node, the first micro-batch is executing its backward computation task on the computation stream, i.e. Attention.backward corresponding to the first computation node, and the two again achieve parallelism.

[0088] In another embodiment, as shown in Figure 4 When the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, the backward 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 is complementary to the above case, ensuring that hardware resources can be fully utilized at each stage of the pipeline.

[0089] Specifically, referring to the second row of overlapping parts of Figure 4 When the second micro-batch is executing its forward communication task, i.e. Dispatch.forward corresponding to the first communication node, the task is running on the communication stream and occupies the communication bandwidth. Meanwhile, the first micro-batch is executing its backward computation task on the computation stream, i.e. MLP.backward corresponding to the second computation node. The communication process of Dispatch.forward and the computation process of MLP.backward achieve time overlap.

[0090] Similarly, referring to the fourth row of overlapping parts of Figure 4 When the second micro-batch is executing its forward communication task, i.e. Combine.forward corresponding to the second communication node, the first micro-batch is executing its backward computation task on the computation stream, i.e. Attention.backward corresponding to the first computation node, and the two again achieve parallelism.

[0091] The method provided by the embodiment of the present application greatly reduces the idle waiting time of a computing device such as a GPU by ingeniously overlapping a forward stage task of one micro batch with a reverse stage task of another micro batch, whether the overlap is between a computing task and a communication task or between a communication task and a computing task. This strategy of hiding the execution time of one task within the execution time of another task is the fundamental reason why the present application can break through the bottleneck of the prior art and significantly improve hardware utilization and end-to-end training efficiency.

[0092] The model distributed training optimization device provided by the present application is described below. The model distributed training optimization device described below can be referred to in correspondence with the model distributed training optimization method described above.

[0093] Based on any of the above embodiments, Figure 5 is a structural schematic diagram of the model distributed training optimization device provided by the present application, as Figure 5 shown, the device comprises: The construction unit 510 is configured to construct a computing flow and a communication flow based on a computing operation and a communication operation in a model, the computing flow being configured to execute a computing task, and the communication flow being configured to execute a communication task. The execution unit 520 is configured to execute an interleaved pipeline scheduling for tasks of at least two micro batches in a training iteration of the model. The interleaved pipeline scheduling comprises scheduling a computing task of a first micro batch to the computing flow for execution, while scheduling a communication task of a second micro batch to the communication flow for execution, or scheduling the communication task of the first micro batch to the communication flow for execution, while scheduling a computing task of the second micro batch to the computing flow 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 being any two different micro batches in the at least two micro batches.

[0094] The device provided by the embodiment of the application fundamentally creates conditions for parallel processing of tasks by constructing independent computing flow and communication flow based on the computing operation and the communication operation in the model. On this basis, the originally serially waiting computing task and the communication task are overlapped by performing staggered pipeline scheduling on different micro-batch tasks, which effectively hides the communication time in the computing time, thereby breaking the serial execution bottleneck of computing and communication. Since the computing task and the communication task can be executed in parallel, the computing unit of the computing device and the communication link responsible for data transmission can be in a busy state at the same time, greatly reducing the idle time of the device and improving the hardware utilization, thereby effectively improving the model training efficiency. In addition, compared with the algorithm optimization scheme relying on an auxiliary loss function, the application is an optimization at the execution scheduling level, and the performance improvement effect can be immediately realized after the start of the training iteration, without waiting for multiple iterations as in algorithm optimization, solving the problem of insufficient real-time performance. More importantly, the application improves efficiency by optimizing the task scheduling process without introducing any new calculation amount, such as calculating auxiliary loss, avoiding additional calculation overhead.

[0095] Based on any of the above embodiments, the construction unit 510 comprises: a node reconstruction subunit configured to reconstruct the computing operation in the model into at least one computing node and reconstruct the communication operation in the model into at least one communication node; an execution flow construction subunit configured 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 being used to execute the computing task corresponding to the at least one computing node, and the communication flow being used to execute the communication task corresponding to the at least one communication node.

[0096] Based on any of the above embodiments, the node reconstruction subunit is specifically configured to: fuse the attention calculation, the residual calculation, the routing calculation and the rearrangement operation in the model into a first computing node; fuse the fully connected layer calculation and the inverse rearrangement operation in the model into a second computing node.

[0097] Based on any of the above embodiments, the node reconstruction subunit is specifically configured to: reconstruct the distribution operation in the model into a first communication node and reconstruct the merging operation in the model into a second communication node, the first communication node and the second communication node both performing global exchange communication.

[0098] Based on any of the above embodiments, the device further comprises a synchronization unit, and the synchronization unit is configured to: obtain an asynchronous communication handle based on the communication flow; Before performing a target computing task on the computing stream, a waiting operation is inserted for waiting for a target communication task associated with the asynchronous communication handle to be performed, and an input of the target computing task depends on an output of the target communication task.

[0099] Based on any of the above embodiments, when the target communication task is a distribution communication task in a forward stage, the target computing task is a computing task corresponding to the second computing node in the forward stage. When the target communication task is a merge communication task in a reverse stage, the target computing task is a computing task corresponding to the second computing node in the reverse stage. When the target communication task is a distribution communication task in a reverse stage, the target computing task is a computing task corresponding to the first computing node in the reverse stage.

[0100] Based on any of the above embodiments, the apparatus further comprises: A buffer allocation unit is configured to pre-allocate an independent communication buffer for a communication task associated with the communication stream before performing the interleaved pipeline scheduling. A scheduling adjustment unit is configured to, if the target communication task associated with the asynchronous communication handle is not performed within a preset time length of the waiting operation, adjust a scheduling of a task of one or more subsequent micro-batches.

[0101] Based on any of the above embodiments, the parallel processing comprises at least one of: Parallel processing the reverse communication task of the first micro-batch and the forward computing task of the second micro-batch; Parallel processing the reverse computing task of the first micro-batch and the forward communication task of the second micro-batch.

[0102] Figure 6 An example of a schematic diagram of an entity structure of an electronic device is shown in FIG. 1. Figure 6As shown, the electronic device can 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 can communicate with each other through the communication bus 640. The processor 610 can invoke the logical instructions in the memory 630 to execute the model distributed training optimization method, which includes: based on the computing operations and the communication operations in the model, constructing a computing flow and a communication flow respectively, the computing flow is used to execute the computing task, and the communication flow is used to execute the communication task; in the training iteration of the model, for at least two micro-batch tasks, an interleaved pipeline scheduling is performed; wherein the interleaved pipeline scheduling includes: scheduling the computing task of a first micro-batch to the computing flow for execution, while scheduling the communication task of a second micro-batch to the communication flow for execution, or scheduling the communication task of the first micro-batch to the communication flow for execution, while scheduling the computing task of the second micro-batch to the computing flow for execution, so that the tasks of the first micro-batch and the tasks of the second micro-batch are processed in parallel, and the first micro-batch and the second micro-batch are any two different micro-batches in the at least two micro-batches.

[0103] In addition, the logical instructions in the memory 630 described above can be implemented in the form of a software functional unit and sold or used as an independent product, which can be stored in a computer readable storage medium. Based on such understanding, the technical solutions of the present application or the part that contributes to the related art or part of the technical solutions can be embodied in the form of a software product. The computer software product is stored in a storage medium and includes a plurality of instructions for causing a computer device (which can be a personal computer, a server, or a network device, etc.) to execute all or part of the steps of the method described in various embodiments of the present application. The aforementioned storage medium includes: a U disk, a mobile hard disk, a read-only memory (ROM, Read-Only Memory), a random access memory (RAM, Random Access Memory), a magnetic disk or an optical disk, and various program code storage media.

[0104] In another aspect, the present application also provides a computer program product, which comprises a computer program, the computer program being stored in a non-transitory computer readable storage medium, and the computer program being executable by a processor to enable a computer to perform the model distributed training optimization method provided by any of the above methods, the method comprising: based on the computing operation and the communication operation in the model, constructing a computing flow and a communication flow respectively, the computing flow being used for performing the computing task, and the communication flow being used for performing the communication task; in the training iteration of the model, performing an interleaved pipeline scheduling for the tasks of at least two micro-batches; wherein the interleaved pipeline scheduling comprises: scheduling the computing task of a first micro-batch to the computing flow for execution, while scheduling the communication task of a second micro-batch to the communication flow for execution, or scheduling the communication task of the first micro-batch to the communication flow for execution, while scheduling the computing task of the second micro-batch to the computing flow for execution, so as to process the tasks of the first micro-batch and the second micro-batch in parallel, the first micro-batch and the second micro-batch being any two different micro-batches in the at least two micro-batches.

[0105] In yet another aspect, the present application also provides a non-transitory computer readable storage medium, which stores a computer program, the computer program being executable by a processor to implement the model distributed training optimization method provided by any of the above methods, the method comprising: based on the computing operation and the communication operation in the model, constructing a computing flow and a communication flow respectively, the computing flow being used for performing the computing task, and the communication flow being used for performing the communication task; in the training iteration of the model, performing an interleaved pipeline scheduling for the tasks of at least two micro-batches; wherein the interleaved pipeline scheduling comprises: scheduling the computing task of a first micro-batch to the computing flow for execution, while scheduling the communication task of a second micro-batch to the communication flow for execution, or scheduling the communication task of the first micro-batch to the communication flow for execution, while scheduling the computing task of the second micro-batch to the computing flow for execution, so as to process the tasks of the first micro-batch and the second micro-batch in parallel, the first micro-batch and the second micro-batch being any two different micro-batches in the at least two micro-batches.

[0106] The device embodiments described above are merely illustrative, wherein the units described as separate components can or can not be physically separate, and the components displayed as units can or can not be physical units, i.e., can be located in one place, or can be distributed on multiple network units. Part or all of the modules can be selected to achieve the purpose of the present embodiment scheme according to actual needs. Those skilled in the art can understand and implement without creative labor.

[0107] Those skilled in the art can clearly understand the technical solutions of the various embodiments from the above description of the embodiments, and the various embodiments can be implemented by means of software with the necessary general hardware platforms, and of course can also be implemented by hardware. Based on such understanding, the above technical solutions, or the essential part or contribution to the related art, can be embodied in the form of a software product, which can be stored in a computer readable storage medium, such as a ROM / RAM, a magnetic disk, an optical disk, and the like, and includes a plurality of instructions to make a computer device (which can be a personal computer, a server, or a network device, etc.) execute the methods described in the various embodiments or some parts of the embodiments.

[0108] Finally, it should be noted that: the above embodiments are only used to illustrate the technical solutions of the present application, and not to limit them; although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that: it can still modify the technical solutions recorded in the foregoing embodiments, or make equivalent replacement for some technical features; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present application.

Claims

1. A model distributed training optimization method, characterized in that: include: Based on the computing operations and communication operations in the model, a computing flow and a communication flow are constructed respectively, wherein the computing flow is used to execute computing tasks, and the communication flow is used to execute communication tasks; During a training iteration of the model, performing staggered pipeline scheduling for at least two mini-batches of tasks; The staggered pipeline scheduling includes: scheduling computing tasks of a first micro-batch to the computing flow for execution, and scheduling communication tasks of a second micro-batch to the communication flow for execution, or scheduling communication tasks of the first micro-batch to the communication flow for execution, and scheduling computing tasks of the second micro-batch to the computing flow for execution, so that tasks of the first micro-batch and tasks of the second micro-batch are processed in parallel, and the first micro-batch and the second micro-batch are any two different micro-batches among the at least two micro-batches.

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

3. The model distributed training optimization method according to claim 2, characterized in that: The step of reconstructing the computing operation in the model into at least one computing node includes: Merging the attention calculation, residual calculation, routing calculation, and reordering operation in the model into a first calculation node; The fully connected layer calculation and the inverse shuffling operation in the model are integrated into a second computing node.

4. The model distributed training optimization method according to claim 2, characterized in that: The reconstructing the communication operation 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 of which perform global switching communication.

5. The model distributed training optimization method according to claim 3, characterized in that: Also includes: Based on the communication flow, obtaining an asynchronous communication handle; Before executing the target computing task on the computing flow, a wait operation is inserted, wherein the wait operation is used to wait for the target communication task associated with the asynchronous communication handle to be completed, and the input of the target computing task depends on the output of the target communication task.

6. The model distributed training optimization method according to claim 5, characterized in that: When the target communication task is a distributed communication task in the forward phase, the target computing task is a 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 computing task is a computing task corresponding to the second computing node in the reverse phase; When the target communication task is a distributed communication task in the reverse phase, the target computing task is a computing task corresponding to the first computing node in the reverse phase.

7. The model distributed training optimization method according to claim 5, characterized in that: Also includes: Before executing the staggered pipeline scheduling, pre-allocating independent communication buffers for communication tasks associated with the communication flow; Within the preset duration of the waiting operation, if the target communication task associated with the asynchronous communication handle is not completed, the scheduling of subsequent one or more micro-batches of tasks is adjusted.

8. The model distributed training optimization method according to any one of claims 1 to 7, characterized in that: The parallel processing includes at least one of the following: Processing the reverse communication task of the first micro-batch and the forward computation task of 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 are processed in parallel.

9. 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, the model distributed training optimization method as described in any one of claims 1 to 8 is implemented.

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

Images