Source-linked AI summary
Don't be lazy: CompleteP enables compute-efficient deep transformers
Nolan Dey, Bin Claire Zhang, Lorenzo Noci, Mufan Li, Blake Bordelon, Shane Bergsma, Cengiz Pehlevan, Boris Hanin, Joel Hestness
TL;DR
LLM training must allocate limited compute across model architecture and other hyperparameters, but parameterizations can fail to transfer optimal settings across depth. This paper evaluates parameterizations and identifies CompleteP, which provides depth-wise hyperparameter transfer and achieves compute savings in large pre-LN transformers.
Problem
LLM training requires choosing model architecture and hyperparameters to minimize pretraining loss under a fixed compute budget.
Method
The paper compares parameterizations across depth and width, and extends CompleteP with scaling prescriptions for LayerNorm parameters and AdamW’s epsilon.
Results
11.8% and 34.4% FLOP savings over µP were achieved by 1.9B-parameter models trained with α = 1 for optimally-shaped and 179-layer models, respectively.
Takeaways & Limitations
CompleteP provides hyperparameter transfer when varying depth and width and supports compute-efficient transformer shape selection.
Takeaways & Limitations
The theoretical analysis studies width and depth scaling at fixed token count, limiting its direct applicability to the fixed-TPP compute-optimal regime.
Abstract
from arXiv · showhide
We study compute efficiency of LLM training when using different parameterizations, i.e., rules for adjusting model and optimizer hyperparameters (HPs) as model size changes. Some parameterizations fail to transfer optimal base HPs (such as learning rate) across changes in model depth, requiring practitioners to either re-tune these HPs as they scale up (expensive), or accept sub-optimal training when re-tuning is prohibitive. Even when they achieve HP transfer, we develop theory to show parameterizations may still exist in the lazy learning regime where layers learn only features close to their linearization, preventing effective use of depth and nonlinearity. Finally, we identify and adopt the parameterization we call CompleteP that achieves both depth-wise HP transfer and non-lazy learning in all layers. CompleteP enables a wider range of model width/depth ratios to remain compute-efficient, unlocking shapes better suited for different hardware settings and operational contexts. Moreover, CompleteP enables 12-34% compute efficiency improvements over the prior state-of-the-art. All experiments were run on Cerebras CS-3 systems. A minimal implementation is available at https://github.com/EleutherAI/nanoGPT-mup/tree/completep.
1 Introduction
The paper studies parameterizations for scaling transformer width and depth while preserving hyperparameter transfer and avoiding lazy learning. It identifies CompleteP (α = 1), with extensions for optimizer and normalization hyperparameters, as the parameterization that meets these goals and improves compute efficiency.
- Larger language models improve performance but require costly choices over architecture, data, batching, training, and hyperparameters under fixed compute budgets.
- µP supports approximately constant optimal hyperparameters when scaling width, enabling a “tune small and train large” strategy.
- Depth-aware parameterizations differ in how they rescale transformer residual-block outputs before addition to the residual stream, using α ∈ [0.5, 1].
- CompleteP adds principled scaling rules for LayerNorm and bias learning rates, AdamW weight decay, and AdamW ϵ.
- Only α = 1 enables depth-wise transfer of learning rate and initialization standard deviation across 2–128 layers.
- 11.8% and 34.4% FLOP savings over µP are achieved by 1.9B-parameter α = 1 models with optimal shapes and 179 layers, respectively.
2 Related work
Related work develops parameterizations for stable training and hyperparameter transfer across width and depth, while prior conclusions about optimal transformer shapes depend on the parameterization used. CompleteP is presented as combining transfer with compute savings and supporting deeper or narrower efficient models.
- NTK and mean-field/µP parameterizations provide infinite-width training dynamics and proposals for transferring hyperparameters across model scale.
- Depth-scaling work studies a family indexed by α ∈ [0.5, 1], with prior disagreement over whether hyperparameter transfer is possible and which α works best.
- Prior empirical approaches normalize layer outputs or learning rates for width transfer, while depth transfer is often more ad hoc and resembles α = 0.5 or α = 1.
- Kaplan et al. reported N : L ≈ 100 as an optimal transformer ratio under standard parameterization, but standard parameterization does not fairly support stable width and depth scaling.
- CompleteP revisits shape optimality with controlled width and depth scaling and finds that even N : L ≈ 10 remains close to compute optimality.
3 Methodology
The experiments use pre-normalized decoder-only Transformer language models and compare parameterizations through width and depth multipliers. The methodology extends infinite-depth prescriptions with practical scaling rules for optimizer and normalization parameters.
- Experiments train pre-normalized decoder-only Transformers with untied embeddings, ALiBi, ReLU2, AdamW, and an autoregressive next-token prediction loss.
- Models use width N and depth L, with scaling expressed through multipliers m_N = N/N_base and m_L = L/L_base relative to N_base = 256 and L_base = 2.
- The tested parameterizations are standard parameterization, µP, and α ∈ {0.5, 1}, which are equivalent for the base model.
- The implementation extends prior prescriptions with corrections for bias and LayerNorm learning rates, AdamW ϵ, and weight decay.
4 Depth-wise HP transfer and α
Depth-wise transfer tests show that standard parameterization, µP, and α = 0.5 do not preserve optimal hyperparameters as depth changes, whereas CompleteP does. Under compute-optimal training, larger tokens-per-parameter reduce learning-rate sensitivity.
- Traditional HP transfer: With 300M tokens, batch size B = 128, and λ_base = 0, standard parameterization, µP, and α = 0.5 lack stable optimal learning rates and initialization standard deviations across depth.
- Traditional HP transfer: With SP, µP, and α = 0.5, models at different depths do not share the same optimal hyperparameters.
- CompleteP: CompleteP keeps optimal η_base and σ_base stable across depths and consistently improves deeper-model loss without hyperparameter tuning.
- CompleteP: CompleteP demonstrates hyperparameter transfer from 2 to 128 layers, exceeding the depths of LLaMA-70B and LLaMA-405B cited in the paper.
- Compute-efficient HP transfer: Under compute-optimal training with 20 TPP, compute-efficient batch sizes, and tuned λ_base, learning-rate transfer is less sensitive than in the 300M-token setting.
- Compute-efficient HP transfer: Larger TPP empirically reduces hyperparameter sensitivity.
5 Re-examining compute-optimal N:L ratio
The paper revisits compute-optimal transformer width-to-depth ratios under parameterizations with simultaneous width and depth control. CompleteP preserves near-optimal performance across a wider ratio range and yields increasing FLOP savings over µP as depth grows.
- CompleteP enables N:L ≈10 to remain close to compute-optimal when width and depth are properly controlled.This addresses confounding from hyperparameter detuning and instability in deeper SP/µP models.
- The study evaluates models at 50M, 300M, and 1.5B non-embedding parameter scales while varying N and L under compute-optimal training.Models approximately maintain Pnon-emb = 12N^2L and use 20 tokens per parameter.
- Compute-optimal N:L trends larger as model scale increases across all tested parameterizations.
- 34.4% FLOP savings occur for the deepest 1.5B models with CompleteP versus µP.CompleteP’s advantage grows with depth because µP becomes more detuned from lacking depth-wise hyperparameter transfer.
- For Pnon-emb=1.5B, N:L = 11.8 remains within 1% of compute-optimal with α = 1, compared with N:L = 38.7 for µP.Shallow-wide models suit latency-sensitive hardware, while narrow-deep models can benefit low-memory hardware through layer-by-layer weight streaming.
6 Desiderata for Hyperparameter Transfer
The paper proposes three desiderata for parameterizations that scale width and depth: stable initialization, maximal residual-stream updates, and complete feature learning. It argues that only α = 1 satisfies the complete-feature-learning criterion as depth grows.
- Stable Initialization: Parameterizations are designed to keep hidden representations and network outputs stable, meaningful, and consistent across model sizes.
- Maximal Residual Stream Update: Depth-aware residual networks scale each block output by L^-α, with α constrained to [0.5, 1].
- Maximal Residual Stream Update: Each residual block should contribute order 1/L to feature movements, while non-residual blocks contribute constant-order changes.
- Complete Feature Learning: Complete feature learning requires that every hidden layer and the model output remain non-lazy with respect to any subset of model parameters.
- Complete Feature Learning: For α < 1, nonlinear contributions diminish as depth increases because the two update terms have matching order only when α = 1.
- Complete Feature Learning: Only CompleteP, corresponding to α = 1, ensures stable training, maximal updates, and complete feature learning while scaling width and depth.
7 Limitations
The theoretical analysis is conducted at fixed token count, limiting its direct applicability to the fixed-TPP compute-optimal regime.
- The theoretical analysis scales width and depth at fixed token count, limiting direct applicability to the fixed TPP compute-optimal regime.
8 Conclusion
The paper finds that CompleteP transfers hyperparameters across depth and width while reducing pre-training FLOPs in compute-optimal settings. It formalizes complete feature learning as a distinguishing property and reports experiments up to 1.5B non-embedding parameters.
- CompleteP (α = 1) provides hyperparameter transfer when varying transformer depth and width.
- CompleteP achieves significant FLOP savings during pre-training, including compute-optimal settings with jointly scaled batch, dataset, and model sizes.
- Complete feature learning distinguishes CompleteP and formalizes a desideratum in which hidden layers and outputs avoid lazy behavior.
- Empirical tests used models with up to 1.5B non-embedding parameters, while future work targets larger state-of-the-art LLMs.
- The methods aim to reduce compute, financial, environmental, and hyperparameter-tuning costs associated with large-model training.
B Coordinate check test for verification of stable training
The coordinate check tests whether activation magnitudes remain stable as transformer depth increases. After correcting bias and LayerNorm learning-rate scaling, both α = 0.5 and α = 1 achieve stability at any depth.
- The coordinate check measures activation-size changes after the final residual addition hL while scaling model depth.
- The literature’s α = 0.5 prescription initially failed the transformer stability check, unlike the expected theoretical behavior.
- Adjusting bias and LayerNorm learning rates restored stability for α = 0.5.
- Both α ∈ {0.5, 1} achieve stability for any depth, but the test is necessary rather than sufficient for complete feature-learning desiderata.
C.1 Desideratum 1: Stable Initialization
The parameterizations are designed to keep signal propagation and parameter updates stable as width and depth scale. The derivations prescribe scaling rules across transformer components, while noting that some architectural extensions still require further analysis.
- Stable signal propagation: For α = 1, signal propagation retains dependence on activation and initialization rather than becoming trivial at large depth.
- Stable updates: The learning-rate scaling is chosen so weight updates produce Θ(1) changes to residual-stream variables.The same width-and-depth scaling applies to AdamW up to momentum and the separately handled ϵ parameter.
- Generalizability: Scaling rules for MoE, long context, batch size, gradient clipping, LAMB, and momentum require more theoretical derivation.The approach currently derives scaling for finitely many data points and training steps, then tests hyperparameter transfer empirically.
- Architecture and optimizer components: The parameterization extends beyond weights to biases, LayerNorm parameters, weight decay, and AdamW ϵ.The prescribed rules aim to maintain the correct update scale in each layer.
F Depth Scaling Training Dynamics
Training dynamics favor CompleteP (α = 1) as depth increases, while α = 0.5 exhibits loss-crossing behavior in which deeper models take longer to overtake shallower ones.
- Depth scaling: CompleteP (α = 1) consistently achieves lower training loss when model depth scales up.The comparison uses loss curves at CompleteP’s optimal learning rate η = 2^-8.
- Depth scaling: In α = 0.5, deeper models continually take longer steps before their loss scales better than shallower models.This loss-crossing behavior aligns with α = 0.5 requiring a larger learning rate as depth increases.
- Experimental setup: The experiments visualize training dynamics with loss curves for all tested parameterizations.
G.1 Compute-efficient setup
The compute-efficient study trains models under a 20-tokens-per-parameter setup while varying width and depth, then measures loss and FLOP efficiency across aspect ratios. It also documents the fitting procedure used to estimate compute-efficient ranges.
- Compute-efficient setup: Models are trained for 20 tokens per parameter to represent a compute-optimal tradeoff between parameters and training tokens.Batch sizes are also scaled according to total training FLOPs.
- Model shapes: The aspect-ratio study varies width and depth while approximately holding non-embedding parameters at 12N^2L.The reported FLOPs are measured granularly because embedding and attention costs make them unequal despite matched non-embedding parameter counts.
- Analysis procedure: FLOP savings are estimated by fitting power laws for loss versus FLOPs and inverting them to obtain FLOPs as a function of loss.Cubic fits to log aspect-ratio data are used to determine intersections for the reported range analysis.