SlowGuess/NeurIPS_2023_OCR huggingface dataset:
H3T: Efficient Integration of Memory Optimization and Parallelism for High-Throughput Transformer Training
Yuzhong Wang\*, Xu Han\*, Weilin Zhao\*
Guoyang Zeng $^{2}$ , Zhiyuan Liu $^{\dagger, 1, 3}$ , Maosong Sun $^{\dagger, 1, 3}$
$^{1}$  NLP Group, DCST, IAI, BNRIST, Tsinghua University, Beijing, China
2 ModelBest Inc., Beijing, China
$^{3}$  Collaborative Innovation Center for Language Ability, Jiangsu Normal University, China
{yz-wang21,zwl23}@mails.tsinghua.edu.cn
zengguoyang@modelbest.cn {hanxu2022,liuzy,sms}@tsinghua.edu.cn
In recent years, big models based on Transformers have achieved state-of-the-art performance on many artificial intelligence (AI) tasks. Despite the success of these Transformer-based models, their huge parameter size poses a serious challenge to their training, both from the storage and computation perspectives. To this end, memory optimization (e.g., rematerialization and offloading) and parallelism (e.g., data parallelism and model parallelism) are widely explored to make training Transformers more efficient. In this paper, we propose a framework to automatically find an efficient integration of memory optimization and parallelism for HighThroughput Transformer Training (named H3T), which is rarely considered by existing efforts for training big Transformer-based models. Specifically, we design search algorithms to combine appropriate memory optimization strategies and parallelism schemes to achieve a balance between memory overhead and training efficiency. We implement H3T based on an open-source toolkit BMTrain and then use H3T to train the Transformers of different sizes to evaluate the efficiency of H3T. The experimental results show that H3T outperforms the most popular deep learning (DL) toolkit Megatron-DeepSpeed by  $1.2 \times 4.3 \times$  training speed while reducing  $34.6\% \sim 80.5\%$  of memory overhead. Moreover, H3T can use only 64 NVIDIA A100 GPUs to train GPT-3-175B, which is very difficult for existing DL toolkits. The source code is available at https://github.com/OpenBMB/BMTrain/tree/h3t.
In recent years, the emergence of Transformers [47] has significantly advanced the development of the AI field. Owing to the strong abilities of sequence modeling and transduction brought by the attention mechanisms, Transformer-based models have achieved state-of-the-art performance on many tasks and have become the foundation architecture for various AI directions, such as natural language processing [33, 9, 34, 4], computer vision [10, 23, 22], and multimodal processing [32, 39, 38].
Despite the outstanding performance, with the trend of continuously growing model size, the computation and storage costs become a severe challenge for training most big Transformer-based models. For example, the popular big model GPT-3 [4] has 175 billion parameters, and all these parameters require at least 700 GB of memory to store in f1oat32 format, not to mention the forward activations,
* Indicates equal contribution.
Indicates corresponding authors.
backward gradients, and optimizer states. Even ignoring the memory overhead, according to the existing analysis [28], it will take about 288 years to train GPT-3 with a single NVIDIA V100 GPU.
To tackle these issues, researchers explore various memory optimization strategies and parallelism techniques. Memory optimization strategies such as offloading [42, 6, 16, 2, 41, 37, 45] and rematerialization [12, 5, 20, 18, 19] aim to reduce the memory overhead of one (or each) single GPU. Parallelism techniques such as data parallelism [53, 30, 11] and parameter parallelism [8, 17, 27, 51, 36] mainly focus on improving both training efficiency and memory utilization across multiple GPUs. Some hybrid memory optimization strategies [49, 31, 3, 15] and automatic parallelism techniques [52, 46, 26] have also been proposed to further improve the storage and computation. Although promising results have been achieved on both tracks of memory optimization and parallelism, respectively, very limited efforts have been devoted to combining them both in an easy and efficient way.
In this paper, we propose a framework to automatically find an efficient integration of memory optimization and parallelism for High-Throughput Transformer Training (named H3T), which is rarely considered by existing efforts for training big models. In H3T, we formalize memory optimization strategies (including offloading and rematerialization) and parallelism techniques (including data parallelism, parameter parallelism, and zero redundancy optimizer (ZeRO)) into multiple optimization switches. We then arrange an automatic solver that selects appropriate switches at each Transformer layer, while the core purpose of the solver is to achieve better training efficiency under the memory constraints. Since each transformer layer has a set of optimization switches independent of the other layers, the giant search space poses a challenge for the solver. Toward this end, we introduce greedy and dynamic programming (DP) algorithms to the searching process, enabling the solver to consider both search performance and search efficiency.
In experiments, our simulation studies show that H3T with our designed solvers outperforms traditional manual optimization without H3T. We also implement H3T based on BMTrain, an open-source toolkit that accelerates the training of big Transformers with various parallelism techniques. We conduct actual-running experiments to compare H3T and Megatron-DeepSpeed [40], one of the most popular distributed learning toolkits for DL. The results show H3T can train big Transformers  $1.2 \times \sim 4.3 \times$  faster than Megatron-DeepSpeed while reducing the memory usage by  $34.6\% \sim 80.5\%$ . Moreover, supported by our implementation of H3T, it is possible to train GPT-3 with 175 billion parameters by using 64 NVIDIA A100 GPUs.
We summarize our contributions as follows: (1) We design H3T, which is, to our best knowledge, the first to automatically integrate memory optimization and parallelism for training big Transformers. (2) We conduct adequate experiments that show H3T outperforms conventional manual optimization as well as the most popular DL toolkit Megatron-DeepSpeed. (3) Supported by our implementation of H3T, we can train GPT-3 with 175 billion parameters on 64 NVIDIA A100 GPUs.
2.1 Transformer-based Model
Transformer [47] is a widely-used deep neural network architecture. Due to the powerful capability of modeling data, the Transformer architecture is used to build various pre-trained models. These pre-trained models can capture rich knowledge from massive unlabeled data and then transfer the captured knowledge to handle various complex AI tasks. Transformer-based pre-trained models have become one of the cornerstones of recent AI developments, especially in the NLP domain. Currently, Transformer-based models achieve state-of-the-art performance on almost all NLP tasks [9, 21, 35, 33, 34, 4]. Inspired by the success of these models in the NLP domain, Transformers are also explored in other important AI domains like CV [10, 23, 22] and Multi-Modal (MM) [32, 39, 38, 29].
Despite the remarkable achievements of Transformers, it is a big challenge to train big Transformer-based models. On the one hand, the parameter size of popular models is increasing in an explosive way and has already exceeded Moore's Law. Take OpenAI's GPT as an example: the parameter sizes of GPT-1 [33], GPT-2 [34], and GPT-3 [4] released in 2018, 2019, and 2020 are 117 million, 1.5 billion, and 175 billion, respectively. On the other hand, Transformers contain heavy dense algebraic operations, such as matrix multiplication, making the time and memory complexities of these operations much higher than the parameter size. The combination of rapid parameter growth and huge computational overhead challenges the DL community.
2.2 Memory Optimization
Rematerialization. Rematerialization, also known as checkpointing, is the process of discarding intermediate results within the GPU and recomputing these discarded results when needed. Rematerialization is first proposed by Grimm et al. [13] and first implemented by Griewank et al. [12]. Chen et al. [5] reduce the memory complexity of the conventional rematerialization method to  $O(\sqrt{N})$  without changing the computational time complexity, which is a milestone that makes rematerialization known to researchers and implemented in many DL toolkits. However, Chen's work is based on artificial rules and cannot precisely balance memory overhead and training efficiency for different models and environments, limiting its performance and scalability. Many researchers devote their efforts to developing automatic rematerialization methods to tackle this issue. Dynamic programming (DP) [14, 20] and mixed-integer linear programming (MILP) [18] are both adopted to find better rematerialization schemes by automatically arranging which neural network layers need to be rematerialized to reduce the memory overhead while minimizing the efficiency loss.
Offloading. Offloading aims to transfer appropriate tensors from GPUs to CPUs and prefetch them back when they are required for computation. Preliminary offloading methods [42, 6, 44] mainly focus on using empirical rules to select the tensors for offloading. However, since offloading introduces complicated communication overhead for data transferring, it is not easy to empirically find the global optimal solution. For this reason, researchers attempt to offload tensors more intelligently. Huang et al. [16] perform tensor offloading based on a custom-designed genetic algorithm, while Beaumont et al. [2] use a DP algorithm. Besides the above offloading methods that mainly focus on activation tensors and convolutional neural networks (CNNs), some other researchers [41, 45] also explore offloading parameters and optimizer states, which can effectively reduce GPU memory usage for other types of model architectures like Transformers.
Hybrid Memory Optimization. Rematerialization saves memory by bringing additional tensor computation, while offloading saves memory by introducing extra tensor movement. These two optimization strategies can coexist, so integrating them and saving more memory is a natural idea. Early works [49, 31] attempt to intelligently combine them using heuristics and achieve remarkable performance. Afterward, some scholars explore using better algorithms, such as DP [3] and MILP [50], to model the integration of rematerialization and offloading and thus to find better hybrid optimization schemes. Although these hybrid methods have achieved promising results, existing works ignore incorporating them with parallelism techniques that are crucial for training big models. In view of this, we explore finding an efficient integration of memory optimization strategies and parallelism techniques for training big Transformers.
2.3 Parallelism
Data Parallelism. Data parallelism [53, 30, 11] is a typical approach for distributed DL. Its main idea is to divide the batch of training data into several mini-batches and process them on different devices (GPUs) in parallel. To synchronize the model parameters distributed on multiple devices, the gradients of parameters must be averaged among all devices before updating parameters in the parameter optimizer. Based on data parallelism, zero redundancy optimizer (ZeRO) [36] is proposed to partition the optimizer states and parameter gradients, which can significantly reduce the memory overhead of data parallelism without changing the computational complexity.
Model Parallelism. In recent years, the parameter size of DL models is continuously growing, and the conventional training scheme that allocates the entire model on all devices can easily run out of memory. To tackle this issue, researchers explore various model parallelism methods that distribute model parameters over multiple devices. There are generally two branches of model parallelism. One is pipeline parallelism [8, 17, 27, 51], which partitions model parameters by layer level and lets devices be responsible for the computation of different layers. The other is parameter parallelism [28, 40], which partitions parameter matrices within a layer to enable distributed computing in multiple devices. Like data parallelism, model parallelism can also be incorporated into ZeRO by partitioning the model at the tensor level and gathering the required parameters before computation. Most recent distributed DL toolkits for training big Transformers, such as BMTrain and Megatron-DeepSpeed, are designed based on data parallelism, model parallelism, and ZeRO.
Automatic Parallelism. Although many efforts have attempted to combine model parallelism and data parallelism to improve efficiency further, finding the optimal combination based on empirical
Figure 1: The overall architecture of H3T.
rules is not easy. To this end, some scholars explore automatic parallelism methods. Galvatron [26] uses a decision tree and a DP algorithm to find a hybrid parallelism solution among multiple parallelism dimensions. Alpa [52] and Unity [46] focus on finding the parallelism scheme at the computational graph level and use DP algorithms to search for the scheme based on their defined parallel operators. However, most existing works of automatic parallelism pay more attention to execution efficiency and rarely involve memory optimizations, which makes the memory overhead suboptimal for training big models. Therefore, we introduce memory optimization strategies to the existing parallelism schemes for better results.
3 H3T Framework
3.1 Formalization of Training Transformers
We first formalize the process of training Transformers. Generally, training DL models involves three stages: forward propagation, backward propagation, and parameter optimization. As a typical sequential model, Transformers can be formalized as an ordered sequence of neural layers, with each layer only connected to adjacent ones. So we go deep into the layer aspect and define  $\text{Forward}_i$ ,  $\text{Backward}_i$ ,  $\text{Optimize}_i$  as the three stages at the  $i$ -th layer, respectively.
In the forward propagation,  $\mathbf{Forward}_i$  takes  $\mathbf{H}_i$  and  $\mathbf{P}_i$  to calculate  $\mathbf{H}_{i+1}$ , where  $\mathbf{H}_i$  and  $\mathbf{H}_{i+1}$  are the hidden states, and  $\mathbf{P}_i$  is the parameters of the  $i$ -th layer. Along with generating some intermediate activations  $\mathbf{A}_i$ , the forward propagation at the  $i$ -th layer can be formalized as
$$
\mathbf {H} _ {i + 1}, \mathbf {A} _ {i} = \operatorname {F o r w a r d} _ {i} \left(\mathbf {H} _ {i}, \mathbf {P} _ {i}\right). \tag {1}
$$
The final output  $\mathbf{H}_{N + 1}$  is used to calculate the loss function and the initial gradient  $\nabla \mathbf{H}_{N + 1}$ . Then, in the backward propagation,  $\mathrm{Backward}_i$  is to calculate the gradient  $\nabla \mathbf{H}_i$  and  $\nabla \mathbf{P}_i$ , by using  $\nabla \mathbf{H}_{i + 1}$ ,  $\mathbf{H}_i$ ,  $\mathbf{P}_i$ , and the intermediate activations  $\mathbf{A}_i$  generated by  $\mathrm{Forward}_i$ , which can be given as
$$
\nabla \mathbf {H} _ {i}, \nabla \mathbf {P} _ {i} = \text {B a c k w a r d} _ {i} \left(\mathbf {H} _ {i}, \mathbf {P} _ {i}, \mathbf {A} _ {i}, \nabla \mathbf {H} _ {i + 1}\right). \tag {2}
$$
After the backward propagation,  $\mathsf{Optimize}_i$  is to use some parameter optimization functions to update the parameters  $\mathbf{P}_i$  according to the gradients  $\nabla \mathbf{P}_i$ .
In each training step, the conventional training way is first to run Forward $_i$  from  $i = 1$  to  $N$ , then run Backward $_i$  in reverse order, and finally use the parameter optimizer to update parameters according to gradients. For H3T, the training step is similar to the conventional one, but it is allowed to adjust the execution order or add some additional operations under the premise that the gradients are calculated correctly and the parameter optimization operations are performed stably, which can bring us much flexibility and convenience to integrate memory optimization and parallelism.
3.2 Optimization Switches
3.2.1 Activation-aspect Optimization Switches
For the sequential model, activations involve hidden states  $\mathbf{H}_i$  and intermediate activations  $\mathbf{A}_i$ . They play important roles in model training and meanwhile consume a lot of GPU memory, especially intermediate activations [42]. As mentioned in the background, scholars have developed two major approaches, rematerialization and offloading, to reduce the memory overhead caused by activations.
Considering the properties of hidden states and intermediate activations, we adopt rematerialization to optimize the memory of intermediate activations while using offloading to save the memory of
Figure 2: Activation-aspect switches.
hidden states. This setting is also applied by some existing work of hybrid optimization [3]. The two optimization switches are illustrated in Figure 2.
Rematerializing intermediate activations is to evict  $\mathbf{A}_i$  after forward steps and recompute  $\mathbf{A}_i$  before backward steps. To differentiate, we define  $\mathrm{Forward}_i^*$  as the forward computation without saving  $\mathbf{A}_i$ . As Figure 2 shows, rematerialization for the  $i$ -th layer replaces  $\mathrm{Forward}_i$  with  $\mathrm{Forward}_i^*$  in the forward stage, and insert an extra  $\mathrm{Forward}_i$  before  $\mathrm{Backward}_i$  in the backward stage.
Offloading hidden states is to offload  $\mathbf{H}_i$  and prefetch them before the backward propagation steps. Different from intermediate activations, if we similarly try to evict and recompute hidden states, the training process will suffer from excessive recomputations, and the performance will also be affected to a great extent. To this end, offloading hidden states can be more efficient.
These two optimizations can significantly reduce the memory cost of activations. Specifically, switching on both of them for every layer will reduce the resident activation memory to 0.
3.2.2 Parameter-aspect Optimization Switches
Traditional training process store all  $\mathbf{P}_i$  and  $\nabla \mathbf{P}_i$  on GPUs. As described above,  $\mathbf{P}_i$  and  $\nabla \mathbf{P}_i$  are only used for Forward $_i$ , Backward $_i$ , and Optimize $_i$ . Therefore, it is a straightforward idea to reduce this part of memory overhead by offloading optimizer states, parameters, and gradients to CPU RAM, as shown in Figure 3a. Generally, the number of gradients equals the number of parameters, and for most common optimizers, the number of optimizer states equals or even exceeds the number of trainable parameters. Therefore, offloading them is an efficient way to reduce the GPU memory overhead. On this basis, we further implement a high-performance CPU optimizer $^3$  that runs Optimize $_i$  on CPUs.
Moreover, we employ parameter parallelism for H3T. Specifically, we store partitioned parameters  $\mathbf{P}_i^{(j)}$  on each GPU. Before Forward $_i$ , we use an all gather operation to gather them from all GPUs to get  $\mathbf{P}_i$  for computation. After Backward $_i$ , we evict the gathered  $\mathbf{P}_i$  and use a reduce scatter operation to average and scatter the  $\nabla \mathbf{P}_i$  to each GPU.
In this way, the parameters and gradients form a closed loop, as shown in Figure 3a. Due to the holistic nature and outstanding memory-saving performance, we globally turn on the above offloading optimizations and parallelism in this paper.
In addition to the above global optimizations, we have two advanced parameter-aspect switches.
Parameter partition (Figure 3b) is to release gathered parameters after using them and gather them before their next use. Parameter partition mainly focuses on saving the wasted memory between the forward and backward steps of each layer, which is an issue overlooked by model parallelism.
Lazy-prefetching (Figure 3c) is to evict the partitioned parameters after every reduce scatter operation and lazily prefetch them back before every all gather operation. Lazy-prefetching aims to eliminate the memory cost of resident partitioned parameters.
Both parameter partition and lazy-prefetching can avoid the long-term persistence of parameters during the forward-to-backward period, thus effectively saving the memory of parameters. Specifically, turning on both optional switches for every layer will reduce the resident parameter memory cost to 0.
3The CPU optimizer is implemented based on Intel AVX instructions. More details are in the appendix.
(a) Basic parameter optimization scheme based on tensor offloading and parameter parallelism.
(b) Parameter partition.
Figure 3: Parameter-aspect switches.
(c) Lazy-prefetching.
3.2.3 Implementation: Concurrent Streams and Execution Switches
We implement the above optimizations with three concurrent streams for H3T. Specifically, we have (1) A calculation stream that takes charge of forward, backward, and parameter optimization. (2) An NVLink $^4$  stream that covers the GPU-GPU communication operators, i.e., all gather and reduce scatter. (3) A PCI-E stream that covers the GPU-CPU communication operators, i.e., offloading and prefetching. Following most previous works [49, 31, 3, 50], asynchronous communication and computation can overlap different parts of the time overhead, thus minimizing the efficiency loss. As illustrated in Figure 4, the main idea of our implementation is to overlap the post-process (scattering and offloading) of the previous layer and the preparation (gathering and prefetching) of the next layer with the computation of the current layer, and we call these overlapped operations a "segment".
Figure 4: An illustration of backward segment  $i$ . Without loss of generality, we suppose the switches of layer  $i - 1$  and  $i + 1$  are all switched on.
Furthermore, we design three execution switches for H3T. Economical forward and economical backward prioritize the scattering and offloading of the previous layer in order to save more memory. In contrast, non-economical modes prioritize gathering and prefetching of the next layer, which may reduce the delay in the computation of the next layer. Segment synchronization is to force all streams to be synchronized between two segments. More details about the implementation and the three switches are introduced in the appendix.
4NVLink is a high-speed multi-GPU communication link developed by NVIDIA. More details are here.
3.3 Optimization Solver
The optimizations save memory while bringing performance losses. To tackle this, H3T involves a solver to decide the appropriate switches to achieve the trade-off between memory and speed.
3.3.1 Problem Definition
We first state the problem definition for the solver. We define  $\mathbf{s} = (s_1, s_2, \dots, s_N)$  as an optimization sequence for the model, where  $s_i \in S$  represents the optimization scheme for the  $i$ -th layer, while  $S$  is the set of all combinations of optimization switches for one single layer. In this paper,  $S$  involves all combinations of the 7 layer-level switches described above, so we have  $|S| = 2^7 = 128$ .
We define Runtime(s), Memory(s) as the runtime (per step) and the memory overhead (per device) given the model and the optimization sequence s, respectively. Then, the goal of the solver is to find an optimization sequence s with lower Runtime(s) under the memory constraint  $M$  (per device):
$$
\operatorname {m i n i m i z e} \operatorname {R u n t i m e} (\mathbf {s}), \text {s . t . M e m o r y} (\mathbf {s}) \leq M. \tag {3}
$$
A key problem is how to precisely estimate Runtime(s) and Memory(s). Towards this goal, we introduce a profiler that collects the runtime and memory overhead of each operation of each layer at the first several training steps. Then we implement a model simulator to simulate the training process of each layer based on the model architecture and implementation, thus helping us estimate Runtime(s), Memory(s). Besides, the model simulator can also help the solver evaluate any solutions or sub-solutions when exploring different optimization schemes.