Cramming: Training a Language Model on a Single GPU in One Day — part 2/3

Cleaned-text code points [15788, 31071). Bibliography omitted here; source.md includes it.

We start our investigation with a close analogue to the original raw text sources of Devlin et al. (2019) , using a recent dump of the English Wikipedia ( 20220301.en ) and English bookcorpus , noting the commentary of Tan (2019) ; Bandy & Vincent (2021) . We force all text into lower-case, strip accents and non-ascii characters and create an English tokenizer from scratch based only on this data. We choose WordPiece with a vocabulary size of 2 15 = 32768 2^{15}=32768 ( Wu et al., 2016 ) . We found no significant change in performance with BPE ( Sennrich et al., 2016 ) or SentencePiece with Unigrams ( Kudo, 2018 ; Kudo & Richardson, 2019 ) . Smaller vocabulary sizes ( 2 12 , 2 13 , 2 14 2^{12},2^{13},2^{14} ) resulted in worse performance, while larger vocabulary sizes ( 2 16 2^{16} ) we not reliably better . We pack tokenized data into randomized sequences of length 128 128 and separate unrelated fragments by <sep> The performance impact from dropping this separator was minimal. No impact was observed from including a <cls> token in pretraining . The shorter sequence length is sufficient for the downstream applications that we are targeting and simplifies attention computations. Packing data into full sequences limits us to simpler sequence losses, but uses the available compute optimally Liu et al. (2019) ; Izsak et al. (2021) . For the targeted compute settings, this sequence length results in micro-batch sizes of 64 64 to 96 96 for most variations of the base BERT architecture on the gtx2080ti , which we will accumulate into larger batch sizes. With our limited compute budget, this produces enough samples to run single-epoch training ( Komatsuzaki, 2019 ; Hernandez et al., 2022 ) where no data point is revisited.

### 4.2 Modifying the Architecture

The most obvious way to efficiently scale down training is by modifying the model architecture; intuitively, it seems likely that smaller/lower capacity models will be optimal in the cramming regime. In this section, we study the relationship between model type and training efficiency. We see that scaling laws create a strong barrier to scaling down. Per-token efficiency of training depends strongly on model size, but not transformer type. Furthermore, smaller models learn less efficiently, and this largely mitigates any throughput gains. Fortunately, the fact that training efficiency is nearly constant across models of the same size means that we can boost performance by finding architecture modifications that speed up gradient computation while keeping the parameter count nearly constant. This makes architecture selection fairly straightforward as we can make design choices based primarily on how they affect computation time for a single gradient step.

##### Scaling laws hold in the low-resource regime

A large corpus of research in recent years has developed architectural improvements to speed up the original transformer. Many of these methods have not been found to improve training for the large-scale T5 architecture Narang et al. (2021) ; Tay et al. (2022a) . But, in the low compute setting where data throughput is of utmost importance, maybe this is the way forward? Scaling laws have been observed by Kaplan et al. (2020) in the high-resource regime, and seem to hold strongly in the limit as resources grow. Surprisingly, these laws also hold in the limit of extreme compute down-scaling, and they create a barrier to low-cost training.

We exemplify the effect of scaling laws for many transformer variants from the literature in Figure 1 , where we train each architecture variant with optimized training hyperparameters as described below in Section 4.3 . We apply these architecture variants to a shared baseline model that incorporates Pre-Normalization and rotary embedding. Figure 1 visualizes the progress of MLM loss versus the number of tokens ingested in total and all architectures run with the same time budget.

We observe that varying the transformer type and size has only minimal impact on the final loss after 24 hours. Models with more parameters learn more efficiently, as their MLM loss decreases faster on a per-gradient basis. However, smaller architectures make up for their slower learning efficiency by higher throughput, and thus process more tokens over the limited budget. Figure 1 shows that different architectures are unpredictable throughout an initial stage of training (the first 1B tokens), after which the per-token efficiencies differ by only a multiplicative constant (a horizontal shift due to the log axis). This constant depends almost entirely on the model size, not model type, so that all choices reach a MLM loss around 1.9 1.9 at the end of training.

##### Exploiting the scaling law.

The scaling laws seem to bar us from making large gains via major changes to the transformer size and type, as per-token performance is tightly coupled to model size. As a result, we find no improvements when using a funnel-transformer architecture ( Dai et al., 2020 ; Nawrot et al., 2022 ) , when dropping FFN layers ( Sridhar et al., 2022 ) , or when using recurrent layers ( Lan et al., 2019 ) , even when trained with BPTT as in Schwarzschild (2021) . Rescaling architectures to be deep-narrow ( Tay et al., 2021 ; Wies et al., 2021 ) provides no gains .

