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

TL;DR

The paper addresses the high cost and fragility of hyperparameter tuning for large-scale transformers by introducing Complete(d)P, a unified parameterisation that enables transfer of both global and per-module hyperparameters across width, depth, batch size, and token horizon. It extends this framework with QK-normalisation adaptations, embedding-scaled dynamics, and a depth-Kronecker per-module HP scheme, then validates transferability across scales and compute horizons, achieving significant speed-ups (e.g., ~27% in large-scale settings) over global baselines. A key finding is that per-module HPs not only transfer but can be optimised at small scales to yield large-scale benefits, and that a careful SDE-based reparameterisation under batch-size and horizon scaling preserves training dynamics. Together, these contributions offer a practical recipe to reduce compute waste in HP search and to realize faster, more stable large-scale training for transformer models.

Abstract

Hyperparameter tuning can dramatically impact training stability and final performance of large-scale models. Recent works on neural network parameterisations, such as $μ$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.

Figures (18)

• Figure 1: (Left): We optimise hyperparameters at a small 50M parameters/1.6B tokens scale (learning rate, initialisation scale, Adam $\varepsilon, \beta_1, \beta_2$ and weight decay) with an evolutionary strategy. These hyperparameters (HPs) can be either optimised globally with a shared value across the entire model, or per-module (with 13 module types, some additionally tuned per depth). (Middle): The per-module approach leads to better results at the 50M scale -- optimal global HPs require $2.3\times$ longer training to achieve same performance. (Right): Crucially, our new parameterisation, Complete(d)P , enables a direct transfer (without any subsequent tuning) to a $\sim\!600\times$ larger FLOP budget. While this is true for the global HPs, this is also the case for our granular per-module HP setup, which at large scale result in a $1.32\times$ speed-up and improved benchmark performance (see \ref{['fig:7B-eval-benchmarks']}).
• Figure 2: Hyperparameter transfer for global learning rate across depth and width. Each setting is run with three independent seeds.
• Figure 3: Learning rate transfer with batch-size.Left: Learning rates transfer when using the square-root rule in \ref{['eq:multipliers_bs']}Right: Learning rates fail to transfer without adjustment. Each setting is run with three independent seeds.
• Figure 4: Weight decay transfer with batch-size.Left: Weight decay fails to transfer with batch-size without any adjustments. Middle: The rescaled (effective) weight decay $\lambda / \sqrt{\kappa}$ where $\kappa$ is the increase does transfer. Right: The effective weight decay transfers when rescaling all hyperparameters following our AdamW SDE scaling rule. Each setting is run with three independent seeds.
• Figure 5: Learning rate transfer across training horizon -- adjusting the number of tokens by changing the number of training iterations while holding batch-size constant. Left: Break-down of transfer of the global learning rate. Right: Stability of the “effective” learning rate -- one that preserves the AdamW SDE integration horizon. Each setting is run with three independent seeds.