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

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

# Cramming: Training a Language Model on a Single GPU in One Day

###### Abstract

Recent trends in language modeling have focused on increasing performance through scaling, and have resulted in an environment where training language models is out of reach for most researchers and practitioners. While most in the community are asking how to push the limits of extreme computation, we ask the opposite question: How far can we get with a single GPU in just one day?

We investigate the downstream performance achievable with a transformer-based language model trained completely from scratch with masked language modeling for a single day on a single consumer GPU. Aside from re-analyzing nearly all components of the pretraining pipeline for this scenario and providing a modified pipeline with performance close to BERT, we investigate why scaling down is hard, and which modifications actually improve performance in this scenario. We provide evidence that even in this constrained setting, performance closely follows scaling laws observed in large-compute settings. Through the lens of scaling laws, we categorize a range of recent improvements to training and architecture and discuss their merit and practical applicability (or lack thereof) for the limited compute setting.

## 1 Scaling Up and Scaling Down

Large-scale training of machine learning models with transformer architectures has lead to ground-breaking improvements in many sub-fields of natural language processing including language understanding and natural language generation ( Vaswani et al., 2017 ; Dosovitskiy et al., 2021 ; Radford et al., 2019 ) . The nowadays accepted (but historically surprising) key behavior of these systems is that they reliably scale – they continuously improve in performance when the number of model parameters and amount of data grow. These increases in performance are well-described by various power laws as studied by Kaplan et al. (2020) . This sets up a dominant paradigm in which scaling is the key to performance improvement ( Sutton, 2019 ) .

The power of scale has set off a race to produce extremely large models, which in turn has created an environment where few researchers or practitioners feel that they are capable of training a language model. The original BERT model Devlin et al. (2019) , which became a cornerstone transformer for many practical applications in natural language understanding, already required a significant amount of computation to train. Yet, the reproduction and improvements in Liu et al. (2019) further increased its performance by cranking up the level of computation by orders of magnitude. As these pre-trained checkpoints became popular for a range of downstream applications ( Wolf et al., 2020 ) , the competition for the largest language model became a focal point for industrial labs. This led to training runs that improved the performance of pretrained language models at the expense of computation at the zettaFLOP scale ( Raffel et al., 2020 ; Yang et al., 2020 ; Zaheer et al., 2021 ) and later at the extremely large yottaFLOP scale ( Brown et al., 2020 ; Black et al., 2022 ; Chowdhery et al., 2022 ; Rae et al., 2022 ) .

Our goal is to turn this trend on its head and investigate how to best scale down language model training and what trade-offs emerge when doing so: What downstream performance can be achieved by a modest researcher when training from scratch with a single GPU for a single day? The ability to train a language model to the performance level of BERT with such modest resources has several interesting implications. For one, if scaled-down model pretraining is a viable analogue of large-compute pretraining, then this opens up a host of further academic investigations that are currently hard to realize for large-scale models. For example, research questions about the differences between existing and new pre-training tasks, tracing model predictions to data points ( Ilyas et al., 2022 ) , security questions such as membership inference ( Carlini et al., 2022 ) and data poisoning ( Geiping et al., 2021 ) , and a wide range of empirical investigations into topics such as stability or generalization that arise during training ( Nagarajan & Kolter, 2019 ; Jiang et al., 2019 ) . At the same time, we can imagine situations in which legal requirements make it unclear whether models trained on public data with uncertain origin are permissible, and where a practitioner is interested in retraining their language models using a specialized or trustworthy data source ( Wilka et al., 2017 ; Gold & Latonero, 2017 ) .

In addition, we are motivated to benchmark the overall conceptual progress of research in this area over the last years, beyond simply turning the scaling knob. The goal of achieving BERT-like performance with modest training resources would have seemed unthinkable in 2018, and yet with modern advances and transformer training techniques this may now be possible.

To answer these questions, we consider a challenge we call ‘‘Cramming’’ -- learning a whole language model the day before the test. Our studies begin by investigating many facets of the training pipeline to see which modifications actually improve performance in the scaled-down scenario. We provide evidence that even in this constrained setting, performance closely follows scaling laws observed in large-compute settings. An unsurprising consequence of these laws is that scaling down is hard; while smaller model architectures enable speeding up gradient computations, overall rates of model improvement over time remain nearly constant. Nonetheless, we can find changes to the training recipe that exploit scaling laws to yield improvements by improving the effective rate of gradient computations without compromising model size. In the end, we are able to train models that achieve respectable performance -- often close to and sometimes exceeding BERT on GLUE tasks -- on a shoestring budget 1 1 1 We provide code to replicate all experiments at https://github.com/JonasGeiping/cramming . .

