Method and apparatus for training neural network models with zero bubble pipeline parallelization

Abstract

The invention discloses a method and apparatus for training a neural network model with zero bubble pipeline parallelization. Various embodiments relate to a computer-implemented method of training a neural network model with zero bubble pipeline parallelization, the computer-implemented method comprising: performing a plurality of forward propagations through the neural network model, where each of the plurality of forward propagations transforms a corresponding input x to a corresponding output y; executing a plurality of back-propagations through the neural network model, wherein the plurality of back-propagations are divided into a plurality of gradient calculation propagations B and a plurality of parameter calculation propagations W; performing a plurality of gradient calculation propagations B for a corresponding input x and a corresponding output y; and determining pipeline bubbles and performing a plurality of parameter calculation propagations W during the pipeline bubbles.

Description

Technical Field

[0001] Various aspects of the present disclosure relate to computer-implemented methods and apparatus for training neural network models using zero-bubble pipelining parallelization. Background Art

[0002] The area of ​​distributed model training has become a hot topic in the deep learning community, especially with the advent of larger and more complex models. Training these models typically requires a large number of Graphics Processing Units (GPUs) interconnected in various topologies. Over the past few years, various parallelization techniques have been proposed for training Deep Neural Networks (DNNs). Data parallelism (DP) has become the default strategy for small and medium-sized models due to its simplicity. However, beyond a certain model size, it is impossible to fit the model parameters on a single GPU, which is why model parallelism is used.

[0003] There are two main model parallelism schemes: tensor parallelism (TP) and pipeline parallelism (PP). TP divides the matrix multiplications in a layer to several devices, while PP splits the entire model into different stages that can be processed across different devices. There is another alternative form of model parallelism: Zero Redundancy Optimizer (ZeRO), which spreads parameters across devices while maintaining the simplicity of DP.

[0004] Recent studies have shown that achieving optimal performance in large-scale training scenarios requires a non-trivial interaction between DP, TP, and PP policies. In the presence of abundant interconnect resources, such as NVLink between GPUs within a compute node, a mixture of DP, TP, and ZERO policies works effectively.

[0005] Pipeline parallelization is one of the key components of large-scale distributed training, but its efficiency is limited by pipeline bubbles. Delving into the complexity of pipeline parallelization, it is found that the efficiency of its implementation depends heavily on the amount of device idle time (known as pipeline bubbles). Bubbles seem inevitable due to the dependencies between layers. Previous methods have attempted to reduce the bubble rate by increasing the number of concurrent batches in the pipeline. However, a direct consequence of doing so is an increase in peak memory requirements. To alleviate this, GPipe discards some intermediate activations while recalculating them during backpropagation. However, this approach introduces a computational overhead of about 20%. A significant improvement to GPipe's limitations is called one-forward-one-backward (1F1B) scheduling.

[0006] 1F1B provides faster memory cleanup by scheduling backpropagation ahead of time. It produces similar bubble rates for the same number of micro-batches, but has a clear advantage in peak memory. It also reduces bubble size by distributing multiple stages to the same device without requiring additional batches at the expense of more communication.

[0007] Despite various efforts, residual bubbles remain the biggest problem for pipeline parallelization to date. Therefore, efficient methods for pipeline parallelization are needed. Summary of the invention

[0008] Various embodiments relate to a computer-implemented method for training a neural network model using zero-bubble pipelining parallelization, the computer-implemented method comprising: performing a plurality of forward propagations through the neural network model, wherein each of the plurality of forward propagations transforms a corresponding input x into a corresponding output y; performing a plurality of back propagations through the neural network model, wherein the plurality of back propagations are divided into a plurality of gradient computation propagations B and a plurality of parameter computation propagations W; performing the plurality of gradient computation propagations B for the corresponding input x and the corresponding output y; and determining a pipeline bubble and performing the plurality of parameter computation propagations W during the pipeline bubble.

[0009] According to one embodiment, each of the plurality of gradient computation propagations B is performed for the corresponding input x and the corresponding output y after each of the plurality of forward propagations.

[0010] According to one embodiment, each parameter calculation propagation W in the plurality of parameter calculation propagations W is performed for the corresponding input x and the corresponding output y after each gradient calculation propagation B in the plurality of gradient calculation propagations B.

[0011] According to one embodiment, a pipeline bubble is the idle time when multiple forward passes and multiple gradient calculation passes B are not being performed.

[0012] According to one embodiment, a heuristic algorithm is used to determine the best schedule for performing each step of the multiple forward passes, the multiple gradient calculation passes B, and the multiple parameter calculation passes W.

[0013] According to one embodiment, activation memory for each step of multiple forward propagations, multiple gradient calculation propagations B, and multiple parameter calculation propagations W is calculated.

