Automatically Planning Optimal Parallel Strategy for Large Language Models

Each parallel strategy has its own focus and limitations. Therefore, efficient training of LLMs on large scale usually requires a combination of multiple parallelization methods.

Data parallelism (DP) can divides a large batch size of data to multiple workers(Li et al. 2014b). It replicates the entire model on multiple workers and accelerates the learning efficiency of the model through parallel computing. The respective gradient will be accumulated periodically to ensure the consistency of the weights in different workers.

Tensor parallelism (TP) partitions weights and activation tensors of LLMs over multiple devices(Dean et al. 2012), and communicates at specific locations at each layer to aggregate block tensors. The parameters can be split along its row or column dimension to reduce the number of communications. With the segmentation strategy provided by Megatron-LM(Shoeybi et al. 2019), each transformer layer only requires two communications in the forward/backward computation. However, because TP communication occurs at each layer, the communication frequency is relatively high. In the backward stage, a certain overlap algorithm(Rashidi et al. 2021; Wang et al. 2022) can be used to perform communication and calculation at the same time, thus reducing communication time cost.

Pipeline parallelism (PP)(Jia, Zaharia, and Aiken 2019; Yang et al. 2021) splits the layers of LLMs into parts and allocates them to multiple workers. Activetion are communicated point-to-point between workers and are calculated in sequence. Because the execution of different layers of LLMs must wait in line, a certain amount of time waste is inevitable. This extra time is called pipe bubbles. There have been a lot of methods trying to reduce the pipeline bubbles. GPipe (Huang et al. 2019) divided a training batch into multiple micro batches and queued for computation to reduce the waiting time of subsequent NPUs. PipeDream-Flush (Narayanan et al. 2021a) is an improved version of GPipe. It chooses the strategy that one forward pass followed by one backward pass (i.e., 1F1B), which effectively reduces the memory usage compared to GPipe. By dividing the reverse calculation into two parts, the new work(Qi et al. 2023) achieves lower bubble rate and effectively improves the throughput of PP.

Others parallelism. The combination of DP TP PP is called 3D parallelism. In addition, there are many other parallel ways. Some methods solve the memory overhead introduced by LLM long sequences through sequence parallelism(Li et al. 2021; Korthikanti et al. 2023). Optimizer parallelism (Rajbhandari et al. 2021)divides optimizer states, gradients, and parameters to reduce redundant memory usage. This makes it possible to train larger models with limited memory resources. Parallel method for Mixture-of-Experts (MoE) model is also proposed(Kim et al. 2021), and the model parallelism is carried out by splitting experts. Most of these parallel approaches are independent of 3D parallelism and can be optimized separately. In addition, most distributed frameworks are based on 3D parallelism. Therefore, in this paper, we only consider the optimal strategy search of 3D parallelism, and will continue to incorporate other parallel methods in the future.

Before our work, there was some other projects tried to solve the problem of optimal parallel strategy searching.

Mindspeed(Ascend 2024a) provides a strategy search algorithm based on profile, which can perform probabilistic search in the pruned strategy space. The profiling-based algorithm provides precise order-preserving estimation, but it takes a long time to search. Gavatron(Miao et al. 2022) uses decision tree and dynamic programming to search for optimal strategy, and supports asymmetric model parallelism across devices. It only uses profiling for computing power estimation, training time is obtained by simulation, which greatly reduces the time of searching algorithm. However, it ignores the optimization of micro-batch. We demonstrate that ignoring this parameter may result in missing the global optimum.

As an analytical performance model, Calculon(Isaev et al. 2023) provides a training time estimation method, which can also be used to select the optimal parallel strategy. In contrast to Gavatron, Calculon sets the micro batch size to the maximum size that the memory can hold for the sake of computing density. However, our experiments indicate that there is an optimal value for micro batch size. When the optimal value is exceeded, additional increase will only lead to an increase in bubbles, which in turn reduces training efficiency. InternEvo(Chen et al. 2024) is a new parallel training framework with specific optimizations for long sequence transformers. The training time estimation and automatic parallelism are included in the framework, but part of the communication is ignored, which may affect the order preservation.

