Packing: Towards 2x NLP BERT Acceleration

1 Introduction . Since its introduction in 2019 , BERT [ 5 ] 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 [ 20 ] . Available sizes range from tiny ( 2 layers with hidden size 128 ) to large ( 24 layers with hidden size 1024 ) [ 6 , 5 ] . The introduction of ALBERT [ 14 ] 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 some methods to enable researchers to accelerate the pre-training of BERT by as much as 2x . The de-facto pre-training dataset Wikipedia , as well as many other NLP datasets , show a positively 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 . Overall , the sample lengths range between 5 tokens up to 512 ( see Figure 1 ) . Samples of length 512 represent only 23.5 % of the dataset . While processing the padding tokens wastes compute , it is 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 . The most obvious way to reduce the extent of padding Preprint . Under review . in the dataset is to group samples by size before batching , i.e. , process the shorter samples together and longer samples together . Typically such an approach would still involve padding but less than if padding all sequences to the same maximum length . For example BERT [ 5 ] 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 ) ) . Another example of this approach is Faster Transformer [ 18 ] which groups samples of similar size together in one batch and fills up with padding only to the maximum length in this batch . 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 [ 4 ] , 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 intensive 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 . Third , each batch now contains a consistent number of real tokens . While we demonstrate the effectiveness of packing specifically on the Wikipedia dataset , packing SQUaD [ 19 ] or GLUE datasets [ 22 , 23 ] for BERT also leads to significant speed-ups ( some in excess of 9x ) [ 1 ] ( sections A and B ) . 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 [ 7 ] . 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 . 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 instance , GPT-3 [ 3 ] is trained with a maximum sequence length of 2048 where a single sequence may contain multiple segments separated by a special end of segment token . The last segment in each sequence is simply cut to meet the sequence length requirement . 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 ) to prevent cross-contamination between examples in a pack . This ensures that the characteristics of the original dataset and model are matched as closely as possible . 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 packing algorithms which easily pack datasets with millions of sequences in a matter of seconds ( or less ) . We empirically show that the proposed packing algorithms produce a nearly-optimal packing scheme for Wikipedia pre-training dataset . We show how to compute the per-sequence loss by inexpensively un-packing the loss . We provide code for building an attention mask which prevents attention between tokens of different sequences in the pack . We demonstrate that the convergence of the BERT large model on the packed dataset is equivalent to that on the un-packed dataset . We show that with the packed dataset , we are able to achieve a nearly 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 [ 5 ] . Each sequence is composed of one < CLS > token followed by the first part of sentences , followed by a < SEP > token , and then finally the second part of sentences . Because parts are created in sentence-level increments there is no token-level control of sequence length . Together with already short parts , empirically , 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 . 3 Methods . Our approach consists of three distinct components . Firstly , we pack the data efficiently during preprocessing to make full use of the sequence length ( Sections 3.1.1 and 3.1.2 , see also [ 1 ] Section D ) . Secondly , we adapt the self-attention mask to prevent the model from attending between different sequences in the same pack ( Section 3.2 ) . Other components of the model , such as the feed-forward layer [ 21 ] , operate on a per-token basis and do not require any modification . Thirdly , we compute the loss and accuracy on a per-sequence basis to match the canonical BERT implementation ( Section 3.3 ) . This is achieved by unpacking the per-pack loss at the framework level , without the use of custom kernels . Additionally , we provide suggestions for hyperparameter adjustment ( Section 3.4 ) that lead to analogous convergence behavior between the packed and un-packed BERT implementations . 3.1 Packing algorithms . The problem of optimally concatenating multiple sequences of different length until a maximum combined length is reached can be directly framed as a bin-packing or stock cutting problem . Since an exact solution is strongly NP-complete [ 13 ] , we propose two new heuristic algorithms that are tailored to this particular instance . A detailed introduction to packing is provided in [ 1 ] ( Section D ) . 3.1.1 Shortest-pack-first histogram-packing ( SPFHP ) . Shortest-pack-first histogram-packing ( SPFHP ) consists of three main components . First , the packing algorithm works on the bins in the sequence length histogram ( with bin size 1 ) rather than the individual samples . Second , we operate on the sorted data from longest to shortest sequences . This comes basically for free due to the use of histograms . Third , we apply the worst-fit algorithm [ 11 , 26 ] onto this histogram , where the currently observed sample goes to the pack1 that has the most space left to reach maximum packing depth ( “ shortest-pack-first ” ) . If the sample does not fit , a new pack is created . A variant is to limit the packing depth , in other words the maximum number of sequences that are allowed in a pack . Therefore , we only extend an existing pack if it is not already at maximum packing depth . The detailed code for the algorithm is provided in [ 1 ] ( listing 3 ) . 3.1.2 Non-negative least squares histogram-packing ( NNLSHP ) . The proposed NNLSHP algorithm is based on re-stating the packing problem as a ( weighted ) nonnegative least squares problem ( NNLS ) [ 2 ] of the form wAx = wb where x 0 . The vector b is the histogram containing the counts of all the sequence lengths in the dataset . Next , we define the A matrix ( the “ packing matrix “ ) by first generating a list of all possible sequence length combinations ( “ strategies ” ) that add up exactly to the maximum sequence length . We focus specifically on strategies that consist of at most 3 sequences per pack ( independent of b ) and encode each strategy as a column of the sparse matrix A . For example , a strategy consisting of the sequence length 128 , 128 , and 256 in represented a column vector that has the value 2 at the 128th row , the value 1 at the 256th row , and zero at all other rows . The variable x describes the non-negative repetition count for each strategy . So a 24 in the ith row of x means that the strategy represented by the ith column of A should repeat 24 times . Moreover , in the un-weighted setting , Ax = b states that we would like to “ mix ” the pre-defined strategies ( columns of A ) such that the number of samples matches the histogram b , and where each strategy is used x 0 times . We use the residual weight w to control the penalization of the Ax b residual on different sequence lengths ( different rows of b ) . Heuristically , we set the weight of 0.09 for all sequences of length 8 or smaller because they are considered acceptable padding sequences . All other sequence lengths get weight 1 . After solving wAx = wb for x 0 using an off-the-shelf solver we obtain a floating point solution , which means that the repetition counts are not necessarily integers . Since we can not use a non-natural number of strategies , we round the solution x̂ to the nearest integer . The error introduced by this rounding is found to be negligible . Further details are provided in [ 1 ] ( Section D.4 ) .

Summary: This paper proposes a simple and effective approach to accelerate the standard BERT pretraining by packing multiple sentences to a single training sample with fixed length, in order to reduce padding tokens and fully utilize the computation. The packing algorithms itself are very efficient and the overhead introduced in training (attention mask and loss calculation) is limited (below 5%). Experiments show that with hyperparameter adjustment, the BERT pretrained with packed data can result in nearly 2x speed up. Compared to the unpadding approach with custom computational kernels, the packing approach is inherently load-balanced and thus scale better to multiple distributed accelerators. Overall, I believe this is a quite general approach that can be applied to large-scale training with data having variable lengths.