SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 3 DP Solver (basic version)
Require:  $N,M,m,$  ModelSimulator  $\mathbf{\Phi}_i(s_{i - 1},s_i,s_{i + 1})$    
1: sol(.,,.,.)  $\leftarrow$  null   
2: runtime(.,.,.,.)  $\leftarrow +\infty$    
3: for  $s_1\in S$  do   
4: sol(0,0,0,null,  $s_1)\gets (s_1)$    
5: runtime(0,0,0,null,  $s_1)\gets 0$    
6: end for   
7: for  $i = 1$  to  $N$  do   
8: for memcurr,mempeak,  $s_{i - 1}$  ,si s.t. sol(i-1,memcurr,mempeak,  $s_{i - 1}$  ,si)  $\neq$  null do   
9: for  $s_{i + 1}\in S$  do   
10: if  $i = N$  then   
11:  $s_{i + 1}\gets$  null   
12: end if   
13:  $\Delta \mathrm{mem}_{i}$  ,residentMemi,peakMemi,runtime  $\leftarrow$  ModelSimulator  $\mathbf{\Phi}_i(s_{i - 1},s_i,s_{i + 1})$    
14: mem'curr  $\leftarrow$  memcurr+residentMemi+  $\Delta$  mem;   
15: mem'peak  $\leftarrow$  max{mempeak,memcurr+peakMem}+residentMem;   
16: runtime'  $\leftarrow$  runtime(i-1,memcurr,mempeak,  $s_{i - 1}$  ,si)+runtime;   
17: if mem'peak  $<  m$  and runtime'  $<$  runtime(i,mem'curr,mem'peak,  $s_i,s_{i + 1})$  then   
18: sol(i,mem'curr,mem'peak,  $s_i,s_{i + 1})\gets$  sol(i-1,memcurr,mempeak,  $s_{i - 1}$  ,si)+(si+1)   
19: runtime(i,mem'curr,mem'peak,  $s_i,s_{i + 1})\gets$  runtime'   
20: end if   
21: end for   
22: end for   
23: end for   
24: soptimal  $\leftarrow$  null   
25: runtimeoptimal  $\leftarrow +\infty$    
26: for memcurr,mempeak,  $S_N$  s.t. sol(N,memcurr,mempeak,  $S_N$  ,null)  $\neq$  null do   
27: if runtime(N,memcurr,mempeak,  $S_N$  ,null)  $<$  runtimeoptimal then   
28: soptimal  $\leftarrow$  sol(N,memcurr,mempeak,  $S_N$  ,null)   
29: runtimeoptimal  $\leftarrow$  runtime(N,memcurr,mempeak,  $S_N$  ,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$  ( $s_{i-1}, s_i, s_{i+1}$ ) to represent the simulation of the  $i$ -th segment, while the result  $\Delta \text{mem}_i$ , residentMem $_i$ , peakMem $_i$ , runtime $_i$  denotes the delta memory after forward segment  $i$ , the resident memory of layer  $i$ , the peak memory during the forward and backward of segment  $i$ , and the total runtime of segment  $i$ , respectively.
For this basic version of DP solver, the state complexity is  $\mathrm{O}\left(m^{2}N|S|^{2}\right)$ , while the transfer complexity is  $\Theta (|S|)$ . On this basis, we attempt to do some pruning to improve its efficiency.
We first notice that the optimization scheme  $s_i, s_{i+1}, s_{i+2}$  are only used in the running segment simulation during transfer. However, the simulation for segment  $i+1$  only requires  $s_{i+1}$  and part of  $s_i, s_{i+2}$ . Specifically, for the previous layer  $(i)$  and the next layer  $(i+2)$ , we only care about three communication-related switches: hidden offloading, lazy-prefetching, and parameter partition. Therefore, we change  $s_i, s_{i+1} \in S$  to  $r_i, r_{i+1} \in R$  in the DP state, where  $R$  denotes the set of combinations of these three switches. Similarly, we only enumerate  $r_{i+2} \in R$  for transfer because it is enough for the simulation of segment  $i+1$ . In addition, we need to extend  $r_{i+1}$  to  $s_{i+1}$  by another enumeration when transferring. In this way, the state complexity can be optimized to  $\Theta(m^2 N|R|^2)$ , while the transfer complexity remains  $\Theta(|S|)$  because the time complexity of enumerating  $r_{i+2}$  and enumerating  $s_{i+1}$  extended by  $r_{i+1}$  are  $\Theta(|R|)$  and  $\Theta\left(\frac{|S|}{|R|}\right)$ , respectively.
Besides, we can prune the non-optimal sub-solutions to improve efficiency further. In particular, for any two sub-solutions with the same  $i, r_i, r_{i+1}$ , if one has a lower  $\text{mem}_{\text{curr}}$ , a lower  $\text{mem}_{\text{peak}}$ , and a shorter runtime, then the other one will be pruned because it is definitely not an optimal substructure of the global optimal solution. In practice, this pruning can be easily implemented by a 2-dimensional scanning with the time complexity of  $\Theta(m^2)$  for any given  $i, r_i, r_{i+1}$ .
Under the first pruning, the time complexity of the DP solver is reduced to  $\mathrm{O}\left(m^{2}N|S||R|^{2}\right)$ , where  $|S| = 64$ ,  $|R| = 8$ . For the experiments in this paper, we always take  $m = 128$ . Although the second pruning does not improve the time complexity, it can reduce more than  $99\%$  of the DP states in actual runs, which greatly improves the solver's efficiency.
Based on the basic DP algorithm and the two prunings described above, we have the pseudo-code of our pruned DP solver shown in Algorithm 4.
C.1 Model Size
The model sizes used in the experiment are shown in Table 2. The parameters  $(N, d_{\mathrm{hidden}}, d_{\mathrm{FFN}}$ , and  $N_{\mathrm{head}})$  of BERT-1.8B, BERT-6B, and BERT-13B are referenced from T5-3B and T5-11B [35]. The parameters of BERT-100B are referenced from GPT3-175B [4].
Table 2: Different sizes of BERT in this paper.
Size	#Parameter	N	dhidden	dFFN	Nhead
1.8B	1.84 × 109	48	1024	16384	32
6B	6.68 × 109	48	1024	65536	128
13B	1.38 × 1010	48	2048	65536	128
100B	1.02 × 1011	56	12288	49152	96
C.2 Hardware Configurations of Environment
8 × 2080Ti: The GPU is NVIDIA GeForce RTX 2080 Ti (10 GB). The CPU is Intel(R) Xeon(R) CPU E5-2680 v4 @ 2.40GHz. The GPU number is 8, and the CPU core number is 56. CPU RAM size is 126GB. The PCI-E version is 3, and the bandwidth is 15.8GB/s. NVLink is not available.
$\mathbf{8} \times \mathbf{A100}$ : The GPU is NVIDIA A100-SXM4-40GB. The CPU is Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz. The GPU number is 8, and the CPU core number is 128. CPU RAM size is 503GB. The PCI-E version is 4, and the bandwidth is 31.5GB/s. There are 12 fully connected NVLink 3.0.
$\mathbf{64}\times \mathbf{A100}$  : 8 nodes of  $^{\prime \prime}8\times \mathrm{A100}^{\prime \prime}$
C.3 Training Settings
As we mentioned in the main text, we implement CPU AdamW [24] based on Intel AVX instructions and use it to train our BERT models. More details of the implementation are stated in Section A.3.
In this paper, our profiler collects the runtime information in the first  $S$  steps ( $S \geq 10$ ). Then we run our solver to determine the optimization sequence of the model. After that, we skip  $S$  steps to wait
Algorithm 4 DP Solver (pruned version)
Require:  $N,M,m,$  ModelSimulator  $\mathbf{\Phi}_i(r_{\mathrm{i - 1}},s_i,r_{\mathrm{i + 1}})$    
1: sol(.,.,.,.)  $\leftarrow$  null   
2: runtime(.,.,.,.)  $\leftarrow +\infty$    
3: optimState  $\leftarrow \{\}$    
4: for  $r_1\in R$  do   
5: sol(0,0,0,null,  $r_1)\gets ()$    
6: runtime(0,0,0,null,  $r_1)\gets 0$    
7: optimState  $\leftarrow$  optimState  $\cup \{(0,0,null,r_1)\}$    
8: end for   
9: for  $i = 1$  to  $N$  do   
10: for (memcurr, mempeak,  $r_{i - 1},r_i)\in$  optimState do   
11: for  $s_i\in$  extended  $(r_i)$  do   
12: for  $r_{i + 1}\in R$  do   
13: if  $i = N$  then   
14:  $r_{i + 1}\gets$  null   
15: end if   
16:  $\Delta \mathtt{mem}_{i}$  ,residentMemi,peakMemi,runetimei  $\leftarrow$  ModelSimulator  $(r_{i - 1},s_i,r_{i + 1})$    
17: mem'curr  $\leftarrow$  memcurr+residentMemi  $^+$ $\Delta$  memi   
18: mem'peak  $\leftarrow$  max{mempeak,memcurr+peakMemi}  $^+$  residentMemi   
19: runtime'  $\leftarrow$  runtime(i-1,memcurr,mempeak,  $r_{i - 1},r_i)$  +runtimei   
20: if mem'peak  $<  m$  and runtime'  $<$  runtime(i,mem'curr,mempeak,ri,ri+1) then   
21: sol(i,mem'curr,mem'peak,ri,ri+1)  $\leftarrow$  sol(i-1,memcurr,mempeak,ri-1,ri)+ (si)   
22: runtime(i,mem'curr,mem'peak,ri,ri+1)  $\leftarrow$  runtime'   
23: end if   
24: end for   
25: end for   
26: end for   
27: optimState  $\leftarrow \{\}$    
28: for  $r_i\in R$  do   
29: for  $r_{i + 1}\in R$  do   
30: for memcurr  $= 0$  to m do   
31: for mempeak  $=$  memcurr to m do   
32: runtimeA  $\leftarrow$  runtime(i,memcurr-1,mempeak,ri,ri+1)   
33: runtimeB  $\leftarrow$  runtime(i,memcurr,mempeak-1,ri,ri+1)   
34: if runtime(i,memcurr,mempeak,ri,ri+1)  $<  \min \{]$  runtimeA',runtimeB}then   
35: optimState  $\leftarrow$  optimState U { (memcurr,mempeak,ri,ri+1)}   
36: else   
37: runtime(i,memcurr,mempeak,ri,ri+1)  $\leftarrow$  min{runtimeA',runtimeB}   
38: end if   
39: end for   
40: end for   
41: end for   
42: end for   
43: end for   
44: Soptimal  $\leftarrow$  null   
45: runtime optimal  $\leftarrow +\infty$    
46: for memcurr, mempeak,  $r_N$  s.t. sol(N,memcurr, mempeak,  $r_N$  ,null)  $\neq$  null do   
47: if runtime(N,memcurr, mempeak,  $r_N$  ,null)  $<$  runtime optimal then   
48: Soptimal  $\leftarrow$  sol(N,memcurr, mempeak,  $r_N$  ,null)   
49: runtime optimal  $\leftarrow$  runtime(N,memcurr, mempeak,  $r_N$  ,null)   
50: end if   
51: end for   
Ensure: Soptimal
for the model warm-up. Finally, we measure the runtime for at least  $S$  steps and reported the average result of the step time.
In the actual-run experiment, we set the memory constraint of the 2080Ti environment as  $M = 8000\mathrm{MB}$ , and  $M = 28\mathrm{GB}$  for A100. The memory constraints do not equal the GPU memory capacity because the actual memory occupation will be more than the memory usage due to memory fragmentation, memory cache, and other factors. We have a discussion about this problem in Section E.2 For all experiments in this paper, we set  $m = 128$  as the default.
We use the half-precision float for all parameters and input tensors. Initial parameters and all input tensors are randomly generated. We fix the random seed and check the loss is invariant to ensure the correctness of our implementation. The sequence length is set to 512, which is maximal for BERT.
We adopt BMTrain 0.1.8 and DeepSpeed 0.7.6 for our experiment. For a fair comparison, we turn on optimizer offloading and rematerialization for Megatron-DeepSpeed. In contrast, we do NOT introduce other optimizations such as tensor parallelism, pipeline parallelism, or FlashAttention for BOTH H3T and Megatron-DeepSpeed. We will leave the implementation of these optimizations and the comparison as our future work.
D Supplemental Experiments
D.1 Memory Test
Table 3: The minimal memory cost of our framework and Megatron-DeepSpeed (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 to their flexibility.
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	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	6.8	-	-	-	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  $4 \times$  larger than Megatron-DeepSpeed, which provides users more flexibility and convenience in training their models. We can also find that H3T can save  $34.6\% \sim 70.7\%$  ( $70.3\% \sim 80.5\%$ ) GPU memory compared with Megatron-DeepSpeed ZeRO-3 (ZeRO-2) with the same batch size and other settings. Furthermore, we are able to train larger models under limited hardware resources, such as training BERT-6B on  $8 \times 2080\mathrm{Ti}$ .
Not only that, based on our implementation, we can train GPT-3-175B [4] on  $64 \times \mathrm{A}100$  (sequence length is 512, and batch size is 512) while using only about 11.1 GB of memory per device at a minimum. This is really an exciting result because we would never have dreamed that a 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, so if we can set the batch size and the sequence length larger, we will have a great increase in the training throughput. We will incorporate CPU RAM optimizations to solve this problem in our future work.
D.2 End-to-end Experiment on SuperGlue
We conduct end-to-end training on four actual tasks from SuperGLUE [48] to check the soundness of our implementation. Specifically, we train BERT-large-cased [9] using BMTrain with H3T and PyTorch with APEX $^5$ , respectively. When using PyTorch with APEX, we adopt mixed precision
<sup>5</sup>https://github.com/NVIDIA/apex
Table 4: The experimental results on SuperGLUE.
Toolkits	BoolQ	CB (acc. / F1)	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 the fairness of comparison. The experimental result is shown in Table 4. From the result, we can find that our system can achieve comparable results with PyTorch on all four tasks, which greatly confirms the correctness of our implementation.
D.3 Energy Test
Figure 10: The energy cost (kWh) per sequence of H3T and Megatron-DeepSpeed. Here we use our DP solver for H3T and use the ZeRO-2 mode for Megatron-DeepSpeed.
We evaluate the energy cost of our framework and the baseline, Megatron-DeepSpeed, and the result is shown in Table 10. From the table, we find that our framework achieves lower energy cost with Megatron-DeepSpeed in most cases, except for training the 6B model on  $8 \times \mathrm{A}100$ , where we also achieve a comparable result. This result uncovers that our framework is more energy-saving than Megatron-DeepSpeed.
D.4 Solver Efficiency Test
Table 5: The solver runtime within different settings. The environment is  $64 \times \mathrm{A}100$ . 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 engineers sometimes change the model size and architecture in their works. For most people, this change does not always happen, so we do not need the solver to be as fast as possible, but the efficiency of the solver is still in demand. Therefore, we test the efficiency of the three solvers, and the result is shown in Table 5.
From the result, we can find that the random solver and greedy solver run significantly faster than the DP solver, although the minute-level runtime of the DP solver is quite acceptable compared to the training time of the large-scale model. The DP solver owes its acceptable efficiency to the two prunings, especially the second one, which reduces more than  $99\%$  of the DP states. In contrast, DP without these two prunings will take over a day or even a month to get the solution, which is greatly unacceptable.
However, this is not the best for the DP solver. For deadline reasons, we use Python to implement the serial version of the three solvers. In the future, we will port them to  $\mathrm{C}++$  and optimize them using multithreading and parallelism techniques. In this way, we expect the DP solver to gain at least another  $100\times$  improvement in efficiency and be 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, 8×2080Ti, M = 8, DP solver.
Simulated case	parameter partition	Switch-on ratio of	offloading hidden state
lazy-prefetching	Rematerialiaztion
/	90%	8%	92%	21%
PCI-E speed ×0.1	0%	0%	98%	0%
FPU speed ×0.1	88%	90%	88%	81%
M←25	50%	50%	65%	56%
Solver←Greedy	0%	0%	100%	0%
We show a case study in Table 6 to intuitively show the ability of our automatic solvers. From the result, we find that if we manually slow down the PCI-E communication, the solver can adaptively switch on fewer communication-aware optimizations and more computation-aware optimizations. The opposite situation occurs if we manually slow down the computation. When the memory constraint becomes relaxed, the solver reduces the switch-on ratio of all four switches. These results show that the DP solver is well adapted to the actual environment to customize the optimization scheme. Besides, we find the greedy solver does not offer as much flexibility as the DP solver in provisioning multiple optimization switches, although it still performs well in most cases of the simulation and actual runs.
E.1 Scalability of H3T
In this part, we discuss the scalability of H3T. Although H3T is based on the solvers and optimization switches introduced above and implemented on BMTrain for Transformer training in this paper, the framework and the idea of H3T can be applied in many other scenes. An illustration is shown in Figure 11.
Optimization techniques and the solver. Scholars and programmers can introduce any optimization or parallelism techniques and design their own solver for the specified optimization switches. The switches do not have to be binary and layer-level as they are in this paper, as long as they can be formalized as several variables. For example, we can have a global switch of pipeline parallelism, or we can also have some continuous optimization hyper-parameters, e.g., the ratio of tensor offloading. Then, the design of the solver highly depends on the optimization switches. To support different switches and hyper-parameters, we may need a different solver from that we proposed in this paper. Despite this, we argue that the three solvers designed in this paper work well for most discrete binary layer-level optimization switches.
Models. H3T is designed for Transformer-based models but can also be applied to other neural models, which can be formalized into the sequential model. Since the structure of sequential models and transformer-based models is becoming more and more widely used in recent years, we argue H3T has good generality and compatibility in the present and future.
Toolkits. In this paper, we implement H3T on BMTrain. Scholars and programmers can also implement H3T on other toolkits like DeepSpeed, etc.
Figure 11: The scalability of H3T.
E.2 Limitation of H3T
Although H3T performs well in our experiment and has good scalability, it still has some limitations. This section discusses the limitations of H3T from two perspectives: memory and model structure.
Memory. In this paper, we follow most existing works [2, 3, 18, 20] and defaultly use 'memory' to represent the memory allocation on GPU. Although this setting works well in most cases, it still brings two limitations.
(1) The memory allocation does not equal the actual memory usage. In actual runs, the GPU memory is usually managed by the underlying DL frameworks like PyTorch, TensorFlow, etc. And memory usage is usually greater than memory allocation due to many complex reasons, such as cache and memory fragmentations. Therefore, the user must set an  $M$  lower than the actual GPU memory capacity to ensure H3T's proper running. Moreover, this also means H3T cannot perfectly estimate the memory usage, which leaves space for further improvement of our solver. One solution is to implement the memory pool management module in-house to gain absolute control over memory management [15]. To tackle this, we will develop a memory management module for absolute control over the GPU memory in the futures.
(2) We only focus on the GPU memory and neglect the CPU memory. This does not matter in most cases because the CPU RAM is usually large enough for all parameters, gradients, and optimizer states. However, when we train GPT-3-175B on  $64 \times \mathrm{A}100$ , the CPU RAM runs out of memory if we set the batch size as 1024 or larger. This result reveals that CPU memory management is also in high demand for training extremely big models. Some researchers attempt to offload tensors from CPU RAM to the secondary storage (e.g., NVMe SSD) [37, 45], which is also our future direction to resolve this problem.
Model structure. As discussed in the scalability part, H3T supports any models that can be converted into sequential models. Nevertheless, there is one exception, namely, the models with dynamic structures, because we cannot search for a fixed optimization scheme for the dynamic model by pre-computing. This is a common problem with this kind of pre-computed optimization [3, 18, 20, 16, 31], and some other real-time optimization frameworks may be a better choice for these dynamic models [19, 15]. However, it is almost impossible to introduce approaches with heavy computations (e.g., DP and MILP) for real-time optimization, so most of these works employ heuristics to achieve adequate performance. In this paper, we mainly focus on Transformer-based models that are mostly fixed structures, so it is a better choice to design a pre-computed optimization framework like H3T.