Unlike previous work, our autoparallel algorithm takes into account more comprehensive variables, such as global batch size and micro batch size, which are often assumed to be constants in other work. The newly introduced variables will cause the search space to become larger, especially global batch size will expand the search space to infinity. In order to reduce the complexity caused by the increase of variables, we set up a white-box simulation system. The mathematical proof based on the white-box system helps us to prun the search space a lot, even if more variables are introduced, the search strategy space becomes very limited.

To ensure the effectiveness of optimal strategy search, the simulation of training duration must be accurate and stable, and the order-preserving property must be provided. In this section, we introduce our training duration simulation model. Table LABEL:notation1 lists symbols to be used in this paper.

In order to make an accurate simulation, we split the training duration into several items and model them separately. Specifically, the total training duration is represented as equation 1,

In addition to the main computation and communication, we also consider the computing duration of all-reduce, which is a proportion of the communication duration. This item takes a relatively small proportion of time. However, for the accuracy of subsequent analysis, this item is still included. This part of the calculation is not provided in the paper, but can be obtained by multiplying the corresponding communication time by a computing power/bandwidth coefficient.

The parameters required for modeling include hardware configuration, model configuration, and parallel settings. The hardware parameters are obtained based on the information of the training cluster. The communicate slow down rate $q$ refers to a bandwidth decrease caused by communication jamming when the communication members increases. As analyzed by InternEvo (Chen et al. 2024), the increase of intra-node and inter-node communication members reduces the efficiency of communication operators such as all-reduce, reduce-scatter, and all-gather. Computing power utilization refers the actual computing power compare to its maximum, which change with the computing intensity of model. We use $\rho$ to represent its reciprocal for easy publicity.Parameters that affect the tensor shape, such as $t$ and $b$, affect the computing power utilization.

The model parameters are obtained based on the large language model that the user needs to train. The parallel parameters are objects that need to be optimized in this work and are regarded as the unknown. Compared with other work focusing on automatic parallelism, we do not assume that data parameters such as global batch size and micro batch size are constant. In the experimental section, we will prove that these parts affect the solution of the optimal.

The computation duration is usually the largest part of the total parallel training, so the accuracy of the estimation is important. The process we do the simulation can be summarized as follows:

• •

  Computation duration = Number of floating point operations per NPU $\times$ Maximum computing power per NPU $\times$ Computing power utilization.

For the first part of the above equation, given that the dense large language model is a variant of the Transformer model, we estimate the number of floating point per NPU by analyzing the Transformer architecture. The compute of the transformer-based large language model is mainly from tensor multiplication, which happen mainly at the attention, feed-forward, and vocabulary output layer. The computation of the vocabulary output layer is proved to be unnegligible, because other devices have to wait for this part in pipeline parallel, which introduces a long extra bubble time.

The computing power utilization of the NPU is related to the operator implementation and hardware. Here, we treat it as a black box and modelling it by small-scale profiling. Specifically, computing power utilization $\rho$ is considered as a function of $b,s,h,t$. The simulation function of $T_{F}$ is shown in equation 2,

The specific modeling process is not given in the paper, but it is worth noting that we are only simulating the last NPU of the PP stage, which contains the vocabulary output layer. If we do not simulate the last stage of the PP, the previous NPU must wait for the last stage to complete the computation of the output layer, which makes the bubble time more complex.

Similar to the simulation of the computation duration, the simulation of the communication duration is carried out by the following basic ideas:

• •

  Communication duration = Total communication volume / (Full speed bandwidth $\times$ Slow down rate).

The total communication volume depends on the model and the parallel approach, and this part of the simulation is static and can be obtained through direct analysis of the parallel training framework. On the other hand, the full speed bandwidth and bandwidth slow down rate are dynamic during the training process, depending on the physical connection form and communication algorithm. For example, if a meshed connection and communication algorithm are used within the server, the bandwidth of communication will not decrease with the increase of members. On the other hand, in a ring communication method, when there are fewer communication members, the bandwidth will be high, but as the number of members increases, the bandwidth will linearly decrease. Same as the estimation of computation duration, the dynamic part of communication is obtained through profiling.

In our training framework, all-reduce occurs in both the multi head attention layer and feed-forward layer. This communication can overlap with computation during back propagation. So the basic idea of communication overlap is that, the communication of the activation and the computation of weight can be synchronized during the backward stage.

Figure 2 shows the stages of two all-reduce operations, with the direction of the arrows indicating forward propagation. During backward propagation, the communication between $X_{1}$ and $X_{2}$ can be synchronized with the computations and updates of $W_{A1}$, $W_{A2}$, $W_{P1}$, and $W_{P2}$.

