A Recipe for Training Large Models

Introduction

help you train large models (>1B parameters)
avoid instabilities
save experiments that started to fail without restarting at 0

Table of Contents:

When do these training guidelines apply?

Your model is over 1 billion parameters
Your model is going to train for at least multiple days
You don't think that you are limited by your dataset
Admittedly, this is much harder to estimate. As an example dalle-mini started overfitting with 400M parameters on a dataset of 17M images after 3 days on a TPU v3-8 (50 epochs)
It is easier to think in term of number of epochs you expect to iterate versus the dataset
You are training from scratch, or on a very different task
if not, you may be able to freeze your model and just fine-tune a few layers, way easier!
You don't have the budget to perform 50-100 experiments

The Recipe To Train Large Models

Experimentation phase

Always start with a small version to ensure your entire pipeline works

Scale your model in 3-10x increments before training the final version

Never try new things on a large model unless it at least seems to work on your own smaller model (and don't trust blindly a paper! Anyone can make mistakes, inflate results, or their tricks may not apply to your own problem). When training your large model you will mainly be busy fighting instabilities so you don't want to add another reason for issues.
Start with a simple baseline that is known to work so you have a good reference model
Avoid making conclusions too fast, some runs may initially look promising but end up diverging later or just reach a plateau earlier than others

Experiment fast, try quick ideas, and make sure to keep notes as you go: you'll revisit your reports and conclusions frequently.
Add new configuration parameters for your experiments (such as norm_type, use_alibi, and even use_idea_1 is better than nothing) and log them all so you can quickly compare experiments and display the difference between your experiments.
If you log your runs with Weights & Biases, it's easy to select a graph that compares a few runs, pull it into a report, and add a quick note of what insight it brings.
See An Evaluation of Transformer Variants as an example: I was just adding tiny sections over time as I ran experiments to record my conclusions, and I keep coming back to it.

Stable Training Of Large Models

Optimizer
I recommend Distributed Shampoo as a very stable optimizer. The memory overhead (vs. Adam) disappears when using multiple nodes (due to sharding)
Precision
Some instabilities are just solved by increasing precision: float32 > bfloat16 > float16
You typically will want to keep your model parameters in full precision and only decrease the precision of computations. Be mindful of certain operations that may require higher precision such as attention layers and losses (softmax, contrastive loss especially with large batch size, etc).
If you don't get significant advantages with lower precision (speed or memory constraints), you should just use a higher precision.
Weight decay
You don't need weight decay with Shampoo
Weight decay may help at the start of training but you should decrease it to 0 during training (if your optimizer allows it)

Weight decay helps the start of training but leads to a plateau

Dropout
I never use dropout for large models with enough data.
Model architecture
I have been having success with NormFormer + GLU, another good one if you want fewer LayerNorm's is the cosine similarity from Swin v2.
I use no bias.
See An Evaluation of Transformer Variants for more details.
Weight initialization
Larger layers often require smaller weights.
Initializing residual connection weights to 0 can improve stability if needed.
Other stabilization methods: if you don't use Shampoo you may benefit from gradient clipping (in that case I'd suggest clipping by global norm).
EMA parameters
If you have enough memory, you can compute EMA of parameters during training.
Otherwise, you can keep a copy of model checkpoints and average over them.
It is mainly beneficial for diffusion models and has a minimal effect on other models.

Your Model Size

Typically the bigger you can handle the better.
Limitations may be based on your compute budget. Additionally:
See what model size your hardware can handle.
Find the most similar model known to have been trained and note their model size, dataset, and training budget.
Refer to the closest applicable scaling laws available.
Finally confirm your training budget based on batch size, number of steps/samples to be seen, and your actual training time per step (measured on a short experiment).
Size your model based on the intended application:
What are the hardware limitations of the target host device?
What is the inference speed requirement?
Test actual inference performance on a dummy untrained model.

The Batch Size