## 2 Tying our hands behind our back: A setup with limited compute

Before we start this investigation, we want to outline the extent of limitations we are interested in. The rules for cramming are as follows: • A transformer-based language model of arbitrary size is trained with masked-language modeling, completely from scratch.

• Existing pretrained models cannot be included in any part of the pipeline.

• Any raw text (excluding downstream data) can be included for training. This means that one can achieve speedups by making judicious choices about how and when to sample data, provided the sampling mechanism does not require a pre-trained model.

• The downloading and pre-processing of raw data is exempted from the total compute budget. Pre-processing may include CPU-based tokenizer construction, tokenization, and filtering, but cannot include representation learning (e.g. pre-training a word embedding is not allowed, unless it is counted towards the final runtime).

• Training proceeds on a single GPU for 24 hours.

• Downstream performance is evaluated on GLUE ( Wang et al., 2018 ) . Downstream finetuning on GLUE is limited to brief training with only the training data of the downstream task (we consider 5 epochs or less) and needs to work with hyperparameters set globally for all GLUE tasks. Downstream finetuning is excluded from the total compute budget.

In our implementation, we analyze both a setup with a classical rtx2080ti GPU (released September 2018) and separate setups with a more modern rtxa4000 or rtxa6000 GPU (released October 2020). We pair each unit with 4 CPU cores and 32GB of RAM.

Why these limitations? We are principally interested in re-investigating the original BERT setup of Devlin et al. (2019) with limited compute. The optimal architecture of the transformer is not fixed, as the optimal size and shape depends on scaling laws ( Kaplan et al., 2020 ) . The limitations on usage of existing models rule out distillation from an existing model ( Turc et al., 2019 ; Jiao et al., 2020 ; Sun et al., 2020 ; Wang et al., 2020b ; Kaliamoorthi et al., 2021 ) and data filtering based on existing large models ( Golchin et al., 2022 ) , both of which ultimately answer questions about compression and transfer of already processed information. Further, we do not want to limit data to the original dataset used to train BERT, wanting to allow for possible improvements through better data curation and quality. The rtx2080ti GPU is a natural candidate for this experiment, given that it was released before Devlin et al. (2019) , but the more recent rtxa4000 is also interesting, as a more recent consumer-grade workstation variant. Finally we also test the rtxa6000 , being arguably the upper limit of a single-user workstation. At the finetuning stage we want to mimic the original BERT finetuning and evaluation setup, but provide additional limits to prevent gains based on tuning of only the downstream procedure, for example via computationally extensive downstream training ( Bahri et al., 2021a ) , use of multiple downstream datasets (for example continued pretraining with MNLI before finetuning other tasks ( Izsak et al., 2021 ) ), and extended hyperparameter optimization for each GLUE task ( Devlin et al., 2019 ; Liu et al., 2019 ; Lan et al., 2019 ) .

## 3 Related Work on Efficient Transformers

##### How long does it take to train BERT?

In general, this question is hard to answer, due to wildly varying hardware and software setups and differing measures of efficiency ( Dehghani et al., 2021 ) . An upper bound on the compute of a training run can be established by finding the total number of (low-precision) floating point operations available over the wallclock budget of the run. This peak of total FLOPs in a given time interval is generally not reached in actual compute, even for highly optimized models ( Chowdhery et al., 2022 ) , but represents the paid budget required to realize a training run. We summarize budgets for a few select training runs in Table 1 . After the original training run for BERT on TPUs, initial reactions estimated up to 11 days of compute for comparable results on GPUs ( Dettmers, 2018 ) . However, sustained improvements, especially in software, have reduced the upper limit significantly ( You et al., 2019 ; Narasimhan, 2019 ) . Yet, recipes and implementations generally require entire server nodes (for GPUs) or TPU slices and target larger BERT architectures.