The time estimation of overlap can be summarize as equation 11,

The all-reduce communication time of activation has been described in the TP communication duration, and the computation time of the weight can be expressed as:

There are two formulas, mainly due to the different weights that need to be updated in the multi-headed attention (MHA) layer and the multi-layer perceptron (MLP) layer.

By simulating the training duration items, we be able to obtain the specific expression of Formula 1, thereby whitening the original black box training duration model. Next, we analyze the constituents of our simulation function and establish an integer programming problem as target:

$M_{NPU},M_{m}$ refers to the memory capacity of each NPU and the memory required by the model. This constraint is used to indicate that the NPU memory meets the requirements of the model. This is a very naive representation, and we’ll discuss this part later.

The complete objective function consists of the items in the simulation model. We will not expand this formula in the text, since it will take up too much space. This programming function cannot be solved by brute-force, as it contain global batch size $G$ and micro batch size $b$, which is infinite in the search space. So in this chapter we prun the parameter space by memory restriction and mathematical analysis. With appropriate constraints, most of the search space, include $G$ and $m$ can be pruned by analysis.

Unlike the degree of $p$ and $t$, which are limited by the number of NPUs, global batch size $G$ and micro batch size $b$ have unlimited values, which is the primary difficulty of the programming problem. We first prove that the total time is monotonic with respect to $G$, by rearrange function 13 as:

Item $\phi_{1},\phi_{2}$ is the rearrangement of the simulation function and is strictly positive. It is easy to see that $G$ has a multiplier effect on training time. When the $G$ is large, the time consumption items of $\phi_{2}$ is effectively reduced, and the $\phi_{1}$ part remains static. Therefore, from a perspective of throughput and training efficiency, a large $G$ is preferred. However, large $G$ also means less randomness for gradient, which can lead to a loss in model performance and impair the overall training efficiency. Based on this trade-off, the training of large language models should determine a global batch size according to the requirement of gradient randomness, but a larger global batch size can accelerate the training, by reducing the computation time, bubble time, and DP communication time in $\phi_{2}$

The pruning of micro batch size $b$ is generally based on two considerations: computing efficiency and memory limitation. In terms of computing efficiency, the increasing of $b$ expands the input tensor shape of each operator, which boost computing power utilization, but the increase in computing power is marginally decreasing. On the other hand, an increase in $b$ leads to an increase in bubbles. Based on this trade-off, we think that there is an optimal value for $b$, which is easy to verify by the partial derivative of $T_{T}$, as Formula 20.

Where $\omega_{1},\omega_{2},\omega_{3},\omega_{4}$ are strictly positive. Also, since $\rho$ is inversely proportional to $b$, the sign of the partial derivative in the right expression of Formula 16 is negative. Therefore, there exists an optimal value for $b$. By addressing the relaxation problem of formula 16, we can determine a possible range of $b$. This inference greatly pruns the search space and makes enumeration possible.

As we increasing $b$ for better computing power utilization, the size of the activation also increases, which greatly increases the memory pressure of parallel training. Moreover, a larger activation memory footprint means that we need to use more NPUs for model parallelism, reducing the resources available for data parallelism, which increases training time. Next, let’s limit the size of b by analyzing the memory.

One approach to pruning the search space is to prun the strategy of significantly exceeding the limit of the memory capacity. In part of the work, memory capacity is considered a gray-box model. In actual scenarios, memory fragmentation and memory reclamation time points are unknown. As a result, the actual memory utilization may not be as expected. However, a basic memory simulation model can still guide our work to some extent, pruning strategies that exceed the memory limit.

In object function 18, we use naive formulas to describe memory limits. Based on the naive idea, that the pipeline parallel and tensor parallel should meet the total memory required for model. Here we provide a more accurate description. During the training duration, the NPU memory usage includes the model weight, optimizer status, and activation. The total memory usage during the training process is as follows:

$M_{w},M_{o},M_{a}$ indicate the memory usage of weight, optimizer and activation respectively. Here, the 1F1B scheduling strategy is used to reduce the activation memory usage. The peak activation size that needs to be stored in the pipeline queue is reduced from $bm$ to $bp$. Besides, note that the input activation are not tensor parallelized(Korthikanti et al. 2023), which further increases memory overhead.

