Scaling Optimal LR Across Token Horizons

TL;DR

This paper investigates how the optimal learning rate for training large language models varies with token horizon, demonstrating that it follows a predictable scaling law and providing practical transfer rules.

Contribution

It introduces the first large-scale empirical study on hyperparameter transfer across dataset size (token horizon), revealing a scaling law for optimal learning rate.

Findings

Optimal LR decreases with longer token horizons.

Optimal LR follows a predictable scaling law.

Applying the scaling law improves training efficiency.

Abstract

State-of-the-art LLMs are powered by scaling -- scaling model size, dataset size and cluster size. It is economically infeasible to extensively tune hyperparameter for the largest runs. Instead, approximately optimal hyperparameters must be inferred or \textit{transferred} from smaller experiments. Hyperparameter transfer across model sizes has been studied in Yang et al. However, hyperparameter transfer across dataset size -- or token horizon -- has not been studied yet. To remedy this we conduct a large scale empirical study on how optimal learning rate (LR) depends on token horizon in LLM training. We first demonstrate that the optimal LR changes significantly with token horizon -- longer training necessitates smaller LR. Secondly we demonstrate the the optimal LR follows a scaling law, and that the optimal LR for longer horizons can be accurately estimated from shorter horizons via such scaling laws. We also provide a rule-of-thumb for transferring LR across token horizons with zero overhead over current practices. Lastly we provide evidence that LLama-1 used too high LR, and estimate the performance hit from this. We thus argue that hyperparameter transfer across data size is an important and overlooked component of LLM training.