[0014] According to one embodiment, the heuristic algorithm uses the computed activation memory for each step of multiple forward propagations, multiple gradient computation propagations B, and multiple parameter computation propagations W to determine the optimal schedule.

[0015] According to one embodiment, the computed activation memory is used to schedule as many forward passes as possible before the gradient computation propagation B to minimize pipeline bubbles.

[0016] According to one embodiment, the neural network model is a feed-forward neural network.

[0017] Various embodiments relate to a system for training a neural network model using zero-bubble pipeline parallelization, comprising: a processor, a memory, the memory storing at least one program code, the at least one program code being loaded and executed by the processor to: perform a plurality of forward propagations through the neural network model, wherein each of the plurality of forward propagations transforms a corresponding input x into a corresponding output y; perform a plurality of back propagations through the neural network model, wherein the plurality of back propagations are divided into a plurality of gradient calculation propagations B and a plurality of parameter calculation propagations W; perform the plurality of gradient calculation propagations B for corresponding inputs x and corresponding outputs y; and determine pipeline bubbles and perform a plurality of parameter calculation propagations W during pipeline bubbles.

[0018] According to one embodiment, each of the plurality of gradient computation propagations B is performed for the corresponding input x and the corresponding output y after each of the plurality of forward propagations.

[0019] According to one embodiment, each parameter calculation propagation W in the plurality of parameter calculation propagations W is performed for the corresponding input x and the corresponding output y after each gradient calculation propagation B in the plurality of gradient calculation propagations B.

[0020] According to one embodiment, a pipeline bubble is the idle time when multiple forward passes and multiple gradient calculation passes B are not being performed.

[0021] According to one embodiment, a heuristic algorithm is used to determine the best schedule for performing each step of the multiple forward passes, the multiple gradient calculation passes B, and the multiple parameter calculation passes W.

[0022] According to one embodiment, activation memory for each step of multiple forward propagations, multiple gradient calculation propagations B, and multiple parameter calculation propagations W is calculated.

[0023] According to one embodiment, the heuristic algorithm uses the computed activation memory for each step of multiple forward propagations, multiple gradient computation propagations B, and multiple parameter computation propagations W to determine the optimal schedule.

[0024] According to one embodiment, the computed activation memory is used to schedule as many forward passes as possible before the gradient computation propagation B to minimize pipeline bubbles.

[0025] According to one embodiment, the neural network model is a feed-forward neural network.

[0026] According to one embodiment, a computer-readable storage medium is provided, wherein the storage medium stores at least one program code, and the program code is executed by a processor to implement the above training method.

[0027] According to one embodiment, a computer program product includes computer instructions stored in a computer-readable storage medium; a processor of a computer device reads the computer instructions from the computer-readable storage medium, and the processor executes the computer instructions, thereby causing the computer device to implement the above-mentioned training method.

[0028] It should be noted that the embodiments described in the context of a method for training a neural network model using zero-bubble pipelining parallelization are equally valid for the system, and vice versa. BRIEF DESCRIPTION OF THE DRAWINGS

[0029] The invention will be better understood by reference to the detailed description considered in conjunction with the non-limiting examples and accompanying drawings, in which:

[0030] Figure 1 An exemplary method for training a neural network model using zero-bubble pipelining parallelization according to an embodiment is shown.

[0031] Figure 2 A computational graph showing a multi-layer perceptron (MLP).

[0032] Figure 3 An exemplary diagram of 1F1B pipeline scheduling according to an embodiment is shown.

[0033] Figure 4A-4B An exemplary handcrafted pipeline schedule is shown in accordance with an embodiment.

[0034] Figure 5 An exemplary post-authentication strategy according to an embodiment is shown.

[0035] Figure 6 Shown is a graph comparing throughput of different pipeline schedules according to an embodiment.

[0036] Fig. 7A and Figure 7B An exemplary post-verification strategy of a pipeline schedule generated by ZB-2p and its profiled execution process is shown according to an embodiment.

[0037] Figure 8 The relationship between memory limit and bubble rate according to an embodiment is shown.

[0038] Fig. 9A and Fig. 9B Scheduling by W grouping and scheduling by parameter grouping according to an embodiment are shown. DETAILED DESCRIPTION

[0039] The following detailed description refers to the accompanying drawings, which illustrate specific details and embodiments in which the present disclosure can be practiced. These embodiments are described in sufficient detail to enable those skilled in the art to practice the present disclosure. Other embodiments may also be utilized, and structural and logical changes may be made without departing from the scope of the present disclosure. The various embodiments are not necessarily mutually exclusive, as some embodiments may be combined with one or more other embodiments to form new embodiments.

