DOCSEARCHbyAlgolia
Submit Docs

microsoft/mup

repository·main·Indexed 23 days ago

https://github.com/microsoft/mup

A tool for implementing Maximal Update Parametrization (μP) in PyTorch models. It enables Hyperparameter Transfer (μTransfer), allowing users to tune hyperparameters on small models and apply them to larger models with zero-shot stability. The library includes utilities for saving and loading base shapes, as well as the mup.coord_check module to verify implementation correctness by analyzing activation coordinates as model width increases.

Tokens
4.6K
Snippets
14
Records
20
Agent score
82%

What's inside mup

  1. Understand the core desiderata of μP

    main

    Maximal Update Parametrization (μP) is designed to satisfy three specific requirements during training of wide networks:

    1. Stable Activations: Every (pre)activation vector in the network should have $\Theta(1)$-sized coordinates.
    2. Stable Output: The neural network output should be $O(1)$.
    3. Maximal Updates: All parameters should be updated as much as possible (in terms of scaling in width) without leading to divergence.

    These requirements ensure that as you scale the width of your model, the internal dynamics and the learning signal remain stable, enabling effective hyperparameter transfer (μTransfer) from smaller models to larger ones.

  2. How μP scaling works under the hood

    main

    When set_base_shapes(model, ...) is called, each parameter tensor p in the model receives a p.infshape attribute. This attribute tracks which dimensions are infinite (scaling dimensions like d_model) and which are finite (fixed dimensions like vocabulary size).

    This metadata is used by mup initializers and optimizers to automatically scale parameters or learning rates. For example, the Adam learning rate for a weight p is calculated as globalLR / p.infshape.width_mult(), where width_mult() is the ratio of fan_in / base_fan_in.

  3. Install mup from source

    main

    To install from source, clone the repository and run the following commands in the repository directory:

    pip install -r requirements.txt
pip install -e .

  4. Verify μP implementation with coordinate check

    main

    Before scaling up to larger models, verify that your μP implementation is correct by checking the size of activation coordinates as model width increases. Using the --coord_check flag will generate plots in the ./coord_checks directory.

    In a correct μP implementation, coordinate sizes should remain stable as width increases, whereas in Standard Parametrization (SP), they will grow. Use the --load_base_shapes flag to provide the base shapes saved during the initial setup.

    # for SGD
python main.py --load_base_shapes width256.bsh --optimizer sgd --lr 0.5 --cuda --coord_check

# for Adam
python main.py --load_base_shapes width256.bsh --optimizer adam --lr 0.01 --cuda --coord_check

  5. Start μP training with hyperparameter transfer

    main

    Once base shapes are saved and the implementation is verified, you can scale up the model width. By providing the base shapes via --load_base_shapes, the script uses μP, allowing you to use the same hyperparameters from the small model for the wider models.

    Warning: If you do not specify the --load_base_shapes flag, the script will default to training using Standard Parametrization (SP).

    python main.py --load_base_shapes width64.bsh

  6. Train a scaled-up μP Transformer

    main

    Once the implementation is verified, you can scale up the model width (e.g., increasing --d_model) and train using the same hyperparameters used for the small model. To ensure μP is used instead of Standard Parametrization (SP), you must provide the base shapes using the --load_base_shapes flag.

    If --load_base_shapes is omitted, the script defaults to training an SP model.

    # for SGD
python main.py --d_model 4096 --load_base_shapes width256.bsh --optimizer musgd --lr 0.5 --cuda

# for Adam
python main.py --d_model 4096 --load_base_shapes width256.bsh --optimizer muadam --lr 0.01 --cuda

  7. Start μP Training with Scaled Width

    main

    Once the implementation is verified, you can scale up the model width using the --width_mult flag. When using μP, you can use the same hyperparameters (like learning rate) from the small base model, and they should transfer effectively to the wider model.

    Important: If you do not specify --load_base_shapes, the script will default to training a Standard Parametrization (SP) model instead of a μP model.

    # for SGD
