u-$\mu$P: The Unit-Scaled Maximal Update Parametrization

TL;DR

u-$$P combines Maximal Update Parametrization with Unit Scaling to create a training scheme that is size-independent, easy to train in low precision, and enables efficient hyperparameter tuning.

Contribution

The paper introduces u-$$P, a novel scheme that integrates $$P with Unit Scaling, simplifying training and hyperparameter optimization across model sizes.

Findings

u-$$P models achieve equal or lower loss than $$P models.

u-$$P works effectively in FP8 precision.

The scheme simplifies hyperparameter sweeping and training procedures.

Abstract

The Maximal Update Parametrization ($\mu$P) aims to make the optimal hyperparameters (HPs) of a model independent of its size, allowing them to be swept using a cheap proxy model rather than the full-size target model. We present a new scheme, u-$\mu$P, which improves upon $\mu$P by combining it with Unit Scaling, a method for designing models that makes them easy to train in low-precision. The two techniques have a natural affinity: $\mu$P ensures that the scale of activations is independent of model size, and Unit Scaling ensures that activations, weights and gradients begin training with a scale of one. This synthesis opens the door to a simpler scheme, whose default values are near-optimal. This in turn facilitates a more efficient sweeping strategy, with u-$\mu$P models reaching a loss that is equal to or lower than comparable $\mu$P models and working out-of-the-box in FP8.