[0040] Embodiments described in the context of one of the systems or methods are also valid for the other systems or methods. Likewise, embodiments described in the context of a system are also valid for the method, and vice versa.

[0041] Features described in the context of one embodiment may be applicable to the same or similar features in other embodiments accordingly. Features described in the context of one embodiment may also be applicable to other embodiments accordingly, even if not explicitly described in these other embodiments. In addition, additions and / or combinations and / or alternatives described for features in the context of one embodiment may be applicable to the same or similar features in other embodiments accordingly.

[0042] In the context of various embodiments, the articles “a”, “an”, or “the” used with respect to features or elements include reference to one or more of the features or elements.

[0043] As used herein, the term "and / or" includes any and all combinations of one or more of the associated listed items.

[0044] Figure 1 An exemplary method for training a neural network model using zero-bubble pipelining parallelization according to an embodiment is shown.

[0045] like Figure 1 As shown, there may be a method 100 for training a neural network model using zero-bubble pipeline parallelization. In the method 100, a first step 102 may include performing a plurality of forward propagations through the neural network model, wherein each of the plurality of forward propagations transforms a corresponding input x into a corresponding output y. A second step 104 may include performing a plurality of back propagations through the neural network model, wherein the plurality of back propagations are divided into a plurality of gradient calculation propagations B and a plurality of parameter calculation propagations W. A third step 106 may include performing a plurality of gradient calculation propagations B for corresponding inputs x and corresponding outputs y. A fourth step 108 may include determining a pipeline bubble and performing a plurality of parameter calculation propagations W during a pipeline bubble.

[0046] It is understandable that the above Figure 1The operations described are not limited to this particular order. Any suitable modified order of operations may be used.

[0047] According to one embodiment, each of the plurality of gradient computation propagations B is performed for the corresponding input x and the corresponding output y after each of the plurality of forward propagations.

[0048] According to one embodiment, each parameter calculation propagation W in the plurality of parameter calculation propagations W is performed for the corresponding input x and the corresponding output y after each gradient calculation propagation B in the plurality of gradient calculation propagations B.

[0049] According to one embodiment, a pipeline bubble is the idle time when multiple forward passes and multiple gradient calculation passes B are not being performed.

[0050] According to one embodiment, a heuristic algorithm is used to determine the best schedule for performing each step of the multiple forward passes, the multiple gradient calculation passes B, and the multiple parameter calculation passes W.

[0051] According to one embodiment, activation memory for each step of multiple forward propagations, multiple gradient calculation propagations B, and multiple parameter calculation propagations W is calculated.

[0052] According to one embodiment, the heuristic algorithm uses the computed activation memory for each step of multiple forward propagations, multiple gradient computation propagations B, and multiple parameter computation propagations W to determine the optimal schedule.

[0053] According to one embodiment, the computed activation memory is used to schedule as many forward passes as possible before the gradient computation propagation B to minimize pipeline bubbles.

[0054] According to one embodiment, the neural network model is a feed-forward neural network, such as a multi-layer perceptron (MLP).

[0055] A detailed description of methods and systems for training neural network models using zero-bubble pipelining parallelization is discussed below.

[0056] Figure 2 A computational graph showing a multi-layer perceptron (MLP).

[0057] Multilayer Perceptron (MLP) is a type of feed-forward artificial neural network.

[0058] According to various embodiments, a neural network is broken down into stacked layers. There are two functions associated with each layer, forward and reverse. Figure 2 As shown, MLP 100 has a forward propagation 102 and a backward propagation 104 .

[0059] According to various embodiments, in the forward propagation F, the input x is transformed into the output y using the parameterized mapping f(x,W). Backward propagation (critical for training) involves two computations: and

[0060] Thus, the two computations compute the gradient B with respect to the input x and the parameters W of the layer. Traditionally, B and W are grouped and provided as a single backward function. This design is conceptually user-friendly and also works well for DP because the propagation of the gradient of the weights at layer i can overlap with the backward computation at layer i-1. However, in pipeline parallelization, this design unnecessarily increases the computation of sequential dependencies, i.e., B at layer i-1 depends on W at layer i, which is generally not conducive to pipeline efficiency because it creates pipeline bubbles.

[0061] According to various embodiments, back propagation 104 may be divided into two parts, namely, gradient calculation propagation B 104A and parameter calculation propagation W 104B. This will greatly improve pipeline efficiency by reducing sequential dependencies and reducing pipeline bubbles.

[0062] Figure 3 An exemplary diagram of 1F1B pipeline scheduling according to an embodiment is shown.

[0063] 1F1B is a pipeline schedule, where a forward propagation is followed by a backward propagation.