Other work discussing improvements to BERT targets compute settings closer to the original BERT, for example SqueezeBERT ( Iandola et al., 2020 ) employs 8 Titan RTX cards for four days. Sellam et al. (2022) note that the original BERT training run is an outlier and doubling its training time more reliably reproduces the original results.

Our central point of comparison for BERT training with limited resources is the work of Izsak et al. (2021) who also attempt the goal of training BERT within 24 hours with overall similar limitations, but use a full server node with 8 V100 GPUs. Izsak et al. (2021) choose a BERT LARGE \text{BERT}_{\text{LARGE}} architecture variant and train with sequence length of 128 128 , including a range of tweaks such as modified learning rates schedules, large batch sizes, sparse prediction and packed sequences. We re-evaluate this setup as a baseline setting for our own compute budget (which is about 15x smaller).

##### Studies of Efficient Transformers

Recent years have seen a flurry of research working to improve and modify the transformer architecture proposed in Vaswani et al. (2017) and we refer to Treviso et al. (2022) for a recent categorization and review of research in this area. Several meta-studies have investigated proposed improvements and modifications: Narang et al. (2021) evaluate a large range of architectural modifications applied to the T5 model pipeline of Raffel et al. (2020) on tasks in both language understanding and translation. The encoder-decoder structure of T5 is closer in spirit to the original transformer setup, but is understood to behave similarly to BERT when using the encoder component ( Liu et al., 2021a ) . Evaluating modifications with 1.75 days of compute on TPU slices they find that most improvements do not reliably materialize gains in final accuracy. Tay et al. (2021) work in the same setting and evaluate the optimal shape of T5 derived architectures and its relative effects on downstream performance as models are scaled. Further exploration of the scaling behavior of various architectural improvements in Tay et al. (2022a) find that only few modifications outperform the original architecture of Vaswani et al. (2017) at all scales, especially when evaluating downstream accuracy. The meta-study investigating improvements in preparation for extreme-scale training in Scao et al. (2022) focuses on minor modifications to layout, positional embeddings and data sources for autoregressive models, and other extremely-large scale training runs have so far been similarly conservative in their settings ( Brown et al., 2020 ; Black et al., 2022 ; Rae et al., 2022 ) .

In general though, these evaluations target larger compute settings than we intend to use, and are concerned with whether improvements (often from academic sources and proposed with evaluations on small scales) translate to larger scales. In this work, we set aside the question of (up)scaling and focus only on the limited compute.

##### Scaling Laws

The difficulty in finding tangible improvements is echoed in the scaling laws of Kaplan et al. (2020) . Over a wide range of transformer model shapes, Kaplan et al. (2020) find only model size (as number of parameters in non-embedding layers) strongly predicts performance. Further, for a fixed compute budget, an optimal model size can be derived, but performance is only mildly connected to model size - larger models processes less data per unit of compute, but improve faster by almost the same margin. While the precise coefficients and shape of these scaling laws continue to be iterated on ( Hoffmann et al., 2022 ) and adapted for related settings ( Bansal et al., 2022 ; Clark et al., 2022 ; Bahri et al., 2021b ) , their overall logic appears hard to escape, even if power laws fit observations somewhat less well on small scales.

## 4 Investigations

For our experimental evaluation we implement and test a considerable number of proposed modifications to the setup of Devlin et al. (2019) for their merits in our limited compute setting as described in Section 2 . We first clarify the common implementation and initial data setup, and then investigate architectural, training and dataset improvements.

### 4.1 Implementation Details

We implement everything in PyTorch ( Paszke et al., 2017 ) and to limit our gains from the ”software lottery” ( Hooker, 2021 ) we do not use specialized implementations, which would further bias results towards well-established components. We keep everything on the implementation level of the PyTorch framework, allowing only automated operator fusion ( Sarofeen et al., 2022 ) that can be applied to all components. Only after choosing a final architecture variant, we then re-enable the efficient attention kernel described in Dao et al. (2022) . We run all experiments and ablation studies with the same setup of automated mixed precision ( Micikevicius et al., 2018 ) for standard 16- and 32-bit floating point precision (over full 32-bit float, scaled 16-bit ( Rasley et al., 2020 ) and pure bfloat16 ( Wang & Kanwar, 2019 ) . We find no benefit from offloading ( Ren et al., 2021 ; Rasley et al., 2020 ) in our setting.) .

##### Initial Data Setup