Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Duration

TL;DR

This paper introduces a unified parameterisation for scaling and transferring hyperparameters across model modules, sizes, batch, and duration, enabling efficient hyperparameter transfer and improved training speed in large models.

Contribution

It extends previous hyperparameter transfer methods by unifying scaling axes and enabling per-module hyperparameter transfer with practical guidelines.

Findings

Hyperparameter transfer is effective across modules and scales.

Significant training speed improvements in large language models.

Unified parameterisation simplifies high-dimensional hyperparameter optimization.

Abstract

Hyperparameter tuning can dramatically impact training stability and final performance of large-scale models. Recent works on neural network parameterisations, such as $\mu$P, have enabled transfer of optimal global hyperparameters across model sizes. These works propose an empirical practice of search for optimal global base hyperparameters at a small model size, and transfer to a large size. We extend these works in two key ways. To handle scaling along most important scaling axes, we propose the Complete$^{(d)}$ Parameterisation that unifies scaling in width and depth -- using an adaptation of CompleteP -- as well as in batch-size and training duration. Secondly, with our parameterisation, we investigate per-module hyperparameter optimisation and transfer. We characterise the empirical challenges of navigating the high-dimensional hyperparameter landscape, and propose practical guidelines for tackling this optimisation problem. We demonstrate that, with the right parameterisation, hyperparameter transfer holds even in the per-module hyperparameter regime. Our study covers an extensive range of optimisation hyperparameters of modern models: learning rates, AdamW parameters, weight decay, initialisation scales, and residual block multipliers. Our experiments demonstrate significant training speed improvements in Large Language Models with the transferred per-module hyperparameters.