python main.py --width_mult 2 --optimizer musgd

# for Adam
python main.py --width_mult 2 --optimizer muadam

  9. Verify μP implementation with a Coord Check

    main

    A coordinate check (or coord check) is a method to verify that your Maximal Update Parametrization (μP) implementation is correct. It involves calculating the average size (the l1 norm, or x.abs().mean()) of activation vectors and model outputs across different model widths and training steps.

    Correct Implementation Behavior:

    • The l1 values should remain stable (horizontal curves) as width increases.
    • Performance (training loss) should consistently improve as the model gets wider.

    Incorrect Implementation Behavior:

    • l1 values blow up or shrink to 0 as width increases.
    • Performance gets worse as the model gets wider.

    Exceptions to watch for: In a correct μP implementation, the following may shrink to 0 at initialization (at a $1/\sqrt{\text{width}}$ rate) but should become roughly flat after a few training steps:

    1. The network output.
    2. The attention logits in a Transformer.

    To resolve these transient discrepancies at initialization, it is recommended to:

    • Initialize the output layer (using MuReadout) with readout_zero_init=True.
    • Manually initialize the query matrix in a Transformer to 0.
    from mup.coord_check import get_coord_data, plot_coord_data

# construct a dictionary of lazy μP models with differing widths
def lazy_model(width):
    # `set_base_shapes` returns the model
    return lambda: set_base_shapes(MyMuModel(width), 'my/base/shape/path.bsh')
    # Note: any custom initialization with `mup.init` would need to
    # be done inside the lambda as well

models = {64: lazy_model(64), ..., 1024: lazy_model(1024)}

# make a dataloader with small batch size/seq len
# just for testing
dataloader = ...

# record data from the model activations over a few steps of training
# this returns a pandas dataframe
df = get_coord_data(models, dataloader)

# This saves the coord check plots to filename.
plot_coord_data(df, save_to=filename)

# If you are in jupyter notebook, you can also do
# `plt.show()`
# to show the plot

  10. Configure learning rate schedulers with mup optimizers

    main

    The mup optimizers (MuAdam, MuSGD) create refined parameter groups to scale learning rates according to μP. This is compatible with standard PyTorch learning rate schedulers, but if you implement a custom scheduler, you must update the learning rate relatively rather than setting it absolutely.

    Correct (Relative):

    pg['lr'] *= 2

    Incorrect (Absolute):

    pg['lr'] = 1e-3 * 2
    # OK: setting learning rate relatively
optimizer = mup.MuAdam(model.parameters(), lr=1e-3)
for pg in optimizer.param_groups:
    pg['lr'] *= 2

  11. Implement μP in a PyTorch model

    main

    To use Maximal Update Parametrization (μP), you must modify your model definition and training loop to use mup components instead of standard PyTorch ones.

    1. Model Definition Changes

    • Output Layer: Replace nn.Linear with MuReadout.
    • Shared Weights: If tying weights with an input nn.Embedding layer, use MuSharedReadout(input_layer.weight).
    • Attention Scaling: For Transformers, use 1/d scaling for attention scores instead of 1/sqrt(d). For backward compatibility with common head dimensions (e.g., $d=64$), use 8/d.

    2. Setting Base Shapes

    μP requires a base_model (small width) and a delta_model (to define which dimensions scale) to calculate the correct scaling for your target model.

    • Call set_base_shapes(model, base_model, delta=delta_model) as soon as possible, before re-initialization and optimizer definition.
    • Alternatively, save shapes to a file using make_base_shapes and load them later with set_base_shapes(model, filename) to save memory.

    3. Initialization and Optimization

    • Initialization: Replace torch.nn.init functions (e.g., uniform_, xavier_uniform_) with their mup.init equivalents.
    • Optimizer: Use MuSGD or MuAdam from mup.optim instead of torch.optim.