[0064] like Figure 3 As depicted, 1F1B 300 begins with a warm-up phase. During this phase, workers (also referred to as devices 302A-302D) perform different numbers of forward propagations, with each phase typically performing one more forward propagation than the phase immediately following it. After warm-up, each device 302A-302D enters a steady state in which they alternately perform one forward propagation and one backward propagation, thereby ensuring a uniform workload distribution between phases. In the final phase, each device processes the backward propagation of the unfinished in-flight micro-batches, thereby completing the batch.

[0065] Figure 4A-4B An exemplary hand-crafted pipeline schedule is shown in accordance with an embodiment.

[0066] In the system disclosed in this article, back propagation is divided into gradient propagation B and parameter propagation W.

[0067] Parameter propagation W can be flexibly scheduled anywhere after the corresponding gradient propagation B in the same stage. This allows strategic placement of parameter propagation W to fill pipeline bubbles. There are many possible schedules that improve over IFIB, which make different trade-offs in terms of bubble size, communication cost, and memory footprint.

[0068] Figure 4A A memory efficient schedule named ZB-H1 is shown. This memory efficient schedule ensures that the peak memory usage does not exceed the peak memory usage of 1F1B. ZB-H1 generally follows the 1F1B schedule, but adjusts the starting point of parameter propagation W according to the number of warm-up micro-batches. This ensures that all devices maintain the same number of in-flight micro-batches. Therefore, the pipeline bubble size is reduced to one-third of the size of 1F1B. This reduction is because gradient propagation B is started earlier across all devices compared to 1F1B, and the tail bubble is filled by the later started parameter propagation W. Since parameter propagation W generally uses less memory than gradient propagation B (see Table 1), the first device has a peak memory usage consistent with 1F1B.

[0069] Table 1: FLOPs per propagation and activation memory required per transformer layer

[0070]

[0071] Figure 4B A zero-bubble schedule named ZB-H2 is shown. When there is a memory footprint larger than 1F1B and there are a sufficient number of micro-batches, a zero-bubble schedule can be achieved, which is labeled ZB-H2. This zero-bubble schedule introduces more F propagations in the warm-up phase to fill the bubbles before the initial gradient propagation B. Parameter propagations W are reordered at the tail, which changes the layout from a trapezoid to a parallelogram, eliminating all bubbles in the pipeline. In addition, the synchronization between optimizer steps is also canceled here.

[0072] In various embodiments, p is used to denote the number of stages, and b is used to denote the size of each micro-batch. For the transformer architecture, the number of attention heads is denoted as a, the sequence length is denoted as s, and the hidden dimension size is denoted as h. The notation MB / MW is used to represent the activation memory required for one B / W propagation, and TF / TB / TW is used to represent the running time of one F / B / W propagation. For simplicity, only quantitative analysis is performed on the transformer architecture using a typical setting similar to GPT-3, where the hidden dimension size inside the feedforward is 4h and the dimension size of each attention head is h / a.

[0073] In various embodiments, only matrix multiplication (matmul) operations are considered when calculating floating point operations (FLOPs) because they contribute most of the computation in the transformer layer. For each matrix multiplication operation in the forward pass, there are two matrix multiplication operations with the same FLOPs in the corresponding backward pass (see Figure 1), where each belongs to B or W. The approximate formula for calculating the FLOP of the transformer layer is shown in Table 1. It is shown that TW < TF < TB and TB + TW = 2TF, and the activation memory required for B is estimated. After B is completed, it releases some activations that are no longer used, but retains some extra gradients for W As shown in Table 1 above, the total memory required for W is less than that for B.

