Title: H3T: Efficient Integration of Memory Optimization and Parallelism for Large-scale Transformer Training
Venue: NeurIPS 2023
Accepted: yes
Keywords: ML System

Parallelism Learning

Memory Optimization

Data Parallelism

Model Parallelism

Parameter Parallelism

ZeRO

Rematerialization

Checkpointing

Tensor Offloading

Dynamic Programming
Abstract: 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 High-Throughput 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 \sim 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.

H3T: Efficient Integration of Memory Optimization and Parallelism for High-Throughput Transformer Training

Abstract

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 H ighT hroughput T ransformer T raining (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 × ∼ 4 . 3 × training speed while reducing 34 . 6% ∼ 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 .

1 Introduction

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 GBof memory to store in float32 format, not to mention the forward activations,

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 H ighT hroughput T ransformer T raining (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 × ∼ 4 . 3 × faster than Megatron-DeepSpeed while reducing the memory usage by 34 . 6% ∼ 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 Background and Related Work

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 Transformerbased 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( √ 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

Profiling the runtime and memory cost for each layer.

Transformer

-based

Model

Layer

𝑁

Switches

…

…

Forward

Backward

Solver

§

3.3.2)

(

§

3.3.1)

Layer 1

Activation-aspect

Switches (

§

3.2.1)

Rematerialization

Offloading

Parameter-aspect

Switches (

§

3.2.2)

Partition

Offloading (global)

Lazy-prefetching

Execution

Switches (

§

3.2.3)

Economical Forward

Economical Backward

Segment Sync.

Profiler

Search for a solution of appropriate optimization switches for each layer.

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 Forward i , Backward i , Optimize i as the three stages at the i -th layer, respectively.

In the forward propagation, Forward i takes H i and P i to calculate H i +1 , where H i and H i +1 are the hidden states, and P i is the parameters of the i -th layer. Along with generating some intermediate activations A i , the forward propagation at the i -th layer can be formalized as

The final output H N +1 is used to calculate the loss function and the initial gradient ∇ H N +1 . Then, in the backward propagation, Backward i is to calculate the gradient ∇ H i and ∇ P i , by using ∇ H i +1 , H i , P i , and the intermediate activations A i generated by Forward i , which can be given as

After the backward propagation, Optimize i is to use some parameter optimization functions to update the parameters P i according to the gradients ∇ 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 H i and intermediate activations 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

Offloading

Rematerialization

Forward 𝑖

∗

𝐀𝑖

(Evict after

Forward)

⋯⋯

𝐀𝑖

⋯⋯

⋯⋯

Forward𝑖

(Re-compute)

Backward𝑖

𝐇𝑖

RAM

Offload

after

Forward

Prefetch

before

Re-forward

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 A i after forward steps and recompute A i before backward steps. To differentiate, we define Forward ∗ i as the forward computation without saving A i . As Figure 2 shows, rematerialization for the i -th layer replaces Forward i with Forward ∗ i in the forward stage, and insert an extra Forward i before Backward i in the backward stage.

Offloading hidden states is to offload 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 P i and ∇ P i on GPUs. As described above, P i and ∇ 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 P ( j ) i on each GPU. Before Forward i , we use an all gather operation to gather them from all GPUs to get P i for computation. After Backward i , we evict the gathered P i and use a reduce scatter operation to average and scatter the ∇ 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 .

GPU 1

∇ Parameter i

ℒ𝑖

Parameter𝑖

∇

𝑃 𝑖

0

𝑃 𝑖

0

RAM

∇𝐏𝑖

𝐏𝑖

GPU 0

∇𝐏𝑖

𝐏𝑖

𝐏 𝑖

0

∇𝐏 𝑖

0

All Gather

……

Offload

Prefetch

CPU Optimizer

Based on SIMD

Optimize

Lazy-prefetching

Partition

After Optimize

Evict

Before Forward

After

Backward

Offloading (global)

(a) Basic parameter optimization scheme based on tensor offloading and parameter parallelism.

GPU 1

∇ Parameter i

ℒ𝑖

Parameter𝑖

∇

𝑃 𝑖

0

𝑃 𝑖

0

All Gather

……

Evict

Before Forward

Forward & Backward

After

Backward

Forward & Backward

GPU 0

∇𝐏 𝑖

0

Partition

∇𝐏𝑖

𝐏𝑖

𝐏 𝑖

0

(b) Parameter partition.

GPU 1

∇

𝑃 𝑖

0

𝑃 𝑖

0

GPU 0

𝐏 𝑖

0

∇𝐏 𝑖

0

……

Prefetch

After Optimize

Before All Gather

Lazy-prefetching

Evict

After

Reduce

Offloading (global)

(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'.

Post-process of the

previous layer (

𝒊 + 𝟏

)

and

Preparation of the

next layer (

𝒊 - 𝟏

)

Calculation of the

current layer (

𝒊

)

Calc. Stream

NVLink Stream

PCI-E Stream

Forward𝑖

(re-compute)

Backward𝑖

AllGather

𝑖-1

Param

ReduceSatter 𝑖+1

∇

Offload

𝑖+1

∇

Prefetch

𝑖-1

Hidden

Prefetch

𝑖-1

Param

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.

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 s = ( s 1 , s 2 , · · · , s N ) as an optimization sequence for the model, where s i ∈ 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):

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.

How to pick appropriate switches is another big challenge. The optimization of different layers can be in any combination, so the number of global solutions can be up to | S | N . It is not easy to find a good solution from such exponential search space with the memory constraint and toward the goal of minimizing the training runtime. Toward this goal, we design three search algorithms in this paper.

3.3.2 Search Algorithms

Random solver . Brute force does not work for our task because of the exponential search space, so we design a simple method based on random sampling to find a suboptimal solution more efficiently. Specifically, we randomly sample several optimization sequences from S N and pick the best one.

Greedy solver . The random solver is simple but not effective for large-scale models. The main reason is that the search space grows exponentially with N , but the sampling time is constant, which means the probability we hit a good solution becomes much lower with the growth of the model size. To address these issues, we design a greedy solver to search for better solutions. The core idea of our greedy solver is first assuming all optimization switches are turned on, then greedily turning off some switches to increase memory utilization and improve training efficiency. Due to the space limitation, detailed descriptions of the greedy solver are in the appendix.

Dynamic Programming (DP) Solver . Our task strives for better runtime under the memory constraint, similar to the knapsack problem [25, 7] that pursues higher value under the weight constraint. Inspired by the classic dynamic programming (DP) algorithm for the knapsack problem [1, 43], we design a search algorithm based on DP to solve our task better. Specifically, our DP algorithm considers the optimal sub-solutions for specified prefixes and memory limits. When transferring, we need to enumerate the next layer's switches and estimate the next segment's running status.

Current memory usage (

memcurr

)

Peak memory usage (

mempeak

)

𝑖

𝑖

……

1

……

1

Forward

……

……

Backward

The process of DP

Memory

Segment

Formally, the DP state can be represented by the following variables: (1) i : The prefix length, which means this state is considering segment 1 ∼ i (including forward segment 1 ∼ i and backward segment i ∼ 1 , similarly hereinafter). (2) mem curr: The memory usage after forward segment 1 ∼ i .

(The memory usage before backward segment i ∼ 1 is essentially equal and can be easily calculated.) (3) mem peak: The peak memory usage during segment 1 ∼ i . (We have to record mem peak because the resident memory may increase due to non-lazy-prefetching, leading to an increase in the peak memory from the previous segments.) (4) s i , s i +1 : The switches of layer i and i +1 , which is used to simulate segment i +1 .

We use sol ( i, mem curr , mem peak , s i , s i +1 ) to represent the optimal sub-solution of the state. Then we can enumerate s i +2 to transfer to the next state ( i +1 , mem ′ curr , mem ′ peak , s i +1 , s i +2 ) . During the transfer, we need to: (1) Simulate the forward segment i +1 and backward segment i +1 based on s i , s i +1 , s i +2 . (2) Calculate mem ′ curr , mem ′ peak and the total runtime based on mem curr , mem peak, the previous runtime, and the simulation result. (3) Update sol ( i +1 , mem ′ curr , mem ′ peak , s i +1 , s i +2 ) if the peak memory mem ′ peak does not exceed M and the runtime is shorter than the existing solution's.

There are two noteworthy details for our DP algorithm. First, considering the overlapping subproblems condition, we must ensure the state is finite, so we discretize the memory limit into m +1 values that are evenly distributed in [0 , M ] . Second, following the condition of optimal substructure, the problem for each state must be reducible to the corresponding sub-problem for the other states. In this paper, only cross-segment communications violate this condition, so we force segment synchronization switched on and remove it from the switch set (thus | S | reduce to 64 ). These two assumptions may affect the result but make our algorithm design much simpler and more efficient. We leave the design of the more comprehensive DP algorithm as our future work.

In addition to the basic version of the DP algorithm, we have two prunings to improve its efficiency. Please refer to the appendix for them along with the pseudo-code and time complexity analysis.

4 Experiments

4.1 Experimental Settings

Models . We use BERT[9], one of the most common Transformer-based language models, as an example for our experiments. We customize different BERT sizes for our experiments, including 1.8B, 6B, 13B, and 100B. Traditional small models, such as BERT-base (0.1B) or BERT-large (0.3B), are not considered in this paper because they can be easily trained using conventional approaches.

Environments . We test the performance of H3T in three different experimental environments: (1) 1 node of 8 × NVIDIA GeForce RTX 2080 Ti 10 GB (8 × 2080Ti); (2) 1 node of 8 × NVIDIA A100-SXM4-40GB with NVLink 3.0 (8 × A100); and (3) 8 nodes of 8 × A100 (64 × A100).

Implementations . We implement H3T based on BMTrain, an open-source toolkit that accelerates the training of big Transformers with various parallelism techniques.

Baselines . Except for H3T with random solver and greedy solver, we introduce manual optimization strategy without H3T and another toolkit Megatron-DeepSpeed for comparison. (1) In the simulation study, we test the performance of manually enumerating the combination of global optimizations and picking the best one. In other words, we simulate arbitrarily and globally turning on or off all the switches and picking the best-performing solution that does not exceed the memory limit. (2) In the actual-running experiment, we attempt to turn on all optimization switches to achieve the lowest memory overhead since testing all manual optimization schemes is time-consuming and resource-intensive. (3) We introduce one of the most popular distributed deep learning toolkits, Megatron-DeepSpeed [40], for comparison in our actual-running experiment. For a fair comparison, we test Megatron-DeepSpeed with ZeRO-2 and ZeRO-3 respectively (ZeRO-2 distributes optimizer states and gradients; ZeRO-3 distributes optimizer states, gradients, and parameters [36]). Please refer to the appendix and the code for detailed configurations of Megatron-DeepSpeed and H3T in the actual-running experiment.

We do not test POFO [3], which also uses DP to solve the memory optimization problem for big model training. We consider it unfair to compare POFO with H3T because POFO does not support any parameter-aspect optimization, such as parameter offloading and parallelism. Considering the large model size, it is almost impossible for a vanilla POFO to train the large-scale language models in this paper. If we arbitrarily implement parameter optimizations for POFO, additional communication will affect the well-designed offloading solution of POFO and seriously hurt its performance.

Other settings . Please refer to the appendix for other detailed experimental settings, including detailed model size, hardware configurations, training settings, etc.

4.2 Simulation Studies





















































































































































































































































































In this section, we report our simulation studies to evaluate the ideal performance of H3T. We profile the models in different environments and use the model simulator to evaluate the performance of different solvers and the manual strategy. The result is shown in Figure 6.

The result shows that our DP solver outperforms the other three baseline approaches in most cases. The greedy solver also performs well and is usually close to DP. Both DP and greedy solver have the flexibility to adjust the optimization solution based on variant memory constraints. This conclusion is applicable in all three environments, showing H3T has good adaptability to GPUs with different numbers, performance, and connections. In contrast, the random solver is less stable and less effective. Although the performance gap is not significant when memory is sufficient, the random solver, more often than not, struggles even to find an available solution.

Compared to three automatic solvers, manual optimization without H3T is often the worst, especially when memory is sufficient. This indicates that our layer-level automatic solver is able to precisely decide the optimization scheme for large-scale Transformer-based models given the memory constraint. In contrast, the traditional manual strategy performs worse and is heavy for users as well.

However, the above rules do not always stand. For BERT-1.8B with a batch size of 128 on 8 × 2080Ti, the greedy solver underperforms the random solver. For BERT-6B on 8 × 2080Ti, the DP solver continuously underperforms the greedy solver and even underperforms the manual strategy in some cases. This is because the local optimality may let the greedy algorithm fall into a suboptimum, and the segment synchronization assumption and the memory discretization may influence the DP algorithm. We argue that it is normal for such problems to occur in a few cases.

4.3 Actual-running Experiments

We conduct our actual-running experiment in the three environments, and the result is shown in Table 1. The results demonstrate our great advantage over Megatron-DeepSpeed. Compared to the Megatron-DeepSpeed with ZeRO-2 (ZeRO-3), H3T improves the training speed by 1 . 2 × ∼ 1 . 8 × ( 1 . 9 × ∼ 4 . 3 × ) in our experiments. Comparing the results of BERT-6B and BERT-13B, both with the batch size of 128 and both on 8 × A100, we find that H3T outperforms Megatron-DeepSpeed more significantly for larger model size, which means H3T is not only faster but also more adapted to bigger models.

Focusing on the results of BMTrain, we find that H3T with DP solver and greedy solver both outperform BMTrain without H3T by up to 1 . 2 × speed up when memory is sufficient. In contrast, when memory is tight, the two automatic solvers can also enable more memory optimization and trade time performance for keeping model training. This result is consistent with the findings from the simulation study, which further verifies the flexibility of our automatic solvers. Also, similar to the simulations, the performance of DP and greedy are close, but DP still outperforms greedy a little in most cases. Besides, the random solver fails to solve for all of the settings, so we do not report its results in the table.

The memory cost of our implementation is also better than Megatron-DeepSpeed. From Table 1, we find Megatron-DeepSpeed often runs out of memory for larger models and batch sizes, which are better supported by H3T. Hence, users can train larger models with more limited hardware, e.g., train BERT-6B on 8 × 2080Ti. Not only that, H3T can train GPT-3-175B [4] on 64 × A100 with a batch size of 512 and using only about 11 GB of memory per GPU. This is really an exciting result because we would never have dreamed that a 175B-parameter model could be trained with such limited hardware, not to mention with such low GPU memory overhead. We conduct a comprehensive memory test in the appendix and find our implementation can save 34 . 6% ∼ 80 . 5% GPU memory compared with Megatron-DeepSpeed under various settings.

Besides, we verify the correctness of our implementation. (1) We fix the input data and check whether the output loss equals the release version of BMTrain during the whole training process, which justifies our incremental development based on BMTrain. (2) We do an end-to-end training experiment on 4 actual tasks of SuperGLUE [48]. The results show that H3T can achieve comparable results with PyTorch, which confirms the soundness of the whole system. The detailed experimental results on SuperGLUE are in the appendix.

We conduct three extra experiments in this paper. Due to the space limitation, we briefly introduce them here and elaborate on the details in the appendix. (1) We conduct an energy test and find that H3T is more energy-efficient than Megatron-DeepSpeed for training big Transformer-based models. (2) We test the efficiency of the solver to confirm it is not too slow. (3) We conduct a case study to show the automatically generated optimization scheme makes sense.

5 Conclusion

In this paper, we propose a novel framework for automatic integration of memory optimization and parallelism for H ighT hroughput T ransformer T raining, named H3T . Despite the challenge of large search space, H3T significantly outperforms some manual strategies and one of the most popular DL toolkits, Megatron-DeepSpeed, in our experiments. More excitingly, H3T can use only 64 NVIDIA A100 GPUs to train Transformer-based models with more than 100 billion parameters like GPT-3-175B. We hope that the release of H3T will lower the barrier to training big Transformer-based models and promote the development of the AI community.

Acknowledgments

This work is supported by the National Key R&D Program of China (No.2022ZD0116312), Major Project of the National Social Science Foundation of China (No. 22&ZD298), the National Key R&D Program of China (No. 2022ZD0160501), and Institute Guo Qiang at Tsinghua University. Weilin Zhao is supported by Tsinghua University Initiative Scientific Research Program.

References

Maram Assi and Ramzi A. Haraty. A survey of the knapsack problem. In Proceedings of ACIT , pages 1-6, 2018. https://ieeexplore.ieee.org/document/8672677 .

Olivier Beaumont, Lionel Eyraud-Dubois, and Alena Shilova. Optimal gpu-cpu offloading strategies for deep neural network training. In Proceedings of Euro-Par , pages 151-166, 2020. https://link.springer.com/chapter/10.1007/978-3-030-57675-2_10 .

Olivier Beaumont, Lionel Eyraud-Dubois, and Alena Shilova. Efficient combination of rematerialization and offloading for training dnns. In Proceedings of NeurIPS , volume 34, pages 23844-23857, 2021. https://papers.nips.cc/paper_files/paper/2021/file/ c8461bf13fca8a2b9912ab2eb1668e4b-Paper.pdf .

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Proceedings of NeurIPS , volume 33, pages 1877-1901, 2020. https: //arxiv.org/pdf/2005.14165.pdf .

Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 , 2016. https://arxiv.org/pdf/1604. 06174.pdf .

Xiaoming Chen, Danny Z. Chen, and Xiaobo Sharon Hu. modnn: Memory optimal dnn training on gpus. In Proceedings of DATE , pages 13-18, 2018. https://ieeexplore.ieee.org/ document/8341972 .

Tobias Dantzig and Joseph Mazur. Number: The language of science . 2007. https:// philpapers.org/rec/DANNTL .

Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc'aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. In Proceedings of NeurIPS , volume 25, 2012. https://proceedings.neurips.cc/paper/ 2012/file/6aca97005c68f1206823815f66102863-Paper.pdf .

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL , pages 4171-4186, 2019. https://www.aclweb.org/anthology/N19-1423.pdf .

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020. https://arxiv.org/pdf/2010.11929.pdf .

Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 , 2017. https://arxiv.org/pdf/1706. 02677.pdf .

Andreas Griewank and Andrea Walther. Algorithm 799: Revolve: An implementation of checkpointing for the reverse or adjoint mode of computational differentiation. ACM Transactions on Mathematical Software (TOMS) , 26(1):19-45, 2000. https://dl.acm.org/doi/pdf/10. 1145/347837.347846 .

José Grimm, Loïc Pottier, and Nicole Rostaing-Schmidt. Optimal time and minimum spacetime product for reversing a certain class of programs . PhD thesis, INRIA, 1996. https: //inria.hal.science/inria-00073896/file/RR-2794.pdf .

Julien Herrmann, Olivier Beaumont, Lionel Eyraud-Dubois, Julien Hermann, Alexis Joly, and Alena Shilova. Optimal checkpointing for heterogeneous chains: How to train deep neural networks with limited memory. arXiv preprint arXiv:1911.13214 , 2019. https://arxiv. org/pdf/1911.13214.pdf .

Zhongzhe Hu, Junmin Xiao, Zheye Deng, Mingyi Li, Kewei Zhang, Xiaoyang Zhang, Ke Meng, Ninghui Sun, and Guangming Tan. Megtaichi: Dynamic tensor-based memory management optimization for dnn training. In Proceedings of ICS , pages 1-13, 2022. https://dl.acm. org/doi/pdf/10.1145/3524059.3532394 .

Chien-Chin Huang, Gu Jin, and Jinyang Li. Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. In Proceedings of ASPLOS , pages 1341-1355, 2020. https://dl.acm.org/doi/pdf/10.1145/3373376.3378530 .

Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Proceedings of NeurIPS , volume 32, 2019. https://proceedings.neurips.cc/paper/2019/file/ 093f65e080a295f8076b1c5722a46aa2-Paper.pdf .

Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and Ion Stoica. Checkmate: Breaking the memory wall with optimal tensor rematerialization. In Proceedings of MLSys , volume 2, pages 497-511, 2020. https://proceedings. mlsys.org/paper/2020/file/084b6fbb10729ed4da8c3d3f5a3ae7c9-Paper.pdf .

Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. Dynamic tensor rematerialization. In Proceedings of ICLR , 2021. https://arxiv.org/pdf/2006.09616.pdf .

Mitsuru Kusumoto, Takuya Inoue, Gentaro Watanabe, Takuya Akiba, and Masanori Koyama. A graph theoretic framework of recomputation algorithms for memory-efficient backpropagation. In Proceedings of NeurIPS , volume 32, 2019. https://arxiv.org/pdf/1905.11722.pdf .

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019. https://proceedings.neurips.cc/ paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf .

Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of CVPR , pages 12009-12019, 2022. https://arxiv.org/pdf/2111.09883.pdf .

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of ICCV , pages 9992-10002, 2021. https://arxiv.org/pdf/2103.14030.pdf .

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proceedings of ICLR , 2019. https://arxiv.org/pdf/1711.05101.pdf .

George B Mathews. On the partition of numbers. Proceedings of the London Mathematical Society , 1(1):486-490, 1896. https://academic.oup.com/plms/article-abstract/ s1-28/1/486/1495829 .

Xupeng Miao, Yujie Wang, Youhe Jiang, Chunan Shi, Xiaonan Nie, Hailin Zhang, and Bin Cui. Galvatron: Efficient transformer training over multiple gpus using automatic parallelism. arXiv preprint arXiv:2211.13878 , 2022. https://arxiv.org/pdf/2211.13878.pdf .

Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: generalized pipeline parallelism for dnn training. In Proceedings of SOSP , pages 1-15, 2019. https://dl.acm. org/doi/pdf/10.1145/3341301.3359646 .

Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient large-scale language model training on gpu clusters using megatron-lm. In Proceedings of SC , pages 1-15, 2021. https://arxiv.org/pdf/2104.04473.pdf .

OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023. https://arxiv. org/pdf/2303.08774.pdf .

Thomas Paine, Hailin Jin, Jianchao Yang, Zhe Lin, and Thomas Huang. Gpu asynchronous stochastic gradient descent to speed up neural network training. arXiv preprint arXiv:1312.6186 , 2013. https://arxiv.org/pdf/1312.6186.pdf .

Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma, Qian Xiong, Fan Yang, and Xuehai Qian. Capuchin: Tensor-based gpu memory management for deep learning. In Proceedings of ASPLOS , pages 891-905, 2020. http://alchem.usc.edu/portal/static/download/ capuchin.pdf .

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In Proceedings of ICML , pages 8748-8763, 2021. https://arxiv.org/pdf/2103.00020.pdf .

Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training. OpenAI blog , 2018. https://www.cs.ubc.ca/ ~amuham01/LING530/papers/radford2018improving.pdf .

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019. https://cdn.openai.com/better-language-models/language_models_are_ unsupervised_multitask_learners.pdf .

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research , 21(140):1-67, 2020. http: //jmlr.org/papers/v21/20-074.html .

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In Proceedings of SC , pages 1-16, 2020. https://arxiv.org/pdf/1910.02054.pdf .

Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zeroinfinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of SC , pages 1-14, 2021. https://arxiv.org/pdf/2104.07857.pdf .

Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022. https://arxiv.org/pdf/2204.06125.pdf .

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In Proceedings of ICML , pages 8821-8831, 2021. https://arxiv.org/pdf/2102.12092.pdf .

Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of KDD , pages 3505-3506, 2020. https://dl.acm.org/doi/10.1145/3394486.3406703 .

Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training. In Proceedings of ATC , pages 551-564, 2021. https://arxiv.org/pdf/2101. 06840v1.pdf .

Minsoo Rhu, Natalia Gimelshein, Jason Clemons, Arslan Zulfiqar, and Stephen W Keckler. vdnn: Virtualized deep neural networks for scalable, memory-efficient neural network design. In Proceedings of MICRO , pages 1-13, 2016. https://arxiv.org/pdf/1602.08124.pdf .

Harvey M Salkin and Cornelis A De Kluyver. The knapsack problem: A survey. Naval Research Logistics Quarterly , 22(1):127-144, 1975. https://onlinelibrary.wiley.com/ doi/abs/10.1002/nav.3800220110 .

SB Shriram, Anshuj Garg, and Purushottam Kulkarni. Dynamic memory management for gpu-based training of deep neural networks. In Proceedings of IPDPS , pages 200-209, 2019. https://ieeexplore.ieee.org/abstract/document/8820980 .

Xiaoyang Sun, Wei Wang, Shenghao Qiu, Renyu Yang, Songfang Huang, Jie Xu, and Zheng Wang. Stronghold: fast and affordable billion-scale deep learning model training. In Proceedings of SC , pages 1-17, 2022. https://zwang4.github.io/publications/sc22-2.pdf .

Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandeep Baines, Carlos Efrain Quintero Narvaez, Vinay Ramakrishnaiah, Nirmal Prajapati, Pat McCormick, Jamaludin Mohd-Yusof, et al. Unity: Accelerating dnn training through joint optimization of algebraic transformations and parallelization. In Proceedings of OSDI , pages 267-284, 2022. https: //www.usenix.org/system/files/osdi22-unger.pdf .

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of NeurIPS , volume 30, 2017. https://proceedings.neurips.cc/paper/2017/file/ 3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf .

Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Superglue: a stickier benchmark for generalpurpose language understanding systems. In Proceedings of NeurIPS , pages 3266-3280, 2019. https://w4ngatang.github.io/static/papers/superglue.pdf .

Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song, Zenglin Xu, and Tim Kraska. Superneurons: Dynamic gpu memory management for training deep neural networks. In Proceedings of PPoPP , pages 41-53, 2018. https://arxiv.org/pdf/1801. 04380.pdf .

Lijie Wen, Zan Zong, Li Lin, and Leilei Lin. A swap dominated tensor re-generation strategy for training deep learning models. In Proceedings of IPDPS , pages 996-1006, 2022. https: //ieeexplore.ieee.org/document/9820619 .

Jun Zhan and Jinghui Zhang. Pipe-torch: Pipeline-based distributed deep learning in a gpu cluster with heterogeneous networking. In Proceedings of CBD , pages 55-60, 2019. https: //ieeexplore.ieee.org/abstract/document/8916305 .

Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Joseph E Gonzalez, et al. Alpa: Automating inter-and intra-operator parallelism for distributed deep learning. In Proceedings of OSDI , pages 559-578, 2022. https://www.usenix.org/system/files/osdi22-zheng-lianmin.pdf .

Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient descent. In Proceedings of NeurIPS , volume 23, 2010. https://proceedings.neurips. cc/paper/2010/file/abea47ba24142ed16b7d8fbf2c740e0d-Paper.pdf .

A Implementation Details

A.1 Timing Diagram of Multiple Streams with Optimization Switches

Calc.

NVLink

PCI-E

F𝑖

⋯

B𝑖

⋯

AG𝑖

P

RS 𝑖

∇

O𝑖

∇

⋯

⋯

⋯

O𝑖

P 𝑖

P

⋯

⋯

Forward

Backward

Optimizer

Evict Gathered Parameter

(a) Default workflow based on parameter offloading, optimizer offloading, and param-

eter parallelism.

F 𝑖

∗

Calc.

NVLink

PCI-E

⋯

B𝑖

⋯

AG𝑖

P

RS 𝑖

∇

O𝑖

∇

⋯

⋯

⋯

Forward

Backward

Evict Gathered Parameter

F𝑖

(b) Rematerialization.

F 𝑖

∗

Calc.

NVLink

PCI-E

⋯

B𝑖

⋯

AG𝑖

P

RS 𝑖

∇

O𝑖

∇

⋯

⋯

⋯

Forward

Backward

Evict Gathered Parameter

P 𝑖

H

⋯

O𝑖

H

(c) Offloading hidden states.

Calc.

NVLink

PCI-E

F𝑖

⋯

B𝑖

⋯

AG𝑖

P

RS 𝑖

∇

O𝑖

∇

⋯

⋯

⋯

Forward

Backward

Evict Gathered Parameter

Evict Gathered Parameter

AG𝑖

P

⋯

(d) Parameter partition.

Calc.

NVLink

PCI-E

F𝑖

⋯

B𝑖

⋯

AG𝑖

P

RS 𝑖

∇

O𝑖

∇

⋯

⋯

⋯

O𝑖

P 𝑖

P

⋯

⋯

Forward

Backward

Optimizer

Evict Gathered and Partitioned Parameter

Allocate Memory for Partitioned Parameter

(e) Lazy-prefetching.

F 𝑖

∗

Forward ∗

F𝑖

Forward

B𝑖

Backward

RS 𝑖

∇

O𝑖

∇

Offload Gradient

AG𝑖

P

All Gather

P 𝑖

H

O𝑖

Optimize

P 𝑖

P

Prefetch Parameter

O𝑖

H

Offload Hidden State

(f) The legend.

Figure 7 illustrate the timing diagram of our implementation of H3T. Figure 7a is the default workflow based on parameter offloading, optimizer offloading, and parameter parallelism. Figure 7b and 7c are the workflow with two activation-aspect switches, respectively. Figure 7d and 7e are the workflow with two parameter-aspect switches, respectively.

A.2 Execution-aspect Switches

In this part, we introduce the detail of our three execution switches.

Economical forward/backward . These two switches are to prioritize the scattering and offloading of the previous layer in order to save more memory. In contrast, non-economical forward/backward prioritize gathering and prefetching of the next layer, which may reduce the delay in the computation of the next layer. Since NVLink is not available for all GPUs, e.g., the NVIDIA GeForce series, we design the workflow with and without NVLink, respectively, and the timing diagram is shown in Figure Figure 8. It is worth mentioning that there are still asynchronous operations in the economical mode, so it is not memory-optimal but balances the efficiency.

Segment synchronization . Let the computation of one layer and the concurrent communications (post-process of the previous layer and preparation of the next layer) be a 'segment'. Then segment

Economical Forward

Calc.

NVLink

PCI-E

F 𝑖

∗

⋯

AG𝑖+1

P

⋯

P 𝑖+1

P

⋯

⋯

O𝑖

H

⋯

⋯

Forward Segment

𝒊

F 𝑗

∗

⋯

AG𝑗+1

P

⋯

P 𝑗+1

P

⋯

⋯

O𝑗

H

⋯

⋯

Economical Forward

Forward Segment

𝒋

(a) Non-economical and economical forward with NVLink.

⋯

RS 𝑗+1

∇

O𝑗+1

∇

⋯

⋯

⋯

F𝑗

AG𝑗-1

P

P 𝑗-1

P

P 𝑗-1

H

⋯

⋯

Economical Backward

Economical Backward

B𝑖

⋯

RS 𝑖+1

∇

O𝑖+1

∇

⋯

⋯

Calc.

NVLink

PCI-E

⋯

F𝑖

AG𝑖-1

P

P 𝑖-1

P

P 𝑖-1

H

⋯

⋯

𝒊

𝒋

(b) Non-economical and economical backward with NVLink.

F 𝑗

∗

⋯

AG𝑗+1

P

⋯

P 𝑗+1

P

⋯

O𝑗

H

⋯

Forward Segment

𝒋

Economical Forward

Economical Forward

Calc.

PCI-E

F 𝑖

∗

⋯

AG𝑖+1

P

⋯

P 𝑖+1

P

⋯

O𝑖

H

⋯

Forward Segment

𝒊

(c) Economical and non-economical forward without NVLink.

⋯

RS 𝑗+1

∇

O𝑗+1

∇

⋯

⋯

F𝑗

AG𝑗-1

P

P 𝑗-1

P

P 𝑗-1

H

⋯

𝒋

Economical Backward

Economical Backward

B𝑖

⋯

RS 𝑖+1

∇

O𝑖+1

∇

⋯

Calc.

PCI-E

⋯

F𝑖

AG𝑖-1

P

P 𝑖-1

P

P 𝑖-1

H

⋯

𝒊

(d) Economical and non-economical backward without NVLink.

synchronization can be defined as: all streams need to be synchronized between two segments. In other words, if segment synchronization is on, the next segment must keep waiting until all tasks of the current segment end. In contrast, if segment synchronization is off, the next segment can start before the current stage ends. An illustration of segment synchronization is shown in Figure 9. Similar to the economical mode, segment synchronization can reduce the memory overhead, but may sacrifice a little performance. Segment synchronization is also a condition of our DP solver.

A.3 CPU Optimizer

As we mentioned in the main text, we implement our CPU optimizer based on Intel AVX instructions. Intel A VX (Advanced Vector Extensions) instructions are SIMD (Single Instruction Multiple Data) instructions proposed by Intel for high-performance parallel computing. With the help of Intel A VX instructions, our CPU optimizer can achieve comparable performance to the conventional GPU optimizer, which greatly ensures training performance.

Segment Synchronization

Calc.

PCI-E

F 𝑖

∗

⋯

⋯

O𝑖

H

⋯

F 𝑖+1

∗

⋯

Forward Segment

𝒊

Start before

the previous

segment

ends.

Forward Segment

𝒊 + 𝟏

Calc.

PCI-E

F 𝑖

∗

⋯

⋯

O𝑖

H

⋯

F 𝑖+1

∗

⋯

Forward Segment

𝒊

Wait until

the previous

segment

ends.

Segment Synchronization

Forward Segment

𝒊 + 𝟏

B Solver Details

In this part, we mainly show the pseudo-code of the three solver algorithms and analyze their time complexity. We also describe details of the greedy solver and the pruning of the DP solver, which are not introduced in the main text due to the space limitation.

B.1 Random Solver

The pseudo-code of our random solver is shown in Algorithm 1. Here we use a hyper-parameter T rand to represent the sampling times. Since the time complexity of the model simulator is Θ( N ) , the time complexity of the random solver is Θ( NT rand ) .

B.2 Greedy Solver

The pseudo-code of our greedy solver is shown in Algorithm 2. We first assume all optimization switches are turned on, which reduces the memory overhead to the lowest (in our framework). Then, in order to increase memory utilization and improve training efficiency, we want to turn off some switches for part of the layers. Here we have a greedy assumption that these "partial layers" have

the same switches turned off, which can be easily enumerated as s off ∈ S . Then we enumerate n off , representing how many layers we will turn off s off for. Finally, we greedily put these s off into the initial sequence evenly.

The greedy solver introduces two greedy assumptions, where the former is centered on the idea that the best is always the best, and the latter on the idea that an even distribution may block communications less. We cannot prove that either of them is optimal, but they both make sense, so we expect our greedy solver to outperform the random solver.

The complexity of the enumeration is Θ( N | S | ) , so taking the model simulator into account, we have the global time complexity of Θ ( N 2 | S | ) .

B.3 DP Solver

Algorithm 3 DP Solver (basic version)

̸

The basic version of the DP solver is introduced in the main paper, while the pseudo-code is shown in Algorithm 3. Here we use ModelSimulator i ( s i-1 , s i , s i+1 ) to represent the simulation of the i -th segment, while the result ∆ mem i , residentMem i , peakMem i , runtime i denotes the delta memory after forward segment i , the resident memory of layer i , the peak memory during the forward and backward of segment i , and the total runtime of segment i , respectively.

For this basic version of DP solver, the state complexity is O ( m 2 N | S | 2 ) , while the transfer complexity is Θ( | S | ) . On this basis, we attempt to do some pruning to improve its efficiency.

̸

We first notice that the optimization scheme s i , s i +1 , s i +2 are only used in the running segment simulation during transfer. However, the simulation for segment i +1 only requires s i +1 and part of s i , s i +2 . Specifically, for the previous layer ( i ) and the next layer ( i + 2 ), we only care about three communication-related switches: hidden offloading, lazy-prefetching, and parameter partition. Therefore, we change s i , s i +1 ∈ S to r i , r i +1 ∈ R in the DP state, where R denotes the set of combinations of these three switches. Similarly, we only enumerate r i +2 ∈ R for transfer because it is enough for the simulation of segment i +1 . In addition, we need to extend r i +1 to s i +1 by another enumeration when transferring. In this way, the state complexity can be optimized to Θ ( m 2 N | R | 2 ) , while the transfer complexity remains Θ( | S | ) because the time complexity of enumerating r i +2 and enumerating s i +1 extended by r i +1 are Θ( | R | ) and Θ ( | S | | R | ) , respectively.

Besides, we can prune the non-optimal sub-solutions to improve efficiency further. In particular, for any two sub-solutions with the same i, r i , r i +1 , if one has a lower mem curr, a lower mem peak, and a shorter runtime, then the other one will be pruned because it is definitely not an optimal substructure of the global optimal solution. In practice, this pruning can be easily implemented by a 2-dimensional scanning with the time complexity of Θ ( m 2 ) for any given i, r i , r i +1 .

Under the first pruning, the time complexity of the DP solver is reduced to O ( m 2 N | S || R | 2 ) , where | S | = 64 , | R | = 8 . For the experiments in this paper, we always take m = 128 . Although the second pruning does not improve the time complexity, it can reduce more than 99% of the DP states in actual runs, which greatly improves the solver's efficiency.

Based on the basic DP algorithm and the two prunings described above, we have the pseudo-code of our pruned DP solver shown in Algorithm 4.

C Experimental Settings

C.1 Model Size

The model sizes used in the experiment are shown in Table 2. The parameters ( N , d hidden , d FFN, and N head) of BERT-1.8B, BERT-6B, and BERT-13B are referenced from T5-3B and T5-11B [35]. The parameters of BERT-100B are referenced from GPT3-175B [4].

C.2 Hardware Configurations of Environment

8 × 2080Ti : The GPU is NVIDIA GeForce RTX 2080 Ti (10 GB). The CPU is Intel(R) Xeon(R) CPU E5-2680 v4 @ 2.40GHz. The GPU number is 8, and the CPU core number is 56. CPU RAM size is 126GB. The PCI-E version is 3, and the bandwidth is 15.8GB/s. NVLink is not available.

8 × A100 : The GPU is NVIDIA A100-SXM4-40GB. The CPU is Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz. The GPU number is 8, and the CPU core number is 128. CPU RAM size is 503GB. The PCI-E version is 4, and the bandwidth is 31.5GB/s. There are 12 fully connected NVLink 3.0.

64 × A100 : 8 nodes of '8 × A100'.

C.3 Training Settings

As we mentioned in the main text, we implement CPU AdamW [24] based on Intel AVX instructions and use it to train our BERT models. More details of the implementation are stated in Section A.3.

In this paper, our profiler collects the runtime information in the first S steps ( S ≥ 10 ). Then we run our solver to determine the optimization sequence of the model. After that, we skip S steps to wait

Algorithm 4 DP Solver (pruned version)

̸

for the model warm-up. Finally, we measure the runtime for at least S steps and reported the average result of the step time.

In the actual-run experiment, we set the memory constraint of the 2080Ti environment as M = 8000 MB, and M = 28 GB for A100. The memory constraints do not equal the GPU memory capacity because the actual memory occupation will be more than the memory usage due to memory fragmentation, memory cache, and other factors. We have a discussion about this problem in Section E.2 For all experiments in this paper, we set m = 128 as the default.

We use the half-precision float for all parameters and input tensors. Initial parameters and all input tensors are randomly generated. We fix the random seed and check the loss is invariant to ensure the correctness of our implementation. The sequence length is set to 512 , which is maximal for BERT.

We adopt BMTrain 0.1.8 and DeepSpeed 0.7.6 for our experiment. For a fair comparison, we turn on optimizer offloading and rematerialization for Megatron-DeepSpeed. In contrast, we do NOT introduce other optimizations such as tensor parallelism, pipeline parallelism, or FlashAttention for BOTH H3T and Megatron-DeepSpeed. We will leave the implementation of these optimizations and the comparison as our future work.

D Supplemental Experiments

D.1 Memory Test

We test the minimal memory cost of the two frameworks under different settings, and the result is shown in Table 3. From the result, we can find that in most of our settings, our implementation supports the batch size at least 4 × larger than Megatron-DeepSpeed, which provides users more flexibility and convenience in training their models. We can also find that H3T can save 34 . 6% ∼ 70 . 7% ( 70 . 3% ∼ 80 . 5% ) GPU memory compared with Megatron-DeepSpeed ZeRO-3 (ZeRO-2) with the same batch size and other settings. Furthermore, we are able to train larger models under limited hardware resources, such as training BERT-6B on 8 × 2080Ti.

Not only that, based on our implementation, we can train GPT-3-175B [4] on 64 × A100 (sequence length is 512, and batch size is 512) while using only about 11.1 GB of memory per device at a minimum. This is really an exciting result because we would never have dreamed that a 175Bparameter model could be trained with such limited hardware, not to mention with such low GPU memory overhead. Unfortunately, we cannot train GPT-3-175B with a larger batch size or a sequence length due to the CPU RAM memory limitation of 512 GB. Currently, the step time of H3T training GPT-3 within this setting is about 40 . 8 second. According to our observation, the bottleneck is now in communication but not computation, so if we can set the batch size and the sequence length larger, we will have a great increase in the training throughput. We will incorporate CPU RAM optimizations to solve this problem in our future work.

D.2 End-to-end Experiment on SuperGlue

We conduct end-to-end training on four actual tasks from SuperGLUE [48] to check the soundness of our implementation. Specifically, we train BERT-large-cased [9] using BMTrain with H3T and PyTorch with APEX 5 , respectively. When using PyTorch with APEX, we adopt mixed precision

training to ensure the fairness of comparison. The experimental result is shown in Table 4. From the result, we can find that our system can achieve comparable results with PyTorch on all four tasks, which greatly confirms the correctness of our implementation.

D.3 Energy Test























































































We evaluate the energy cost of our framework and the baseline, Megatron-DeepSpeed, and the result is shown in Table 10. From the table, we find that our framework achieves lower energy cost with Megatron-DeepSpeed in most cases, except for training the 6B model on 8 × A100, where we also achieve a comparable result. This result uncovers that our framework is more energy-saving than Megatron-DeepSpeed.

D.4 Solver Efficiency Test

AI scientists and engineers sometimes change the model size and architecture in their works. For most people, this change does not always happen, so we do not need the solver to be as fast as possible, but the efficiency of the solver is still in demand. Therefore, we test the efficiency of the three solvers, and the result is shown in Table 5.

From the result, we can find that the random solver and greedy solver run significantly faster than the DP solver, although the minute-level runtime of the DP solver is quite acceptable compared to the training time of the large-scale model. The DP solver owes its acceptable efficiency to the two prunings, especially the second one, which reduces more than 99% of the DP states. In contrast, DP without these two prunings will take over a day or even a month to get the solution, which is greatly unacceptable.

However, this is not the best for the DP solver. For deadline reasons, we use Python to implement the serial version of the three solvers. In the future, we will port them to C++ and optimize them using multithreading and parallelism techniques. In this way, we expect the DP solver to gain at least another 100 × improvement in efficiency and be much more senseless in the real application.

D.5 Case Study

We show a case study in Table 6 to intuitively show the ability of our automatic solvers. From the result, we find that if we manually slow down the PCI-E communication, the solver can adaptively switch on fewer communication-aware optimizations and more computation-aware optimizations. The opposite situation occurs if we manually slow down the computation. When the memory constraint becomes relaxed, the solver reduces the switch-on ratio of all four switches. These results show that the DP solver is well adapted to the actual environment to customize the optimization scheme. Besides, we find the greedy solver does not offer as much flexibility as the DP solver in provisioning multiple optimization switches, although it still performs well in most cases of the simulation and actual runs.

E Discussion

E.1 Scalability of H3T

In this part, we discuss the scalability of H3T. Although H3T is based on the solvers and optimization switches introduced above and implemented on BMTrain for Transformer training in this paper, the framework and the idea of H3T can be applied in many other scenes. An illustration is shown in Figure 11.

Optimization techniques and the solver . Scholars and programmers can introduce any optimization or parallelism techniques and design their own solver for the specified optimization switches. The switches do not have to be binary and layer-level as they are in this paper, as long as they can be formalized as several variables. For example, we can have a global switch of pipeline parallelism, or we can also have some continuous optimization hyper-parameters, e.g., the ratio of tensor offloading. Then, the design of the solver highly depends on the optimization switches. To support different switches and hyper-parameters, we may need a different solver from that we proposed in this paper. Despite this, we argue that the three solvers designed in this paper work well for most discrete binary layer-level optimization switches.

Models . H3T is designed for Transformer-based models but can also be applied to other neural models, which can be formalized into the sequential model. Since the structure of sequential models and transformer-based models is becoming more and more widely used in recent years, we argue H3T has good generality and compatibility in the present and future.

Toolkits . In this paper, we implement H3T on BMTrain. Scholars and programmers can also implement H3T on other toolkits like DeepSpeed, etc.

Switches in this paper

DL Toolkits

Pytorch

BMTrain

……

Distributed Learning Toolkits

…

Optimization

Techniques

Parameter Partition

Tensor Offloading

Rematerialization

……

Overlapping

DeepSpeed

Switches

……

Solvers

Greedy Solver

DP Solver

……

Formalize into

Design for

Big Models

Transformer Models

Other

Models

H3T implemented in this paper

The scalability of H3T

E.2 Limitation of H3T

Although H3T performs well in our experiment and has good scalability, it still has some limitations. This section discusses the limitations of H3T from two perspectives: memory and model structure.

Memory . In this paper, we follow most existing works [2, 3, 18, 20] and defaultly use 'memory' to represent the memory allocation on GPU. Although this setting works well in most cases, it still brings two limitations.

(1) The memory allocation does not equal the actual memory usage. In actual runs, the GPU memory is usually managed by the underlying DL frameworks like PyTorch, TensorFlow, etc. And memory usage is usually greater than memory allocation due to many complex reasons, such as cache and memory fragmentations. Therefore, the user must set an M lower than the actual GPU memory capacity to ensure H3T's proper running. Moreover, this also means H3T cannot perfectly estimate the memory usage, which leaves space for further improvement of our solver. One solution is to implement the memory pool management module in-house to gain absolute control over memory management [15]. To tackle this, we will develop a memory management module for absolute control over the GPU memory in the futures.

(2) We only focus on the GPU memory and neglect the CPU memory. This does not matter in most cases because the CPU RAM is usually large enough for all parameters, gradients, and optimizer states. However, when we train GPT-3-175B on 64 × A100, the CPU RAM runs out of memory if we set the batch size as 1024 or larger. This result reveals that CPU memory management is also in high demand for training extremely big models. Some researchers attempt to offload tensors from CPU RAM to the secondary storage (e.g., NVMe SSD) [37, 45], which is also our future direction to resolve this problem.

Model structure . As discussed in the scalability part, H3T supports any models that can be converted into sequential models. Nevertheless, there is one exception, namely, the models with dynamic structures, because we cannot search for a fixed optimization scheme for the dynamic model by pre-computing. This is a common problem with this kind of pre-computed optimization [3, 18, 20, 16, 31], and some other real-time optimization frameworks may be a better choice for these dynamic models [19, 15]. However, it is almost impossible to introduce approaches with heavy computations (e.g., DP and MILP) for real-time optimization, so most of these works employ heuristics to achieve adequate performance. In this paper, we mainly focus on Transformer-based models that are mostly fixed structures, so it is a better choice to design a pre-computed optimization framework like H3T.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nLarge language models have made significant advancements in terms of accuracy and representativeness across various applications. However, their large memory consumption and intensive computation size pose critical challenges for existing hardware platforms.\n\nThis paper proposes a memory-latency co-optimization training system that enables more efficient DNN training by reducing both memory consumption and latency. Given that the optimization is an NP-hard problem, this paper studies three optimization schemes using different solver searching algorithms. The evaluation demonstrates that the proposed solver-based optimizations significantly outperform other state-of-the-art frameworks, achieving up to 4.3x speedup and 80.5% less memory overhead.\n\nSTRENGTHS:\n- Training transformer is an important and hot research topic.\n- This paper is well-organized and easy to follow. It is also well-motivated.\n- The proposed techniques generally make sense.\n- The evaluations are thorough and show a significant reduction in memory usage and latency.\n\nWEAKNESSES:\n- The paper lacks clarity on the contribution and requires further clarification on some technical details. The main contribution proposed in this paper is in Section 3.3.2, while there are limited details about these optimizations in the main paper, especially for the DP solver.\n- To improve the understanding of the optimizations, it would be helpful to show a breakdown analysis of which specific optimization (e.g., re-computation, parallelism) contribute more to memory and latency reduction.\n- Some figures have font sizes that are too small to read, such as Fig. 2.\n- It would be clearer to explain the major difference between ZeRO-2 and ZeRO-3.\n\nQUESTIONS:\n- Further clarification of the memory reduction ratio would be helpful. Does the claim of 80.5% less memory consumption include the model size?\n- What is the actual overhead (in terms of time) for the Brute-Force solution?\n- Is the proposed solution compatible with layer reordering [1]?\n\n[1] MCUNetV2: Memory-Efficient Patch-based Inference for Tiny Deep Learning\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFor efficient training of large language models, how parallelism and memory optimization are jointly configured is essential. Since the search space is too broad to explore, each has been studied separately. However, the separate search for memory optimization and parallelism could lead to a suboptimal solution. Therefore, efficiently exploring the joint search space and finding the optimal solution is a significant research problem.\n\nThis paper addresses it by proposing a novel system: H3T (High-Throughput Transformer Training). H3T incorporates memory optimizations (i.e., rematerialization and offloading) and parallelism strategies (data, model, and pipeline) into a single abstraction, which the authors defined as an optimization switch. Then, H3T automatically finds the best solution under the constraints (e.g., memory capacity) with a dynamic programming solver. The evaluation results show that the H3T solver can find better solutions than manual search results. Even more, their training performances (step time) are faster than Megatron-DeepSpeed.\n\nSTRENGTHS:\nAs the authors pointed out, H3T is the first attempt to automatically integrate memory optimization and parallelism strategy for large Transformer-based models.\n\nMoreover, the authors conducted various experiments, including memory tests, e2e evaluation, energy tests, and solver efficiency tests. The solve efficiency results show that H3T can find solutions in less than 20 minutes even for the DP solver.\n\nTwo typos found:\n- Alph => Alpa in line 137\n- $P_i$ and $ \\nabla P_{i} $ is => $P_i$ and $ \\nabla P_{i} $ are in line 191\n\nWEAKNESSES:\nAlthough the authors say that H3T covers model parallelism including pipeline parallelism, the paper does not explain how pipeline parallelism is incorporated. It looks like H3T is an automatic optimizer for tensor (parameter) parallelism and memory optimizations.\n\nFurthermore, I think the experiments could be enhanced.\n- Since H3T is an automatic optimizer, it is more plausible when adding a comparison with other automatic optimizers. For example, compared to Alpa or Unity, can H3T operate better?\n- Did the Megatron-DeepSpeed results use the pipeline and tensor parallelism that Megatron provides? I expect that Megatron can perform much better for the 13B model without using the ZeRO optimizer (or only ZeRO-1 optimizer).\n\nQUESTIONS:\nQ) The paper highlights that **H3T can train a GPT-3 175B on 64xA100 40GB with a batch size of 512 and using only about 11 GB of memory per GPU**. Could you also give the step time of training the GPT-3 175B model? According to a Megatron paper [1], we can train the GPT-3 175B model on 64xA100 80GB with a notable step time of 13.75 seconds with a batch size of 64. Although it uses a different A100 cluster, each of which has 80GB memory capacity, comparing the values would be beneficial because the prices of A100 40GB and A100 80GB do not differ much (e.g., AWS p4d.xlarge is 32.77$ per hour and p4de.xlarge is 40.96 per hour). It could be much cheaper when using A100 80GB if possible.\n\nQ) Does Megatron-DeepSpeed includes the latest Megatron [1] that supports selective recomputation and sequence parallelism?\n\nQ) For the BERT-13B result with 64xA100 GPUs, why ZeRO-3 fails though ZeRO-2 succeeds? To my knowledge, ZeRO-3 is a more aggressive memory-saving approach than ZeRO-2.\n\nQ) Do BMTrain and Megatron-DeepSpeed use FlashAttention?\n\nQ) Is there any reason why $d_{FFN}$ is much bigger than $d_{hidden}$? Only the 100B model follows the common scaling ($d_{FFN} = 4  d_{hidden}$). For the others, it seems unnatural that $d_{FFN}$ is larger than $16  d_{hidden}$.\n\nReference \n\n[1] Korthikanti, Vijay Anand, et al. \"Reducing activation recomputation in large transformer models.\" Proceedings of Machine Learning and Systems 5 (2023).\n\nLIMITATIONS:\nNo potential negative societal impact exists in the paper. The authors adequately addressed the system's limitations in Appendix.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a framework to automatically find an optimized integration of memory and parallelism for training large language models (LLMs). To this end, this paper first formulates the memory optimization strategies and parallelism techniques, and then solves the equation to find the best configuration for the training of LLMs. Using three different solvers, this paper shows a significant boost in terms of training speed and memory compared to the DeepSpeed framework across different models.\n\nSTRENGTHS:\n-- The main strength of this paper is its results on training LLMs. It shows a significant improvement compared to DeepSpeed.\n-- The main idea of finding the best memory optimization strategy and parallelism technique automatically is interesting and of significant importance since the models are becoming larger.\n\nWEAKNESSES:\n-- The main concern of mine is related to the configuration used for training using DeepSpeed. DeepSpeed supports different memory optimization strategies and parallelism techniques which can be set manually via its configuration file. However, Table 1 of this paper shows that DeepSpeed fails to fit some models. I don't understand what could be the reason. My understanding of DeepSpeed is that with a proper config, it can fit the model at the cost of speed. Moreover, increasing the batch size shouldn't increase the memory as long as it reflects as gradient accumulation steps as discussed in PipeDream [25]. Therefore, my suggestion to the authors is to include the details of memory strategy and parallelism technique used in Table 1. Without such details, it is hard to comment on the significance of this paper.\n\nQUESTIONS:\n1- My main question is about the config file of DeepSpeed for the results in Table 1 (see https://huggingface.co/docs/accelerate/usage_guides/deepspeed). Without the details of DeepSpeed configuration, it is hard for me to evaluate the results of this paper. As a result, I will go with a midpoint score for now. Once I see the config file, I'll decide to increase or decrease my score.\n2- Have you also compared your work with FSDP? How do you perform compared to this framework?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers large model training and proposes a search algorithm for various memory optimizations.\n\nSTRENGTHS:\nPlease see the \"Questions\" section.\n\nWEAKNESSES:\nPlease see the \"Questions\" section.\n\nQUESTIONS:\n- Lines 89-91: In what sense are you using the terms \"dynamic\" and \"static\"? I think, according to the most common way these terms are used, the methods referenced in the paragraph starting in line 89 are all static.\n\n- Line 105: \"These two optimization strategies are independent\". I think this statement may need to be changed or removed. Rematerialization and offloading are not necessarily independent in the sense that one can offload a tensor for predecessor node (in a computation graph) and then later use that tensor to rematerialize another tensor (e.g. output of a successor node). There is actually recent work that tries to optimizes the integrated search space for these two memory optimization techniques. Please correct me if I misunderstand what is meant by \"independent\".\n\n- I'm confused about one thing related to the proposed algorithm. Does the algorithm optimize for parallelism? That is, does your search include different parallelism schemes or are you assuming a fixed parallelism scheme and search for other things given this fixed parallelism? \n\n- The Megatron-DeepSpeed baseline in Table 1: Do you have a sense of why ZeRO-2 beats ZeRo-3? \n\n- For ZeRO, are you running mixed-precision training? Perhaps more importantly, what precision is used for your method vs ZeRO?\n\nMinor:\n- no space between analysis and [26] in line 30\n- Line 74, MM is short for?\n- Is it convention to capitalize \"Transformers\"?\n- Is this statement in line 85 correct: \"Rematerialization is first proposed by Griewank et al. [12]\"? I believe there are older references for rematerialization.\n\nLIMITATIONS:\nPlease see the \"Questions\" section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}