Scaling with Collapse: Efficient and Predictable Training of LLM Families

ListenWhen Loss Curves Collapse Into One11:43

Show notes

Cerebras researchers show that whole LLM training loss curves collapse onto a single universal trajectory across scales — from 100M to 3.9B parameters — but only when three controls are held fixed: the AdamW timescale tau, tokens-per-parameter, and the LR schedule. That collapse becomes a live diagnostic: it flagged a numerical bug 30% earlier than raw loss, and it enables reliable hyperparameter early-stopping at 10–30% of training. We unpack the mechanism, the Celerity family that lands on the compute-efficiency frontier, and what it means for anyone training big generative models. **Sources** - Scaling with Collapse: Efficient and Predictable Training of LLM Families

Subscribe to the channel: https://podcast.lumalabs.link/27a25f51ea81d41a8253f46c6133a35e/feed.xml

TL;DR

Bergsma et al. (Cerebras) take Qiu et al.’s scaling-collapse result to LLM scale under practical joint scaling of width, depth, batch size, and weight decay. They identify three controls that together govern normalized loss-curve shape under µP: the AdamW normalized timescale $\tau = \eta\lambda B_s/(D_s)$ , the tokens-per-parameter ratio TPP, and the LR schedule. When all three are held fixed across sizes and $\tau$ is set optimally for the given TPP, loss curves collapse from 100M to 3.9B — a signature of compute-efficient training. They release Celerity, a 300M–3.9B family trained in three TPP bands (20/80/234) that lands on the accuracy/compute Pareto frontier for open models, and demonstrate two practical uses: collapse residuals surface a numerical instability 30% of training earlier than raw loss trends, and a small-scale-fit parametric predictor of normalized curves enables reliable HPO early-stopping at 10–30% of training.

Key claims

Under µP, three controls govern normalized-TLC shape: the AdamW normalized timescale $\tau = \eta\lambda B_s/D_s$ , the tokens-per-parameter ratio TPP, and the LR schedule; fixing all three across sizes with $\tau$ optimal for the TPP produces collapse [§3, Fig. 1 middle].
$\tau$ modulates TLC shape as an implicit time-averaging window over AdamW updates: small $\tau$ gives short EMA memory (fast early bias reduction, higher variance floor), large $\tau$ averages more (slower start, lower floor) — the “fast-then-flatten” pattern under constant LR [§3, Fig. 3, App. B.3, Eq. 3].
TPP modulates TLC shape independently and scale-invariantly: higher-TPP normalized curves drop earlier and flatten longer, with a shape almost identical across a 30× parameter range at fixed TPP [§3, Fig. 4].
Llama-2’s family (7B/13B/34B/70B trained at decreasing TPP as size grows) fails to collapse under the paper’s normalization, whereas Celerity at fixed TPP does — collapse is a signature of compute-efficient scaling, not merely a mathematical artifact [Fig. 1 left vs middle].
LR-decay schedules invert the ordering of TLCs across $\tau$ : at constant LR, small- $\tau$ curves stay lower throughout; with decay, small- $\tau$ curves gain a late variance-suppression drop and finish lower [App. B.3, Fig. 10]. Normalization by final loss removes the curvature scale, leaving normalized shape a function of $(\tau, \text{schedule fraction})$ only — hence collapse across scales at matched $\tau$ .
Celerity 300M/500M/900M/1.8B/3.9B family, trained in three TPP bands (20, 80, 234) with fixed $\tau$ -optimal HPs, achieves the accuracy/FLOPs Pareto frontier for open pre-training-only models on 7 downstream tasks; against BTLM (Celerity’s non-annealing peer) reaches parity with 75% fewer FLOPs [§4, Fig. 2, Table 9].
Collapse residuals detected a 1.8B numerical instability at ~60% of training that only became visible in raw loss at ~90%; ablations against the collapse reference isolated the bug to a specific microbatch-size-triggered loss-kernel numerical issue [Fig. 6 right, Fig. 1 right].
A four-parameter functional form $L(f) = f^{-A}\cdot G(f,\text{schedule};w)$ with $A(\tau) = c_A\tau^{\alpha}$ and $w(\text{TPP}) = c_w\text{TPP}^\gamma$ (fit on 111M-scale runs) predicts normalized 3.3B curves with mean absolute error 2× the per-curve oracle floor — 1000× cheaper than direct-scale sweeps [§5, Eq. 4–5, Fig. 8, Table 10].
Aligning partial in-progress curves against the fitted prediction identifies the winning HP configuration at 10–30% of training with negligible final-loss regret, whereas “current-best” partial-run selection succeeds at 3.3B but fails at 1.7B — collapse-based selection is the more general solution [§5, Fig. 9].
When tuning LR/batch/weight-decay for LLMs, previous work has commonly fixed $\lambda$ (unintentionally varying $\tau$ ), which makes mid-training loss a poor predictor of final ordering; fixing $\tau$ instead preserves the ranking of curves throughout training [Fig. 7].
CompleteP (µP + depth) transfers HPs more reliably than plain µP for Celerity’s joint width–depth scaling ladder [§4, Fig. 15].