Based on this, we can update the boundary of $p$ and $t$ to be:

where

Different from naive thinking, $p$ and $t$ do not distribute model memory evenly, part of activated peak size can only split by tensor parallelism, and the input activation of each layer must be fully loaded based on the tensor parallelism strategy(Narayanan et al. 2021b). Therefore, each server must provide a basic memory size to prevent activation memory overflow, which limits the minimum value of $t$.

As analyzed above, we impose the following three types of restrictions on the search space:

• •

  $G$ can accelerate training, but is specified by the user based on the consideration of the model performance.

• •

  $b$ affects the computing power utilization and bubble time simultaneously, therefore there exists an optimal value. Based on the white box simulation and memory, we can set an upper limit for the search space of $b$.

• •

  The choice of $p$ and $t$ is limited by the memory.

After pruning, the overall optimal strategy searching algorithm is as Algorithm 1.

In this section, we conducted experiments to verify the accuracy and robustness of our algorithm. The main purpose of the experiment was to examine three aspects:

• •

  Simulation precision: How accurately does the simulation algorithm estimate training time under different parallel settings.

• •

  Rank preservation: Whether the algorithm can find the global optimal parallel setting and whether the ranking of sub-optimal strategies is accurate.

• •

  Inference correctness: Whether our inferences about $G,b$, and memory hold.

Our experiment was conducted using 16 Ascend 910b NPUs on 2 servers, with each NPU having a maximum computing power of 313T and a memory capacity of 64GB. Internal communication within the server is done through mesh architecture, while communication between servers is done through ring. The total data volume and global batch size of the experiment is fix to 256. The distributed training framework is ModelLink, a large language models solution that is well adapted to the NPU. The software environment used is python3.8, pytorch2.1.0.

To verify the simulation accuracy of the algorithm for training time, we estimate the training time under different models and parallel configurations, and compare the results with the real profiling. Since most of the profiling tools cannot measure the granularity of our formula 1, we combine those item to align with the profiling tools.What needs to be compared is computation duration, communication duration, and the overlap. Among them, bubbles time and reduce computation time are both included in the communication duration.

Table 2 shows the simulation precision of our algorithm. The experimental models include Baichuan2-7b(Yang et al. 2023), Qwen-14b(Bai et al. 2023) and Aquila2-7b(Zhang et al. 2024). The number of layers of the model is fixed at 32 to reduce the memory overhead and allow us to compare more combinations. We not only measured the optimal parallel strategy, but also the top 5 suboptimal strategies to verify the reliability of the algorithm. Strategies are sorted by their throughput.

It can be seen that our estimation accuracy of the overall computation time is quite good, with an average estimation precision of 96.89% and a minimum overall accuracy over 91.44%. This indicates that our modeling of training duration is reliable and provides a solid foundation for determining the optimal strategy. Observing the fitting effect of molecular terms, it can be found that the estimation of computational time is the most accurate. Based on the estimation of computing power utilization, the estimation precision of computation duration can reach an average of 97.45%. Estimates of the duration of communication are not very accurate and generally low. This is because the slight deviation in computing efficiency of each worker introduces random waiting times. In our algorithm, the waiting time is about 1 second, and we do not deal with this random term specifically because it has limited influence on optimal strategy planning.

In the test, our algorithm successfully suggests the global optimal strategy for all the tested models. The optimal parallel strategy for Baichuan2-7b and Qwen-14b is $(2,4,2,2)$. Among the remaining strategies, $t=4$ is optimal, slightly better than $t=8$ suggested by Megatron-LM.

In addition, our algorithm not only accurately finds the global optimal parallel strategy, but also correctly estimates and sorts the top five sub-optimal strategies. In the experiment, all the strategies are sorted correctly, which proves that our algorithm has reliable rank-preserving property. The deviation between the estimated time and the actual training time is stable.

In previous section, we propose several pruning theorems based on our white box model. These theorems constrain our searching space, such as $G$ and $b$, enabling enumeration based algorithms to conduct. Previous work often ignores these two hyperparameters, especially in the setting of $b$, where most of the work was controversial. For example, Gavatron sets $b$ to 1, while Calculon sets it to the maximum value within the memory capacity. Here, we set up experiments to verify the validity of the inference of $G$ and $b$.