Always use the biggest batch size you can.
Find a minimum requirement for batch size based on similar models that have been trained successfully.

Use gradient accumulation to increase your effective batch size.
Warning: not applicable on certain models such as CLIP where the loss is based on calculations over the entire batch at a time.
Use gradient checkpointing (evaluate the impact on training speed).
Use model sharding.
Reduce your model size if you have to, a batch size too low can cause low performance of your model.

At the start of training, a lower batch size can help you train faster. However, it adds noise to the gradient so can potentially lead to instability, and it can be hard to detect at which point you should switch back to full batch size.
Later during training, an increase in batch size has similar effects as a decrease in the learning rate, with an immediate decrease in loss. Based on the impact it has on training speed, it can be preferable to increase the batch size (through gradient accumulation) vs. lowering the learning rate.

Model Sharding

Finding the smallest number of devices over which to shard the model (the device mesh can be defined as an array sharded over model and data), experimenting with various configurations of partitioning strategies.
Making the batch size as large as possible.
Experimenting with model size adjustments to evaluate the impact on training speed.
Testing FSDP strategy and fully shard data and model across all devices

The Learning Rate

Start on a smaller model, learning rate ports pretty well if your model size increases in x5-10 increments, with minor adjustments.
Perform a quick search:
Find the best learning rate that goes in about x3 increments (1e-3, 3e-3, 1e-2…).
Start with x10 increments initially to go faster.
Run the experiments long enough to get some confidence.
It is essential to optimize at least the learning rate before an expensive training run as it has a significant impact on training speed (and even training success).

At the start of training, use a warmup
It can only be beneficial
set the warmup duration based on your training budget (should be less than 5-10%)
During training, keep the learning rate constant:
It let you decouple its effect on metrics (lowering learning rate typically induces a sharp decrease of loss)
Try from time to time to lower/increase the learning rate (I use x3 increments) and observe the effect it has on the training speed (slope), you can always readjust it without having to go back in time
If you are really lazy, you can use a cosine or linear decay but keep in mind that it's only because you're really lazy 😬
At the end of training, use a final cosine or linear decay to 0 for a little boost

When is Training Complete?

Large Model Infrastructure

Logging of training/validation loss + relevant metrics (accuracy…).
Logging of parameters and gradients:
The norm of parameters and gradients should regularly be logged so you can refer to it during instabilities
Histograms can be logged optionally (it will impact training speed but is nice to enable when debugging instabilities)
At start of training, manually verify that all gradients are flowing properly in your model and that parameters are being updated



Saving/resuming from the checkpoint.
You may want to handle separate model and optimizer states.
Logging of prediction samples.
There is nothing better than visualizing actual sample predictions.
If it requires too much compute, you should create a script that computes them against your latest model checkpoint and start it on a different instance.
A small demo so you can quickly experiment with your model:
It is essential to be able to test your model during training.
A simple notebook or a Gradio demo (better to share with other users) are very fast to implement and will let you detect potential issues early on.
Keep a training log of your model with:
Your training configuration, hardware, and hyperparameters.
Your loss, metrics, etc. over time with all relevant runs.
A small note on every single change you made between runs so you can understand your graphs later.
Sample predictions over time if possible.
See the dalle-mini training journal as an example.

Conclusion

Resources

NeurIPS 2022, Workshop "Has it Trained Yet", "The Road to Craiyon" related to the scaling of dalle-mini
Another cool guide: "Deep Learning Tuning Playbook" from Varun Godbole, George E. Dahl, Justin Gilmer, Christopher J. Shallue, Zachary Nado
T5X repository for the cool resources on model sharding with JAX

Acknowledgements

Awesome review and suggestions from Zachary Nado
Google TPU Research Cloud (TRC) program for providing computing resources
Weights & Biases for providing the infrastructure for experiment tracking and model management
🤗 Hugging Face for the models and training scripts used for these experiments