Tempo: Accelerating Transformer-Based Model Training through Memory Footprint Reduction
Authors: Muralidhar Andoorveedu
Zhanda Zhu
Bojian Zheng
Gennady Pekhimenko
Venue: NeurIPS 2022 Accept (2022)
Keywords: Implementation
Systems
Transformers

Abstract: Training deep learning models can be computationally expensive. Prior works have shown that increasing the batch size can potentially lead to better overall throughput. However, the batch size is frequently limited by the accelerator memory capacity due to the activations/feature maps stored for the training backward pass, as larger batch sizes require larger feature maps to be stored. Transformer-based models, which have recently seen a surge in popularity due to their good performance and applicability to a variety of tasks, have a similar problem. To remedy this issue, we propose Tempo, a new approach to efficiently use accelerator (e.g., GPU) memory resources for training Transformer-based models. Our approach provides drop-in replacements for the GELU, LayerNorm, and Attention layers, reducing the memory usage and ultimately leading to more efficient training. We implement Tempo and evaluate the throughput, memory usage, and accuracy/loss on the BERT Large pre-training task. We demonstrate that Tempo enables up to 2× higher batch sizes and 16% higher training throughput over the state-of-the-art baseline. We also evaluate Tempo on GPT2 and RoBERTa models, showing 19% and 26% speedup over the baseline.

