Efficient Packing: Towards 2x NLP Speed-Up without Loss of Accuracy for BERT

1 INTRODUCTION . Since its introduction in 2019 , BERT ( Devlin et al. , 2019a ) has been the backbone driving the most exciting advances in Natural Language Processing ( NLP ) . Pre-training BERT from scratch requires substantial computational resources which may be out of reach for researchers and industry professionals . To some extent this has been addressed by the public release of pre-trained models of different sizes and depths ( Turc et al. , 2019 ) . The introduction of ALBERT ( Lan et al. , 2019 ) and Switch transformers ( Fedus et al. , 2021 ) further improved the accessibility of larger models . However , the dependence on pre-trained models limits the ability of researchers to explore new backbone architectures . Furthermore , it limits the extent to which practitioners in industry can leverage internal datasets and adapt the model to their particular needs . Hence , any approach that speeds up the pre-training process is desirable from an economical as well as environmental perspective . In this paper , we present efficient methods to enable researchers to accelerate the pre-training of BERT by as much as 2x without loss of accuracy . The de-facto pre-training dataset Wikipedia , as well as many other NLP datasets , have a skewed distribution of sequence lengths . We show that padding tokens ( wasted compute ) represent 50 % of all tokens of the Wikipedia pre-training dataset at sequence length 512 . Thus , by avoiding processing the padding tokens one can get a 2x speed-up . Overall , the lengths range between 5 tokens up to 512 ( see Figure 1 ) . Samples of length 512 represent only 23.5 % of the dataset , a surprising result given that the pre-processing in BERT attempts to “ pack ” together sentences so as to fill the sequence length as completely as possible ( Devlin et al. , 2019c ) . While processing the padding tokens wastes compute , it is still the most standard approach for leveraging modern massively-parallel compute especially on GPUs . These are most efficient when applying the same operation to each sequence in a batch . By padding all sequences to the same maximum sequence length , they can easily be batched . We note that this naive batching is the most widely used and provided in the vanilla BERT implementation as well as the Hugging Face framework ( Wolf et al. , 2020 ) and thus considered as our baseline for comparison . The most obvious way to reduce the extent of padding in the dataset is to group samples by size before batching ( SORT ) , i.e. , process the shorter samples together and longer samples together . BERT is pre-trained in two phases , where the first phase uses sequence length 128 for 900K steps and the second phase uses sequence length 512 for 100K steps . However even by splitting the training in this way , the wasted compute due to padding is approximately 20 % ( see Figure 1 ) . Other examples of this “ sorted batching ” approach can be found in Faster Transformer ( NVIDIA , 2021 ) , lingvo ( Shen et al. , 2019 ) fairseq ( Ott et al. , 2019 ) , and RoBERTa ( Liu et al. , 2019 ) , which group samples of similar size together in one batch and fill up with padding only to the maximum length in this batch . This approach can be highly efficient in cases where the dataset length is multiple orders of magnitude larger than the batch size and the number of different sequence lengths . Despite its high computational efficiency , this approach has multiple drawbacks . We outline these below and propose an alternative which maintains the high efficiency , while also circumventing the downsides . Firstly , sorting the data can reduce the overall convergence speed when the batch size is large because it violates the i.i.d . assumption on the data distribution ( Bottou et al. , 2018 ; Meng et al. , 2019 ) . Secondly , processing batches with shorter sequence lengths under-utilizes the compute compared to running the same batch size with a longer sequence length . For GPUs , a common heuristic to mitigate this effect is to adjust the batch size to keep the number of processed tokens near constant ( Ott et al. , 2019 ; Liu et al. , 2019 ) . In general however , the relationship between the sequence length and the optimum batch size is more complex and maximizing compute utilization can require the model to be sharded differently across multiple accelerators . Avoiding this , often manual process , is important for ease of use and the portability of methods across different hardware architectures . Thirdly , modern NLP applications are optimized and compiled for fixed tensor sizes using tools such as XLA ( XLA , 2021 ; Fedus et al. , 2021 ) , which provides a ≈ 7x acceleration for BERT in MLPerf™ ( Mattson et al. , 2020 ) compared to the non-XLA baseline ( XLA , 2021 ) . Changing the sequence length or batch size requires re-optimization of the computational graph and recompilation of the program for the new tensor shapes . For complex models such as BERT , optimization and recompilation take a non-negligible amount of time . Even if one pre-compiled and cached all combinations of batch size and sequence length , the kernels would still need to be re-uploaded to the device every time the shapes change . Depending on how frequently the tensor shapes change , the overhead from switching kernels adds up . To avoid these issues , it is preferable ( and common ) to work with fixed tensor shapes for the entire duration of the training run . More advanced approaches for reducing the padding overhead rely on custom computational kernels . Loosely these are referred to as “ un-padding ” approaches . In Effective Transformer ( ByteDance Inc. , 2021 ) , the input batch is provided as a padded matrix but padding values are dynamically removed and restored during different calculation stages . While un-padding implementations are highly sophisticated and are able to completely circumvent the processing of padding tokens , they introduce a significant overhead due to the multiple GPU kernel launches ( i.e. , one kernel per sequence rather than one kernel per batch ) . Additionally the time to process each batch will fluctuate depending on the sequence lengths in each batch , i.e. , batches with only shorter sequences will typically be processed faster . When working with more than one accelerator , this variability in throughput results in all devices in the cluster waiting for the device with the most compute inten- sive batch to finish processing . As such , un-padding approaches are not appropriate for deployment on large clusters . The “ packing ” based approach introduced in this paper offers significant advantages over un-padding approaches . Firstly , packing is implemented directly at the framework level and requires no additional custom kernel implementations . Secondly , the processing time for each batch is independent of the content of the batch , allowing the packing based approach to maintain the same speed-up whether running on a single device or thousands . While we demonstrate the effectiveness of packing specifically on the Wikipedia dataset , packing SQuAD ( Rajpurkar et al. , 2016 ) or GLUE datasets ( Warstadt et al. , 2018 ; Wang et al. , 2018 ) for BERT also leads to significant speed-ups ( some in excess of 9x ) ( Sections H and I ) . The effectiveness of packing is a result of both the length distribution of the documents in the source datasets as well as the different text preprocessing steps for BERT ( Devlin et al. , 2019c ) . The use of bi-directional self-attention in BERT implies that the input sequences should contain complete sentences . If a sentence is abruptly cut short , the hidden state on other ( preceding ) tokens in the sequence will be affected . Language models with causal attention ( only attending to previous tokens in the input ) do not have this issue to the same degree . For such models , if a sequence is cut short at an arbitrary token , the other tokens ( which occur earlier in the sequence ) will not be affected . This ability to cut sequences arbitrarily completely trivializes the packing problem for models based on causal attention . For instance , GPT-3 ( Brown et al. , 2020 ) is trained with a maximum sequence length of 2048 where a single sequence may contain multiple segments of sentences separated by a special end of segment token . The last segment in each sequence is simply cut to meet the sequence length requirement making the packing problem trivial and avoiding any padding . In the interest of computational efficiency GPT-3 does not mask the attention between different segments in a sequence . In contrast , the packing approach presented in this paper introduces a mask in the attention layer ( see Section 3.2.2 ) to prevent cross-contamination between examples in a pack . Note , we mask the interaction between different sequences and not between different sentences or segments in the same sequence . This ensures that the characteristics of the original dataset and model are matched as closely as possible . RoBERTa and many other models in production like T5 ( Raffel et al. , 2019 ) use a similar packing approach as GPT-3 , combining full sentences/sequences with GREEDY packing ( first come first concatenate ) and also separation tokens or additional padding . The RoBERTa ablation study shows that mixing of sentences from different documents reduces accuracy , but it is used nonetheless for load balancing reasons which indicates that sorted batching is not sufficient . In summary , the contributions of the paper are as follows . In Section 2 , we produce histograms of the Wikipedia pre-training dataset showing the high percentage of padding tokens . We present two new deterministic and efficient packing algorithms which efficiently pack datasets with millions of sequences in a matter of seconds ( or less ) in Section 3.1 . We empirically show that the proposed packing algorithms produce a nearly-optimal packing scheme In Section 3.2 and Section 3.3 , we explain how the BERT model can be adjusted to show the same convergence behavior on packed and unpacked sequences . In Section 4.2 , we demonstrate that the convergence of the BERT large model on the packed dataset is equivalent to that on the un-packed dataset with 2x throughput increase on the Wikipedia sequence length 512 pre-training dataset . 2 WIKIPEDIA BERT PRE-TRAINING DATASET . BERT is pre-trained using masked-language modelling and next-sentence prediction on a large corpus of Wikipedia articles . Each sequence is composed of one < CLS > token followed by the first “ segment ” of sentences , followed by a < SEP > token , and then finally the second “ segment ” of sentences . Because these “ segments ” are created in sentence-level increments there is no token-level control of sequence length . Furthermore 10 % ( default value , ( Devlin et al. , 2019b ) ) of sequences are intentionally cut short . This leads to significant levels of padding , especially for longer maximum sequence lengths ( see Figure 1 ) . At sequence length 128 ( commonly used in phase 1 of pre-training ) the theoretical speed-up is around 1.2 , at sequence length 384 this increases to 1.7 , and finally at sequence length 512 ( commonly used for phase 2 of pre-training ) it is 2.0 . Despite the widespread use of the Wikipedia dataset for pre-training BERT such histograms have , to the best of our knowledge , not been published previously . This has perhaps lead to the underestimation of the speed-up opportunity available . To put things into perspective , the sequence length 512 dataset contains 8.33 billion tokens , of which 4.17 billion are padding tokens .

Summary: This paper proposes efficient packing methods for training sequences of BERT, such that the 50% of the padding tokens in the Wikipedia dataset is avoided to speed up the training. These methods include shortest-pack-first histogram-packing (SPFHP) and non-negative least-squares histogram-packing (NNLSHP) algorithms, which are shown to be straightforward to implement and have little impact on the performance. Empirical studies show that a near 2x speedup over the vanilla BERT training is achieved by the proposed methods.