MiCS: Near-linear Scaling for Training Gigantic Model on Public Cloud

Deep learning model training process mainly consists of three phases, i.e., forward computation, backward computation, and parameter updating. In order to train the model faster, we can harness the computing power of multiple machines. A gradient synchronization step is performed before updating the model parameters to ensure all workers will use the same set of parameters to evaluate the incoming new training samples.

Deep learning model training is memory consuming as it needs to hold the model states including model parameters, gradients from backward computation, and optimizer states for parameter updating. Because of the limited on-device memory resource, activation checkpointing and gradient accumulation are typically enabled. Activation checkpointing discards the activation outputs from the forward phase and requires activation recomputation in the backward phase. Gradient accumulation divides one large data batch into multiple small micro-batches to reduce the memory footprint of storing activation outputs. However, for models with billions of parameters, these two techniques alone are not sufficient. Many solutions targeting at gigantic model training are thus proposed.

In this paper, we use the term “gigantic model” to refer to those Deep Neural Network models that consist of billions of densely connected parameters, which means both the size of the model and the per-sample computation of the model, i.e., floating-point operations (FLOPs), are “gigantic”. Nowadays, the commonly adopted models that fall into this category are transformer-based models (Narayanan et al., 2021; Rajbhandari et al., 2019; Radford et al., 2019; Brown et al., 2020; Zhang et al., 2022; Smith et al., 2022) and latest wide computer vision models (Zagoruyko and Komodakis, 2017).

Traditionally, developers use model-parallel (MP) distributed training for gigantic model training. The basic idea is to distribute the model parameters and computations across multiple devices for each training sample. Thus, the memory for storing model states is also distributed across devices. This way of distributing computations comes with issues. Tensor model parallelism as one MP method requires lots of communications during computation (Shoeybi et al., 2020). On the other hand, pipeline MP strategy is advocated with smaller communication overheads, but it suffers from pipeline bubbles and causes under-utilization. Besides, MP solutions are not directly compatible with common frameworks like PyTorch or Tensorflow, and they require non-trivial engineering effort from the user side. Lastly, some of the MP designs (Narayanan et al., 2021; Naumov et al., 2019) are model-specific solutions, making them hard to generalize.

Compared to the MP solutions, ZeRO (Rajbhandari et al., 2019) powered data-parallel (DP) solutions are general to various models and do not require model refactoring. ZeRO partitions the model states onto all devices on the cluster to reduce the memory consumption on each device. ZeRO has three different stages, corresponding to three different levels of memory reduction: ZeRO-1 partitions optimizer states only; ZeRO-2 partitions gradients and optimizer states; ZeRO-3 partitions all three states, i.e., parameters, gradients and optimizer states, evenly across all devices on the training cluster. The full-fledged ZeRO allows us to train the extremely large models when we have a large enough cluster. However, we have to pay communication costs for gathering model parameters during both forward and backward. Figure 1 illustrates the forward and backward passes in ZeRO-3 powered DP, in which the parameters of each layer are partitioned across all the ranks. Here we use the same convention in the high-performance computing (HPC) community where we use a rank number to identify a computing device. Before computing the activations or gradient for a layer, all parameters of this layer are gathered back by all-gather communication. After computing the gradients on each rank with its own part of the data, the gradients are synchronized and partitioned across all ranks using reduce-scatter communication, which aggregates gradients among all ranks and partitions the gradients at the same time. The gradient partition is necessary for gigantic models with billions of parameters due to the limited memory on each rank.

ZeRO-Offload (Ren et al., 2021) and ZeRO-Infinity (Rajbhandari et al., 2021) are two extensions to ZeRO-3. These two systems offload model parameters, gradients, and optimizer states to CPU memory and NVMe SSDs. Both systems suffer from the same communication overheads as ZeRO-3, which will be discussed in the next subsection.

