Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training

Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, Joel Hestness

TL;DR

The paper addresses how hyperparameters should be scaled during large-language-model pre-training by proposing a unified AdamW-based framework that connects the optimal weight decay and batch-size strategies to model size, data size, and compute. It reveals a power-law relation between the AdamW timescale and tokens-per-parameter, enabling prediction of λ_opt across N, D, and B, and shows that B_opt and B_crit scale with data size D independently of N, with a practical method to estimate Pareto-optimal N-D-B configurations for time vs compute. The findings offer actionable guidance for hardware-aware HP tuning and training planning, suggesting that λ tuning is often more robust than η tuning when scaling batch size and that modest overtraining can improve time and parallelism outcomes. Together, these results help practitioners balance wall-clock time and compute costs in real-world LLM pre-training while informing remote hardware and scheduling decisions.

Abstract

Efficient LLM pre-training requires well-tuned hyperparameters (HPs), including learning rate $η$ and weight decay $λ$. We study scaling laws for HPs: formulas for how to scale HPs as we scale model size N, dataset size D, and batch size B. Recent work suggests the AdamW timescale, $τ= B/(ηλD)$, should remain constant across training settings, and we verify the implication that optimal $λ$ scales linearly with B, for a fixed N and D. However, as N and D scale, we show optimal $τ$ obeys a precise power law in the tokens-per-parameter ratio, D/N. This law thus provides a method to accurately predict $λ$opt in advance of large-scale training. We also study scaling laws for optimal batch size Bopt (the B enabling lowest loss at a given N,D) and critical batch size Bcrit (the B beyond which further data parallelism becomes ineffective). In contrast to prior work, we find both Bopt and Bcrit scale as power laws in D, independent of model size, N. Finally, we analyze how these findings inform the real-world selection of Pareto-optimal N and D under dual training time and compute objectives. All experiments were run on Cerebras CS-3 systems.

Power Lines: Scaling Laws for Weight Decay and Batch Size in LLM Pre-training

TL;DR

The paper addresses how hyperparameters should be scaled during large-language-model pre-training by proposing a unified AdamW-based framework that connects the optimal weight decay and batch-size strategies to model size, data size, and compute. It reveals a power-law relation between the AdamW timescale and tokens-per-parameter, enabling prediction of λ_opt across N, D, and B, and shows that B_opt and B_crit scale with data size D independently of N, with a practical method to estimate Pareto-optimal N-D-B configurations for time vs compute. The findings offer actionable guidance for hardware-aware HP tuning and training planning, suggesting that λ tuning is often more robust than η tuning when scaling batch size and that modest overtraining can improve time and parallelism outcomes. Together, these results help practitioners balance wall-clock time and compute costs in real-world LLM pre-training while informing remote hardware and scheduling decisions.

Abstract

Efficient LLM pre-training requires well-tuned hyperparameters (HPs), including learning rate and weight decay . We study scaling laws for HPs: formulas for how to scale HPs as we scale model size N, dataset size D, and batch size B. Recent work suggests the AdamW timescale, , should remain constant across training settings, and we verify the implication that optimal scales linearly with B, for a fixed N and D. However, as N and D scale, we show optimal obeys a precise power law in the tokens-per-parameter ratio, D/N. This law thus provides a method to accurately predict opt in advance of large-scale training. We also study scaling laws for optimal batch size Bopt (the B enabling lowest loss at a given N,D) and critical batch size Bcrit (the B beyond which further data parallelism becomes ineffective). In contrast to prior work, we find both Bopt and Bcrit scale as power laws in D, independent of model size, N. Finally, we analyze how these findings inform the real-world selection of Pareto-optimal N and D under dual training time and compute objectives. All experiments were run on Cerebras CS-3 systems.

Paper Structure

This paper contains 54 sections, 24 equations, 10 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Scaling of the AdamW timescale $\tau$, and optimal weight decay $\lambda_{\text{opt}}$
  3. Background: $\hbox{$\mu$P}$, AdamW, and $\tau_\text{epoch}$
  4. $\hbox{$\mu$P}$
  5. The EMA view of AdamW
  6. Methods: The AdamW timescale for LLMs, $\tau$, and its scaling
  7. Experimental details
  8. Results: $\tau$ and scaling $\lambda$
  9. Discussion
  10. Scaling of optimal batch size $B_{\text{opt}}$ and critical batch size $B_{\text{crit}}$
  11. Background: $B_{\text{opt}}$ and $B_{\text{crit}}$
  12. $B_{\text{opt}}$
  13. $B_{\text{crit}}$
  14. Methods: estimating $B_{\text{opt}}$ and $B_{\text{crit}}$, and their scaling
  15. Estimating $B_{\text{opt}}$
  16. ...and 39 more sections

Figures (10)

  • Figure 1: Hyperparameters and their power lines: Optimal $\tau$ obeys a power law in tokens-per-parameter ($\text{left}$), while optimal batch size ($\text{middle}$) and critical batch size ($\text{right}$) obey power laws in $D$. Faded markers indicate points not used in fitting; all fits generalize well to larger-scale runs.
  • Figure 2: (610M 20TPP): For each $B$, we sweep $\lambda$ and find ${\tau}_{\text{opt}}$ ($\text{left}$). ${\tau}_{\text{opt}}$ is stable around 0.21 for $B \in [63, 2016]$, meaning $\lambda_{\text{opt}}$ scales linearly with $B$ over this range ($\text{middle}$). When sweeping $\eta$ ($\text{right}$), the lower boundary over all curves is a bowl with a minimum at 0.21; the smallest $B$ settings have ${\tau}_{\text{opt}}$ within 2$\times$ of this value, but as $B$ increases, ${\tau}_{\text{opt}}$ quickly drifts higher.
  • Figure 3: (111M): Fitted $B$-specific power laws, $L_B(D)$, for inferring steps to reach target loss $\hat{L}$ (arrowed blue line).
  • Figure 4: \ref{['eq:tradeoff']} fits observed data well: Increasing $B$ (moving leftward along curves) decreases optimization steps (x-axis), but requires more tokens (y-axis) to reach target loss (in color). $B_{\text{crit}}$ is the transition point along each fitted curve.
  • Figure 5: Prior work suggests $B_{\text{opt}}$ scales in $C$ ($\text{left}$) and $B_{\text{crit}}$ in loss ($\text{middle}$), but this only holds at a fixed $N$/TPP (same data as \ref{['fig:hook']}$\text{middle}$/$\text{right}$); \ref{['fig:hook']} shows scaling in $D$ is the fundamental relationship. Plotting $B_{\text{deepseek}}(C)$ law from bi2024deepseek ($\text{right}$), but over $D$ (using $C \approx 6ND$ to obtain $C$ for the spurious dependence on $N$), we see bi2024deepseek used generally efficient $B$ values (i.e., within the $B_{\text{opt}} < B < B_{\text{crit}}$ regime) despite fitting $C$ rather than $D$ ($B_{\text{opt}}$ and $B_{\text{crit}}$ lines from \ref{['fig:hook']}).
  • ...and 5 more figures

Theorems & Definitions (1)

  • Definition 3.1