SlowGuess/NeurIPS_2023_OCR huggingface dataset:
How to pick appropriate switches is another big challenge. The optimization of different layers can be in any combination, so the number of global solutions can be up to  $|S|^N$ . It is not easy to find a good solution from such exponential search space with the memory constraint and toward the goal of minimizing the training runtime. Toward this goal, we design three search algorithms in this paper.
3.3.2 Search Algorithms
Random solver. Brute force does not work for our task because of the exponential search space, so we design a simple method based on random sampling to find a suboptimal solution more efficiently. Specifically, we randomly sample several optimization sequences from  $S^{N}$  and pick the best one.
Greedy solver. The random solver is simple but not effective for large-scale models. The main reason is that the search space grows exponentially with  $N$ , but the sampling time is constant, which means the probability we hit a good solution becomes much lower with the growth of the model size. To address these issues, we design a greedy solver to search for better solutions. The core idea of our greedy solver is first assuming all optimization switches are turned on, then greedily turning off some switches to increase memory utilization and improve training efficiency. Due to the space limitation, detailed descriptions of the greedy solver are in the appendix.
Dynamic Programming (DP) Solver. Our task strives for better runtime under the memory constraint, similar to the knapsack problem [25, 7] that pursues higher value under the weight constraint. Inspired by the classic dynamic programming (DP) algorithm for the knapsack problem [1, 43], we design a search algorithm based on DP to solve our task better. Specifically, our DP algorithm considers the optimal sub-solutions for specified prefixes and memory limits. When transferring, we need to enumerate the next layer's switches and estimate the next segment's running status.
Figure 5: An illustration for  $i$ ,  $\text{mem}_{\text{curr}}$ ,  $\text{mem}_{\text{peak}}$ .
Formally, the DP state can be represented by the following variables: (1)  $i$ : The prefix length, which means this state is considering segment  $1 \sim i$  (including forward segment  $1 \sim i$  and backward segment  $i \sim 1$ , similarly hereinafter). (2)  $\mathsf{mem}_{\mathrm{curr}}$ : The memory usage after forward segment  $1 \sim i$ .
(The memory usage before backward segment  $i \sim 1$  is essentially equal and can be easily calculated.) (3) mem_peak: The peak memory usage during segment  $1 \sim i$ . (We have to record mem_peak because the resident memory may increase due to non-lazy-prefetching, leading to an increase in the peak memory from the previous segments.) (4)  $s_i$ ,  $s_{i+1}$ : The switches of layer  $i$  and  $i+1$ , which is used to simulate segment  $i+1$ .
We use  $\mathsf{sol}(i,\mathsf{mem}_{\mathrm{curr}},\mathsf{mem}_{\mathrm{peak}},s_i,s_{i + 1})$  to represent the optimal sub-solution of the state. Then we can enumerate  $s_{i + 2}$  to transfer to the next state  $(i + 1,\mathsf{mem}_{\mathrm{curr}}^{\prime},\mathsf{mem}_{\mathrm{peak}}^{\prime},s_{i + 1},s_{i + 2})$ . During the transfer, we need to: (1) Simulate the forward segment  $i + 1$  and backward segment  $i + 1$  based on  $s_i,s_{i + 1},s_{i + 2}$ . (2) Calculate  $\mathsf{mem}_{\mathrm{curr}}^{\prime},\mathsf{mem}_{\mathrm{peak}}^{\prime}$  and the total runtime based on  $\mathsf{mem}_{\mathrm{curr}},\mathsf{mem}_{\mathrm{peak}}$ , the previous runtime, and the simulation result. (3) Update  $\mathsf{sol}(i + 1,\mathsf{mem}_{\mathrm{curr}}^{\prime},\mathsf{mem}_{\mathrm{peak}}^{\prime},s_{i + 1},s_{i + 2})$  if the peak memory  $\mathsf{mem}_{\mathrm{peak}}^{\prime}$  does not exceed  $M$  and the runtime is shorter than the existing solution's.
There are two noteworthy details for our DP algorithm. First, considering the overlapping subproblems condition, we must ensure the state is finite, so we discretize the memory limit into  $m + 1$  values that are evenly distributed in  $[0, M]$ . Second, following the condition of optimal substructure, the problem for each state must be reducible to the corresponding sub-problem for the other states. In this paper, only cross-segment communications violate this condition, so we force segment synchronization switched on and remove it from the switch set (thus  $|S|$  reduce to 64). These two assumptions may affect the result but make our algorithm design much simpler and more efficient. We leave the design of the more comprehensive DP algorithm as our future work.
In addition to the basic version of the DP algorithm, we have two prunings to improve its efficiency. Please refer to the appendix for them along with the pseudo-code and time complexity analysis.
Models. We use BERT[9], one of the most common Transformer-based language models, as an example for our experiments. We customize different BERT sizes for our experiments, including 1.8B, 6B, 13B, and 100B. Traditional small models, such as BERT-base (0.1B) or BERT-large (0.3B), are not considered in this paper because they can be easily trained using conventional approaches.
**Environments.** We test the performance of H3T in three different experimental environments: (1) 1 node of  $8 \times$  NVIDIA GeForce RTX 2080 Ti 10 GB ( $8 \times$  2080Ti); (2) 1 node of  $8 \times$  NVIDIA A100-SXM4-40GB with NVLink 3.0 ( $8 \times$  A100); and (3) 8 nodes of  $8 \times$  A100 ( $64 \times$  A100).
**Implementations.** We implement H3T based on BMTrain, an open-source toolkit that accelerates the training of big Transformers with various parallelism techniques.
Baselines. Except for H3T with random solver and greedy solver, we introduce manual optimization strategy without H3T and another toolkit Megatron-DeepSpeed for comparison. (1) In the simulation study, we test the performance of manually enumerating the combination of global optimizations and picking the best one. In other words, we simulate arbitrarily and globally turning on or off all the switches and picking the best-performing solution that does not exceed the memory limit. (2) In the actual-running experiment, we attempt to turn on all optimization switches to achieve the lowest memory overhead since testing all manual optimization schemes is time-consuming and resource-intensive. (3) We introduce one of the most popular distributed deep learning toolkits, Megatron-DeepSpeed [40], for comparison in our actual-running experiment. For a fair comparison, we test Megatron-DeepSpeed with ZeRO-2 and ZeRO-3 respectively (ZeRO-2 distributes optimizer states and gradients; ZeRO-3 distributes optimizer states, gradients, and parameters [36]). Please refer to the appendix and the code for detailed configurations of Megatron-DeepSpeed and H3T in the actual-running experiment.
We do not test POFO [3], which also uses DP to solve the memory optimization problem for big model training. We consider it unfair to compare POFO with H3T because POFO does not support any parameter-aspect optimization, such as parameter offloading and parallelism. Considering the large model size, it is almost impossible for a vanilla POFO to train the large-scale language models in this paper. If we arbitrarily implement parameter optimizations for POFO, additional communication will affect the well-designed offloading solution of POFO and seriously hurt its performance.
Other settings. Please refer to the appendix for other detailed experimental settings, including detailed model size, hardware configurations, training settings, etc.
4.2 Simulation Studies
Figure 6: The results of the simulation study.
In this section, we report our simulation studies to evaluate the ideal performance of H3T. We profile the models in different environments and use the model simulator to evaluate the performance of different solvers and the manual strategy. The result is shown in Figure 6.
The result shows that our DP solver outperforms the other three baseline approaches in most cases. The greedy solver also performs well and is usually close to DP. Both DP and greedy solver have the flexibility to adjust the optimization solution based on variant memory constraints. This conclusion is applicable in all three environments, showing H3T has good adaptability to GPUs with different numbers, performance, and connections. In contrast, the random solver is less stable and less effective. Although the performance gap is not significant when memory is sufficient, the random solver, more often than not, struggles even to find an available solution.
Compared to three automatic solvers, manual optimization without H3T is often the worst, especially when memory is sufficient. This indicates that our layer-level automatic solver is able to precisely decide the optimization scheme for large-scale Transformer-based models given the memory constraint. In contrast, the traditional manual strategy performs worse and is heavy for users as well.
However, the above rules do not always stand. For BERT-1.8B with a batch size of 128 on  $8 \times 2080\mathrm{Ti}$ , the greedy solver underperforms the random solver. For BERT-6B on  $8 \times 2080\mathrm{Ti}$ , the DP solver continuously underperforms the greedy solver and even underperforms the manual strategy in some cases. This is because the local optimality may let the greedy algorithm fall into a suboptimum, and the segment synchronization assumption and the memory discretization may influence the DP algorithm. We argue that it is normal for such problems to occur in a few cases.
4.3 Actual-running Experiments
We conduct our actual-running experiment in the three environments, and the result is shown in Table 1. The results demonstrate our great advantage over Megatron-DeepSpeed. Compared to the Megatron-DeepSpeed with ZeRO-2 (ZeRO-3), H3T improves the training speed by  $1.2 \times \sim 1.8 \times$ $(1.9 \times \sim 4.3 \times)$  in our experiments. Comparing the results of BERT-6B and BERT-13B, both with the batch size of 128 and both on  $8 \times \mathrm{A}100$ , we find that H3T outperforms Megatron-DeepSpeed more significantly for larger model size, which means H3T is not only faster but also more adapted to bigger models.
Table 1: The time per step of the actual-running experiment (unit: second). We do not report the random solver because it fails on all these settings. - means the program runs out of memory.
Environment	8×2080Ti	8×A100	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	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
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	7.1	17.7	22.1	23.0	5.7	19.8	9.9	28.6	7.8	26.9	36.3	61.8
Megatron	ZeRO-3	11.3	-	-	-	15.7	-	43.2	-	-	-	-	-
-DeepSpeed	ZeRO-2	9.3	-	-	-	6.8	-	17.7	-	10.6	-	-	-
Focusing on the results of BMTrain, we find that H3T with DP solver and greedy solver both outperform BMTrain without H3T by up to  $1.2 \times$  speed up when memory is sufficient. In contrast, when memory is tight, the two automatic solvers can also enable more memory optimization and trade time performance for keeping model training. This result is consistent with the findings from the simulation study, which further verifies the flexibility of our automatic solvers. Also, similar to the simulations, the performance of DP and greedy are close, but DP still outperforms greedy a little in most cases. Besides, the random solver fails to solve for all of the settings, so we do not report its results in the table.
The memory cost of our implementation is also better than Megatron-DeepSpeed. From Table 1, we find Megatron-DeepSpeed often runs out of memory for larger models and batch sizes, which are better supported by H3T. Hence, users can train larger models with more limited hardware, e.g., train BERT-6B on  $8 \times 2080$ Ti. Not only that, H3T can train GPT-3-175B [4] on  $64 \times$  A100 with a batch size of 512 and using only about 11 GB of memory per GPU. This is really an exciting result because we would never have dreamed that a 175B-parameter model could be trained with such limited hardware, not to mention with such low GPU memory overhead. We conduct a comprehensive memory test in the appendix and find our implementation can save  $34.6\% \sim 80.5\%$  GPU memory compared with Megatron-DeepSpeed under various settings.
Besides, we verify the correctness of our implementation. (1) We fix the input data and check whether the output loss equals the release version of BMTrain during the whole training process, which justifies our incremental development based on BMTrain. (2) We do an end-to-end training experiment on 4 actual tasks of SuperGLUE [48]. The results show that H3T can achieve comparable results with PyTorch, which confirms the soundness of the whole system. The detailed experimental results on SuperGLUE are in the appendix.
We conduct three extra experiments in this paper. Due to the space limitation, we briefly introduce them here and elaborate on the details in the appendix. (1) We conduct an energy test and find that H3T is more energy-efficient than Megatron-DeepSpeed for training big Transformer-based models. (2) We test the efficiency of the solver to confirm it is not too slow. (3) We conduct a case study to show the automatically generated optimization scheme makes sense.
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-DeepSpeed, in our experiments. More excitingly, H3T can use only 64 NVIDIA A100 GPUs to train Transformer-based models with more than 100 billion parameters like GPT-3-175B. We hope that the release of H3T will lower the barrier to training big Transformer-based models and promote the development of the AI community.
This work is supported by the National Key R&D Program of China (No.2022ZD0116312), Major Project of the National Social Science Foundation of China (No. 22&ZD298), the National Key R&D Program of China (No. 2022ZD0160501), and Institute Guo Qiang at Tsinghua University. Weilin Zhao is supported by Tsinghua University Initiative Scientific Research Program.
[45] Xiaoyang Sun, Wei Wang, Shenghao Qiu, Renyu Yang, Songfang Huang, Jie Xu, and Zheng Wang. Stronghold: fast and affordable billion-scale deep learning model training. In Proceedings of SC, pages 1-17, 2022. https://zwang4.github.io/publications/sc22-2.pdf.
[46] Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandep Baines, Carlos Efrain Quintero Narvaez, Vinay Ramakrishnaiah, Nirmal Prajapati, Pat McCormick, Jamaludin Mohd-Yusof, et al. Unity: Accelerating dnn training through joint optimization of algebraic transformations and parallelization. In Proceedings of OSDI, pages 267-284, 2022. https://www.usenix.org/system/files/osdi22-unger.pdf.
[47] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of NeurIPS, volume 30, 2017. https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
[48] Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Superglue: a stickier benchmark for general-purpose language understanding systems. In Proceedings of NeurIPS, pages 3266-3280, 2019. https://w4ngatang.github.io/static/papers/superglue.pdf.
[49] Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song, Zenglin Xu, and Tim Kraska. Superneurons: Dynamicgpu memory management for training deep neural networks. In Proceedings of PPOPP, pages 41-53, 2018. https://arxiv.org/pdf/1801.04380.pdf.
[50] Lijie Wen, Zan Zong, Li Lin, and Leilei Lin. A swap dominated tensor re-generation strategy for training deep learning models. In Proceedings of IPDPS, pages 996-1006, 2022. https://ieeexplore.ieee.org/document/9820619.
[51] Jun Zhan and Jinghui Zhang. Pipe-torch: Pipeline-based distributed deep learning in agpu cluster with heterogeneous networking. In Proceedings of CBD, pages 55-60, 2019. https://ieeexplore.ieee.org/abstract/document/8916305.
[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://www.usenix.org/system/files/osdi22-zheng-lianmin.pdf.
[53] Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient descent. In Proceedings of NeurIPS, volume 23, 2010. https://proceedings.neurips.cc/paper/2010/file/abea47ba24142ed16b7d8fbf2c740e0d-Paper.pdf.
A.1 Timing Diagram of Multiple Streams with Optimization Switches
(a) Default workflow based on parameter offloading, optimizer offloading, and parameter parallelism.
(b) Rematerialization.
(c) Offloading hidden states.
(d) Parameter partition.
(e) Lazy-prefetching.
(f) The legend.
Figure 7: Timing diagram of H3Twith different optimization switches.
Figure 7 illustrate the timing diagram of our implementation of H3T. Figure 7a is the default workflow based on parameter offloading, optimizer offloading, and parameter parallelism. Figure 7b and 7c are the workflow with two activation-aspect switches, respectively. Figure 7d and 7e are the workflow with two parameter-aspect switches, respectively.
A.2 Execution-aspect Switches
In this part, we introduce the detail of our three execution switches.
Economical forward/backward. These two switches are to prioritize the scattering and offloading of the previous layer in order to save more memory. In contrast, non-economical forward/backward prioritize gathering and prefetching of the next layer, which may reduce the delay in the computation of the next layer. Since NVLink is not available for all GPUs, e.g., the NVIDIA GeForce series, we design the workflow with and without NVLink, respectively, and the timing diagram is shown in Figure Figure 8. It is worth mentioning that there are still asynchronous operations in the economical mode, so it is not memory-optimal but balances the efficiency.
Segment synchronization. Let the computation of one layer and the concurrent communications (post-process of the previous layer and preparation of the next layer) be a "segment". Then segment
(a) Non-economical and economical forward with NVLink.
(b) Non-economical and economical backward with NVLink.
(c) Economical and non-economical forward without NVLink.
(d) Economical and non-economical backward without NVLink.
Figure 8: Timing diagrams of H3T's workflow in 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 is on, the next segment must keep waiting until all tasks of the current segment end. In contrast, if segment synchronization is off, the next segment can start before the current stage ends. An illustration of segment synchronization is shown in Figure 9. Similar to the economical mode, segment synchronization can reduce the memory overhead, but may sacrifice a little performance. Segment synchronization is also a condition of our DP solver.
A.3 CPU Optimizer
As we mentioned in the main text, we implement our CPU optimizer based on Intel AVX instructions. Intel 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 to the conventional GPU optimizer, which greatly ensures training performance.
Figure 9: A schematic diagram of segment synchronization. The legend is the same as in Figure 7f.
B Solver Details
In this part, we mainly show the pseudo-code of the three solver algorithms and analyze their time complexity. We also describe details of the greedy solver and the pruning of the DP solver, which are not introduced in the main text due to the space limitation.
B.1 Random Solver
Algorithm 1 Random Solver
Require:  $N,M$  ,Runtime (s),Memory (s),  $T_{\mathrm{rand}}$  1: soptimal  $\leftarrow$  null 2: for  $i = 1$  to  $T_{\mathrm{rand}}$  do 3: Randomly sample an optimization sequence  $\mathbf{s}\in S^N$  4: if Memory (s)  $\leq M$  and Runtime (s)  $<$  Runtime (soptimal) then 5: soptimal  $\leftarrow$  s 6: end if 7: end for   
Ensure: soptimal
The pseudo-code of our random solver is shown in Algorithm 1. Here we use a hyper-parameter  $T_{\mathrm{rand}}$  to represent the sampling times. Since the time complexity of the model simulator is  $\Theta(N)$ , the time complexity of the random solver is  $\Theta(NT_{\mathrm{rand}})$ .
B.2 Greedy Solver
Algorithm 2 Greedy Solver
Require:  $N,M$  ,Runtime (s),Memory (s)   
1:soptimal  $\leftarrow$  null   
2:for  $n_{\mathrm{flexible}} = 1$  to  $N$  do   
3: for  $s_\mathrm{flexible}\in S$  do   
4: Let s be the sequence that turns on all optimization switches.   
5: for  $i = 1$  to  $n_{\mathrm{flexible}}$  do   
6:  $\mathbf{s}\left[1 + (i - 1)\lfloor \frac{N}{n_{\mathrm{flexible}}}\rfloor \right]\gets s_{\mathrm{flexible}}$    
7: end for   
8: if Memory (s)  $\leq M$  and Runtime (s)  $<$  Runtime (soptimal) then   
9: Soptimal  $\leftarrow$  s   
10: end if   
11: end for   
12: end for   
Ensure: soptimal
The pseudo-code of our greedy solver is shown in Algorithm 2. We first assume all optimization switches are turned on, which reduces the memory overhead to the lowest (in our framework). Then, in order to increase memory utilization and improve training efficiency, we want to turn off some switches for part of the layers. Here we have a greedy assumption that these "partial layers" have
the same switches turned off, which can be easily enumerated as  $s_{\mathrm{off}} \in S$ . Then we enumerate  $n_{\mathrm{off}}$  representing how many layers we will turn off  $s_{\mathrm{off}}$  for. Finally, we greedily put these  $s_{\mathrm{off}}$  into the initial sequence evenly.
The greedy solver introduces two greedy assumptions, where the former is centered on the idea that the best is always the best, and the latter on the idea that an even distribution may block communications less. We cannot prove that either of them is optimal, but they both make sense, so we expect our greedy solver to outperform the random solver.
The complexity of the enumeration is  $\Theta(N|S|)$ , so taking the model simulator into account, we have the global time complexity of  $\Theta(N^2|S|)$ .
B.3 DP Solver