ZeRO’s model state partitioning mechanism results in the heavy use of collective communication for gathering model states , which is demonstrated in Figure 1. Specifically, ZeRO-3 transmits $3(n-1)M/n$ bytes (Rajbhandari et al., 2019) in forward and backward passes, where $M$ denotes the size of the parameters of the model in bytes, $n$ denotes the number of devices. The transmitted data volume is as large as tens to hundreds of gigabytes for models with tens to hundreds of billions of parameters. The cost of transmitting these data crossing the entire cluster cannot be easily hidden via pipelining the communication and computation. Our timeline measurement shows that for a BERT model with 10B parameters, parameter gathering could take 2.85$\times$ more time than computation in forward pass. Similar expensive communications also exist in the backward computation and gradient synchronization, which hurts the performance of ZeRO especially when the network bandwidth between devices is less preferable.

There are two main factors that contribute to the costly communications when using ZeRO-3 on the cloud. Firstly, at the hardware level, many of the available internode network bandwidths and latency of cloud-based GPU clusters are not as good as DGX systems (gcloud-gpu-bandwidths, 2022; AWS-P3-Instances, 2022; azure-gpu-ndv2-series, 2022; Ziegler et al., 2022). Moreover, unlike on-premise clusters, the network topology of cloud-based clusters is out of users’ control, which could negatively impact the network performance (Luo et al., 2020; Luo et al., 2021). Secondly, at the algorithmic side, the latency of collective algorithms for communication has a positive correlation with the communication scale and the startup time for transmission (Chan et al., 2007)¹¹1The latency of tree algorithms is bounded with $\left\lceil\log_{2}(p)\right\rceil\alpha$, and the ring algorithms have a latency term $2p\alpha$, where $p$ denotes the number of participants and $\alpha$ denotes the startup time for transmission, §7.1.7 in (Chan et al., 2007). Therefore, as the scale grows, the latency becomes more significant and hurts the performance of communication at a large scale. In addition, previous studies (Zhang et al., 2020a) suggest that network bandwidth may not be the performance bottleneck of distributed model training. Based on our measurements, as the cluster size grows, we need larger message sizes to saturate the bandwidth. Figure 2 shows that small-size message such as 128MB obtains poor bandwidth utilization on 16 and 32 nodes.

In practice, we may not be able to always communicate large messages due to the memory constraint. Instead, it is better to control the communication scale to improve the bandwidth utilization, especially on the cloud with less favorable network conditions.

As mentioned in §2.3, the frequent communication among all devices significantly hampers the training performance of ZeRO powered DP solutions. This motivates us to design a new system that reduces the cost of communications while preserving the generality and usability advantages. We found the communication overhead can be effectively reduced by shrinking the communication scale, i.e., reducing the number of participants in a collective communication. With the reduced communication scale, the majority of communications are restricted to a smaller group of devices. This allows us to maintain high bandwidth utilization in communications for various sized messages, as shown in Figure 2. In addition, because the transmitted data volume is positively correlated to the number of participants, reducing the communication scale reduces the data volume. We give detailed descriptions of our methodology in the next section.

We define the notations used in this section as follows:

• •

  $n$: Number of devices or ranks in the cluster.

• •

  $k$: Number of devices on each computational node.

• •

  $M$: Size of a model in bytes.

• •

  $p$: Number of devices for holding a model replica.

• •

  $s$: Number of micro-steps.

• •

  $B_{g}$: Effective communication bandwidth among devices belonging to the group $g$. We define the effective communication bandwidth as the bandwidth measured using collective communication. Effective communication bandwidth takes algorithm latency into account. Thus it is smaller than the theoretical bandwidth of hardware specification. For a fixed message size, when the number of nodes increases, the effective bandwidth shrinks.

• •

  $C$: Time cost.

Partitioning model states across all the devices causes significant communications overheads during training. Such communication overheads scale with the number of participants in a single collective communication (§2.3). To reduce the communication overheads, we consider distributing the model states over a subset of devices to reduce the scale of the communication. A modern single computing device typically has tens of gigabytes of memory, and tens of them provide sufficient memory for a model with tens of billions of parameters. For example, a model with 10 billion parameters takes about 160GB of memory when training with Adam optimizer using mixed-precision. Partitioning the model states across 8 V100 (32GB) GPUs is already more than enough. By using 8 V100 GPUs instead of all the devices for holding one model states replica, we can effectively reduce the scale of communication. If 8 GPUs are located on a single node, then we can leverage high-speed intra-node connections such as NVLink/NVSwitch to perform the most communications. Next, we give a general form of model states partitioning in our system and provide an analysis of the benefits.

