H3T: Efficient Integration of Memory Optimization and Parallelism for Large-scale Transformer Training
Authors: Yuzhong Wang
Xu Han
Weilin Zhao
Guoyang Zeng
Zhiyuan Liu
Maosong Sun
Venue: NeurIPS 2023 poster (2023)
Keywords: ML System
Parallelism Learning
Memory Optimization
Data Parallelism
Model Parallelism
Parameter Parallelism
ZeRO
Rematerialization
Checkpointing
Tensor Offloading
Dynamic Programming
TL;DR: A framework for High-Throughput Transformer Training (H3T) that can train GPT-3-175B using only 64 NVIDIA A100 GPUs. The implementation can find efficient integration of memory optimization and parallelism for actual training of Transformers.

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
Transformer Training
Yuzhong Wang*!, Xu Han Weilin Zhao
Guoyang Zeng', Zhiyuan Liu 1...1...,a.ao Maosong Sun
NLP Group, DCST, IAI, BNRIST, Tsinghua University, Beijing, China
ModelBest Inc., Beijing, China
Collaborative Innovation Center for Language Ability, Jiangsu Normal University, China
{yz-wang21 wwlzzwlll2l@. Smails tsinghua. edu. cn
zengguoyangdmodelbess.cco cn (hanxu2022, 10e,,li22,, liuzy, sms sms)@tsinghua, edu. cn
Abstract
In recent years, big models based on Transformers have achieved state-of-the-art performance many artificial intelligence (AI) tasks. Despite the success of these Transformer-based models, their huge parameter size poses a serious challenge 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 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 Trranfformer-bare based models. Specifically, we design search algorithms to combine appropriate memory optimization strategies and parallelism schemes to achieve balance between memory overhead and training efficiency. We implement H3T based use H3T to train the Transformers of different sizes to evaluate the efficiency H3T. The experimental results show that H3T outperforms the most popular deep learning (DL) toolkit Megatron- Sreaatoon-eeepSped by 1.2x ~ 4.3x training speed while reducing 34.6% 80.5% of memory overhead. Moreover, H3T can use only 64 NVIDIA A100 GPUs to train GPT-3-1757 which is very difficult for existing DL toolkits. The source code is available at /httss://ggtmaeie com OpenBBB/ BMTrain/trne ttree///t
open-source toolkit BMTrain and then
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 models have achieved state-of- performance on many tasks
mechanisms, Transformer- and have become the foundation architecture for various AI directions, such as natural language processing
computer vision 23 22 and multimodal processing 32
Despite the outstanding performance, with the trend of continuously growing model size, the compu- tation and storage costs become severe challenge for training most big Transformer-bas 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 float32 format, mention the forward activations,
Indicates equal contribution.
Indicates corresponding authors.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

backward gradients, and optimizer states. Even ignoring the memory overhead, according to the existing analysis [28]. will take about 288 years to train GPT-3 with 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
T6
2
41 [37 45] and rematerial-
ization 122 5 20
18 [9] aim to reduce the memory overhead of one (or each) single GPU. Parallelism
techniques such as data parallelism [53] 30
and parameter parallelism 8 I7 27 51
36 mainly
focus on improving both training efficiency and memory utilization across multiple GPUs, Some hy-
brid memory optimization strategies 49 61 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 framework to automatically find efficient integration of memory opti- mization 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 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 tradi- tional 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 44), one of the most popular distributed learning toolkits for DL. The results show H3T can train big Transformers 1.2x ~ 4.3x faster than Megatron- DeepoSped 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 automatically integrate memory optimization and parallelism for training big Transformers (2) We conduct adequate experiments that show H3T outperforms conventional manual optimization as well the most popular DL toolkit Regatron-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
a widely-used deep neural network architecture. Due to the powerful capability of modeling data, the Transformer architecture is used to build various preitrained models. These pre-trained models can capture rich knowledge from massive unlabeled data and then transfer the captured knowledge to handle various complex Al tasks. Transformer-basee pre-trained models have become one of the cornerstones of recent Al developments, especially in the NLP domain. Currently, Transformer-basee models achieve state-of-the-an performance on almost all NLP tasks 21 35
Transformer 47
33 44 Inspired by the success of these models in the NLP domain, Transformers are also explored in other important AI domains like CV 110 23 22 and Multi- (MM) [2 39 38 29
Despite the remarkable achievements of Transformers, it is 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 OpenAl's GPT as an example: the parameter sizes of GPT-1 [3]. GPT-2 [4), and GPT-3 4 released in 2018, 2019, and 2020 are 1 17 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. Rema- terialization 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(VN) without changing the computational time complexity. which is milestone that makes rematerialization known to researchers and implemented in many DL toolkits. However, Chen's 's work 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 developing automatic rematerialization methods tackle this issue. Dynamic programming (DP) [4 20 and mixed-intege linear programming (MILP) [8] 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] 66 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 custom custom-designes designed genetic algorithm, while Beaumont al. 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 natural idea. attempt to intelligently combine them using heuristics and achieve remark-
Early works 49 31 able performance. Afterward, some scholars explore using better algorithms, such as DP 33 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
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) (6) 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.
Data Parallelism. Data parallelism 55 30
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 51) which partitions model parameters by layer level and lets devices be responsible for the computation of differentt layers. The other is parameter parallelism [28] 40). which partitions parameter matrices within 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 designed based on data parallelism, model parallelism, and ZeRO.
Automatic Parallelism. Although many efforts have attempted to combine model parallelism and data parallclism to improve efficiency further, finding the optimal combination based on empirical
3

---

Profiling the nuntime and memory cost. for each layer
Solver
Layer
Transformer
based Model
Model Simulator
Activation-aspect
Parameteraaspet
Execution
(53.3.1)
Profiler
Focward
Switches
Switches ($3.22)
Switches
Layer
Farward Segment Sync.
Offleaddnng
Ofthoding (global)
Search Algorithm ($3.3.2)
Laay-greftcinng
Baclovard
Search for solution appropriate optimization switches for each layer.
Figure I: The overall architecture of H3T.
rules not easy. To this end, some scholars explore automatic parallelism methods. Galvatron 26 uses a decision tree and DP algorithm to find 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 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.
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 typ- ical 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,, Backward,, Optimize, as the three stages at the i-th layer, respectively.
In the forward propagation, Forward; takes H, and P. to calculate H,+1, where H; and H,+1 the hidden states, and P; is the parameters of the i-th layer. Along with generating some intermediate activations A, the forward propagation the i-th layer can be formalized as
H.+1, A, Forward, (H., P.).
(I)
The final output HN+1 is used to calculate the loss function and the initial gradient VH Then, in the backward propagation, Backward, is to calculate the gradient VH, and VP, by using VH H. P,, and the intermediate activations A; generated by Forward; which can be given as
VH. VP: Backward, (H,, P., A,, VH.+).
(2)
After the backward propagation, Optimize; to use some parameter optimization functions to update the parameters P; according to the gradients VP;-
In each training step, the conventional training way first to run Forward, from i = 1 N, then run Backward; 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 is allowed to adjust the execution order add some additional operations under the premise that the gradients 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-aspeet Optimization Switches
For the sequential model, activations involve hidden states H, and intermediate activations A, They play important roles in model training and meanwhile consume lot of GPU memory, especially intermediate activations [42.. As mentioned in the background, scholars have developed two major approaches, rematerialization and offloading, 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 save the memory of

---

Forward/
H
A Evitt Forwand)
Offload after
Prefetch before
Forward; (Re-compute)
Forward
Re-forward
A
Rematerfalization
RA M
fflooding
Backward,
Figure 2: Activation-aspect switches.
hidden states. This setting is also applied by some existing work of hybrid optimization optimization switches are illustrated in Figure 2
3 The two
Rematerializing intermediate activations is to evict A, after forward steps and recompute A; before backward steps. To differentiate, we define Forward; as the forward computation without saving A, As Figure [ shows, rematerialization for the i-th layer replaces Forward, with Forward; in the forward stage, and insert an extra Forward, before Backward, the backward stage.
Offlodding hidden states to offload H; and prefetch them before the backward propagation steps. Different from intermediate activations, if we similarly try evict and recompute hidden states, the training process will suffer from excessive recomputations, and the performance will also be affected 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 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. and VP, on GPUs. As described above, P, and VP, are only used for Forward,, Backward,, and Optimize, Therefore, it is 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 efficient way reduce the GPU memory overhead. On this basis, we further implement a high-performance CPU optimizer that runs Opt Optimize, on CPUs.
Moreover, employ parameter parallelism for H3T. Specifically, store partitioned parameters p(o) on each GPU. Before Forward,, we use an all gather operation to gather them from all GPUs to get P. for computation. After Backward, we evict the gathered P; and use a reduce scatter operation to average and scatter the VP; 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, globally the above offloading optimizations and parallelism this paper.
In addition the above global optimizations, 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ße) is to evict the partitioned parameters after every reduce scatter operation and lazily prefetch them back before every all gather operation. Lazy-prefetching aims eliminate the memory cost of resident partitioned parameters.
Both parameter partition and lazy- prefetching can avoid the long-tern persistence of parameters dur- ing the forward-to-bee ba-kaac 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.
TThe CPU optimizer is implemented based Intel AVX instructions. More details are in the appendix.
5

---

Olloaadng (global)
All Gather
Belore Forward
GPU 1
GPU
RAM
Partition
Lacy-prelelching Prefetch
P:
P (0)
P
Evict
Aftee Optimize
After
VP;
VP (0)
VP:
Backward
After Redune Scaller
Optimize
CPU Optimier
Redace Scatter
Hased SIMD
Aftor Backward
(a) Basic parameter optimization scheme based on tensor offloading and parameter paral- lelism.
Gather
Fueward
GPU
G
Officadingg (glokal)
GPUO
GPUO
Partition
Prefecch
p
P()))
Evict
Before All Gather
AStinnptimize
Aftac
Evict
Afte
Reduce
VP:
Badwwul
Scatte
Femerd Bacowad
(b) Parameter partition.
sazy-prefetching.
Figure 3: Parameter-aspect switches.
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 stream that covers the GPU-GPU communication operators, i.e., all gather and reduce scatter. (3) A PCI-E stream that covers the GPU-CPI communication operators, i.e., offloading and prefetching. Following most previous works 49 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 implementation is 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 "segment".
Calc. Stream
Forward, cosputt
Backward,
Calculation the current layer
Post-process of the previous layer (1
NVLink Stream
ReduceSattor
AllGather?Pram
PCI-E Stream
PrefetchParam
Prefetch
Offload
Preparation of next layer
Figure 4: An illustration of backward segment i. Without loss of generality, suppose the switches of layer 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 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 streams to be synchronized between two segments. More details about the implementation and the three switches are introduced in the appendix.
force all
NNVLin is high-speed multi-GPi GPU communication link developed NVIDIA. More details here

---

3.3 Optimization Solver
The optimizations save memory while bringing performance losses. To tackle this, H3T involves 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 = ($1, $2, ... SNN as an optimization sequence for the model, where s, E 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. this paper, S involves all combinations of the layer-level switches described above,
have [S] = 2"
128.
We define Runtime(s) Memory(s) the runtime (per step) and the memory overhead (per device) given the model and the optimization sequence 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):
minimize Runtime(s), s.t. Memory(s) S M.
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 implement model simulator to simulate the training process of each layer based on the model architecture and implementation, thus helping us estimate Runtime(s), (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 any combination, so the number of global solutions can be up to |S|AN. 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 this paper.
3.3.2 Search Algorithms
Random solver. Brute force does not work for 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 and pick the best one.
Greedy solver. The random solver 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 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 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 con- straint, 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 [ 43 we design a search algorithm based on DP to solve our task better. Specifically, our DP algorithm considers the optimal nol-solutions for specified prefixes and memory limits. When transferring, need to enumerate the next layer's 's switches and estimate the next segment's 's running status.
Memory
Peak memory usage peak) Current memory usage (memeure)
Segment
Forward
The of Backward
Figure 5: An illustration for i, mem meMgeak.
Formally, the DP state can be represented by the following variables: (1) i: The prefix length, which means this state is considering segment 1 (including forward segment ~ and backward segment i ~ 1, similarly hereinafter). (2) mem The memory usage after forward segment 1 ~ i.
7

---

(The memory usage before backward segment i ~ is essentially equal and can be easily calculated.) (3) pompeak: The peak memory usage during segment i. (We have to record mempeak because the resident memory may increase due to non-lazy www-lazy-prefecchngg. leading to an increase in the peak memory from the previous segments.) (4) 8i, 81+1: The switches of layer and + 1, which used to simulate segment + 1.
We use sol mem mempeak S;, 1i+1) to represent the optimal sub-solution of the state. Then we can enumerate S142 to transfer to the next state (i + 1, 1,mem(ur: mem S(+2)) During the transfer, we need to: (1) Simulate the forward segment i + and backward segment i + 1 based on 82, 88+ 1 2-+22 (2) Calculate mem'ur: memp and the total runtime based mem mempeak the previous runtime, and the simulation result. (3) Update sol(i + I, mem pem 81+1, 9t42) if the peak memory mem peak does not exceed M and the runtime is shorter than the existing solution's. 's.
There are two noteworthy details for our DP algorithm. First, considering the overlapping sub- problems 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 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 pscudo-code and time complexity analysis.
4 Experiments
4.1 Experimental Settings
Models. We use BERT()) 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 (O.IB) 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 x NVIDIA GeForee RTX 2080 Ti 10 GB (8 x 2080Ti); (2) 1 node of 8 x NVIDIA 100-SXXM4400GB with NVLink 3.0 (8 x A100); and (3) 8 nodes of 8 x A100 (64 x A100).
Implementations. We implement H3T based on BMTrain an open-source toolkit that accelerates the training of big Transformers with various parallelism techniques.
Except for H3T with random solver and greedy solver, we introduce manual optimization strategy without H3T and another toolkit Dregatron-DeepSpee 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 best-performing solution that does not exceed the memory limit. (2) 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 fair comparison, we test states and gradients; ZeRO-3 distributes optimizer states, gradients, and parameters [36)). Please refer to the appendix and the code for detailed configurations of the actual-running experiment.
Baselines.
with ZeRO-2 and ZeRO-3 respectively (ZcRO-2 distributes optimizer
and H3T in
We do not test POFO B. 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, is almost impossible for 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 well-designed offloading solution of POFO and scriously hurt its performance.
8

---

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
BERT-L.8B on xx2080Ti
BERT-L88 8 8~2080Ti
BERT-6B 8>2080Ti
BERT-6B 8=2080Ti
Seep Time (toc)
Global Batch Size - 128
Global Batch Size 512
Step Time (coc)
Global Batch Size 8
Global Batch Size - 32
20
20
Time
17:00
Time
16.75
5.0
15
16
24
32
16
24
32
16
24
32
24
32
Memory Constraint (GB)
Memory Constraint (GB)
Memory Constraint (GB)
Memory Constmint (GB)
BERT-6B on KxA100
BERT-6B 8*A100
BERT-138 on 8*A100
BERT-13B on
Step Time (68c)
Global Butch Size - 128
Global Batch Size - 512
Step Time (8a))
Global Batch Size
128
Step Time (Ra0)
Global Batch Size
512
5.50
Step Time
19.75
13.0
5.25
31.5
12.5
32
64
96
128
32
128
32
128
96
128
Memory Constraint (GB)
Memory Constriint (GB)
Memory Consraint (GB)
Memory Consrmint (GB)
BERT-13B on 64*A100
BERT-13B on 64*A100
BERT-100B on 64xA100
BERT-100B on 64~A100
Step Tine (865)
Global Batch Size
1024
Step Time (oec)
Global Batch Size
4096
Step Timn (sec)
Global Batch Size
1024
Global Batch Size
2048
10.0
70
29.0
Step Tine
28.5
68
32
64
96
128
32
64
128
32
128
96
128
Memory Constraint (GB)
Memory Constraint (GB)
Memory Constraini (GB)
Memory Constraitt (GB)
DP Solver
Groedy Solver
Ransom Solver
Marual Optimization (wol HST)
Figure 6: The results of the simulation study.
In this section, report 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. contrast, the random solver 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 sufficient. This indicates that our layer-level automatic solver able to precisely de- cide 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-I,8B with a batch size of 128 on 8 80208iTi, the greedy solver underperforms the random solver. For BERT-6B on 8x2080Ti, 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-ruming experiment in the three environments, and the result is shown in Table [ The results demonstrate our great advantage over Dggatron-DeepSpeed. Compared the MegaatronDeeppSpee
with ZeRO-2 (ZeRO-3), H3T improves the training speed by 1.2x ~ 1.8x (1.9x ~ 4.3x) in our experiments. Comparing the results of BERT-6B and BERT-13B, both with the batch size of 128 and both on 8x A000, we find that H3T outperforms more significantly for larger model size, which means H3T not only faster but also more adapted to bigger models.

---

Table 1: The time per step of the actual-running experiment (unit: second). We do not report the random solver because fails on all these settings. means the program runs out of memory.
Environment
8x 2080Ti
8 A100
64xA10
Model Size
1.8B
6B
6B
13B
13B
100B
Global Butch Size
128
512
8
32
128
512
128
512
1024
4096
1024
2048
DP
6.0
17.6
19.4
19.8
5,5
19.5
10.0
28.4
7.4
26.6
35.6
61.6
BMTrain
Greedy
5.8
17.5
19.6
20.5
5.5
19.7
9.7
28.7
7.6
26.8
36.2
62.0
w/o H3T
17.7
22.1
23.0
5.7
19.8
9.9
28.6
7.8
26.9
36.3
61.8
Megatron
ZcRO-3
11.3
15.7
43.2
-DeepSpeed
ZeRO-2
9.3
6.8
17.7
10.6
Focusing on the results BMTrain, we find that H3T with DP solver and greedy solver both outperform BMTrain without H3T by up to 1.2x 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 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 Deeirronn-DeSppe.e... From Table [] we find Negatron-DeepSS better supported by H3T. Hence, users can train larger models with more limited hardware, c.g. train BERT-6B on 8 2080Ti. Not only that, H3T can train GPT-3-175B [4] on 64xA100 A100 with a batch size of 512 and using only about 11 GB of memory per GPU. This really an exciting result because we would never have dreamed that 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.
often runs out of memory for larger models and batch sizes, which are
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 SuperGL.UE [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 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 for training big Transformer-bases models.
H3T is more energy-efficient than Hegatron-DeepSppp (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 High-Throughput Transformer Training, named H3T. Despite the challenge of large search space, H3T significantly outperforms some manual strategies and one of the most popular DL toolkits, Megatron-DeepSSppe NVIDIA A 1100 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 training big Transformer-bast models and promote the development of the Al community.
in our experiments. More excitingly, H3T can use only 64
Acknowledgments
This work supported by the National Key R&D Program of China NN..20222D01163 16312), Major Project of the National Social Science Foundation of China (No. 22&ZD298), the National Key R&D Program of China (No. 20227D0160501), and Institute Guo Qiang at Tsinghua University. Weilin Zhao is supported by Tsinghua University Initiative Scientific Research Program.
10

---

References
[1] Maram Assi and Ramzi A. Haraty. A survey of the knapsack problem. In Proceedings of ACTT, pages 1-6, 2018. https://ieeexpllore ieeee.
ieee. wtw..e/...gg/occmmmtt//////////
[2] Olivier Beaumont, Lionel Eyraud-Dubois. and Alena Shilova. Optimal gpu-epu offloading strategies for deep neural network training. In Proceedings of Euro-Par, pages 151- 2020.
https://link springer ctt wwwpo///wwt.....oooooooooo 1ttp:/1000.0007777----00800 -00---000---77222100
[3] Olivier Beaumont, Lionel Eyraud-Dubois, and Alena Shilova. Efficient combination of re- materialization and offloading for training dnns, In Proceedings of NeurlPS, volume 34, pages 23844-23857, 2021. https: bttpss:/ppaperss http://ppppp..pp......p. wwtpppprrffllsaaa////ppppp http:/ppppp/2222// /2///////////////lle pdf
[4] 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 few-shot learners. In Proceedings of NewrIPS, volume 33, pages 1877-1901, 2020. https
www.xiv.ogg/df////// pttp:/dff/20500 14165.1114 pdf
[5] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint vrXiv:160.00112. 2016. http://///xxv..o//d......
06174.
[6] Xiaoming Chen, Danny Z. Chen, and Xiaobo Sharon Hu. modan: Memory optimal dnn training on gpus. Proceedings of DATE, pages 13-18, 2018. https::/ieeeepplore ieee org/ document/8841977
[7] Tobias Dantzig and Joseph Mazur. Number: The language of science. 2007. https:/
phi philpapers owg/rrg//e..........e. wttp://///err////////T/T. DANNTLL
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large scale distributed deep networks. In Proceedings of NeurlPS, volume 25, 2012. https://proceddangs meurips..c/papees Paper, pdf
[9] Jacob Devlin, Ming-Way Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, In Proceedings of NAACL, pages 4171-4186, 2019. httpp::/www acwweb..lwbbb.. wwwp:/awlle....gaaannnaaeaeeeeee
hwt.h//......r..aooooooo///////e wwthaaoooooooyaaaaae 1423. pdf
[10] 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.1 1929, 2020. http:://axxivoorg/pfffff /210 11929. .df.
[11] 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 hour. arXiv preprint arXiv: (ele::0066.07777 2017. https::/aaxxvvorggpppaa. 02677.pd
[12] Andreas Griewank and Andrea Walther. Algorithm 799: Revolve: An implementation of check- pointing for the reverse or adjoint mode of computational differentiation. ACM Transactions on Mathematical Software (TOMS), 26(1): (6611:19945 2000.
hhtps: //dl acm. org/doi/pdf/1
1145/347837 347846
[13] José Grimm, Loic Pottier, and Nicole Rostaing-Schmine Optimal time and minimum space- time product for reversing a certain class of programs. PhD thesis, INRIA, 1996. https
//inria hal siwieenee/inriaa-007770000
21--2222222------------ .df
[14] 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 133114 2019. https: httts://arxiss arxivv org /9911 19211991..1111111111 13214.pdd

---

[15] 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 dan training. In Proceedings of ICS, pages I 1-13, 2022. htttps//dl.co
org/doi/pdf/10n 10. 1145/3524055 3532394
[16] 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://dll acm. www.aa..............///e..eeoote org/doi/pdf 10. 11145/3373376.3378565
[17] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V Le. Yonghui Wu, et al. Gpipe: Effi- cient training of giant neural networks using pipeline parallelism. In Proceedings of NeurlPS, volume 32, 2019. https://prooeednss..oe newrips 09566600022210007777777222222222
www.eeeeeeennnnn...n...... cc/ cateriis.ccppppppee pappr/2019/ff1ll
pdf
[18] 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 remate- rialization. In Proceedings of MLSys, volume 2. pages 497- 2020. https.tttstt// httpps::pproceedinngg Paper popprrrpff
[19] 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.
http://p/6/0000066666.66 pdf
[20] Mitsuru Kusumoto, Takuya Inoue. Gentaro Watanabe, Takuya Akiba, and Masanori Koyama. A graph theoretic framework of recomputation algorithms for memory-cfficient backpropagation. In Proceedings of NeurIPS, volume 32, 2019. hhttps://arxis
httg:/df/19990 11722. pdf
[21] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint /rXiv:1907.116922 2019. https: http:://pocceedeeee
neurips. cc/ Paper paaperppff. pdf
[22] 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.
/2111. 09883.pp8..
[23] 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 ICCV, pages 9992-10002, 2021.
http:///axxvvv.oggg///
/2103. 14030.pdl
[24] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proceedings of IČLR, 2019. hattps://arxi:
httg/pdf/1711.050000
[25] George B Mathews. On the partition of numbers. Proceedings of the London Mathematical Society, (1)::48-490. 1896. https::/ oup. com/ www.o////m//////e.e//eeeeeeeeeet wtwpm///aatiiclle abstraccc s1-28 11-28/1/446/149582)
[26] Xupeng Miao, Yujie Wang, Youhe Jiang, Chuman Shi, Xiaonan Nie, Hailin Zhang, and Bin Cui Galvatron: Efficient transformer training over multiple gpus using automatic parallelism. arXiv preprint 2rXiv:2211.13875 2022. https: fttpss://axxiie /arxivxivvvv..i org/ http/ddf/221111 2211 13878. pdf
[27] 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://ddd acm org/doi/pdf/ 10. 34455/3341111.11 33599466
[28] 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 gpu clusters using megatron-Im. In Proceedings of SC, pages 1-15, 2021. /tttps://arxivvoor//pa////
04473. pdf
[29] OpenAl. Gpt-4 technical report. arXiv preprint arXiv:2303.0077744 2023. https : //arxiv
org/ http/ddf22300 08774.
12

---

[30] 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:
2013. https://arxiv. http:/iwwwwrrppdd//////eee// 1112. 2186.6618..00 pdf
[31] Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma, Qian Xiong, Fan Yang, and Xuehai Qian. Capuchin: Tensor-base based gpu memory management for deep learning. In Proceedings of ASPLOS, pages 891-905, 2020. http:/ //al://acccm capuchin. pdf
usc. www.:uuuucuuuuuuuuou cet/pooraa/staaaee http:/ssttcccc/oooaaata
[32] 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. Proceedings of ICML, pages 8748 0163,,773,,-0 2021.
https://arxiv. org/ org/pdf////// 2203. 00020. pd
[33] Alec Radford, Karthik Narasimhan, Tim Salimans, llya Sutskever, et al. Improving language understanding by generative pre-training, OpenAl blog, 2018. https://wnen ubc. ca/
Simuham01/LLNN6553/paaaaa pttps////ppppeeeraaaaaaaaaaaaaee mrd2df11m0111mmiiim pdf
[34] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9. 2019. httpp::/ccnn.oo htwp://ppppppppoooooeeeeeaaaaaaa mmsupprieee__mlliaae...eeem
saaeees_mollss_eeee are_
pdf
[35] 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 unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. http:
//jmlr.orz wwanllr.orgpppprrs..ee /t/////1122222//// 074. html
[36] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory opti- mizations toward training trillion parameter models. In Proceedings of SC, pages 1-16, 2020.
https://arxiy: https://axxiv.orr//pff 1910. 02054.0055.
[37] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero- infinity: Breaking the gpu memory wall for extreme scale deep learning. In Proceedings of SC, pages 1-14, 2021. https: https:///rxivv .rg.l....rrgggooo arg//df/// htfp//df//211 220.f//214..7888 07857.pd8
[38] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical lext-conditional image generation with clip latents. arXiv preprint 2rXiv:2204.06125, 2022.
https://arxiv http:://arrri.orr///a/////00
06125.pdf
[39] Aditya Ramesh, Mikhail Paylov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In Proceedings of ICMI,, pages 8821-8831, 2021.
12092. pdf
[40] Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System opti- mizations enable training deep learning models with over 100 billion parameters. In Proceedings of KDD, pages 3505-3506, 2020. /https://ll acm www.aaam....ooo.//. www.//gg///oooooooo /ooi 1145/ /0445//339486.... 3406703
[41] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang. Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-seea model training. In Proceedings of ATC, pages 551-564, 2021. https://arxiv orw. htgg/dff2101 06840v1.pdf
[42] Minsoo Rhu, Natalia Gimelshein, Jason Clemons, Arslan Zultiqar, and Stephen W Keckler. vdnn: Virtualized deep neural networks for scalable, memory-effiiien neural network design. In Proceedings of MICRO, pages 1-13, 2016. https://arxii: org httpppff/11f 160.. 08124. pdf
[43] Harvey M Salkin and Cornelis A De Kluyver. The knapsack problem: A survey. Naval Research Logistics Quarterly, 22(1):127-144, 1975. httts://oollinelibraay wtwpeeeiiiirrrrrlllrllleeeeeee wiley com/ doi/abs/10. 1002/nav ww000//aa.380022111110 3800220110
[44] 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.
/tt////tttmmmmmmmea//// attps://1eeexploree ieee, /
13

---

[45] Xiaoyang Sun, Wei Wang, Shenghao Qiu, Renyu Yang. Songfang Huang, Jie Xu, and Zheng Wang. Stronghold: fast and affordable billion-scall scale deep learning model training. In Proceedings of SC, pages 1-17, 2022. https://zwang4 github. 1 httpppullaiiinssssppeococ
pdf
[46] Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandeep Baines, Carlos Efrain Quintero Nar- vaez, Vinay Ramakrishnaiah, Nirmal Prajapati, Pat McCormick, Jamaludin Moso-Yusof, et al. Unity: Accelerating dnn training through joint optimization of algebraic transfor- mations and parallelization. In Proceedings of OSDI, pages 267-284, 2022. https
//www. usenix. http://stee//eeee//diiiiie Sonddi222uggerrr/ pdf
[47] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Proceedings of NewrlPS, volume 30, 2017. lttpss//pproceddi hew.eeddddnnnnnnnnn....... htcpilps.cppppppeee httpppppp221//////// $ffo223447ees11tt0sccccccaaaaaa Paper pdf
[48] Alex Wang. Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Superglue: stickier benchmark for general- purpose language understanding systems. In Proceedings of NeurIPS, pages 3266-3280, 2019. pdd.
https://w4ngatarg:
[49] Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song. Zenglin Xu, and Tim Kraska. Supemnuuonn: Dynamic gpu memory management for training deep neural networks. In Proceedings of PPoPP, pages 41-53, 2018. httpe:///xxvvvoog///reo 04380.pat
[50] Lijie Wen, Zan Zong, Li Lin, and Leilei Lin. swap dominated tensor re-generation strategy for training deep learning models. In Proceedings of IPDPS, pages 996-1006, 2022. https //ieeexplore leee. .rgg (olocumen//982000
[51] Jun Zhan and Jinghui Zhang. Pipe-toreh: Pipeline-based distributed deep learning in a gpu cluster with heterogeneous networking. In Proceedings of CBD, pages 55-60, 2019. https
/ieeexplore. ieee. wwwp/battaa////aeeeooo /ttb:tttrrttt/ommmmmm/9/910000
[52] 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: botps::/xxn //www/ usenix. org/ system/ wtwpllee////sd2--2221212e222 zheng- 1i anmin pdf
[53] Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient descent. In Proceedings of NewrlPS, volume 23, 2010. pttpp://rroceeddngggg neurips
14

---

A
Implementation Details
A.1 Timing Diagram of Multiple Streams with Optimization Switches
Cathered Parmmater
Forward
Backward
O pttmizer
Calc.
Fr
B,
NVLink
AG.
RS?
PCI-E
0f
(a) Default workflow based on parameter offloading, optimizer offloading, and param- eter parallelism.
Thrdt
BextGattndd Facendrr
Forward
Bazlovard
Cak:
Calc.
AGT
WiLis
ACE
RS!
PCI-F
0
PCIEE
(b) Rematerialization.
(c) Offloading hidden states.
10rd Gedana Honendarr
Snaet Gitberd ranmatir
Erat
Forward
Haclovard
Forward
Barlwaard
Optimbeer
NVTisk
AG
AG
RS?
Litk
AG
RS!
PCL-E
PCI-E
Parameter partition.
(c) Lazz-prefetching.
Fi
Forward'
Fi
Forward
R
Rackward
Optimize
AG
Gather
P? Prefetch Parameter
pM
Prefetch Hidden State
RS?
R Reduce Scatter
0 Offload Gradient
Ottload Hidden State
(f) The legend.
Figure 7: Timing diagram of H3Twith different optimization switches.
Figure illustrate ZiluuteZlllstratee the timing diagram of our implementation of H3T. Figure 7alis is the default workflow based on parameter offloading, optimizer offloading. and parameter parallelism. Figure 77 and are the workflow with activation-aspect switches, respectively. Figure 7 and 77 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-cconomical Korward/backwardd prioritize gathering and prefetching of the next layer, which may reduce the delay in the computation of the next layer. Since NVLink 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 g It is worth mentioning that there are still asynchronous operations in the economical mode, it memory memory-optimal but balances the efficiency.
Segment synchronization. Let the computation one layer and the concurrent communications (post-process of the previous layer and preparation of the next layer) be "segment". Then segment
15

---

Forward Segment
Forward Segment j
Calc.
Fj
Fj
NVLink
AG
AG
PCI-E
P++1
of"
of
J+1
X
Economical F Forward
Economical Forward
Non- economical and economical forward with NVLink.
Backward Segment
Backward Segment
Calc.
F;
Fj
Bj
NVLink
RS
AG
RSJ+1
AGI-1
PCI-E
PP_
P
0%-1
X
Economical Backward
M Economical Backward
(b) ccn-cconomical and economical backward with NVLink.
Forward Segment
Forward Segment j
Calc.
Fj
Fj
PCI-E
p.P.
AG
Of!
0
Pf11
AG
Economical Forward
Economical Forward
(c) Economical and con-economical forward without NVLink.
Backward Segment
Backward Segment
Calc.
F;
Fj
PCI-E
p.H.
P
AG
RS
RS
0
pll
AGP
1--1
X Economical Backward
M Economical Backward
(d) Economical and non-economical backward without NVLink.
Figure 8: Timing diagrams of H3T's workflow economical and non-economical mode. The legend is the same as in Figure 7f. Without loss of generality, we supposed all optimizations are switched on.
synchronization can be defined as: all streams need to be synchronized between two segments. In other words, if segment synchronization 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 little performance. Segment synchronization 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 AVX (Advanced Vector Extensions) instructions are SIMD (Single Instruction Multiple Data) instructions proposed by Intel for high-performance parallel computing. With the help of Intel AVX instructions, our CPU optimizer can achieve comparable performance the conventional GPU optimizer, which greatly ensures training performance.
16

---

Forward Segmenti
Forward Segment
Forward Segment
Forward Segment
Start before the previous
Wait until
Calc.
Fi
i+1
Calc.
[+1
the provious segment ends.
seginntt ends.
PCI-E
Of!
PCI-E
Of
Segment Synchronization
Segment Synchronization
Figure 9: A schematic diagram of segment synchronization. The legend is the same as in Figure 7f
B
Solver Details
In this part, 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 introduced in the main text due the space limitation.
B.1 Random Solver
Algorithm 1 Random Solver
Require: N, M, Runtime (s). Memory (s) Trand
1: Suptimal mull
2: for ir to Trand do
3: Randomly sample optimization sequence S.
4: if Memory (s) S M and Runtime (s) < Runtime (Soptimal) then
5: Soptimal
6: end if
7:
end for
Ensure: Soptimal
The pseudo-code of our random solver is shown in Algorithm Here we use hyper-parameter Trand to represent the sampling times. Since the time complexity of the model simulator is θ (N), the time complexity of the random solver is (NTrand).
B.2
Greedy Solver
Algorithm 2 Greedy Solver
Require: N, M, Runtime (s) Memory (s)
1: Soptimal f- null
2: for Aitexibbe = to N do
3: for Sttexible S do
4: Let s be the sequence that turns on all optimization switches.
5: for i to Mitexible do
6:
1 (i - 1) Matxiak
<-
8tlexihle
7: end for
8: if Memory (s) < M and Runtime (s) < Runniime (Soptimal) then
9: Soptimal s
10: end if
end for
12:
end for
Ensure: Suptimul
The pseudo-code of our greedy solver 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 turn off some switches for part of the layers. Here we have greedy assumption that these "partial layers" have
17

---

the same switches turned off, which can be easily enumerated as Soft E S. Then we enumerate Roit- representing how many layers we will turn off Son for. Finally, we greedily put these Saff into the initial sequence evenly.
The greedy solver introduces 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, we expect our greedy solver to outperform the random solver.
greedy assumptions, where the former is centered the idea that the
The complexity of the enumeration is @ (N]S)), so taking the model simulator into account, we have the global time complexity of @ (N2)
B.3 DP Solver
Algorithm 3 DP Solver (basic version)
Require: N, M, m, ModeM Simulator, (8/1, 81, 8+1)
1: sol("..',,,) f- mull
2: runtime(;;- (...s)))
+00
3: for $1 do
4: so1(0, 0, 0, mull, 81) - (81)
5: runtime (0, O, null, 81)
6: end for
7: for i = to N do
8: for memear, mempeak, 81 81 s.t. sol(i (i - memeam mempeakk 81 f null do
9:
for 8 E do if i N then
10:
null 9;+1 + null
12:
end if
Amem, resi resideetteee,,,neeeeeeeee peakMem, runtime; ModelSimulator: 8/s 8:41)
14:
mem +- memeur + residentMem; + Amem;
mem peak t (ee mem + peakMem,) ) + residentMem,
runtime' <- runtime(i - memeur, mempenk, 8;-1, 5;) runtime;
17:
if mem < and runtime' runtime (i, mem 84, 85+1) then
18:
sol (i, mem mem 8,, 81+1) sol(i mem mempeak: 8 ) (s1+1)
runtime (i, mem
mem 81, 8141) runtime"
20: end if
21: end for
22: end for
23: end for
24: Soptimal f- null
25: runtimeopimal + +00
for memeur, mempeak. v sol(N, (N, memeur meMpeak. SN null) * null do
27: if runtime(N, memeur memppekk null) < runtimeopimal then
Soptiall f- sol(N, meM mempeak SN, SN, null)
29:
runtimeopiimal +- runtime N, mem mempeak SN null)
30:
end if
31:
end for
Ensure: Soptimal
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-) Si, Si+) to represent the simulation of the i-th segment, while the result Amem,, residentMem; peakMem, runtime; denotes the delta memory after forward segment i, the resident memory of layer 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² N[S))), while the transfer com- plexity is θ (ISI). On this basis, attempt to do pruning improve its efficiency.
18