While this principle closes one door for scaling down efficiently, it opens another; Because per-gradient efficiency remains nearly constant for all models of the same size, we can exploit scaling laws by quickly searching for architectural choices that speed up computation while keeping model size roughly constant. A number of obvious optimizations fall into this category, and we describe them below, in addition to several other tweaks that provide marginal but worthwhile/free gains.

##### Attention Block:

We disable all QKV biases ( Dayma et al., 2021 ) . This exploits the scaling law by removing a layer of computation, making the forward and backward pass somewhat faster, while keeping the model size nearly constant. We find that we could decrease gradient costs by reducing the number of attention heads ( Merity, 2019 ; Araabi & Monz, 2020 ; Liu et al., 2021b ; Javaheripi et al., 2022 ) , as this parallelizes better on the GPU and provides a slight performance boost. Yet, reducing the amount of heads also decreases finetuning performance, so we ultimately keep all 12 12 heads. We find no benefits from replacements to the softmax operation ( Richter & Wattenhofer, 2020 ) . We further keep the original multi-head self-attention mechanism. A large amount of work has been focused on efficient attention ( Sukhbaatar et al., 2019 ; Beltagy et al., 2020 ; Wang et al., 2020a ; Liu et al., 2021c ) and studies of efficient attention ( Tay et al., 2020a ; Tay et al., 2020b ) . But, because we set the maximal sequence length to 128 128 , attention complexity is less of a concern in our setting. To verify this, we implement the recently proposed FLASH mechanism ( Hua et al., 2022 ) , but find no benefits. We further experiment with Fourier attention as proposed in Lee-Thorp et al. (2021) , but find no improvements. We find rotary embeddings ( Su et al., 2021 ; Black et al., 2022 ) , to provide small benefits, but these are evened out by the drop in speed, so we ultimately decide against these.

##### Feedforward Block:

We find empirical gains from disabling all linear layer biases ( Dayma et al., 2021 ) . Just as for the attention layers, this leverages the scaling law by accelerating gradient computation without noticeable impacts on model size. As a result, we get higher throughput without compromising the rate at which the model improves. We keep the original feedforward block largely unchanged, finding no benefits from changing to another activation than GELU. We do see small improvements from re-ordering the block into a gated linear unit ( Dauphin et al., 2017 ) . In contrast to other work, e.g. ( Black et al., 2022 ) , we do not increase the number of parameters in the FFN block to compensate for the halving of the hidden dimensionality due to gating.

##### Embedding:

We implement scaled sinusoidal positional embeddings as described in Hua et al. (2022) , finding incremental benefits over learned or unscaled sinusoidal embeddings. We see no improvements from decoupling the input and output embeddings ( Chung et al., 2020 ) . The suggestion from Lan et al. (2019) to factorize the input embedding provides no gains in our setting. We include a layer normalization at the end of the embedding block.

##### Layer Structure:

As observed in many studies, we find that pre-normalization with Layer Norms is beneficial over post Layer Norms ( Baevski & Auli, 2018 ; Xiong et al., 2020 ) . We see no additional benefit from other variants of this modification, such as ( Liu et al., 2020b ; Shleifer et al., 2021 ) . Further, replacing Layer Normalization with RMS Normalization provides no gains ( Zhang & Sennrich, 2019 ) . We note that the key effect of pre-normalization is to stabilize training and enable larger learning rates and reduced warmup, and we see limited benefits from including it by itself. We see no benefits from stochastic dropping of entire layers as described in ( Zhang & He, 2020 ) .

##### Head Block:

We find that we can remove the nonlinear head without ill effect. We can further drop the decoder bias ( Radford et al., 2019 ) and gain in memory using sparse token prediction ( Liu et al., 2019 ; Izsak et al., 2021 ) . We add a final Layer Norm to stabilize training further.

### 4.3 Modifying the Training Setup

We study the impact of training hyper-parameters on the BERT-base architecture. The original BERT training recipe understandably results is poor model performance in the cramming setting, and so we revisit a number of standard choices.

##### Objective:

We train with only masked language modeling on fully packed blocks of tokens with a masking rate of 15% and the original setup of Devlin et al. (2019) where 10 % 10\% of all masks are filled with random words and 10% unchanged. We see no improvement from masking at larger rates, e.g. at 40% as proposed in ( Wettig et al., 2022 ) , see Appendix. We see no difference enabling or disabling the mentioned 20% rule. We evaluate other functions for the masked-language objective, such as mean-squared error ( Hui & Belkin, 2021 ) or L1 loss, but find no benefits.