In MiCS, we divide all the devices into multiple groups and partition the model states within each group. Every group has the same number of devices and holds a complete replica of the model states in training. We call these groups partition groups. Each device is tagged with a local group rank. Devices with the same local group rank form another type of group, named replication group, and they hold the same part of the model states. In Figure 3, we give an example that the model states are partitioned onto two devices. Thus every two devices with consecutive rank numbers form a partition group. The devices ranked with odd numbers and even numbers form two replication groups separately. During the training, when a parameter tensor is needed for either the forward or backward computation, MiCS invokes all-gather collective to gather the corresponding model parameters distributed within each partition group. After the gradients are computed on each device, MiCS uses all-reduce collective to aggregate the gradients, and then it partitions the gradients within each partition group.

Now we give the performance analysis of our partitioning strategy in terms of the cost of all-gather. We assume the iteration time is bounded by the communications, which is true based on our measurements (§2.3). The time cost of ZeRO-3’s partition-to-all strategy is $C_{\text{all}}=((n-1)M)/(nB_{\text{all}})$, where $B_{\text{all}}$ denotes effective bandwidth among all devices. The time cost of MiCS is $C_{\text{MiCS}}=((p-1)M)/(pB_{\text{part}})$. Here we assume all partition groups have the same intra-group bandwidth, and we denote this bandwidth by $B_{\text{part}}$. Because the value of function $(x-1)/x$ increases when $x\geq 1$ and $p\leq n$, we have the following inequality for the ratio of two costs.

For models that can be partitioned to devices located on a single node, only local high-speed NVLinks connections are used for all-gather. Based on the measurement on 64 GPUs spread across 8 computational nodes, we get $B_{\text{part}}\simeq 128\text{GB/s}$ and $B_{\text{all}}\simeq 11\text{GB/s}$. Thus, the cost ratio can be as large as 11.6. For the case where there are 32 nodes in total and each partition group consists of 4 nodes, the cost ratio is ranging from 2.7 to 4.9 based on our measurements presented in §2.3. For the model states that can be partitioned within 4 nodes, we can expect about 63.6% to 91.3% time reduction for parameter gathering with our partitioning strategy.

When the model size grows, it requires more devices to hold the model states for the training. If the required devices span multiple computational nodes, inter-node communication is needed for parameter gathering. We can reduce the transmitted data volume over inter-node connections by reducing the scale. Assuming we want to all-gather a message with size $M$ among $p$ participants, the data traffic transmitted among participants is determined by $(p-1)M/p$ (Chan et al., 2007). This means we can split $p$ participants into multiple small groups and perform independent communication within each group. We consider GPUs spanning across multiple nodes as a two-dimensional grid, in which we first aggregate the data across nodes in parallel and then merge the local data on each node.

For hierarchical communication to work properly, we first build communication channels for devices. Assuming each computational node has $k$ devices, MiCS builds $k$ communication channels for inter-node communication and a separate communication channel for intra-node communication (Figure 4 (right)). As a comparison, vanilla collective communication uses a single communication channel for devices spanning across nodes (Figure 4 (left)). In Figures 4, we illustrate the idea using two computational nodes, each of which has two devices, i.e., $p=4$ and $k=2$. Next, we introduce how hierarchical communication works for inter-node communication.

MiCS uses a three-stage algorithm for hierarchical communication. In the first stage, each device uses the inter-node communication channels to do all-gather with the devices that have the same local rank on respective nodes. The inter-node all-gather operations are executed in parallel. In the second stage, the data chunks are rearranged to ensure correctness. In the third stage, we invoke batched intra-node all-gather. In general, for the model states partitioned onto $p$ devices spanning $p/k$ nodes, the number of batched all-gather calls is $p/k$ in the third stage. An example of the algorithm running across two nodes is given in Figure 5. In the following, we explain why we have the second and third stages work as we described here.

