Don't be lazy: CompleteP enables compute-efficient deep transformers

Cengiz Pehlevan (Harvard University) · Shane Bergsma (Cerebras Systems) · Nolan Dey (Cerebras Systems) · Joel Hestness (Cerebras Systems) · Bin Zhang (Cerebras Systems) · Lorenzo Noci (ETH Zürich) · Mufan Li (University of Waterloo) · Blake Bordelon (Harvard University) · Boris Hanin (Princeton)

completepcompute efficiencycompute-efficientdepth-wise hp transferhardware settingshp transferlazy learning regimellm trainingmodel sizemodel width/depth ratiosnon-lazy learningoperational contextsoptimizer hyperparametersparameterizations

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.