Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining

TL;DR

This work introduces Step Law, a universal, data-agnostic scaling law for hyperparameters in large language model pretraining. By training 3,700 LLMs on 100T tokens, the authors reveal convex loss landscapes under fixed $N$ and $D$ and derive power-law relations: $\eta(N,D)=1.79\,N^{-0.713}\,D^{0.307}$ and $B(D)=0.58\,D^{0.571}$, with $B$ independent of $N$. The law generalizes across dense and MoE architectures and diverse data recipes, achieving near-global-optimal performance (0.094% gap) and demonstrating robustness to architectural variations, sparsity, and data composition. The authors provide extensive empirical validation, cross-architecture transferability, and open-source tooling to streamline scalable LLM pretraining. This work promises to dramatically reduce hyperparameter search cost while delivering reliable, scalable training guidelines for industry and academia alike.

Abstract

The impressive capabilities of Large Language Models (LLMs) across diverse tasks are now well established, yet their effective deployment necessitates careful hyperparameter optimization. Although existing methods have explored the influence of hyperparameters on model performance, a principled and generalizable framework across model architectures and data recipes remains absent. In this study, we conduct an unprecedented empirical investigation training over 3,700 LLMs from scratch across 100 trillion tokens, consuming nearly one million NVIDIA H800 GPU hours to establish a universal Scaling Law for hyperparameter optimization in LLM Pre-training, called Step Law. We empirically observe that, under fixed model size ($N$) and dataset size ($D$), the hyperparameter landscape exhibits convexity with a broad optimum, substantially reducing the complexity of hyperparameter search. Building on this insight, we formally define and empirically validate the Step Law: The optimal learning rate follows a power-law relationship with $N$ and $D$, while the optimal batch size is primarily influenced by $D$ and remains largely invariant to $N$.Notably, our estimated optima deviate from the global best performance found via exhaustive search by merely 0.094\% on the test set. To our best known, Step Law is the first that unifies different model shapes and structures, such as Mixture-of-Experts models and dense transformers, as well as establishes optimal hyperparameter scaling laws across diverse data recipes. We contribute a universal, plug-and-play optimal hyperparameter tool for the community, which is expected to advance efficient LLM training at scale. All experimental code, data and checkpoints are publicly available at https://github.com/step-law/steplaw

Figures (35)

• Figure 1: This plot shows the hyperparameter space for a 1B model trained on 100B tokens. We trained 120 LLMs from scratch with different LR and BS combinations, obtaining contour lines and global optimal points based on real data. Optimal points represent the lowest training loss for each LR and BS pair, while contour lines depict the relative loss differences from these optima. Step Law predicts the optimum with the highest accuracy compared to other methods, nearly matching the global optimal points.
• Figure 2: Learning Rate vs. Batch Size Loss Landscape Analysis for 1B Model (Trained on 100B Tokens): Scatter Plots and 3D Surface Visualizations of Hyperparameter Sensitivity.
• Figure 3: Comparison of learning rate schedules. Contour plots illustrate two learning rate schedules: the conventional decay schedule (blue), which sets $\mathrm{lr}_{\min} = \mathrm{lr}_{\max}/10$, and our fixed schedule (red), with $\mathrm{lr}_{\min}=10^{-5}$. The visualization reveals that the conventional decay method leads to a discernible leftward bias in the optimal learning rate range, indicated by the shift of the lowest loss region towards lower learning rates in the blue contours compared to the red.
• Figure 4: Contour plots of training loss (left) and validation loss (right) . Both plots share the global minimum (✗) at batch size 393,216 and learning rate of 0.001950.
• Figure 5: Empirical vs. predicted HP scaling. (a) Scatter points indicate empirical optimal learning rate vs. batch size for model scale $N$; (b) Analogous results for dataset scale D. Curves show our hp-scaling law predictions, with shaded regions representing parameter uncertainty bounds from the sampling-based fitting strategy. Both plots use double logarithmic scaling (1,912 training samples).