The second and third stages are designed to fix the memory discontiguous issue. Otherwise, we would get the wrong output. We use the model states partitioned to two nodes with 4 GPUs for an explanation, shown in Figure 5. The final outputs of the hierarchical communication algorithm should place data $C0$ and $C1$ in the adjacent locations. However, the inter-node all-gather will gather $C0$ and $C2$ into a contiguous memory. Thus, if we directly launch an all-gather collective primitive on the output from the first stage, we will get the wrong memory layout $[C0,C2,C1,C3]$, while the correct one is $[C0,C1,C2,C3]$. To fix this, we add a data movement stage before intra-node all-gather to rearrange the data chunks. Then in the third stage, we launch $p/k=2$ intra-node all-gather collectives in a batch, where each intra-node all-gather operation works on a subset of the data chunks. Launching multiple communications in a batch requires new communication API implementation to get good performance, which is detailed in §4.

The performance benefits of the hierarchical communication strategy depend on the scale of the model states partitioning. Assume the model is partitioned onto $p$ devices, and $p$ is divisible by $k$, where $k$ is the number of devices on each computational node. With the vanilla communication strategy, the inter-node data traffic is $(p-1)M/p$. Using the proposed hierarchical communication, the data volume transmitted over inter-node connections is reduced to $(p-k)M/p$. In this way, the communication volume over the slow inter-node links is reduced by

Given that $p\geq k\geq 1$, this ratio decreases monotonically and approaches $1$ when $p$ increases. Thus, the improvement is less when we have to use more devices to hold a model replica. In a typical setup, we would have $k=8$. A $10$B-$50$B parameter model typically requires $8\leq p\leq 64$ number of workers for holding the model states. In this case, we will obtain $11.1$% to $46.6$% data volume reduction with hierarchical communication.

In the typical distributed training setting, we need to aggregate gradients across all the devices (Goyal et al., 2017). Gradient aggregation is an expensive synchronization step and its cost scales with the number of workers, detailed in §2.3. It ensures that all devices work on the same model states.

To improve the training efficiency, more recent works advocate large-batch training (Li et al., 2021; You et al., 2019; Goyal et al., 2017; You et al., 2018; Zheng et al., 2020). However, due to the limited device memory, practitioners have to resort to gradient accumulation that divides a large batch into multiple micro-batches and accumulates the gradient w.r.t. each micro-batch into a shared memory buffer. In the standard data parallel setting, the gradient synchronization is only needed at the accumulation boundary where all the gradients have been computed. However, ZeRO requires additional gradient synchronization within each micro-step because of gradient partitioning. Since each device is only responsible for holding a part of the gradient, the gradient needs to be partitioned once it is computed. In order to avoid losing the gradient information, gradients have to be aggregated before the partitioning. This makes every gradient partitioning step become a global synchronization barrier among all devices. Since MiCS only partitions the model states into a small group of devices, we can restrict the gradient synchronization within the group for each micro-step and delay the global gradient synchronization to the accumulation boundary. This motivates the design of 2-hop gradient synchronization schedule without over-paying communication costs.

2-hop gradient synchronization performs gradient synchronization within each partition group for each micro-step. Only at the gradient accumulation boundary, global synchronization is performed among the devices that possess the same part of the model. Figure 6 gives an example of a model partitioned onto two devices. Every two consecutive ranks form a partition group. Ranks with odd number and even number indices form two different replication groups, respectively. For illustration purposes, we assume the number of gradient accumulation steps is $s=4$. For each micro-step, MiCS uses reduce-scatter to synchronize gradients within each partition group. At the gradient accumulation boundary, an all-reduce operation is used within each replication group for synchronization. An alternative synchronization schedule is to use all-reduce for gradient synchronization at every micro-step and then partition the gradient on each device. During the partitioning, each device only keeps the part of the gradient that it is responsible for while discarding the rest. This alternative schedule is the default one implemented in DeepSpeed. However, this scheme is redundant and overpays the communication costs.

