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: The paper proposed to pack sequences instead of padding to reach the max sequence length for each sample for BERT pretraining. To achieve similar optimization results with the original training process, the proposed method modified positional embedding, additional attention masks and other optimization hyper-parameters. On the Wikipedia dataset, the proposed method can achieve 2x speedup while achieving similar training loss.