---

We first notice that the optimization scheme 8,, $441, $142 are only used in the running segment simulation during transfer. However, the simulation for segment i only requires s and part of 8i, 81+2. Specifically, for the previous layer (i) and the next layer (i + 2), we only care about three communication-relatee switches: hidden offloading. lazy-prefetching. and parameter partition. Therefore, we change 8, 8441 E S to T4, F+411 E R in the DP state, where R denotes the set of combinations of these three switches. Similarly, we only enumerate r142 R for transfer because is enough for the simulation of segment i + 1. In addition, we need to extend Ti+11 to 8,41 by another enumeration when transferring. In this way, the state complexity can be optimized to (m² N [R])). while the transfer complexity remains θ ((S)) because the time complexity of enumerating 1 and enumerating 8;-1 extended by 1 are θ (I)) and +
|8 [
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, T++1, if one has lower mem a lower mempeak- and 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 2-dimensional scanning with the time complexity of (m² for any given i, ra 1-+1-
Under the first pruning, the time complexity of the DP solver is reduced to O [ R|?). where [S] = 64, [R] = 8. For the experiments 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
C Experimental Settings
C.1 Model Size
The model sizes used in the experiment are shown in Table The parameters (N, dnidden- dEEN, and Nhcad) of BERT-1.88 1.BB, 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].
Table 2: Different sizes of BERT in this paper.
Size
#Parameter
N
daidden
deN
Nuead
1.8B
1.84 x 109
48
1024
16384
32
6B
6.68 x 109
48
1024
65536
128
13B
1.38
1000
48
2048
65536
128
100B
1.02
10a1
56
12288
49152
96
C.2 Hardware Configurations of Environment
8 x 2080T:: The GPU is NVIDIA GeForce RTX 2080 Ti (10 GB). The CPU is Intel(R) Xeon(R) CPU H5-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. NVL.ink is not available.
8 x A100: The GPU is NVIDIA A 1100000M440000
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 and the bandwidth is 31.5GB/s. There 12 fully connected NVLink 3.0.
64 x A100: nodes of "8 x A100".
C.3 Training Settings
As we mentioned the main text, we implement CPU Adam (4 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 2 10). Then we run our solver to determine the optimization sequence of the model. After that, skip S steps to wait
19

---

Algorithm 4 DP Solver (pruned version)
Require: N, M, m, ModelSimulaton (Fi-1, 81, Tisl)
1: sol(".''.,) + null
2: "
-- +00
3: optimState + 0
4: for r1 R do
5: sol(0, 0. 0, mull,
6: runtime mull, +
7: optimState optimState U {(0, 0, m1))
8: end for
9: i- N do
10: for (memeurr: memp 1, ) optimState do
11: for E extended(r)) do
12:
for r++1 E do
13:
if N then
14:
7++1 null
15:
end if
16:
Amem; residentMem; p runtime; f- Modell 1
17:
mem mem + residentMem; + Amem;
18:
mem peak + max mem + peakMem, + residentMem,
19:
runtime; if mem m and runtime' < (i, mem T', T(+1) then
runtime' <- runtime(i - memeur memppak, T: +
sol(i, mem mem Pi, T241) sol(i 1, memeum, memppea Ti-1, (2) ) (s;)
22:
runtime(2 mem mempuak: Ft, T(41) + runtime'
end if end for
24:
25:
end for
26: end for
27:
optimState +
28:
for do
29:
for ri+1 E R do
30:
for mem to do
31:
for mempeak = mem to m do
32:
runtime', runtime mem 1, mempeak, Pí, r++1)
33:
runtimeg <- runtime(( i, mem mempeak - 1.
34:
if runtime memeurr, mempeak: Ti, TT4 min munnrimntmmes runtimef] then optimState + optiopttaattt U (memaur meMpeak, ri, 1)) clse
35:
36:
37:
runtime(i, memuurs mempcak: Pi, T(+1) +- min minntumiime', runtime end if end for 39:
end for
41: end for
end for
43: end for
Septimal f- null
45: runtimeopimal +
for memcum; mempenk, I'N s.t. sol(N, (N, memeurr; mempeak: TN, null). A null do if runtime([,, memeur meMpeak. F', null) < runtimeopiimll then
47:
48:
Soptimal f- sol(N,, memuur meMpeak TN, null)
49:
runtimeopinal runtime N, mem mempcak, TN, null)
50:
end if'
51:
end for
Ensure: Suptimal
20

---

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 = 8000MB, and M 28GB 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 discussion about this problem in Section E.2 For all experiments in this paper, we set m = 128 as the default.
We use the salf-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 1.8 and DeepSpeed 0.7.6 for our experiment. For a fair comparison, we turn In contrast, we do NOT
on optimizer offlnn and rematerialization for introduce other optimizations such as tensor parallelism, pipeline parallelism, or FlashAttention for BOTH H3T and Segatron-DeepSpeee We will leave the implementation of these optimizations and the comparison as our future work.
D Supplemental Experiments
D.1 Memory Test
Table 3: The minimal memory of our framework and Megatron-DeeppSpe (unit: GB). We do not report the memory overhead of our automatic solvers because their memory overhead is always close to the upper bound due their flexibility.
Environment
2080TTTi
8x A 100
64 A100
Model Size
1.8B
6B
6B
13B
13B
100B
Global Batch Size
128
512
8
32
128
512
128
512
1024
4096
1024
2048
BMTrain
Min Memory
1.7
6.0
1.0
2.0
6.0
21.5
7.0
23.2
6.7
22.9
14.1
20.2
Megatron
ZeRO-3
5.8
9.4
10.7
-DeepSpeed
ZeRO-2
20.2
34.4
34.4
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 4x larger than Megatron-DeepSSpeed 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 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 20800T..
ZeRO-3 (ZeRO-2)
Not only that, based on our implementation, we can train GPT-3-175B 4 on 64x A100 (sequence length is 512, and batch size is 512) while using only about 11.1 GB of memory per device at minimum. 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. 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, if set the batch size and the sequence length larger, we will have great increase in the training throughput. We will incorporate CPU RAM optimizations 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-case 9 using BMTrain with H3T and Py PyTorch with APEX-] respectively. When using Py Torch with APEX, we adopt mixed precision
21

---

Table 4: The experimental results on SuperGLUE.
Toolkits
BoolQ
CB (acc. / FFI)
RTE
WiC
BMTrain w/ H3T
75.12
94.64 96.17
71.07
72.50
PyTorch w/ APEX
67.34
94.64 95.91
71.48
72.73
training to ensure fairness comparison. The experimental result shown Table 4 From the result, we can find that our system can achieve comparable results with PyTorch all four tasks, which greatly confirms the correctness of our implementation.
D.3 Energy Test
BERT-L88
BERT-6B
BERT-I3B
BERT-
on 8*2080Ti le-5
8xA100
8*A100
64xA100
Energy (WWb)) sequenees
Energy (kWh) sequencess
Energy (kWh) ssqueecss
Enengy (kWh) sequences
Le-5
le-5
1e-5
4
2-
2-
Н3T
Megatron -DeepSpeed
H3T
Megatron -DeepSpoed
HBT
Megatron -DeepSpeed
H3T
Megutron
-DeepSpeed
Figure 10: The energy cost (kWh) per sequence of H3T and Megatron-DeepSppe Here we use our DP solver for H3T and use the ZeRO-2 mode for Sreeatron-DeepSpeee
We evaluate the energy cost of our framework and the baseline, Segaatron-DeepSeee and the result is shown in Table TO From the table, we find that our framework achieves lower energy cost with Deepppen- in most cases, except training the 6B model on 8 x A100, where we also achieve a comparable result. This result uncovers that our framework is more energy-saving than Megatron- Megatron-Deepppeed.
D.4 Solver Efficiency Test
Table 5; The solver runtime within different settings. The environment is A100. The runtime of the DP solver is influenced by many variables because of the pruning, while the other two are only influenced by the model size. Here M denotes the memory constraint.
Setting
Solver Runtime (min)
Model Size
Batch Size
M (GB)
DP
Greedy
Rand.
13B
1024
25
12.0
0.6
0.1
100B
1024
25
4.7
0.9
0.2
13B
4096
25
2.8
0.6
0.1
13B
1024
100
13.9
0.6
0.1
AI scientists and engincers sometimes change the model size and architecture in their works. For do not need the solver be fast as
most people, this change does not always happen, possible, but the efficiency the solver still 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-leved runtime of the DP solver is quite acceptable compared to training time of 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 get the solution, which is greatly unacceptable.
22

---

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 x00x improvement in efficiency and much more senseless in the real application.
D.5
Case Study
Table 6: A case study of H3T's optimization in different simulated cases.
Default settings: BERT-1.8B, batch size = 128, 8x2080Ti, M = 8, DP solver.
Switch-on ratio of
Simulated case
parameter partition
lazy-prefetching
Rematerialiaztion
offloading hidden state
90%
8%
92%
21%
PCI-E speed x0.1
0%
0%
98%
0%
FPU speed x0.1
88%
90%
88%
81%
M+-25
50%
50%
65%
56%
Greedy
0%
0%
100%
0%%
We show a case study in Table intuitively show the ability of automatic solvers. From the result, we find that we manually slow down the PCI-E communication, the solver can adaptively switch 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, discuss the scalability of H3T. Although H3T is based 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 many other scenes. An illustration shown in Figure I
Optimization techniques and the solver. Scholars and programmers 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 they 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 aoper-parameters we may need 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.
this paper, long they can be
Models. H3T is designed for Transformer-bast 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.
23

---

Convert
Other Models
Big Models
Transformer Models
Sequential Models
Solvers
Greedy Solver
DP Solver
Design for
Switches
Switches in this paper
Formelize into
Optimization Techniques
Rematerialization
Tensor O Offloading
Parameter Partition
Overlaping
H3T implemented in this paper
DL Toolkits
Pytorch
BMTrain
DeepSpeed
The scalability of I3T
Distributed Leaming Toolkits
Figure 11: The scalability of H3T.
E.2 Limitation of H3T
Although H3T performs well 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, follow most existing works 2 I8 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. 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 fragmentationss 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, will develop 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 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 common problem with this kind of pre-computed optimization G I8 20 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, mainly focus Transformer based models that are mostly fixed structures, so it is a better choice to design a pre-computed optimization framework like H3T.
24


Reviews:
[{"review_id": "xacESj7IBQ", "reviewer": "Reviewer_qHi5", "summary": "This paper considers large model training and proposes a search algorithm for various memory optimizations.", "questions": "- 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.", "limitations": "Please see the \"Questions\" section.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Please see the \"Questions\" section.", "weaknesses": "Please see the \"Questions\" section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thorough and contributive review. Here are the replies to your questions:\n\n1. The misuse of \"dynamic\" in lines 89-91.\n\nThanks for pointing out this mistake. Our point is that these approaches are based on static rules which cannot dynamically adapt to the changes in model structure and hardware resources. However, \"dynamic\" and \"static\" do have other meanings according to the most common way they are used. We will fix this problem in the revision.\n\n2. Rematerialization and offloading are not independent.\n\nHere \"independent\" means they do not influence each other in general. In other words, their core ideas and logic do not conflict, so switching on either of them does not means we cannot switch on the other one. However, we also noticed this representation is ambiguous at your reminder. Thanks for pointing this out, and we will fix this in revision.\n\n3. Does the algorithm optimize for parallelism?\n\nOur DP and greedy solver can optimize for parallelism, but in this paper, our implementation focuses more on the memory optimization solution. We have a discussion about this problem in the appendix (section E.1, lines 168~176). Please refer to the appendix for it.\n\n4. Do you have a sense of why ZeRO-2 beats ZeRO-3?\n\nZeRO-3 optimize more but brings more communications, which means ZeRO-3 may be slower than ZeRO-2 (at least not faster than ZeRO-2). Moreover, we guess the ZeRO-3 implementation of Megatron-DeepSpeed is sub-optimal, which also makes ZeRO-3 much slower.\n\n5. What precision is used for your method vs ZeRO?\n\nWe use float16 for both BMTrain with H3T and Megatron-DeepSpeed in our experiment. We think it is a fair competition between our proposed framework and the baseline.\n\n6. Line 74, MM is short for?\n\nMM is short for Multi-Modal. Thanks for pointing this problem out, and we will supplement the full name in the revision.\n\n7. There are older references for rematerialization than Griewank et al. [12].\n\nWe apologize for our laxity, and thanks for pointing this out. Jose Grimm et al. do have another work about rematerialization in 1996. We will fix this in the revision.\n\nWe also thank you for pointing out the other typos, and we will fix all of them in the revision. \n\nWe argue this paper is contributive because it not only proposes a novel and scalable framework to solve the difficulty of training LLMs, but also implements an out-of-the-box version based on the open-source toolkit BMTrain, which is helpful for developers and researchers. We believe both of our contributions are valuable for the AI community and can greatly promote the development of AI and LLM techniques. We would be grateful if you could reconsider your score.\n\nWe sincerely thank you again for your careful review.\n"}, {"review_id": "ukkSNGw8xQ", "reviewer": "Reviewer_khhX", "summary": "This 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.", "questions": "1- 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?", "limitations": "NA", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-- 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.", "weaknesses": "-- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your appreciation of our work. We notice that your primary concern is about the configuration of DeepSpeed used in our experiment. We briefly describe it in the appendix (section C.3, line 108): \"For fairness, we turn on optimizer offloading and rematerialization for Megatron-DeepSpeed.\" We show the key part of the DS config as follows, and we will release the detailed config of DeepSpeed in the revision. \n\n```json\n\"zero_optimization\": {\n  \"stage\": 2,\n  \"offload_optimizer\": {\n    \"device\": \"cpu\",\n    \"buffer_count\": 4,\n    \"pipeline_read\": false,\n    \"pipeline_write\": false,\n    \"pin_memory\": true\n  }\n},\n\n\"fp16\": {\n  \"enabled\": true,\n  \"initial_scale_power\": 12\n},\n```\n\nIt is worth mentioning that according to our test, the parameter offloading is automatically switched on if we switch on optimizer offloading. Therefore, the comparison between our framework and Megatron-DeepSpeed is generally fair.\n\nAnother key point we would like to explain is that, we do not switch on tensor parallelism (TP) and pipeline parallelism (PP) in our experiments. The main reason is that our implementation of H3T in this paper does not involve TP and PP, so we think it is fair to switch off TP and PP for Megatron-DeepSpeed. In the future, we will design the solver for TP and PP schemes. We will also implement it and compare it with Megatron-Deepspeed's TP and PP versions.\n\nIt is also a good idea to compare our framework with FSDP. We do not try it, but we believe it will be a great baseline as a supplement and make our experiments more solid. We will seriously consider this suggestion in the revision.\n\nThanks again for your careful and contribute review. \n"}, {"review_id": "M2a2YOOC2l", "reviewer": "Reviewer_xJUr", "summary": "For 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.", "questions": "Q) 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).", "limitations": " No potential negative societal impact exists in the paper. The authors adequately addressed the system's limitations in Appendix.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "As 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", "weaknesses": "Although 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thorough and contributive review. Here are the replies to your questions:\n\nQ1: Although the authors say that H3T covers model parallelism including pipeline parallelism, the paper does not explain how pipeline parallelism is incorporated. \n\nA1: \nWe are sorry for your misunderstanding, but we only mention pipeline parallelism in the appendix (section E.1, line 171). In this paragraph, we mean anyone can introduce any optimization or parallelism techniques and design their own solver for the specified optimization switches, including pipeline parallelism. For the H3T implemented in this paper, we mainly focus on ZeRO, which does not normally coexist with pipeline parallelism. We leave the TP and PP version H3T as our future work.\n\nQ2: There may be a comparison with other automatic optimizers like Alpa or Unity.\n\nA2: \nAlthough H3T, Alpa, and Unity are all automatic optimizers, they are very different types of automatic optimizers. H3T mainly focuses on layer-level optimizations, such as tensor offloading, rematerialization, and ZeRO. In contrast, Alpa and Unity are kernel-level and operator-level optimizers, which means their granularity is finer, but they struggle to do a better job of optimizing for the structure of transformer-based models. Therefore, we argue that this comparison is unfair, and its priority is not high in our paper. We still sincerely thank you for your suggestion, and we will seriously consider it as a supplemental experiment in the revision.\n\nQ3: Did the Megatron-DeepSpeed results use the pipeline and tensor parallelism that Megatron provides?\n\nA3:\nNo. As we mentioned in A1, we mainly focus on ZeRO in this paper, and we will leave the exploration of other parallelism methods in our future work. Specifically, we will design a TP and PP version of H3T and try to beat Megatron-DeepSpeed with TP and PP.\n\nQ4: Could you also give the step time of training the GPT-3 175B model?\n\nA4:\nThe step time of training GPT-3 175B is about 40 seconds with a batch size of 512. \n\nQ5: Megatron 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. \n\nA5: Memory capacity is a very important reason, but the config is more critical. As we mentioned in A2, we want our comparison to be as fair as possible on basic optimization techniques. Therefore, we employ ZeRO as the parallelism scheme for both H3T and Megatron-DeepSpeed. \nBesides, our throughput (512 / 40 = 12.8) is still significantly better than the result reported by Megatron's paper (64 / 13.75 = 4.7). This result may provide a comparison between H3T and a well-configured Megatron.\n\nQ6: Does Megatron-DeepSpeed includes the latest Megatron that supports selective recomputation and sequence parallelism?\n\nA6: No. We will test the selective recomputation and update the result in the revision. Thanks for this advice. Similar to TP and PP, for fairness, we do not test sequence parallelism in our experiment.\n\nQ7: 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\nA7: Because the ZeRO-3 version is stuck in this experiment. We are not sure whether there is a bug or if we do not run it correctly. But referring to the results on the 8 × A100, we believe that ZeRO-3 will perform very poorly even if it is able to run, so this does not affect the overall conclusions of our experiments.\n\nQ8: Do BMTrain and Megatron-DeepSpeed use FlashAttention?\n\nA8: No. Both of them do not use FlashAttention. So the comparison is fair.\n\nQ9: 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\nA9: The parameters of these models are mainly referenced to T5 (3B and 10B). The diversity of settings could also better validate the generalizability of H3T capabilities.\n\nBesides, we thank you for pointing out the typos (and the detail problems, such as Q3, Q4, Q6, Q7, Q8, and Q9). We will fix the typos and supplement these details in the revision.\n\nWe understand that you would like to see more comparative experiments with other toolkits or Megatron-DeepSpeed with other settings. \nWe appreciate your constructive suggestions, but on the one hand, we prefer to present the most critical comparisons (i.e., fair comparisons using the same underlying optimization methods) in the limited paper space, as this is more indicative of the strengths of our underlying implementation and automated optimization solver. On the other hand, the config of Megatron-DeepSpeed is very complex, and we believe this is why we need an automatic solver to decide the appropriate optimization config. Besides, although H3T implemented in this paper only consider ZeRO as the parallelism solution, it still beats a well-configured Megatron-DeepSpeed (according to the throughput comparison in A5), which reflects BMTrain with H3T is already a better toolkit than Megatron-DeepSpeed. As mentioned in A1, we will further develop the H3T to support TP, PP, and other optimization and parallelism techniques and make it a better toolkit.\n\nIn summary, we argue that our experiments are convincing enough to support the contribution of this paper. We would be grateful if you could reconsider your score! (We would like to take the liberty of confirming with you that you know the review is on a 10-point scale :).)\n\nThanks again for your careful review."}, {"review_id": "gcpzIDtbfG", "reviewer": "Reviewer_QoRR", "summary": "Large 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.", "questions": "- 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", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thorough and contributive review. Here are the replies to your questions:\n\n1. There are limited details about the proposed solvers, especially the DP solver.\n\nWe are sorry for the lack of details in the main paper. But actually, we show many details about our solvers in the appendix (section B), especially for the DP solver (section B.3). In section B.3, we use nearly one page of text and nearly two pages of pseudo-code to introduce the details of the DP solver. We have to put them in the appendix due to the space limitation. This also follows some previous works like [1], which also describes the main idea in the main paper and puts the details in the appendix (this work is accepted by NeurIPS 2021).\n\n2. To improve the understanding of the optimizations, it would be helpful to show a breakdown analysis of the optimization techniques.\n\nThank you for your contributive suggestion! And we will seriously consider this suggestion in the revision. In this paper, we also put much effort into improving the understanding of the optimizations. For example, we show the timing diagram of each optimization switches in the appendix (section A). We also conduct a case study in section D.5 to show the characteristics of different optimizations. \n\n3. Some figures have font sizes that are too small.\n\nWe are sorry for the too-small font sizes, but we have to do so due to the space limitation. Once the paper accepts, we will have one more page, and we will make the font sizes larger in the revision.\n\n4. It would be clearer to explain the major difference between ZeRO-2 and ZeRO-3.\n\nThanks for the advice. \nZeRO has three levels, and their main difference is: (1) ZeRO-1 distributes optimizer states; (2) ZeRO-2 distributes optimizer states and gradients; (3) ZeRO-3 distributes optimizer states, gradients, and parameters. \nWe apologize that we misestimated how much people know about the ZeRO, and we will supplement its description in the revision.\n\n5. Further clarification of the memory reduction ratio would be helpful. Does the claim of 80.5% less memory consumption include the model size?\n\nWe are sorry for the ambiguity. The 80.5\\% memory reduction occurs when training the 13B model with the batch size of 1024 on 64 × A100. In this case, Megatron-DeepSpeed allocates 34.4 GB of memory, and our framework only allocates 6.7 GB, so we have the memory reduction ratio equals (34.4-6.7)/34.4 = 80.5\\%. This is the largest reduction ratio we observe in our experiments, and all the comparisons are across the same setting. The memory test result is shown in the appendix (section D.1, Table 2). Please refer to the appendix for more details.\n\n6. What is the actual overhead (in terms of time) for the Brute-Force solution?\n\nThis question is hard to answer because the actual overhead of Brute-Force is too high to estimate. The theoretical time complexity of Brute-Force is Θ(|S|^N × N), which means even for the 1.8B model (with N=48 layers), Brute-Force will take over 1.89 × 10^117 calculations to get the result. Even if 10^12 calculations can be done per second, we still need nearly 6 × 10^97 **years** to get the result. That is why we say this question is hard to answer, and this also demonstrates the value of our well-designed search algorithm.\n\n7. Is the proposed solution compatible with layer reordering?\n\nThe proposed DP solver does not support layer reordering because this solver is only designed for ZeRO, tensor offloading, rematerialization, and overlapping. However, as we discuss in the appendix (section E.1, line 168~176), Scholars and programmers can introduce any optimization or parallelism techniques and design their own solver for the specified optimization switches.\n\nDue to the space limitation, we have to move many important information and details to the appendix, and we are very sorry for your poor reading experience.\nHowever, we believe this paper is contributing, and our work can greatly promote the development of the AI community and LLM techniques. We would be grateful if you could reconsider your score!\n\nThanks again for your careful review.\n\n\n[1] 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."}]

Author rebuttal:

Meta-review: Paper proposes a system, High-Throughput Transformer Training (H3T), that incorporates memory optimizations and parallelism strategies into a single optimization problem that is then yields a solution under the constraints (e.g., memory) with a DP solver. Resulting solution is compared against an off the shelf, (near) state of the art solution and found to be outperforming it (Megatron-DeepSpeed).

Considering the proposal being the first attempt to automatically integrate memory optimization and parallelism strategy for LLM training and convincing results, the paper is ready to be presented to the wider research audience. That said, some of the reviewers raised concerns mostly revolving around the "generalization" of the proposal and lack of experiments that could give convincing evidence on this matter. 

With the current form and content, the papers targeted impact area is limited to the set of experimental configs and tools included in the design, which is quite significant already. For wider adoption and recognition the authors could consider building on top of their current work, but the submission is, to my assessment, over the bar to start discussions along this line of research -- which is quite new and important.
Decision: Accept (poster)