Tempo: Accelerating Transformer-Based Model Training through Memory Footprint Reduction
Muralidhar Andoorveedu', Zhanda Zhu Bojian Zheng!. Gennady Pekhimenko 1.3
'University of Toronto, Toronto, Canada
Shanghai Jiao Tong University, Shanghai, China
Vector Institute, Toronto, Canada (andoorve, zhanda, bojian, pekhimenko)@cs.. toronto.edi
Abstract
Training deep learning models can be computationally expensive. Prior works have shown that increasing the batch size can potentially lead better overall throughput. However, the batch size is frequently limited by the accelerator memory capacity due the activations/feature maps stored for the training backward pass, as larger batch sizes require larger feature maps to be stored. Transformer-basee -basedd models, which have recently seen surge in popularity due to their good performance and applicability to variety of tasks, have similar problem. To remedy this issue, we propose Tempo, a new approach to efficiently use accelerator (e.g.. GPU) memory resources for training Transformer-bast based models. Our approach provides drop-in replacements for the GELU, LayerNorm, and Attention layers, reducing the memory usage and ultimately leading more efficient training. We implement Tempo and evaluate the throughput, memory usage, and accuracylloss on the BERT BERTLARGE LARGE pre-training task. We demonstrate that Tempo enables up 2 higher batch sizes and 16% higher training throughput over the state-of-the-art bascline. We also evaluate Tempo on GPT2 and RoBERTa models, showing 19% and 26% speedup over the baseline.
1 Introduction
Transformer-base based models such as BERT and GPT-2 [49 have found success in numerous general natural language processing tasks including question answering [51], paraphrasing [[3]
natural language inference [68] and even areas outside language tasks such as image recognition [4]
However, training such models can be highly expensive in terms of time, monetary resources and
carbon footprint 24 For instance, the pre-training of BERT, LARGE takes days to complete on
16 Cloud TPUs (64 TPU chips total) [12), which costs about $ $10,000 (56). Training more recent Transformer-baat based model, GPT-3, has an even more astonishing price tag $12 million(66 Hence, even a small decrease in the end-to-end training time of Transformer-baset models matters.
Although there has been significant progress made in accelerating Transformers using specialized hardware (c.g.. Google TPUs (30), NVIDIA Tensor Cores (39)) in the past few years, a fundamental issue with Transformer- accelerators. For example, even batch size of does not fit into modern GPU with 12GB of memory when training BERT with sequence length 512 (15). Reducing memory footprint [88 52 is a viable option to allow larger batch training, leading to better hardware utilization and ultimately improved training throughput (73]
models that they are limited by the memory capacity of hardware
Many existing approaches to memory footprint reduction (e.g., offloading 52
checkpoint- 33 28), and data compression/encoding 226 6 either have high computational overhead
65 48
ing 8
7
or do not apply to Transformer- baaserrmer-bbeedd models directly. Prior approaches fall into two main categories,
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

neither of which are satisfactory for the Transformer-beet based model case. First, these techniques may utilize the specifics of Transformer-busedd models well, such as the multi-headed attention mechanism used in Transformers [63] or optimization opportunities available in specific layers such as the LayerNorm 4 layer. For example, although checkpointing [ 28 can significantly enlarge batch size, it also brings high overhead (e.g., 30% performance degradation observed in some prior works (8). Second, if prior works are specific, they focus on other types of models/layers with ideas not being applicable to Transformers. For example, Gist and In-Place ABN deal with CNNs 26 531.
be too general 448 6 30 B3 28
In our work, we demonstrate that low overhead memory footprint reduction can lead to positive improvement in throughput. In addition, unlike prior works which do not leverage the specifics of Transformer-basee models, we propose a new approach specifically tailored for Transformer based models, called Tempo. This approach includes three new techniques: (i) In-place GELUU, (ii) In-place LayerNorm, and (iii) Sub-Layer Dropout Recomputation. In In-place GELU and In-place LayerNorm both use alternative derivations for the computation of the backward passes of these layers. These derivations allow some activations that are normally retained during the forward pass (to be later used in the backward pass) to be discarded, leading to more memory-efficient implementation. Sub-Layer Dropout Recomputation discards activations within the high memory footprint attention mechanism during the forward pass, then recomputes these during the backward pass without recomputing extra unnecessary tensors. Tempo is able to increase the training throughput with larger batch sizes by reducing the total memory footprint of the models during training. To our best knowledge, this is the first work to explore memory footprint optimizations specifically for Transformer-bast based layers that show not just footprint reduction, but the actual increase in throughput using the extra memory savings. Tempo reduces the memory footprint of training Transformer- based models by targeting a major part of the total footprint the activation memory (74] (the saved feature maps during the forward pass of the model that are required for backpropagation (54)). ). All the proposed techniques provide a large memory footprint reduction with very low throughput degradation (as low as 1%). Our results show up to 2 improvement in batch size for BERT a.r pre-training at a sequence length of 512 on modern GPUs while increasing the training throughput by up to 16%.
2
Background and Motivation
2.1 Memory Footprint of BERT
BERT 12 is a popular natural language processing model that is based on the Transformer ar- chitecture 663] The model has been successfully applied to a variety of tasks such as question answering (SQuAD (I),, paraphrasing (MRPC (13)), natural language inference (MNLI (68 ), and others [57 72 through a two step training process. The training process entails first training on a general unlabelled data set (pre-training) [12]. The faster second part of the training process (fine-tuning) takes the parameter weights produced by the pre-training section and further trains downstream task such as question answering (51) or sentiment analysis [57] which it accomplishes
through the addition of specialized output layer [2].
The BERT architecture allows for multiple different configurations depending on model hyperparam- eters selected, some being derived from the original Transformers paper; these include the hidden layer size (H), sequence length (S), number of attention heads (A) and number of layers (L).
In the context of this work, we point out some of the relevant parts of the model and their activation memory footprint with respect to these hyperparameters referring to Figure [
! At this point, where attention 63 is calculated we observe that the size of each of the feature maps goes as ((S))- there are variety of previous techniques and models that have been explored in the literature to deal with this problem [61]. Additionally, at this point note that we store three feature maps of size [ x A S²). Calculations based on Figure I at the BERT B BASE parameters show that at a sequence length of 512 these three feature maps account for 56% of the encoder layer activation memory.
2 At these two points, we store the input to the two LayerNorm layers of size [B x S x H]
Here GELU (1 layer is used as the activation function for the preceding fully-connected layer of size [B* x S x 4H). The activation memory for this function stores almost 17% of the total layer activation memory of BERT BASE at sequence length of 128.
2

---

5
Liserr
Reshape
Neat
Despout
LaveeSarm
10 S.:.
[Rix
Linear
Trasspoe
Restage
Reshape
Reshape
Liuae
Q.1
Presions Loyer
Figure 1: A diagram of single Transformer encoder the Huggingface implementation of BERT [69]. As in the BERT paper, A represents the number of attention heads, and И represents the hidden size. We represent the batch size by B and the sequence length by S. Sizes of intermediate tensors (both retained activations and unretained intermediates) are annotated.
[63 layer used in BERT This based on
2.2 Why Activation Memory Matters
As iterated previous works [2 26 73 6 there are multiple benefits to reducing the memory footprint of models. First, it allows for larger models which positively affect the model's performance on downstream tasks [1]. Second, memory footprint reduction can allow for larger batch size. This, in turn, could lead better utilization of the GPU hardware [17, increasing the overall throughput 733. In order to verify this possibility for Transformer-baan
models, we conduct own experiments using Huggingface's BERT implementation [69 to train BERTLARGE on the MRPC [13] fine-tuning task. Figure 2 shows the throughput on this task for sequence lengths of 128 and 512. From the figure, we conclude that there is a steady improvement in batch size when the sequence length is 128. This is also the case when the sequence length is 512, however, in this situation the trend ends more abruptly as the memory consumption of the model exceeds the GPU memory capacity, showing clear opportunity to take advantage of memory footprint reduction.
Throughput (Sequences/s)
30
20
S 128
10
S 512
0
2
8
10
12
14
16
Batch Size
Figure 2: Plots of throughput (sequences/s) vs batch size for BERTLARGE
fine-tuning on the MRPC [13] task at sequence lengths 128 and 512 on four 2080Ti 40 GPUs. The maximum batch sizes are respectively 16 and 2.
12]
We note that previous works on Transformer-based models show that although the model parameters contribute memory footprint, the main memory capacity consumer during training actually the activation feature maps 77 28 8 48
In addition, the majority of this activation memory will used in each the BERT Transformer encoder layers. Profiling the Huggingface BERT BASE implementation 669 the MRPC [13) fine-tuning task at batch size of 32 and
26
33
6

---

sequence length of 128 shows that 66% of the total memory is taken up by these encoder activations. More details on this are shown in Appendix A
2.3 Key Prior Works
There are three major prior techniques used in training memory footprint reduction of deep learning models. The first of these Checkpointing 8 33 28 73.. feature maps in the forward pass while retaining others. Later, in the backward pass, these discarded feature maps may be recomputed from the retained feature maps, and thus used in the computation of the gradients. The second technique is Offloading [48 52 65]. In this case, the main idea involves taking feature maps that would be stored the GPU memory, and instead offloading them to the CPU memory. These techniques can also involve pre-fetching tensors from the CPU memory in anticipation of their use. Offloading suffers from a dependence on system variables such as the communication channel bandwidth [22 48 It also requires extensive engineering effort to avoid high overhead [6. Finally. Comppessiooneencoding this can be divided into two different categories, lossless and lossy [26. 6 However, the fundamental idea is to compress, or reduce the space taken up by feature maps in the forward pass, then decompress it for use in the backward pass.
This technique involves discarding certain
These techniques are usually largely orthogonal to one another as was shown in prior works where both offloading and checkpointing are used simultaneously 48 65 We expand on these techniques in Appendix
2.4 Why Tempo?
Although the techniques in the previous section show good performance a variety of models, they suffer from a variety of issues. Checkpointing's 's scope is often too broad to consider certain layer-specific optimizations and alternative derivations that can provide lower overhead [53]. Fur- thermore, overhead can be high (as much as 30%) []. Offloading can be system- dependent and requires significant engineering effort, while compression Transformer case. Hence, there a clear need for a deeper look at activation memory optimizations for Transformer-based neural networks in particular. To our best knowledge, our work is the first to explore such optimizations tuned to improving the throughput of Transformer-basse models. Table [ shows summary comparison of Tempo and various other techniques, with the major points that differentiates our technique from prior work.
be lossy not applicable to the
Capecktn
Checkmatc
Tempu
Giat
Layer-Spacific
x
Tranformer Specific
Lossless
Drup Layer Roplacement
Online
Table Comparison between Tempo and Capuchin [48], Checkmate 228] ActNN
66,
and Gist (26)
Tempo: Key Ideas
We now present the major ideas that lays behind the design of Tempo: (1) In-place GELU, (2) In-place LayerNorm, and (3) Sub-Layer Dropout Recomputation. The major theme behind all of these ideas is to compute the backward pass as normal, while asing less storage do so. To this end, In-place GELU and In-place LayerNorm compute the output of each layer in-place; instead using the output activation to compute the gradient. Sub-Layer Dropout Recomputation also discards the output, and through a closer look at the structure of the Dropout layer is able to recompute the output without excessive recomputation. We strongly suggest reading Appendix E for the implementation details. We also add in this appendix a new optimization of softmax that we use that further reduces memory [8]
Some of the Gist 26 optimizations are lossy.
"Accuracy of our lossy optimization is tunable, offering flexible tradeoff between the accuracy and the hardware cost.

---

3.1 In-place GELU
The GELU layer used as an activation function for the feed-forward section of the BERT layer (3) in Figure I [1]]. A plot of this function is shown in Figure 3a Referring to the baseline in Figure 3b note that both X and Y are stored for the backward pass. Y is needed for the downstream fully connected layer, while X is stored for the GELU layer itself [46]. Prior work has demonstrated that certain activation functions such as ReI.U may be computed in-place [6)]. This can be done without affecting the calculation of the backward pass. If we were able to compute the GELU function in-place, potentially by recovering the input from the output on the backward pass, we could save the storage required for X. However, this is impossible to do directly. A key observation to make with respect to the GELU function is that it is not bijective - hence there is no function that will be able to compute the input from the output without additional information.
However, we observe that the GELU function is both continuous and has only one extremum, minimum value at 2 = -9.75179 as can be seen in Figure Ba Notably, this implies that just one additional piece of information: which side of the minimum the input originates from, allows us to compute the inverse of the GELU. This is because on each side of the minimum the function one-to-one, and hence the input is recoverable from the output in each section. Based on this key observation, we can discard the input, and simply retain the output of the GELU, as well as the additional information on whether the input is greater than or equal to the value at which the minimum occurs. Figure 3b illustrates the difference between our method and the baseline.
In order to execute this efficiently on a real system, we note that the original derivative in terms of the input can be composed with the function inverse in order to create a composite kernel. This kernel consists of polynomial approximation of this composite function, the approximation being necessary since GELU is transcendental, and therefore the inverse cannot be solved in terms of elementary functions 58
Further details are discussed in Appendix H
GELU(x) = x 2111 + erf( x ) 12
Linear
Lineer Live
Linear PRD
Linear
GE)U(x)
2
Y
GELU Pro
GELA lvas
GELU
GELU Dva
X
dX
4X
-0.70179
Baseline
Terpo
0
Figure 3(b): Saved feature maps between the base- line and Tempo. Note that our method only saves 8-bit mask' that denotes whether the input is Figure 3()): A plot of the GELJ function near greater or less than the minimum value, instead of the origin, along with the marked minimum point, the full 32-bit input feature map.
X
3.2 In-place LayerNorm
The LayerNorm layer is used at multiple points in the Transformer encoder layer 633 which we denote by (2 in Figure [ input from the next layer, as well as the input feature map which is stashed for this computation [46).
Usually, the gradient computation of LayerNorm relies on the gradient
Similar to GELU, we are able to derive an expression for the gradient of the LayerNorm layer as function of its output. In this context, the output of LayerNorm must be stashed to compute the gradient of the successive fully connected layer anyways. Using this approach, the memory footprint overhead of LayerNorm is just the intermediate mean and variance computed in the forward pass. The full derivation is presented in Appendix E which is extended from the treatment of BatchNorm in (53]
Comparison with Checkpointing: Note that although In-place GELU requires more memory compared to recomputing Y from X, it will have increased overhead due to the recomputation.
PPytorch boolean masks use b-bits value 46). Masks can also be implemented as 1 1-bit manually but this brings extra overhead due to unpacking and packing bit tensors.
5

---

Additionally, our technique is orthogonal to conventional Checkpointing, as it could take advantage of the fact that no recomputation is required for the imput X for both In-place GELU and In In-place LayerNorm.
3.3 Sub-Layer Dropout Recomputation
In this section we explore the idea of sub-layer granularity checkpointing, or partial recomputation applied to the Dropout layer [59] found in ① in Figure[] The function of a dropout layer is to set the output of p% entries the incoming feature map to zero ("drop" the outputs) and then scale the remaining outputs by the factor I-p' which makes the network less sensitive to any output of the preceding feature map, thereby making it more robust [59].
We define sub-layer recomputation as technique where recomputation of only some of the feature maps is necessary for the backward pass that may be produced by given layer's output. We observe that better recomputation strategies are possible if we carefully deconstruct layers in the case where they store multiple outputs. This observation can be directly applied to the Dropout layer. In the computation Dropout, both mask (which records the entries which are set to zero in implementatoons of Dropout [46] 77) and output are produced. If a layer-based checkpointing implementation [6 was used, would cause both the mask and the output to be recomputed the backward pass if the layer is checkpointed, thus requiring higher overhead. However, nothing precludes us from simply doing only one of these recomputations. Storing the mask would only reduce the recomputation (including memory transfer) time, while the fact that the mask itself only has Boolean values allows us to keep most of the memory benefit of recomputation. can save the storage required for the output at the critical (SS)) Attention section (Q) in Figure I for the cost of a simple mask multiply. This technique is illustrated in Figure
notice that
this way.
Matcha Malliply
Dropout
Scalur Medsaply
Sealue Mishsply
Scoue Multipty
Scalar Muttiply
Scalae Magghel
TX
Madieed Splect
Muxied 2etl
Ma ladd Selaet
Misled Solect
Gen
4x
4X
Su0Maa
SaR3iiu
Figure 4: Comparison of dropout implementation between the baseline, and our method. Note that we only save the mask, and recompute the other output. The representation on the left is not an exact copy of the Py PyTorch implementation, rather it is an illustrative representation.
3.4 Other Engineering Optimizations
We note that PyTorch uses a memory-inefficien implementation of the softmax function which retains both the input and output of the function for the backward pass [46]. Instead, only the output is necessary. This optimization has also previously been implemented as part of some models in the Huggingface library [18). We use this optimization as well in our implementation of the attention mechanism to further reduce the activation memory pressure.
Evaluation
4.1 Methodology
Infrastructure Our main test setup consists of 4 NVIDIA RTX 2080 Ti GPUs 40 each with 11 GB of memory connected over PCIe v3 47). We also use an Amazon Web Services p3. 8xlagee [ instance consisting of 4 NVIDIA Tesla V100 GPUs [39 each with 16 GB of memory connected

---

using NVLink 42 For our ablation studies, we employ system with an NVIDIA A100 GPU 43
with 40 GB of memory. We summarize the detailed setup Appendix G
Applications We evaluate our work using both the BERT pre-training and fine-tuning tasks [2] For pre-training, we employ the NVIDIA DeepLearmingExmmpee library [41] with the English Wikipedia dataset [67]. We perform the training in two phases, the first (i.e., longer) phase at a sequence length of 128, and the second (i.e., shorter) phase at sequence length of 512 [2 41) For throughput and memory experiments, we use the BERT LARGE configuration. For our fine-tuning task, we use the MRPC [3] paraphrasing task on BERT BASE using the Huggingface library 669
For our ablation studies, we also train both RoBERTa 34 and GPT2 For the evaluation of RoBERTa,
the Fairseq library 444, while GPT2 uses the Huggingface GPT2 model 66 Both of these models use the WikiText Dataset for evaluation [35].
Metrics The first metric we focus is the total memory footprint of our method compared to the baselines. There are two ways to look at this metric. First, we compare the maximum batch size possible for each method. We compare this across sequence lengths of 128 and 512 on B for both 2080Ti and V100 GPUs. Second, we compare the total memory used by Py PyTorch given commonly used batch size for the same parameters. The second metric use is the throughput for which we count the total number of sequences per second processed. Finally, we provide comparison between our method and the baseline method on BERT BASE pre-training in order to compare the loss curves and show the change due to our lossy optimizations, We also provide fine-tuning curves the MRPC 13 task, training for 10 epochs to ensure no significant accuracy deviations.
Our ablation studies only use the throughput metric.
4.2
Results
We use two major baselines. The first baseline is the NVIDIA BERT LARGE model 41 with no memory footprint techniques applied which we refer to as the Baseline. The second one is the same model, with the default checkpointing applied, based on the PyTorch implementation, applied the input of each Transformer encoder layer [46 41] and is similar to the Huggingface implementation [69]. We refer to this baseline as Checkpoint. We refer to our method that uses In-Place GELU, In-Place LayerNorm, Sub-Layer Dropout Recomputation, and the softmax engineering optimization as Tempo.
Impact on Memory Footprint Table shows the maximum batch size and memory consumed at fixed batch size for all three methods. Additionally, the total memory used at a batch size of 15 at sequence length of 128 is 11.3 GB, 8.3 GB and 9.2 GB respectively for Baseline, Checkpoint, and Tempo. From this, we conclude that Checkpoint reduces the memory footprint to a much higher degree than both Baseline and Tempo This is expected, as Checkpoint discards most of the feature maps to be recomputed [69 41 no matter the performance cost. Tempo still provides a significant increase in batch size over Baseline at the sequence length of 512 we see 2 and 1.5x larger batches over Baseline for the 2080 Ti and V100 respectively but, as the next section shows, with much better throughput.
Technique
Soquence Length
Batch Size
Technique
Sequence Length
Batch Size
Baseline
128
15
Baseline
128
28
Baseline
512
Baseline
512
Checkpoint
128
Checkpoint
128
96
Checkpoint
512
Checkpaint
512
18
Tempo
128
24
Tempo
128
41
Tempo
512
2
Tempo
512
Table 2: The maximum batch size on both 2080 Ti (left) and V100 (right) for BERT LARGE.
Impact on Throughput Figure illustraes our main results with respect to throughput. From the figure, we can see that Tempo outperforms both Checkpoint and Baseline across both sequence
'These are the sequence lengths of Phase and Phase 2 of pre-training 12 41
7

---

lengths and across different hardware setups. We observe an improvement of 16% over Baseline on the 2080 Ti sequence length of 512. these settings, we also have an improvement of 8% over Checkpoint. We also observe up to 27 % over Checkpoint on the V100 at sequence length of 512, which also corresponds to a 5% improvement over Baseline. This despite the fact that Checkpoint uses the largest batch size at the beginning of each Transformer encoder layer, and recomputes these layers [69 41]. Hence, an increased batch size also means more recomputation. In contrast, Tempo is able to decrease the total memory footprint, and then convert this decrease into a substantial performance improvement over the Baseline due to the use of only low overhead mechanisms.
2 This is because Checkpoint stores feature maps
Throughput (Sequences/s)
200
B Base.
1,03X
1.DXX
0 Chk.
2 Tempo
100
1.04X
108 X
0
128
512
128
512
Sequence Length
Sequence Length
Figure 5(a): 2080 Ti
Figure 5(b): V100
Figure 5: Throughput experiments at the maximum batch size annotated with the speedup over the best baseline.
Impact on Loss and Accuracy We pre-train BERT BASE to ensure that our model's 's loss curve is not affected by approximate optimizations (c.g., In-Place GELU). Figure 6a shows the loss curve of phase I of BERT BASE pre-training [[2]. We observe almost complete overlap in the loss curves with more than 0.5% difference between Tempo and the baseline at the endpoint. We conclude that within that margin of error our method is satisfactory.
0.875
0.850
AccuraCy
0.800
10.0
Baseline/Checkkointing
0.775
Tempo
Loss
7.5
Beseline/Checckoonting
0.725
Tempo
5.0
Epoch
0
1000
2000 3000 4000 5000 6000 7000
Figure 6(b): Accuracy of BERT fine-tuning
on the MRPC 13 task. We run 10 trials of epochs. The solid line represents the median accuracy of these Figure 6(a): Phase 1 BERT BASS pre-training curve on trials, and the maximum and minimum along the train- ing process by the transparent curves' boundaries.
Step
the English Wikipedia dataset (67]
For the fine-tuning accuracy, we use the pre-trained Huggingface 69 implementation. Figure 6b
shows the results of BERT LARGE fine-tuning [12] on the MRPC [33] task. The figure shows consistent overlap between the maximum and minimum accuracy of Tempo and the baseline, so can conclude that Tempo has little impact on the accuracy of the trained model.
4.3 Ablation Studies
Ablation Study With Respect to Larger Model Parameters on Modern Hardware Platforms We also evaluate on other hardware platforms as well model parameters. First, we use an increased hidden layer size for various configurations. These experiments are conducted platform with

---

an NVIDIA A100 GPU 4S across sequence lengths of 128 and 512. We maintain the hidden layer size H to the number of attention heads A ratio of 64 which is in line with prior works [3 12. The The figure demonstrates two important generalizations of Tempo. First, note that even on newer and more advanced GPUs, Tempo continues to provide a tangible benefit. Second, across larger hidden layer sizes Témpo consistently demonstrates clear improvement over the baseline (as shown in the figure, this can be as high as 39% speedup over Baseline which corresponds to a 16% speedup over Checkpoint). The speedup over Checkpoint is high 20% conclude that Tempo will continue to be applicable to new hardware and larger models.
results are shown in Figure 7
1.5
L07x
1.07
1
L11X
L.02x
1.15X
Normalized Throughput
0.5
128
512
128
512
128
512
128
512
Sequence Length
d Baseline JDCheckpoint T Tempo
Figure 7: Normalized throughput at the maximum batch size, with annotated speedup over the best baseline. From left to right the configurations (a) BERT LARGE (H 1024), (b) BERT BASE H 2048, (c) BERTLARGE H 2048, (d) BERT BASE H 3072.
We also conduct experiments BERT LARGE (modified data points) for sequence lengths larger than 512. Figure X shows the results for this experiment, where demonstrate that Tempo outperforms Baseline longer sequence lengths well which can be as high as 27% % speedup over Baseline. At the same settings, we observe 16% speedup Checkpoint. Tempo also outperforms Checkpoint by as much as 20%. We conclude that yet again Tempo will be able to take advantage of modern hardware, as well as remain advantagcous as sequence lengths increase. Note that the largest sequence length 3072 on Baseline does have enough memory to run.
12 I. Layers instead of 24 for more
1.5
1.06X
1.10X
1.16X
1.16X
1.10X
Normalized Throughput
Base. EB
Chk. $ T'Temp
0.5
128
512
1024
2048
3072
Sequence Length
Figure 8: Normalized throughput relative the Baseline across different sequence lengths the NVIDIA A100 GPU for BERT, BERT,ARGE LARGG modified to use 12 layers. We annotate each bar group with Tempo's 's speedup over the best baseline.
Results on Other Models We conduct experiments on other Transformer models as well: RoBERTa (34) and GPT2 [49] For the evaluation of RoBERTa, we use the Fairseq library (44 as well as a sequence length of 512, while GPT2 uses the Huggingface GPT2 model (6). These experiments are conducted on both 2080 Ti and V100 setups. We note that the improvement over the
baseline is substantial (up 19% and 26% % for GPT2 and RoBERTa respectively the 2080 Ti

---

setup. This improvement corresponds to increase in batch size of 3x and 2x. Furthermore, we also see speedups of 5% and 4% on the V100 setup as well. From these results, we conclude that Tempo generalizes well to other Transformer-basee models besides BERT.
5
Extensions
5.1
Extending In-place GELU
The ideas used in section 3 for In- in-place GELU can be extended to general elementwise layers. The generic steps required for this are listed below, from the the high-level mathematics the low-level kernel based accelerator implementation, To the best of our knowledge, ours is the first work that exposes this potential optimization. This is a generic strategy reduce memory footprint in multi-dimensional space.
Consider elementwise layer with n inputs that applies function f inputs such that y f( 22,,.., to each corresponding element of the input tensors and where the output is re- tained for the backward pass of the subsequent layer.
Discard activation Vi without loss of generality. Determine a function g such that T1 U2, ..., ). For bijective functions of one variable this is simply the inverse.
such a function does not exist without ambiguity, construct functions 91,..., 9j that can recover I1 an interval. Construct a function g. such that 21 g y, 22,22..2 ... where is indicator that denotes the interval from which æ1 from and thus the piecewise selection of one 91,...9 Oj- Polynomially approximate each of ... to construct a new piecewise function 9-p in the case that they cannot be expressed analytically.
For the implementation of the forward pass, fold the computation of m into the computation of f. In essence, construct a new function f. such that (y, m) = (_1,,,,.. Fn). This can be done in a single kernel call.
(m,,m, (m, y, 42 -2, Fnn into the computations of 0f af if the computation of these values requires by composing these functions. In Dez .... den essence, fusing the kernels for the inverse and gradient operator. Then, we require kernel calls to calculate the gradient with respect the loss as before.
For the backward pass, fold the calculation of æ = Ixp
We illustrate this strategy in appendix E in more detail. The crux of the idea is that m, if needed at all, can be stored with less memory than æ1, while keeping the number of kernel calls minimum.
5.2 Auto-Tempo
As part of exploratory future work, we consider the application Tempo as an automatic compiler pass. We propose and prototype two different methods of automatically applying Tempo trans- formers which are available at the link in section 6 The first method is a fast method of profiling beforehand to determine whether memory footprint reduction would help, then applying Tempo to all applicable layers. The second method is fine- fine-grained grained method applies Tempo to subset of the applicable layers where the subset is determined through automatic profile and search, analogous to binary search.
6 Conclusion
We propose Tempo, a new mechanism that reduces the memory footprint of Transformer-basee models at low cost. It shows improvement in throughput of up 16% over the sate--of-hhear baseline for BERTLARGE pre-training task and also shows an improvement in maximum batch size of up 2X on both V100 and 2080 Ti GPUs. Our technique also generalizes well to new models, more modern hardware, as well as diverse model parameters in terms of memory footprint and throughput, demonstrating the robustness of our technique. Our hope is that Tempo will be used with other footprint reduction methods to improve training efficiency of Transformer-basee models. We open-source Tempo for an immediate positive impact on both machine learning researchers and practitioners here: https:/ http://ggttttbtte cww.oof.cc//ooffoataoo /tgpg///wb.com///ooTToooootttttt Tempo
10

---

References
[1] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg. Rajat Monga, Sherry Moore, Derek Gordon Murray, Benoit Steiner, Paul A. Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: A system for large-scale machine learning. In 12th USENIX Symposium Operating Systems Design and Implementation, (OSDI 2016). USENIX Association, 2016. https: ktttps://wat h/ww:///wwwww...o usenix.ee............ org
/
sessions wtw.i/onnnneeeeeetttiiioooo abadi
[2] Jorge Albericio, Patrick Judd, Tayler H. Hetherington, Tor M. Aamodt, Natalie D. Enright Jerger, and Andreas Moshovos. Cnvlutin: neeffetual-Neuroo-PPe Deep Neural Network Computing. In 43rd ACM/IEEE Annual International Symposium on Computer Architecture, (ISCA 2016), 1109/ISCAA 2016 11
2016. https://doi.orr
[3] AWS. Amazon EC2 P3 Instance Product Details, 2019. https://aws. http:///aaaaamaz..oooo.o comzee///cm///////
-stance-types/p3
[4] Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization. CoRR, 2016. http://arxiy> wttp:/iii.rg///bba.b/rbbeb
org/ awwgg/abss/eee /etbb///11700000 06450
[5] Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The Long-Document Trans- former. CoRR, abs/2004.051500 150, 2020. https://arxiv: org/ httpg//bs/2004 051150
[6] Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael W Mahoney, and Joseph E Gonzalez. ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training. In International Conference Machine Learning (ICML), 2021.
[7] Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems. CoRR, 2015. http: ftttp://arxiin org/ abs///bs//// 1112. 01274
[8] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training Deep Nets with Sublinear Memory Cost. CORR, 2016. http://arxiv... www./oogg/////oeeeeoo 1abbb//114..6600
[9] Tianshi Chen, Zidong Du, Ninghui Sun, Jia Wang, Chengyong Wu, Yunji Chen, and Olivier Temam. DianNao: small-footprint high-throughput accelerator for ubiquitous machine- learning. In Architectural Support for Programming Languages and Operating Systems, (ASP-
LOS 2014). ACM, 2014. https: /ttps://doi... orgo 10. 01445//241949 2541967
[10] Yu-Hsin Chen, Joel S. Emer, and Vivienne Sze. Eyeriss: A Spatial Architecture for Energy- Efficient Datallow for Convolutional Neural Networks. In 43rd ACM//EEE Annual International Symposium on Computer Architecture, (ISCA 2016), 2016. https://doi. org/ httgg//..00////1..1. 1109/IS
CA. NA.2016.40
[11] Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking Attention with Performers. In 9th International Conference on Learning Representations, (ICLR 2021), 2021. htttp:://openreviex.. wttp:e//oorrmm/duuuubauuoooooo.
[12] Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Association for Computational Linguistics, 2019, "https
//aclanthology- org/ N19- 1423
[13] William B. Dolan and Chris Brockett. Automatically constructing a corpus of sentential paraphrases. In Proceedings of the Third International Workshop on Paraphrasing (IWP2005), 2005. "htps:///aclathooogy org/ 105-5002
11

---

[14] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An Image is Worth 16x 16 Words: Transformers for Image Recognition Scale. 9th International Conference on Learning Representations, (ICLR
2021), 2021. htttps://operreviev/ wttp://www....t.toooooooo///caad
[15]
Google. Out-of-memory issues. http:://ggghh..oooooooogeeeeeeer
f 1-memory-issues
[16] Priya Goyal. [Re-checkpointiing] Autograd container for trading compute for memory, 2018.
[17] Priya Goyal, Piotr Dollár, Ross B. Girshick, Picter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, Large Minibatch SGD: Training ImageNet in Hour. CoRR, 2017.
[18] Sylvain Gugger. PyTorch DeBERTa model, 2020. http:s///gtccub...///gggggggggge hetr://ooelss/cbbertt///oollllld
[19] Song Han, Xingyu Liu, Huizi Mao, Jing Pu, Ardavan Pedram, Mark A. Horowitz, and William J. Dally. EIE: Efficient Inference Engine on Compressed Deep Neural Network. In 43rd ACM/IEEE Anmual International Symposium on Computer Architecture, (ISCA 2016), 2016.
https: http///ii.rr//////// 1109/ ISCA, 2016.30
[20] Song Han, Huizi Mao, and William J. Dally. Deep Compression: Compressing Deep Neural Network with Pruning. Trained Quantization and Huffman Coding. In 4th International Conference on Learning Representations, (ICLR 2016), 2016. http::/arxiv. http//axiiv.gr//aae//// abs/
10. 00149
[21] Dan Hendrycks and Kevin Gimpel. Bridging Nonlinearities and Stochastic Regularizers with Gaussian Error Linear Units. CoRR, Abss/600.08413, 2016.
_08415
[22] Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the Knowledge in Neural Network. CoRR, abs/1503.02531, 2015. http:://axivv org/ aas////ss//11 0tt///113000022.
[23] ISO. ISO/EC 14882:2011 Information technology technology- Programming languages C++. Inter-
national Organization for Standardization, 2012. http:// /www. iso atg/iso/iso_catall
[24] Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data Movement Is All You Need: A Case Study on Optimizing Transformers. Proceedings of Machine Learning and Systems (MLSys), 2021. httt:::://proceednngs.oo wwggg..msss..... wttpppppppp/pp/////////////////a
Paper. pdd
[25] Animesh Jain, Parker Hill, Shih-Chich Lin, Muneeb Khan, Md. Enamul Haque, Michael A. Laurenzano, Scott A. Mahlke, Lingjia Tang, and Jason Mars. Concise loads and stores: The case for asymmetric compute-memory architecture for approximation. In 49th Annual JEEE/ACM International Symposium on Microarchiteecture, (MICRO 2016), pages 41:1 41:11. IEEE Computer Society, 2016. https://doi org/ 10. 1109/MICRO 2016. 7783744
[26] Animesh Jain, Amar Phanishayee, Jason Mars, Lingjia Tang, and Gennady Pekhimenko. Gist: Efficient data encoding for deep neural network training. In International Symposium on Computer Architecture (ISCA 2018), 2018. //htpss // /ww. mi microsoft. com/ en us/rese tra
attp/ppbbiaattiooniiststeeei.fie ining/
data- encoding deep deee-neural--e
[27] Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and Ion Stoica. Checkmate, 2019, https: http//tttttttt/tt///// wwwtttu..compppattttettt com/parasji...cccccc http:/aaaass/chccaaattt checkmatee
12

---

[28] Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and lon Stoica. Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization. In Proceedings of Machine Learning and Systems (MI.Sys), 2020. /https //proceedings... senteeedingg.mmysss http://s.w....a...aa/a////////// oatp/a/ar///////1000a666oeeeeea0 c9 c9-Paper pdff
[29] Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. TinyBERT: Distilling BERT for natural language understanding. In Findings of the Association for Computational Linguistics: EMNLP 2020. Association for Computational Linguistics, 2020. https://aclarre aacpanho/ccl
www./r20.........ooo 2020 findingnnsennnn wwnlp.mnnn..
[30] Norman P. Jouppi, Cliff Young, Nishant Patil, David A. Patterson, Gaurav Agrawal, Ramin- der Singh Bajwa, Sarah Bates, Suresh Bhatia, Nanette J. Boden, Al Borchers, Rick Boyle, Pierre luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert B. Hagmann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Daniel Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Lc, Chris Leary, Zhuyuan Liu, Kyle A. Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, J. W. Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon. In-datacenter performance analysis of a tensor processing unit. 2017 ACM/IEEE 44th Annual International Symposium Computer Architecture (ISCA), 2017.
[31] Patrick Judd, Jorge Albericio, Tayler H. Hetherington, Tor M. Aamodt, and Andreas Moshovos. Stripes: Bit-serial deep neural network computing. In 49th Annual IEEE/ACM International Symposium on Microarchitecture, (MICRO 2016), 2016. httpp: //doi. org/ 10. 1109/MICR
0. 0.2016.10 7783722
[32] Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. Dynamic Tensor Rematerialization (DTR) Prototype, 2020. attpss//eet thub. www/ /wtp:///tt/ttttttooootppp
[33] Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. Dynamic Tensor Rematerialization. In 9th International Conference on Learning Representations, (ICLR 2021), 2021.
[34] 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. CORR, abs/1907. 11692, 2019. http://arxiv. org/abs/ /100 11692
[35] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer Sentinel Mixture Models. In 5th International Conference Learning Representations, (ICLR 2017), 2017.
htttps://penreviiet
[36] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory F. Diamos, Erich Elsen, David García, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed Precision Training. In 6th International Conference on Learning Representations, (ICLR
2018), 2018
[37] John Nickolls, Ian Buck, Michael Garland, and Kevin Skadron. Scalable Parallel Programming with CUDA: Is CUDA the Parallel Programming Model That Application Developers Have Been Waiting For? Quene, 2008.
0r5.://.00.rr/00.11133335540000 1365500
[38] NVIDIA. NVIDIA Driver Downloads. http:://www..aaaaaaaa.////////a.
[39] NVIDIA. Tesla V100 Data Center GPU, 2017.
enter/v100/
13

---

[40] NVIDIA, GEFORCE RTX 2080 Ti, 2018.
http:::/ww..vviiia..ccooooooooaa aphics-cards/rte 2ttp:rxxx200000-00 2080-
[41] NVIDIA. BERT For http://pwgggggaa..ooo//////////a
wwe.naatwee/yooooo/////ggaggagaa 2019.
[42] NVIDIA. NVLink, 2019. https://www.vvidzz. httpa///ooo/e/aeaaaaaatttttaaaaa
[43] NVIDIA. NVIDIA A100 Tensor Core GPU, 2020.
ta-center/a100/
[44] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A Fast, Extensible Toolkit for Sequence Modeling, In Proceedings of NAACI.HHIT 2019: Demonstrations, 2019.
[45] Angshuman Parashar, Minsoo Rhu, Anurag Mukkara, Antonio Puglielli, Rangharajan Venkate- san, Brucek Khailany, Joel S. Emer, Stephen W. Keckler, and William J. Dally. SCNN: An Accelerator for Compressed sparse Convolutional Neural Networks. In Proceedings of the 44th Anmual International Symposium on Computer Architecture, (ISCA 2017). ACM, 2017.
https://det /rgg://////...1//..... 1145/3079882 3080254
[46]
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Z. Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, (NeurIPS 2019), 2019. https://procceeeee 92F22ffaFF012272774000 Attp://4t0-Asstttt...r
[47] PCI-SIG. Specifications, 2005. //tttsss//pcisss: httpp//sss...ee/epp....aaaaaaaaa
[48] Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma. Qian Xiong, Fan Yang. and Xuchai Qian. Capuchin: Tensor Baseed GPU Memory Management for Deep Learning. In Proceedings of the Twenty-Füfth International Conference on Architectural Support for Programming Lan- guages and Operating Systems (ASPLOS), ASPLOS '20. Association for Computing Machinery, 10. 1115//3333///3337 3378505
2020. /https: //doi
[49] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 2019.
[50] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. ZeRO- infinity: breaking the GPU memory wall for extreme scale deep learning. In SC '21: The International Conference for High Performance Computing, Networking, Storage and Analysis, 000111155111111111 3458877 3476205
2021. https://doi.... org/ 10. 1145/
[51] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods Natural Language Processing. Association for Computational Linguistics.
2016.
www.oogg/////// 1264"
[52] Minsoo Rhu, Natalia Gimelshein, Jason Clemons, Arslan Zulfiqar, and Stephen W. Keckler. vDNN: Virtualized Deep Neural Networks for Scalable, Kiemory-Efficiee Neural Network De- sign. In The 49th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), 2016.
[53] Samuel Rota Bulò, Lorenzo Porzi, and Peter Kontschieder. In-Place Activated BatchNorm for Training of DNNs. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2018.
[54] David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning Representations by Back-propan paok-propgggting Errors. Nature, 323(6088), 1986. htttp://www nature. com/ articss// 323583a0
14

---

[55] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. CoRR, abs/1910.01 108, 2019. http
warxiv.org/abbl /tt//a/b////....1 19900001081100001
[56] Or Sharir, Barak Peleg, and Yoav Shoham. The Cost of Training NLP Models: A Concise Overview. CoRR, abs/2004.08900 2020. https://arxiv. org/ abs 2b04. 08900
[57] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, A. Ng, and Christopher Potts. Recursive Deep Models for Semantic Compositionality Over a Sentiment Treebank. In EMNLP, 2013.
[58] Phillip Spencer. Solution to the Transcendental Equation 2 3 5. 1999. https
/ //www.athh wat.t.tttttttttooooooo wttp:///wwcccceeeeeeeeeetrrreeee
html
[59] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A Simple Way to Prevent Neural Networks from Overfitting. Journal of Machine Learning Research, 15(56), 2014. http://jmlr wwgg/aapess/v15. stappers/v1//rriastttvaaa wtt.ssttvaaaaaaaaaaaaaaa html
[60] Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and Policy Considerations for Deep Learning in NLP. In Proceedings of the 57th Conference of the Association for Computational Linguistics, (ACL 2019), 2019. littps: //o:: org/10. 1865388633// 11t.1//////11111555555 1355
[61] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient Transformers: A Survey. CoRR, 2020. https: /ttpsi/////xvvva///a orgg www./ag./////////ee/........eee/ /b09 20000060900202222222
[62] Guido Van Rossum and Fred L. Drake. Python 3 Reference Manual. CreateSpace, 2009.
[63]
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems (NeurIPS), 2017. ptttttss:/ppoceeeiinne wwwpeediinnnnnnuns..n.e. nwwripe..c.... Paper pdf,
[64] Swagath Venkataramani, Ashish Ranjan, Subarno Banerjee, Dipankar Das, Sasikanth Avancha, Ashok Jagannathan, Ajaya Durg, Dheemanth Nagaraj, Bharat Kaul, Pradeep Dubey, and Anand Raghunathan. ScaleDeep: A Scalable Compute Architecture for Learning and Evaluating Deep Networks. In Proceedings of the 44th Annual International Symposium Computer Architecture, (ISCA 2017). ACM, 2017.
https: : ://ooi//////// ht////i..ogg/1// 10. 1145/3079856 1315/3330885663302444
[65] Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song, Zenglin Xu, and Tim Kraska. Supermeurons: Dynamic GPU Memory Management for Training Deep Neural Networks. In Proceedings of the 23rd ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming (PPoPP). Association for Computing Machinery, 2018. 1145/317848//// 3178491
[66] Kyle Wiggers. OpenAl's massive GPT-3 model is impressive, but size isn't 't everything, 2020. wew.ahhhheeeeeeaannnnnoooeeou gpt-3-st siz
e- e-isnt-everything!
[67] Wikipedia. English Wikipedia, 2021.
Wikkppcia.20ll./ttppss///// http:://///n..kkiiddia..a org/
[68]
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers). Association for Computational Linguistics, 2018. Whtt ps://aclanthology www.////-/-.. 1101
[69]
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: Sf the-art natural language processing. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrutions, (EMNLP). Association for Computational Linguistics, 2020. "https: /wwww/wwww////////// aclweb. org/anthology/2020 emm]
P- x-demos.6 6"
15

---

[70] Geoffrey X. Yu. Tovi Grossman, and Gennady Pekhimenko. Skyline: Interactive In-Editor Computational Performance Profiling for Deep Neural Network Training. In Proceedings of the 33rd ACM Symposium on User Interface Software and Technology (UIST'20), 2020.
[71] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big Transformers for Longer Sequences. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, (NeurIPS 2020), 2020. https: https://procceeding nwwrips http:///.....aaoo//////////////1
Bird:
43725531a9a77361ab9 Awstract.ttattaaoaeeeeaaa html
[72] Rowan Zellers, Yonatan Bisk, Roy Schwartz, and Yejin Choi. SWAG: A Large-Scale Adversarial Dataset for Grounded Commonsense Inference. In EMNLP,, 2018.
[73] Bojian Zheng, Nandita Vijaykumar, and Gennady Pekhimenko. Echo: Compiler- GPU Memory Footprint Reduction for LSTM RNN Training. In 47th ACM/IEEE Annual International Symposium on Computer Architecture, (ISCA 2020), 2020. https://doi. . org/ 11.110..../11
CA45697 2020.000000000 00092
[74] Hongyu Zhu, Mohamed Akrout, Bojian Zheng, Andrew Pelegris, Anand Jayarajan, Amar Phanishayee, Bianca Schroeder, and Gennady Pekhimenko. Benchmarking and Analyzing Deep Neural Network Training. In 2018 IEEE International Symposium Workload Characteriza tion (IIWC),, 2018.
16


Reviews:
[{"review_id": "x4g-K52bH0a", "reviewer": "Reviewer_Qdmx", "summary": "This paper proposes an approach to reduce the memory footprint of transformer-based models for faster training without sacrificing accuracy. The proposed approach Tempo uses In-Place GELU, In-Place LayerNorm, Sub-Layer Dropout Re-computation, and softmax engineering optimization. The main advantage is to compute the backward pass as normal while requiring less storage. Evaluation results show 2x improvement in batch size for BERTLARGE pre-training at a sequence length of 512 while increasing the training throughput by up to 16%.", "questions": "If there are any data for the above issue 3, please briefly summarize.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n+ Clever techniques to enable recomputing input from output in GELU and LayerNorm with minimal overhead\n+ Good observation on dropout that avoids recomputing of unnecessary input\n+ Well organized paper and clear writing. Supplemental materials are helpful\n\nWeaknesses:\n1. Tempo requires more memory than checkpointing (although expected, this is still a fact/weakness)\n2. Optimizations in Tempo are limited to transformer-based models while some other alternatives are not (e.g., [48][28][26])\n3. Evaluation-related:\n- In Section 2.1, a sequence length of 512 is used for showing that feature maps account for 56% of the encoder layer activation memory, but a sequence length of 128 is used for showing that the memory function takes 17%. It would be more convincing if both results were given for the same sequence length. \n- In ablation studies, it would be great to show the impact on accuracy and memory footprint when increasing the number of hidden layers. \n\nMinor issues and suggestions:\n- Figures 3 and 6 are missing the main caption.\n- GPU names can be added as a row in Table 2.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "fwjnGoMX3Wy", "reviewer": "Reviewer_WUQt", "summary": "This paper proposes a solution named Tempo that includes three core techniques, in-place GELU, in-place LayerNorm and sub-layer dropout recomputation, to reduce the memory footprint of training transformer-based models. Given the reduction of memory usage, the same GPU is able to conduct gradient-based optimization with larger mini-batches; therefore, the training throughput is improved. ", "questions": "It seems the experiments are conducted on multi-GPU machines, Is there any parallel schema applied in these benchmarks?\n", "limitations": "Not applicable. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- The out of memory issue is a serious problem when training transformer-based models, especially giant foundation models. \n\n- The analysis of the memory consumption of transformer models during the training procedure is informative. \n\n- The performance boost over the state-of-the-art baseline and checkpoint approach is significant.  \n\nWeaknesses:\n\n- The novelty inherited from the proposed methods is relatively limited. These optimizations seem straightforward from the point of view of implementation.   \n\n- There is a lack of theoretical discussion about the lossy compression in the Tempo system. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "aGuAskVVa87", "reviewer": "Reviewer_mcUf", "summary": "The paper proposes optimizations that reduces memory usage when training transformer models. Reducing memory usage for training models allows using larger batch size which improves training throughput.\nPrevious methods that reduce memory footprint in training are: checkpointing, offloading and compression. \nThe changes presented in the paper are: in-place GELU layer, in-place Layer Normalization and Dropout.\nThe proposed method was evaluated on BERT, RoBERTa, GPT2 model. They show improved training throughput and ability to use larger batch sizes when training BERT with their method.\n\n\n", "questions": "Some suggestions and questions that could improve the paper are:\n\nAdding a comparison to check checkpoint+tempo to demonstrate that tempo can be used with other techniques \n\nAdding a comparison with compression techniques to demonstrate the trade-off of accuracy loss with more efficient training from tempo \n\nTry on more models that uses GELU to demonstrate that this method isn't only suited for transformers\n\nHaving a list of layers from various models (CNN, RNN, GANs, etc) that could be done in-place (that arent done yet) to improve memory footprint in training would be a value study\n\nTable2 should also list and highlight throughput achieved and have total memory usage \n\nFigure3 caption format seems confusing\n\n\n", "limitations": "The paper doesn't discuss limitations or future work.\nDiscussion how improving training on transformer models can lead to lower power consumption in servers systems.\nAnd add future work on what other layers could be done in place.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\nThe paper demonstrates techniques to improve training of relevant transformer models using in-place layer compute. The paper tackles an important problem\n\nOverall the paper is clear. But improvements on diagrams could improve clarity of the paper\n\nCons:\nIt is dubious that this is the first paper to analyze transformer training memory footprint\n\nin-place compute is fairly known optimization method for training thus the method lacks originality \n\nMore experiments and comparison could be added to strengthen the paper\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EriyfpBhqTy", "reviewer": "Reviewer_YdHM", "summary": "This paper proposes three simple yet effective optimizations to speedup the pre-training of large language models by reducing the memory footprint. The resulting approach is called Tempo. Tempo is able to reduce the memory requirement for the pre-training of language models such as BERT, RoBERTa, and GPT-2 significantly. In addition, the authors demonstrate that memory savings can be translated into increased throughput.", "questions": "1. Could you elaborate on how in-place GELU and layer normalization are different from in-place ReLU and batch normalization? Are there any new treatments involved in converting these two operations in-place?\n2. Have you evaluated the pre-training performance on other libraries with complier optimization enabled (e.g., Pytorch with JIT or Tensorflow with XLA enabled)? Is it possible that some compiler optimizations, such as operation/kernel fusion, have helped to reduce memory footprint and improve training throughput? If so, are the proposed optimizations complementary to the existing compiler optimizations? For example, Pytorch performance optimization guide [1] uses GELU as an example for kernel fusion as follows:\n```\n@torch.jit.script\ndef fused_gelu(x):\n    return x * 0.5 * (1.0 + torch.erf(x / 1.41421))\n```\n\n[1] https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The three proposed optimizations for reducing the memory footprint are simple and straightforward. I am convinced that the proposed ideas can help reducing the memory requirement for the pre-training of language models significantly. The paper is well written and easy to understand. The experimental evaluation is solid, as the authors demonstrate not only the results of reduced memory footprint, but also the improvement in throughput.\n\nWeaknesses: My main concern with this paper is whether the proposed optimizations are considered as new contributions. In particular, the first two optimizations propose to turn the activation and normalization functions into in place functions, which helps to reduce the required memory for saving the intermediate tensor for backward propagation. The idea of using in-place function to reduce memory footprint is not new. Although not directly applicable to the Transformer architecture, some previous work and in-place ABNs have proposed to perform activation and batch normalization in-place to reduce memory footprint. In addition, attempts have also been made to add support for in situ GELU in pytorch (https://github.com/pytorch/pytorch/pull/74629).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept