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

TL;DR

This paper introduces CompleteP, a new parameterization method that improves compute efficiency and enables effective depth and non-lazy learning in large language model training, reducing the need for re-tuning hyperparameters.

Contribution

CompleteP is a novel parameterization that ensures hyperparameter transferability and non-lazy learning, enhancing compute efficiency across various model sizes and hardware setups.

Findings

CompleteP achieves 12-34% compute efficiency improvements.

It enables depth-wise hyperparameter transfer and non-lazy learning.

Experiments conducted on Cerebras CS-3 systems validate its effectiveness.

Abstract

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.