Method

The paper builds a GPT2-like architecture (ALiBi, SwiGLU, µP or CompleteP) and sweeps a 111M-to-3.9B scaling ladder over SlimPajama, isolating three sources of TLC shape variation: (1) at fixed model size, sweep $\eta, \lambda, B_s$ independently and observe that curves with matched $\tau$ overlap; (2) at fixed $\tau$ , sweep TPP and observe scale-invariant shape shifts; (3) show the LR schedule phases these effects. A noisy-quadratic model (App. B.3) derives Eq. 3 for constant-LR expected loss as variance_floor(τ) + bias_term(τ)·f^power, explaining the fast-then-flatten shape and predicting inversion of $\tau$ -ordering under LR decay — which is what the experiments show.

Celerity is trained in three TPP bands with fixed $\tau$ across sizes. Widths, depths, batch sizes are co-scaled (Table 1); LR/init/α transferred from 111M proxies via CompleteP. The collapse diagnostic is used live during training via two normalization strategies for the still-unknown final loss: extrapolate from lower-scale power-law fits, or “early-align” by choosing the divisor to best fit the smallest-scale curve over the first 25–50%. The 1.8B instability post-mortem is one demonstration; the HPO early-stopping study is the other, using a parametric predictor fit at 111M with just $A(\tau)$ and $w(\text{TPP})$ scaling.

Results

Celerity forms the accuracy/compute Pareto frontier for open pre-training-only models across 7 downstream tasks; strengthens further when teacher-FLOP costs are included for distilled competitors [Fig. 2, Fig. 16].
Collapse holds tightly at 80 TPP across 300M–3.9B; small early deviations at 20 TPP attributed to warmup-fraction mismatch; late-training divergence at 234 TPP for large models — investigated and attributed to training-data overfitting (held-out data stays aligned) [Fig. 6, Fig. 1 middle].
The 111M-fit predictor achieves MAE < 2× oracle on 3.3B curves; MAE actually improves with scale as TLCs smooth out and fewer extreme HPs are tried at large scale [Table 10, Fig. 8].
HPO early-stopping: at 1.7B/20TPP and 3.3B/30TPP, “predicted best” chooses the true optimum at 10–30% of training with negligible loss gap; “current best” partial-run selection fails at 1.7B [Fig. 9].
Fixing $\tau$ during tuning preserves the ordering of loss curves throughout training; letting $\tau$ vary (the common practice of fixing $\lambda$ instead) inverts orderings mid-training and misleads early-stopping [Fig. 7].

Why it’s interesting

This paper answers the extension question left open by Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks: does the small-scale scaling-collapse phenomenon survive when the joint scaling ladder co-scales width, depth, batch size, and weight decay under AdamW at LLM scale? Yes — with the additional condition that the AdamW normalized timescale $\tau$ has to be held fixed across sizes at the value that is optimal for the given TPP. This crystallizes into a live diagnostic Luma teams could actually run: any in-progress training curve can be compared to a small-scale reference of the same $\tau$ and TPP, and residuals will flag issues 30%+ earlier than raw loss trends do. The paper connects Wang & Aitchison’s AdamW-timescale-as-EMA framing to Bergsma et al.’s earlier LR/λ/B scaling work and to Qiu et al.’s collapse framework, giving all three a common surface. It also complements the Fitting paradigm of How to Set the Learning Rate for Large-Scale Pre-training? by supplying a training-side validation criterion — if the fit gives the right $\tau^*(\text{TPP})$ , the curves collapse; if not, they don’t. And its HPO early-stopping methodology is the most concrete practical output from the scaling-collapse line so far.