[0074] In the case without the assumption of TF = TB = TW, the peak memory and bubble size of ZB-H1 and ZB-H2 are quantified in Table 2. Note that for ZB-H1 and ZB-H2, the peak memory of worker i is ((p - i + 1)M B +(i - 1)M W ≤ pM B and (2p - 2i + 1)M B +(2i - 2)M W ≤ (2p - 1)M B

[0075] Table 2: Comparison between IFIB and our handcrafted schedule

[0076]

[0077] In various embodiments, although handcrafted schedules offer simplicity and better understandability, such schedules face several problems in practical applications. First, the schedule under the assumption of TF = TB = TW introduces unwanted bubbles, especially in models where these values vary greatly. Additionally, the communication time (denoted as Tcomm) required to transfer activations / gradients between stages is often ignored in handcrafted schedules, resulting in significant delays in the pipeline flow. Finally, when the available memory is not sufficient to accommodate enough micro-batches for a bubble-free schedule, it becomes particularly challenging to balance minimizing the bubble size and complying with the memory limit.

[0078] To address these challenges and ensure generalization to practical scenarios, the present invention discloses a heuristic algorithm that automatically searches for the best schedule considering the number of pipeline stages p, the number of micro-batches m, the memory limit Mlimit, and the runtime estimates TF, TB, TW, and Tcomm. Using a heuristic strategy, it always generates optimal or near-optimal solutions (especially when m is large enough). When the problem reaches a certain scale, the problem is also systematically formulated as an integer linear programming (ILP) that can be solved using off-the-shelf ILP solvers. These two approaches can be combined: first using a heuristic solution initially and then fine-tuning it with ILP.

[0079] In various embodiments, the heuristic algorithm may have the following steps:

[0080] In the warm-up phase, schedule as many F propagations as possible before the first B to minimize bubbles, within the memory limit. If the memory limit is not reached, the resulting schedule can still have a small bubble (smaller than TF) before the first B, with local search to make careful decisions under memory constraints, since putting in another F may delay later Bs.

[0081] After the warm-up phase, a pattern is followed in which one F and one B are repeatedly scheduled. When there is a gap larger than TW, W is inserted to fill the bubble. When a bubble appears, but the size is smaller than TW, W is still inserted if the current bubble makes the maximum cumulative bubble size in all phases larger. When the memory limit is reached, W is also inserted to reclaim some memory. In general, the heuristic policy enters a stable state following the 1F-1B-1W pattern.

[0082] Throughout the process, it is always guaranteed that pipeline stage i can schedule at least one more F than stage i+1 at any time before it runs out of F. When this difference exceeds one, a local search is performed to decide whether to skip an F in pipeline stage i if it does not result in more bubbles.

[0083] In each stage, when F and B propagations are exhausted, all remaining W propagations are scheduled one by one.

[0084] In various embodiments, for an ILP design approach, any propagation in the pipeline can be uniquely indexed by (i, j, c), where i∈{1, 2, ..., p} indexes the stage, j∈{1, 2, ..., m} indexes the micro-batch, and c∈{F, B, W} represents the specific propagation of the micro-batch. (i，j，c) can be defined as the time cost, and E (i，j，c) can be defined as the end time of propagation. ΔM (i，j，c) It is used to represent the memory increment caused by propagating (i,j,c). For example, ΔM (·，·，F) =M B , because the forward pass results in a net increase in activations M stored for the backward pass B . ΔM (·，·，B) =M W -M B , which removes the memory stored for B while adding those needed by W, and ΔM (·，·，W) -M W Finally, the variables to be searched are based on the order of propagation in the schedule, introducing variable O for the schedule. (i，j，c)→(i，j′，c′) ∈{0,1}. This variable is an indicator of whether the propagation indexed by (i,j,c) is scheduled before (i,j′,c′).

[0085]

[0086] oeLh (i，j，F) ≥E (i-1，j，F) +T comm +T (i，j，F) (2)

[0087] E (i，j，B) ≥E (i+1，j，B) +T comm +T (i，j，B) (3)

[0088] E (i，j，c) ≥E (i，j′，c′ )+T (i，j，c) -O (i，j，c)→(i，j′，c′) ∞ (4)

[0089]

[0090] Overall, the optimization objective (1) is to minimize the time taken by the longest stage. Constraints (2) and (3) add order dependencies on the propagation of F and B for the same batch in adjacent stages. In addition, (4) adds the dependency constraints imposed when we decide the scheduling order. Finally, (5) limits the peak memory to M limit the following.

[0091] Figure 5 An exemplary post-authentication strategy according to an embodiment is shown.

[0092] In most pipeline parallelization practices, synchronization across pipeline stages is usually performed in the optimizer step for numerical robustness. For example, global gradient norms need to be enforced for gradient norm clipping. In mixed precision settings, global checks are performed for NAN values ​​and INF values, and both values ​​require all-reduce communication across all stages. However, synchronization at the optimizer step destroys the parallelogram (Figure 4) and makes zero bubbles impossible. Therefore, an alternative mechanism is used that is used to bypass these synchronizations while still maintaining synchronous optimization semantics.

[0093] In existing implementations, an all-reduce communication is initiated first to collect the global state, followed by the optimizer steps conditioned on the global state. However, most of the time, the global state has no effect, e.g., global checks for NAN and INF are rarely triggered, since in a robust setting, most iterations should not have numerical issues. Empirically, the gradient clipping rate is also low, so it is reasonable to synchronize the global gradient norm at each iteration.

[0094] Based on these observations, post-update verification replaces pre-synchronization. Figure 5The idea is illustrated in Figure 2. At each stage before an optimizer step, the partially reduced global state is received from the previous stage, combined with the local state of the current stage, and propagated to the next stage. The optimizer step of each stage is controlled by the partially reduced state, for example, aborting updates when a NAN is found. Then, during the warmup phase of the next iteration, the fully reduced global state is passed back from the last stage to the first stage. After receiving the global state, each stage performs a validation that determines whether the previous optimizer step is legal. If a modification to the gradient is required, a rollback is issued.

[0095] In the embodiment, the implementation is based on the open source Megatron-LM project and its performance is evaluated using a model similar to GPT-3, see Table 3 for details.

[0096] In an embodiment, the model used may be a large language model such as GPT3. In an embodiment, any suitable model may be used.

[0097] During the experiment, a specific number of iterations were performed to profile and collect empirical measurements of TF, TB, TW, and Tcomm. Once these values ​​were obtained, they were fed to the automatic pipeline scheduling algorithm to determine the best schedule. Note that both the initial and final pipeline stages have one less transformer layer compared to the intermediate stages. This design is to compensate for the additional embedding lookups and loss calculations in the initial and final stages so that they will not become a bottleneck and cause bubbles to other stages.

[0098] Table 3: Models and fixed settings used in the experiments

[0099]

[0100] Compare the following methods:

[0101] ZB-1p: Automatic search schedule with limit pM B of activated memory, which theoretically has the same peak memory as 1F1B.

[0102] ZB-2p: Automatic search schedule with a limit of 2pM B activation memory, which is the minimum amount of memory that is empirically required to achieve near-zero bubbles (see Figure 8 ).

[0103] 1F1B and 1F1B-I: 1F1B and interleaved 1F1B methods using implementations from Megatron-LM. For interleaved 1F1B, the entire model is divided into a series of blocks that are cyclically adopted by each stage, forming an interleaved pipeline. In the interleaved experiments, the maximum number of blocks is used to ensure the minimum number of bubbles, that is, each transformer layer is used as a block.

[0104] The experiments utilized up to 32 NVIDIA A100 SXM 80G GPUs distributed across 4 nodes connected to each other via a RoCE RDMA network. The runtime of each iteration was recorded after several warm-up iterations. Due to the repeatability provided by the Megatron-LM implementation, the correctness of ZB-1p and ZB-2p can be verified without running the models until convergence. A fixed random seed was used to initialize the model and the losses after each iteration of ZB-1p, ZB-2p, and 1F1B were recorded and then verified to be bit-for-bit identical.

[0105] Figure 6 Shown is a graph comparing throughput of different pipeline schedules according to an embodiment.

[0106] Table 4 shows the experimental results. The experiments show that ZB-2p consistently outperforms all other methods across various settings. Note that the throughput of 1F1B, 1F1B-I, and ZB-1p is strongly positively correlated with the number of micro-batches. In contrast, ZB-2p maintains efficiency even with fewer micro-batches. This is because the bubble rate in ZB-2p is almost zero (Table 5) and its throughput is already close to the upper limit. Here, by comparing the throughput of 1F1B with the number of micro-batches, the throughput of 1F1B is significantly higher than that of 1F1B. Multiply to roughly estimate the upper limit.

[0107] As mentioned above, compared with the 1F1B baseline (pM B ) compared to ZB-2p, the improved efficiency comes at the expense of higher memory consumption (2pM B ) at the expense of 1F1B-I. In contrast, ZB-1p is designed to have similar peak memory cost as the baseline. This shows comparable throughput to 1F1B interleaving in an 8-GPU setting. In a multi-node setting where communication bandwidth is more of a bottleneck, ZB-1p significantly outperforms 1F1B-I, highlighting its advantage of reducing pipeline bubbles without incurring additional communication cost.

[0108] Figure 4: Experimental results details

[0109]

[0110] Fig. 7A and Figure 7BAn exemplary post-verification strategy for the pipeline schedule generated by ZB-2p and its profiling execution process according to an embodiment is shown.

[0111] To quantify the efficiency of pipeline scheduling, the bubble rate is used, which is calculated as (cost-m(TF+TB+TW)) / cost. Here "cost" is defined as the maximum execution time of all stages calculated for each schedule using the profiled TF, TB, TW, and Tcomm values. m(TF+TB+TW) is the optimal execution time when all communication overlaps with computation and therefore there are no bubbles in the pipeline.

[0112] Table 5 presents the bubble rates for different schedules. The hand-crafted schedules ZB-H1 and ZB-H2 are included as baselines for the automatically searched schedules. In most settings, ZB-2p produces a bubble rate below 1%, which is the best among all schedules. In contrast, ZB-H2 performs consistently worse than ZB-2p. This provides strong evidence that the automatic scheduling algorithm is better adapted to realistic scenarios by using more accurate estimates of TF, TB, TW, and Tcomm. In contrast, no such improvement is observed in the comparison of ZB-1p with ZB-H1, which may be because memory limitations become the dominant factor. Note that all public methods significantly outperform 1F1B.

[0113] We also plotted actual execution curves for ZB-2p and profiled it on 16 GPUs to provide direct visual evidence that it is a zero-bubble scheduler. Fig. 7A and Figure 7B As shown, the automatically generated ZB-2p schedule is almost bubble-free. The profiled execution has slightly more bubbles, but maintains good overall alignment.

[0114] Figure 5 : Bubble rate of 1FIB, 1FIB-I, ZB-H1, ZB-H2, ZB-1p, ZB-2p under different settings

[0115]

[0116] Figure 8 The relationship between memory limit and bubble rate according to an embodiment is shown.

[0117] To better understand the impact of peak memory limit, we study the relationship between bubble rate and memory limit. We run the automatic scheduling algorithm with a range of memory limits and plot them. Initially, the bubble rate shows a nearly linear downward trend as the memory limit increases. In theory, the curve should be Nearby stabilized.

[0118] From experience, when TF≈TB and Tcomm is relatively small, 2pM Bis a good threshold for achieving a near-zero bubble rate. Beyond the inflection point, while sufficiently large memory limits do lead to a theoretical zero bubble rate, the costs usually outweigh the benefits.

[0119] Fig. 9A and Fig. 9B Scheduling by W grouping and scheduling by parameter grouping according to an embodiment are shown.

[0120] When data parallelism is considered, all-reduce communication is initiated to collect gradients before the optimizer step. Typically, such communication has a poor overlap with computation propagation, which leads to delays, especially when the communication bandwidth is limited.

[0121] like Figure 4A and Figure 4B As shown, usually at the end of an iteration, multiple W are scheduled. For each W propagation, it includes several independent calculations to calculate the gradients for different parameters. Fig. 9A and Fig. 9B As shown, all of these computations can be reordered to group together those that compute the gradients of the same parameters, thereby achieving optimal overlap between computation and communication.

[0122] According to various embodiments, the relationship between memory limit and bubble rate is affected by the height of the bubble before the first B in the initial stage. For the first micro-batch, forward propagation needs to go from the initial stage to the final stage, and backward propagation reverses this process until it finally returns to the initial stage. The total time from the start to the completion of the first micro-batch takes at least p(T F +T B )+2(p-1)T comm , and this time cannot be squeezed due to the dependency chain. Before the first B propagation in the initial stage, the number of F propagations is denoted as k≥1 and the bubble size is denoted as β≥0. Then:

[0123] M limit ≥ kM B (6)

[0124] β=p(T F +T B )+2(p-1)T comm -kT F -T B =(p-1)(T B +2T comm )+(pk)T F (7).

[0125] When increasing k and keeping The size β of the considered bubble decreases linearly.

[0126] If the number of micro-batches is only 1, a pipeline bubble of size (p-1)(TF+TB+2Tcomm) is created. To fill this bubble, multiple additional F propagations need to be scheduled before the B propagation of the first micro-batch. As this number increases until it reaches , the size of the bubble under consideration should decrease linearly.

[0127] In the experiment, the analysis time of TF, TB, TW and Tcomm in ZB-2p across different settings was recorded. These values ​​were then used to calculate the bubble rate for all methods considered above. These values ​​can be seen in Table 6.

[0128] Table 6: T F 、T B 、T W and T comm Analysis time

[0129]

[0130]

[0131] The methods described herein may be performed by one or more circuits and the various processing or computing units and devices and computing entities described herein may be implemented by one or more circuits. In an embodiment, a "circuit" may be understood as any kind of logic implementation entity, which may be hardware, software, firmware, or any combination thereof. Thus, in an embodiment, a "circuit" may be a hard-wired logic circuit or a programmable logic circuit, such as a programmable processor, for example, a microprocessor. A "circuit" may also be software implemented or executed by a processor, such as any kind of computer program, for example, a computer program using a virtual machine code. According to alternative embodiments, any other kind of implementation of the corresponding functions described herein may also be understood as a "circuit".

[0132] Although the present disclosure has been specifically shown and described with reference to specific embodiments, it will be appreciated by those skilled in the art that various changes may be made to the form and details of the present invention without departing from the spirit and scope of the present invention as defined in the appended claims. Therefore, the scope of the present invention is indicated by the appended claims, and all changes within the meaning and equivalent range of the claims should be included therein.

Claims

1. A computer-implemented method for training a neural network model using zero-bubble pipeline parallelization, the computer-implemented method comprising: performing a plurality of forward propagations through the neural network model, wherein each forward propagation in the plurality of forward propagations transforms a corresponding input x into a corresponding output y; Perform multiple back propagations through the neural network model, wherein the multiple back propagations are divided into multiple gradient calculation propagations B and multiple parameter calculation propagations W; Performing a plurality of gradient calculation propagations B for the corresponding input x and the corresponding output y; as well as A pipeline bubble is determined and the plurality of parameter calculation propagations W are performed during the pipeline bubble.

2. The computer-implemented method of claim 1 , wherein: Each gradient computation propagation B of the plurality of gradient computation propagations B is performed for the corresponding input x and the corresponding output y after each forward propagation of the plurality of forward propagations.

3. The computer-implemented method of claim 1 or 2, wherein: Each parameter calculation propagation W of the plurality of parameter calculation propagations W is performed for the corresponding input x and the corresponding output y after each gradient calculation propagation B of the plurality of gradient calculation propagations B.

4. The computer-implemented method according to any one of claims 1 to 3, wherein: The pipeline bubble is an idle time when the plurality of forward passes and the plurality of gradient calculation passes B are not executed.

5. The computer-implemented method according to any one of claims 1 to 4, wherein: A heuristic algorithm is used to determine an optimal schedule for performing each step of the plurality of forward propagations, the plurality of gradient calculation propagations B, and the plurality of parameter calculation propagations W.

6. The computer-implemented method of claim 5, wherein: An activation memory for each step of the plurality of forward propagations, the plurality of gradient calculation propagations B, and the plurality of parameter calculation propagations W is calculated.

7. The computer-implemented method of claim 6, wherein: The heuristic algorithm uses the calculated activation memory of each step of the plurality of forward propagations, the plurality of gradient calculation propagations B, and the plurality of parameter calculation propagations W to determine the optimal schedule.

8. The computer-implemented method of claim 8, wherein: The computed activation memory is used to schedule as many forward passes as possible before the gradient computation propagation B to minimize the pipeline bubbles.

9. The computer-implemented method of any one of claims 1 to 8, wherein: The neural network model is a feedforward neural network.

10. A system for training a neural network model using zero-bubble pipeline parallelization, comprising: A processor, a memory, the memory storing at least one program code, the at least one program code being loaded and executed by the processor to: performing a plurality of forward propagations through the neural network model, wherein each forward propagation in the plurality of forward propagations transforms a corresponding input x into a corresponding output y; Perform multiple back propagations through the neural network model, wherein the multiple back propagations are divided into multiple gradient calculation propagations B and multiple parameter calculation propagations W; Performing a plurality of gradient calculation propagations B for the corresponding input x and the corresponding output y; as well as A pipeline bubble is determined and the plurality of parameter calculation propagations W are performed during the pipeline bubble.

11. The system according to claim 10, wherein: Each gradient computation propagation B of the plurality of gradient computation propagations B is performed for the corresponding input x and the corresponding output y after each forward propagation of the plurality of forward propagations.

12. The system according to claim 10 or 11, wherein: Each parameter calculation propagation W of the plurality of parameter calculation propagations W is performed for the corresponding input x and the corresponding output y after each gradient calculation propagation B of the plurality of gradient calculation propagations B.

13. The system according to any one of claims 10 to 12, wherein: The pipeline bubble is an idle time when the plurality of forward passes and the plurality of gradient calculation passes B are not executed.

14. The system according to any one of claims 10 to 13, wherein: A heuristic algorithm is used to determine an optimal schedule for performing each step of the plurality of forward propagations, the plurality of gradient calculation propagations B, and the plurality of parameter calculation propagations W.

15. The system of claim 14, wherein: An activation memory for each step of the plurality of forward propagations, the plurality of gradient calculation propagations B, and the plurality of parameter calculation propagations W is calculated.

16. The system of claim 15, wherein: The heuristic algorithm uses the calculated activation memory of each step of the plurality of forward propagations, the plurality of gradient calculation propagations B, and the plurality of parameter calculation propagations W to determine the optimal schedule.

17. The system of claim 16, wherein: The computed activation memory is used to schedule as many forward passes as possible before the gradient computation propagation B to minimize the pipeline bubbles.

18. The system according to any one of claims 10 to 17, wherein: The neural network model is a feedforward neural network.

19. A computer-readable storage medium, characterized in that: The storage medium stores at least one program code, and the program code is executed by a processor to implement the training method according to any one of claims 1 to 9.

20. A computer program product, comprising computer instructions stored in a computer-readable storage medium; a processor of a computer device reads the computer instructions from the computer-readable storage medium, and the processor executes the computer instructions, thereby causing the computer device to implement the training method according to any one of claims 1 to 9.

Images