The performance benefits of the 2-hop gradient synchronization schedule depend on the number of micro-steps and the effective communication bandwidths within partition groups and replication groups. For simplicity, we assume that every partition group has the same effective bandwidth $B_{\text{part}}$, and bandwidth within each replication group is $B_{\text{repl}}$. The time cost of 2-hop schedule is $C_{\text{2-hop}}=(sM(p-1))/(pB_{\text{part}})+2M(n-p)/(nB_{\text{repl}})$, while the time cost for the alternative schedule is $C_{\text{alt}}=2sM(n-1)/(nB_{\text{all}})$. We take the ratio of two costs, and simplify the ratio using inequalities $(p-1)/p\leq(n-1)/n$ and $(n-p)/n\leq(n-1)/n$ when $n\geq p\geq 1$. In the following inequality, we can view the right-hand side as the lower bound for the improvement.

Assuming $s=4$, which is a reasonable setup for large batch training (Shoeybi et al., 2020; Rajbhandari et al., 2019), and assuming $B_{\text{all}}=B_{\text{part}}=B_{\text{repl}}$ for simplicity, we get the lower bound of the ratio at $4/3$. This means at least 25% cost reduction by using the 2-hop schedule. Taking heterogeneous bandwidth into consideration would further reduce the denominator on the right-hand side and helps achieve more gains. We notice that when $s=1$, under the assumption that $B_{\text{all}}=B_{\text{part}}=B_{\text{repl}}$, the 2-hop synchronization is sub-optimal compared to the alternative schedule. However, given the heterogeneity of the effective bandwidths in a large cluster, e.g., having $B_{\text{part}}\simeq B_{\text{repl}}>1.5B_{\text{all}}$ (which is reasonable based on our measurement in §2.3), the 2-hop schedule typically costs less. Therefore, even for $s=1$, in training large models with a large cluster, 2-hop is still preferred.

In this section, we demonstrate the performance advantages of MiCS. First, we show the scalability of MiCS against DeepSpeed, which is the state-of-the-art (SOTA) solution using DP with model states partitioning. The TFLOPS numbers are also reported to show the computation utilization of each GPU. In §5.1.2, we evaluate MiCS and DeepSpeed in a different network condition, i.e., 400Gbps network. In §5.1.5 we provide performance numbers of the 100B model training on a large scale. In §5.1.3, we show MiCS can outperform Megatron-LM-3D (Narayanan et al., 2021), which is a SOTA design for transformer-based language models that uses DP and MP.

To understand the performance contribution of each component in MiCS, we conduct ablation tests in this section. We divide our studies into three subsections. Each subsection corresponds to one of the three components in §3. For each experiment, we present the setups followed by detailed results and takeaways.

We conduct experiments to analyze the performance improvements by using the optimization techniques described in §4. We use the BERT 10B model for the evaluation. In the training, we use the default setup as mentioned at the beginning of §5. When we turn off optimizations that are unique to MiCS and let the model states be partitioned over all devices, MiCS reduces to ZeRO-3 with the optimization techniques in §4, We denote it as “MiCS (ZeRO-3)”. For comparisons, we report the throughput of DeepSpeed ZeRO-3.

Figure 15 shows the improvements of using the proposed system optimizations. MiCS (ZeRO-3) achieves 54.1% better system throughput than DeepSpeed ZeRO-3 when the cluster scales up to 128 GPUs, while the scaling efficiency of DeepSpeed ZeRO-3 drops when we scale out the cluster. In addition, MiCS still significantly outperforms MiCS (ZeRO-3), demonstrating the superiority of minimizing the communication scale.

In this section, we show that MiCS achieves consistent convergence as DeepSpeed, which validates the correctness of our system. We provide the training loss curves for training a 1.5B parameter model on the Wikipedia-en dataset. The model has 48 transformer layers, each of which is constructed with the hidden size 1,600 and intermediate size 6,400. The global batch size is 512. And the micro-batch size is 8 (the number of gradient accumulation steps is 4). The loss validation process does not aim to produce exactly the same loss as DeepSpeed but to ensure the convergence behaviours are the same. We report the training losses on 1 million sequences. As shown in Figure 16, MiCS provides the same convergence as DeepSpeed.