What Language Model to Train if You Have One Million GPU Hours?
The BigScience
Group
Abstract
The crystallization of modeling methods around the Transformer architecture has been a boon for practitioners.
Simple, well-motivated architectural variations that transfer across tasks and scale, increasing the impact of modeling research.
However, with the emergence of state-of-the-art 100B+ parameters models, large language models are increasingly expensive to accurately design and train.
Notably, it can be difficult to evaluate how modeling decisions may impact emergent capabilities, given that these capabilities arise mainly from sheer scale.
Targeting a multilingual language model in the 100B+ parameters scale, our goal is to identify an architecture and training setup that makes the best use of our 1,000,000 A100GPU-hours budget.
Specifically, we perform an ablation study comparing different modeling practices and their impact on zero-shot generalization.
We perform all our experiments on 1.3B models, providing a compromise between compute costs and the likelihood that our conclusions will hold for the target 100B+ model.
In addition, we study the impact of various popular pretraining corpora on zero-shot generalization.
We also study the performance of a multilingual model and how it compares to the English-only one.
Finally, we consider the scaling behaviour of Transformers to chose the target model size, shape, and training setup.
Introduction
Recent years have seen the advent of large language models with emergent capabilities (e.g., zero-shot generalization) arising from sheer scale alone (Radford et al., 2019; Brown et al., 2020). Scaling
LLMs produces a predictable increase in performance; simple scaling laws connect the number of parameters, pretraining dataset size, and compute budget (Kaplan et al., 2020; Ganguli et al., 2022),
[Figure]
providing a clear path towards more capable models.
This paradigm shift can be attributed to the wide adoption of the Transformer (Vaswani et al., 2017), providing a scalable basis for practitioners to build upon.
In this paper, we design an architecture and training setup for a multilingual 100B+ parameters model, seeking to best use a fixed 1,000,
A100-hours budget.
Many promising additions and tweaks have been proposed to the Transformer architecture (Press et al., 2022; Shazeer, 2020;
Dettmers et al., 2021)
but there is a lack of systematic investigation to identify which will best transfer to increased scales, and which are idiosyncrasies of specific setups.
Contributions.
We study the impact of pretraining corpora, positional embeddings, activation functions, and embedding norm on zero-shot generalization.
We base our study on the popular GPTarchitecture (Brown et al., 2020), with experiments at the 1.3B parameters scale.
We consider the impact of massive multilinguality, and, finally, we study the scaling of our models, and draft an architecture for a 176B model
Methods
We first justify our choice to base our model on the popular recipe of combining a decoder-only model with an autoregressive language modeling objective, and introduce our experimental setup.
We then discuss our evaluation benchmarks, and motivate our choice of zero-shot generalization as our key metric.
Finally, we introduce the baselines we compare to throughout the paper.
2. Architecture and Pretraining Objective
In this paper, we base all models on a decoder-only
Transformer pretrained with an autoregressive language modeling objective.
This is a popular choice for large language models (Brown et al., 2020;
Rae et al., 2021; Thoppilan et al., 2022), possibly because it lends itself to zero-shot application to many downstream tasks (Radford et al., 2019).
Alternatives include encoder-decoder models trained with a span-corruption objective (e.g., T5 Raffel et al.
(2019)), as well as non-causal decoders models with visibility over a prefix (so-called PrefixLMs,
Liu et al.
Dong et al.
Our decision is motivated by the findings of Anonymous (2022), which showed that decoderonly models combined with an autoregressive language modelling objective provide the best zeroshot generalization abilities right after pretraining.
Although multitask finetuning (Sanh et al., 2021; Wei et al., 2021) will favor an encoder-decoder with span corruption for best zero-shot generalization, they also found a compromise between these two practices.
Following autoregressive pretraining, decoder-only models can be efficiently adapted into non-causal decoders, simply by extending pretraining with span corruption.
This adaptation produces a second model, which can provide excellent zero-shot generalization after multitask finetuning.
Accordingly, we follow their recommendation, and train an autoregressive decoder-only model which we will later consider adapting and finetuning.
2. Experimental Setup
We follow the architecture and hyperparameters of GPT-3 (Brown et al., 2020).
For learning rate, we use a maximum value of 2 × 10 − 4 , with a linear warm-up over 375M tokens, followed by cosine decay to a minimum value of 1 × 10 − 5 .
We use a 1M tokens batch size, with linear ramp-up over the first 4B tokens, and a sequence length of 2,048.
For optimization, we use Adam (Kingma and Ba, 2014), with β 1 = 0 .
9 , β 2 = 0 .
999 , ϵ = 1 × 10 − 8 , weight decay 0.1, and gradient clipping to 1.0.
We also tie the word embedding and softmax matrix (Press and Wolf, 2017).
Unless noted otherwise, we conduct our experiments with 1.3B parameters models, pretraining on 112B tokens.
We picked this size and dataset size as a compromise between compute cost and the likelihood that our conclusions would transfer to the target 100B+ model.
Notably, we needed to be able to reliably measure zero-shot generalization above random chance.
We note that training for 112B tokens our models bring them significantly above the optimality threshold identified by Kaplan et al.
(2020), and was enough to yield non-random chance zero-shot results.
The main architectural difference with GPT-3 is that all our layers use full attention, while GPTuses alternating sparse attention layers (Child et al., 2019).
The main value of sparse attention layers is to save compute with long sequence lengths.
However at the 100B+ scale, sparse attention layers provide negligible compute savings, as the vast majority of the compute is spent on the large feedforward layers.
Kaplan et al.
(2020) estimated the amount of compute per token to be:
[Equation]
where C forward is the cost for the forward pass, n layer is the number of layers, d is the hidden dimension, and n ctx is the sequence length.
This means if 12 d >> n ctx , the second n layer n ctx d term is negligible, which is the case for our final model where d > 10 , 000 and n ctx = 2048 .
FLOPS or FLOPS?
We report throughput per
GPU in FLOPS and total budgets in PF-days (i.e. one PFLOPS sustained for a day).
It is important to highlight that FLOPS are never directly measured, but always estimated, with widely different practices across papers.
We refer to model FLOP the estimates based the C = 6 ND formula from Model Parameters
Pretraining tokens
[Table] [Table] Babbage [Table]
— GPT
[Table]
Kaplan et al.
(2020), where C is the total compute,
N the model size, and D the number of tokens processed.
These are the FLOP actually used to train the model, and which are used for scaling laws.
We refer to hardware FLOP the estimates reported by our codebase, using the formula from
Narayanan et al.
This notably includes gradient checkpointing and a more thorough accounting of operations depending on model shape.
2. Evaluation Benchmarks
We measure upstream performance with the language modeling loss on an held out sample of the pretraining dataset.
For downstream performance, we could use zero/few-shot generalization, with or without task-specific finetuning.
Upstream performance is not always aligned with downstream performance, and they are sometimes even at odds (Tay et al., 2021).
It is also not always possible to compare losses across architectures, objectives, and tokenizers.
Given that we mainly care about downstream task performance, we base our evaluation on a diverse set of tasks.
Specifically, we choose to measure zero-shot generalization on this set of tasks.
Few-shot and zero-shot results are strongly correlated: we found a Pearson correlation coefficient of 0.93 between zero-shot and few-shot performance across model sizes in Brown et al.
We do not rely on finetuning as it is not how the main final model is likely to be used, given its size and the challenges associated with finetuning at the 100B+ scale.
We use the popular EleutherAI Language Model
Evaluation Harness (EAI harness, Gao et al.
(2021)), evaluating models across 27 diverse tasks that are similar to those used in Brown et al.
(2020) (see Appendix A for details).
Prompts used in the
EAI harness reproduce as closely as possible the evaluation setup of GPT-3.
2. Baselines
We use GPT-Neo (Black et al., 2021), a 1.3B decoder-only auto-regressive language model trained on the Pile (Gao et al., 2020), and GPT3 (Brown et al., 2020), which we access via the
OpenAI API.
We evaluate two models, Babbage and Curie.
We believe Babbage is 1.3B while Curie is 6.7B based on how close our computed results are to those reported in the original paper.
Impact of Pretraining Data
We first study the impact of pretraining data on zero-shot generalization.
More diverse pretraining data, ideally curated from a cross-domain collection of high-quality datasets, has been suggested to help with downstream task performance and zeroshot generalization (Rosset, 2020; Gao et al., 2020).
3. Corpora
We evaluate three possible corpora, all commonly used to train large language models:
• OSCAR v1 (Ortiz Suárez et al., 2019) 2 , a multilingual, filtered version of Common Crawl; • C4 (Raffel et al., 2019), specifically its replication by AllenAI, a processed and filtered version of Common Crawl; • The Pile (Gao et al., 2020), a diverse pretraining corpus that contains webscrapes from Common Crawl in addition to high-quality data from cross-domain sources such as academic texts and source code.
For each pretraining corpus, we train a 1.3B parameter model for 112B tokens.
For the Pile specifically, motivated by good early results at 112B tokens, we train up to 300B tokens, to compare with GPT-3 models and validate against GPT-Neo.
3. Results
Evaluation results are outlined in Table 1.
We find that training on the Pile produces models that are better at zero-shot generalization, with C4 a close second, and OSCAR significantly behind.
Importantly, this finding transfers to larger scales: as part of engineering test runs, a 13B model was trained on OSCAR for 300B tokens.
We found this 13B model to underperform the 6.7B model from OpenAI API which we attribute to the low quality of the English data in OSCAR.
We also note that our model trained on The Pile outperforms the 1.3B GPT-Neo trained on the same dataset.
Finally, our 1.3B model still underperforms the 1.3B model from the OpenAI API by by 1.6%.
It seems most likely that the difference is that of data, but we cannot investigate this further as the GPT-3 training dataset is neither publicly available nor reproducible.
Finding 1.
Diverse cross-domain pretraining data combining web crawls with curated highquality sources significantly improves zeroshot generalization over pretraining datasets constructed from Common Crawl only.
Architecture Ablations
We now consider ablation studies to better identify the best positional embedding, activation function, and embedding normalization placement.
We note that considerations around pretraining objectives and architecture were taken from Anonymous (2022) rather than independently replicated (section 2.1).
4. Positional Embeddings
One aspect of the Transformer architecture that has attracted recent significant interest is the way position information is captured within the model.
Positional embeddings are important because without positional embeddings, Transformers cannot order tokens against one another.
Background
The Transformer paper Vaswani et al.
(2017) proposed two options: static sinusoidal position embeddings and learned position embeddings (i.e., the position of each token is associated with a learned embedding vector).
Learned position embeddings are popular in large language models, and are used for GPT-3.
Su et al.
(2021) later proposed the rotary position embedding method, where the query and key representations inside the self-attention mechanism is modified such that the attention computation captures relative distances between keys and queries.
Recently, Press et al.
(2022) proposed a position method which does not use embeddings, and instead directly attenuates the attention scores based on how far away the keys and queries are.
Results
We compare learned, rotary, and ALiBi position embeddings, and include a baseline without position embeddings.
Our results are presented in Table 2.
Although learned positional embeddings outperforms rotary embeddings, ALiBi yield significant better results than all other alternatives.
We also confirm the discovery of Biderman (2021), that the baseline without explicit position information shows competitive performance.
While bidirectional models require positional embeddings to determine the location of tokens, we find autoregressive models can simply leverage the causal attention masking.
Finding 2.
ALiBi positionnal embeddings significantly outperforms other embeddings for zero-shot generalization.
4. Activation Functions
Background.
Large language models by and large still mostly use the GELU activation (Hendrycks and Gimpel, 2016).
We evaluate a recently proposed alternative, SwiGLU (Shazeer, 2020), which combines both Gated Linear Units
Positional Embedding
[Table]
trained on for (Dauphin et al., 2016) with the Swish activation function (Ramachandran et al., 2017).
With x the input of the layer, we have:
[Equation]
SwiGLU uses 50% extra parameters in the feed forward layers.
As suggested in Shazeer (2020), we compensate for this by reducing the hidden size of the feed forward layer.
Results.
We present our results in Table 3.
SwiGLU produces slightly better results than
GELU; however this comes at a cost of reducing the throughput by approximately a third.
This overhead may primarily be associated with the change in the hidden size of the feedforward network.
Indeed, this new size, 5,456, is divisible by neither the warp size of the GPU (Lashgar et al., 2013) nor the number of streaming multiprocessors, resulting in both tile and wave quantization.
Further hyperparameter tuning could yield better throughputs.
In our final model, we use GELU, but recommend further benchmarks around GELU and SwiGLU.
4. Embedding Norm
Dettmers et al.
(2021) suggests that greater stability of training can be achieved by including an extra layer normalization (Ba et al., 2016) after the embedding layer.
We evaluate the performance impact of such a modification in Table 4.
We note that this incurs a significant reduction in the performance of the model.
However, models above 100 billion parameters are notoriously unstable and require considerable engineering efforts in order to be kept stable.
If this addition provides increased stability when training, it may be valuable. [Table]
Finding 3.
Adding layer normalization after the embedding layer incurs a significant penalty on zero-shot generalization.
Multilinguality
The majority of 100B+ language models have been trained exclusively on English data, with notable exceptions in (monolingual) models trained in Chinese (Zeng et al., 2021; Wu et al., 2021) and Korean (Kim et al., 2021).
Smaller massively multilingual models have seen wider adoption (Xue et al., 2020), but these models are not suitable for zero-shot tasks.
Recent zero-shot results on large GPT-like multilingual models are promising, but English-only performance is usually disappointing (Lin et al., 2021).
Training data.
We train a multilingual language model in order to evaluate the effectiveness and potential impacts of this practice.
We use the OSCAR dataset (Ortiz Suárez et al., 2019) for training, but here we include multiple languages, not only English as in the earlier experiments.
The languages we include are Arabic, Basque, Bengali,
Chinese, Catalan, English, French, Hindi, Indonesian, Portuguese, Spanish, Urdu, and Vietnamese.
We sample each language with a different sampling probability that downsamples the most frequent languages and upsamples the least frequent ones, so that all languages are represented.
We estimate the sampling probabilities similar to Xue et al.
English-only evaluation.
We first evaluate our multilingual model on the same set of English benchmarks we have used previously, results are presented in Table 5.
The multilingual model has significantly lower accuracy than the English-only model on the English benchmark, which is in line with the results from Lin et al.
Multilingual evaluation.
[Table] shot [Table] evaluation is more challenging to setup because it
Embedding Norm
[Table]
requires writing new prompts for each new language.
Therefore, instead of manually writing prompts for each language, we follow the strategy proposed by Lin et al.
(2021), using English prompts for non-English examples–this can be viewed as cross-lingual zero-shot generalization.
They validated this strategy by demonstrating its ability to achieve zero-shot performance on par with (and sometimes even better than) humanwritten language-specific prompts.
This is also an indicator of the cross-lingual capabilities of the model.
We evaluate on XNLI (Conneau et al., 2018), a multilingual NLI dataset that covers 8 of the languages we use for training.
The task uses the following English clozestyle prompt template across all languages: [premise] , right?
[MASK] , [hypothesis]
The prompt fields, [premise] and [hypothesis] , are filled with the premise/hypothesis pairs in the target language.
For zero-shot evaluation, the [MASK] token is replaced with “Yes” (for entailment), “No” (for contradiction) and “Also” (for neutral).
The completion with the highest likelihood according to the model is taken as its prediction.
Our evaluation is different from the zero-shot evaluation of the XTREME benchmark (Hu et al., 2020).
XTREME first finetunes the model on the
English training data of each downstream task, then evaluates on the non-English dataset, attempting cross-lingual generalization.
Our evaluation avoids any finetuning, and instead relies entirely on zeroshot generalization from pretraining.
Results
Table 6 shows the XNLI results of our multilingual model and how it compares to
XGLM (Lin et al., 2021).
We were able to reproduce the results of XGLM-7.5B which validates our evaluation setup.
Furthermore, the table shows that the performance of our 1.3B is inline with the
XNLI 1.7B model, validating that our multilingual setup achieves competitive results.
It is worth notPretraining [Table] ing that our 1.3B model is trained on only 112B tokens from 13 languages while XGLM is trained on 500B tokens from 30 languages.
As far as we are aware, this is the first independent replication of the main results of Lin et al.
Scaling to 176B parameters
We are now ready to scale to the final model.
We have established that a curated high-quality crossdomain pretraining dataset similar to (Gao et al., 2020) will help boost zero-shot generalization; adopted from Anonymous (2022) that it is best to use a decoder-only model with an auto-regressive language modeling objective; identified that ALiBi positional embeddings should be used; decided that a default GELU activation should be preferred; and we have motivated the use of layer normalization on the embeddings to help with model stability.
We should now determine which model architecture (e.g., number of parameters, layers, width) will result in the best model out of our compute budget.
Compute budget.
We have been allocatedweeks of dedicated use of a cluster with 52 nodes of 8 80GB A100 GPUs.
We set four nodes aside as spare, as this cluster is brand new and may experience hardware failures.
This amounts to 1,161,216 A100-hours in total.
Assuming a throughput of 100 model TFLOPS, approximately corresponding to state-of-the-art hardware FLOPS of 150 (Narayanan et al., 2021), we have a compute budget of 4,838 PF-days for the model training.
We round this down to 4,500 PF-days, this ∼ 10 % safety margin accounting for potential downtime and inefficiencies (e.g., batch size ramp-up) during training.
To put this number in perspective, this is ∼ 23 % more than the training budget of GPT-3.
6. Parameters, Tokens, and Shapes Fitting scaling laws.
We establish scaling laws (Kaplan et al., 2020) to verify the scaling behaviour of our model and codebase, and to help decide on optimal model size.
We use English data from OSCAR (Ortiz Suárez et al., 2019) for pretraining, and train 125M, 350M, 760M, 1.3B, and 13B models for 100B to 300B tokens.
Figure 1 shows the smooth decrease in loss against the compute budget as the model size increases.
From Figure 1, we observe the following fit
L ( C ) =
( C c /C )
α C with α C ≈ 0 . 046 and C c
≈ 253 × 10 8 PF-days.
This scaling exponent α C is inline with the 0.050 reported in Kaplan et al. [Table]
[Table]
and the 0.048 reported at larger scale in Henighan et al.
This is a remarkably close fit given that we are using different datasets, codebases, and hyperparameters.
Given how close our fits are, for the rest of this section, we are going to use the coefficients derived in Kaplan et al.
(2020), as they are derived from a larger set of runs.
One caveat is that these scaling coefficients are estimated on the OSCAR pretraining data, but we will apply them to the training of a multilingual model on a dataset that was not developed by the time these experiments were performed.
Comparing the loss curves of the English and multilingual 1.3B models, we observe similar scaling and convergence trends.
This suggests that the capacity of the model is the same whether it is trained on English or multilingual text.
Therefore, we expect English-based scaling laws coefficients to still apply for a multilingual model.
We additionally conjecture that they will transfer to our final dataset.
Optimality vs. convergence.
From the scaling laws (Kaplan et al., 2020), it is possible to derive a Pareto frontier describing the optimal allocation of the compute budget between model size and number of tokens seen.
This optimal allocation achieves the lowest possible loss for a given compute budget.
It is notable that this Pareto optimal frontier describes very large models trained on few tokens; we call this training to optimality .
This is in stark contrast with the common practice of training much smaller models on many more tokens to convergence .
For instance, for a 1.3B parameters model, Kaplan et al.
(2020) predicts optimality at 20B tokens; any additional pretraining compute budget would be better allocated to a larger model.
However, this only provides a partial view.
Scaling laws focus on the upstream performance (pretraining loss) as the main evaluation metric, assuming that it directly translates to downstream performance.
However, prior work showed that this is not always the case (Tay et al., 2021).
Additionally, Table 1 shows that zero-shot generalization continues improving significantly past optimality.
Scaling laws also neglect inference cost.
A larger model is necessarily more expensive to serve, and inference costs at scale can rapidly catch up with training costs (Patterson et al., 2022).
Finally, in the context of multilinguality, we note that most multilingual models are also trained for more tokens (e.g. XGLM on 500B tokens Lin et al.
than a similarly sized monolingual model, as the pretraining dataset will otherwise include very few tokens from low-resource languages.
Accordingly, we choose to use the optimality front as an upper bound on model size and a lower bound on number of training tokens.
Given our compute budget, scaling laws predict an optimal allocation for training a 392B parameter model for 165B tokens.
Consequently, we will use these as constraints: the largest model we can afford is 392B parameters, and the minimum number of tokens to use is 165B tokens.
Prior models.
We outline in Table 7 the architectures of all publicly detailed 100B+ parameters models.
We note that most of these models are trained significantly beyond their respective optimality threshold, in the 300-400B tokens range.
PanGu- α is the biggest exception, but this appears to be due to hardware availability constraints rather than actively motivated by a modeling decision.
Model shape.
Kaplan et al.
(2020) studied the dependence of the loss with model shape, and found only a limited impact within a wide range of feed-forward ratios d ff /d model , aspect ratios d model /n layer , and attention head dimensions.
Notably, for a 1.5B parameters model, aspect ratios from 10-250 and head dimensions from 15 toall results in less than 1% change to the final loss.
Levine et al.
(2020) proposed a theoretically motivated and empirically backed law describing the optimal compromise between width and depth.
Notably, for models past 100B+ parameters, they predict that models such as GPT-3 are too deep, while models in the 10B or smaller range are usually too shallow.
For a GPT-3-sized model with 175B parameters, they predict an ideal depth of 80 layers.
6. Final Model Architecture
From the observations above, we identify three main guidelines for our final model:
• 300-400B tokens.
We want to guarantee our model will train on around 300-400B tokens of data.
This is in the upper range for models in the size range we are pursuing, ensuring that lowresource languages will not be allocated too few tokens.
Using the C = 6 ND approximation (Kaplan et al., 2020), with C = 4 , 500 PF-days and D = 300 -400B tokens, this constrains the model size to be around 160-200B parameters.
• 70-80 layers.
From Levine et al.
and the size constraint above, we estimate that our model should have between 70 and 80 layers.
• Maximum throughput.
Finally, we want the final architecture to have as high of a throughput per GPU as possible, as more compute will translate directly into longer pretraining and thus a better model.
Engineering constraints also come into light here: wide shallow models are typically easier to parallelize across nodes, up to a point where excessive tensor paralellism becomes necessary due to memory constraints.
From these guidelines, we benchmark 20 model configurations, detailed in Appendix B. Among these configurations, we select three of particular interest, outlined in Table 8.
They best fit our guidelines above, and offer high throughput.
We discard configuration (1), as its attention heads are much larger than other models in the literature.
Configuration (3) is shallower than recommended by Levine et al.
(2020), but delivers 3% higher throughput compared to (2).
Thus, we choose configuration (3) and its better throughput, and because a shallower model is easier to deal with at inference time by introducing less latency.
Conclusion
Seeking to establish the best possible model architecture that can be accommodated within a fixed 1,000,000 GPU-hours compute budget, we have presented an extensive study on principled modeling decisions for large language models.
First, we have found that complimenting Common Crawl data with high-quality cross-domain curated data can boost zero-shot generalization, validating previous suggestions (Rosset, 2020; Gao et al., 2020).
Through an ablation study, we have identified ALiBi as the position embedding of choice, confirmed the potential of SwiGLU, and highlighted that stabilizing techniques such as embedding normalization sometimes come at the expense of zero-shot generalization.
Exploring multilinguality, we have found that multilingual models significantly underperform their monolingual counterpart on English zero-shot benchmarks.
Finally, we further provided insights into how scaling laws can be used in practice in the design of large language models.
At variance with previous works, we outlined the full reasoning behind every architectural parameters, including model shape.
We hope our work can help practitioners better understand modeling decisions, leading to better language models, and that our transparency will accelerate future similar work. [Table]
Acknowledgements
This work was granted access to the HPC resources of Institut du deéveloppement et des ressources en informatique scientifique (IDRIS) du Centre national de la recherche scientifique (CNRS) under the allocation 2021-A0101012475 made by Grand équipement national de calcul intensif (GENCI).
In particular, all the trainings ran on the Jean-Zay cluster of IDRIS, and we want to thank the IDRIS team for responsive support throughout the project, in particular Rémi Lacroix.
Evaluations of GPT-3 models were provided in part by the Allen Institute for Artificial Intelligence.
We thank Leo Gao for his expertise and advice on language model evaluation.