##### Choice of Optimizer:

We keep Adam ( Kingma & Ba, 2015 ) as the optimizer of choice, with weight decay of 0.01 0.01 as described in ( Loshchilov & Hutter, 2017 ) , β 1 = 0.9 , β 2 = 0.98 \beta_{1}=0.9,\beta_{2}=0.98 and ε = ​ 10 − 12 \varepsilon=${10}^{-12}$ . To stablize training at no extra cost, we include gradient clipping at a clip value of 0.5 0.5 . We find no noticeable change in varying these parameters in reasonable amounts, e.g. ε = ​ 10 − 6 \varepsilon=${10}^{-6}$ , β 1 = 0.9 , β 2 = 0.999 \beta_{1}=0.9,\beta_{2}=0.999 . We test other first-order adaptive optimizers ( Shazeer & Stern, 2018 ; Liu et al., 2020a ) but find no advantages in our setting. We further find no advantages using higher-order optimizers ( Yadav, 2020 ; Anil et al., 2021 ) , but note that especially for higher-order optimizers there is a greater amount of variability in implementation.

##### Learning Rate Schedule and Peak:

Following the advice of Izsak et al. (2021) , we re-scale the learning rate schedule so that it is tied to our budget and the learning rate decays as the budget reduces to zero. Interestingly, we observe in Figure 2 that while globally a large number of learning rate shapes lead to similar reductions in loss, we find that we can make some gains through the choice of schedule. We find that a simple one-cycle learning rate ( Smith & Topin, 2018 ) with a peak learning rate of ​ 10 − 3 {10}^{-3} leads to minimal pretraining loss within our budget.

##### Batch Size Schedule:

A particularity of our setting is that, due to being limited to a single GPU, the micro-batch size that finds its way onto this GPU ( 96 96 for most experiments) is several times smaller than the optimal batch size. We find that the optimal batch size in this setting is around 1536 1536 for minimal pretraining loss, but 4032 4032 for maximal downstream performance for the 2080ti, i.e. we accumulate gradients and only perform an update every 16 and 42 forward/backward passes, respectively. For the larger A4000 and A6000 cards, this corresponds to a micro-batch size of 128 / 256 128/256 and final batch size of 4096 4096 , which we again accumulate.

Fortunately, we can find small speedups by using an aggressive batch size schedule; we increase the number of averaged micro-batches linearly over the course of training. This results in more progress earlier in training, and leads to a small benefit to performance. We also experiment with automatic and adaptive batching rules ( De et al., 2017 ; Bollapragada et al., 2018a ; Bollapragada et al., 2018b ) , but find that the best results from these adaptive schedules resemble the fixed linear schedule. For simplicity we just stick to the simpler linear schedule.

##### Dropping Dropout

The original BERT model of Devlin et al. (2019) includes dropout as in Vaswani et al. (2017) , which prevents overfitting when training data is small relative to total compute budget. While it can be helpful as a regularizer, dropout effectively reduces the number of gradient updates seen by each parameter, as updates do not occur when the associated feature is dropped. At the same time, update runtime is not strongly effected by the presence of dropout, and so dropout results in a net reduction in updates per second.

In the cramming setting, training data is large compared to compute. Overfitting is not possible due to the single epoch schedule, and we disable dropout during pretraining ( Brown et al., 2020 ) to maximize the number of parameter updates. We re-enable dropout during downstream fine-tuning with a dropout value of 0.1 0.1 . Further, we experiment with length curricula ( Li et al., 2022 ) (see appendix) and token dropping ( Hou et al., 2022 ) , but find no gains in our setting.

### 4.4 Optimizing the Dataset

We found above that scaling laws create a barrier to making major gains (beyond computational efficiencies) with architectural modifications. However, scaling laws do not preclude us from training on better data. Once we have exhausted our ability to train on more tokens per second, we should seek to train on better tokens.

We consider two data based pathways to better down-scaling. First, we can filter, process, or sort the existing data in various ways. Second, we can swap our data source. To this end, we experiment with several subsets of The Pile ( Gao et al., 2020 ) , containing raw text from only Gutenberg , Books3 and Wikipedia (en) . From these Pile datasets we tokenize the first 4 × 10 6 4\text{\times}{10}^{6} entries to generate enough tokens for our single pass. Another popular source of data is C4, the colossal, cleaned version of Common Crawl ( Raffel et al., 2020 ) , from which we stream the first 20 × 10 6 20\text{\times}{10}^{6} entries. For each data source we regenerate its own WordPiece tokenizer as described in Section 4.1 .