Add Maximal Update Parametrization (μP)
Author: qgallouedec
Pull request: yes

Adds opt-in μP via two fields on `PretrainedConfig` (`mup: bool`, `mup_base_width: int`) plus a single base-class hook in `PreTrainedModel`. `Trainer.create_optimizer` picks up `config.mup` automatically and applies the AdamW μP learning-rate split (matrix-like → `lr/m`, vector-like → `lr`) following `mup.MuAdam`; users building their own optimizer can call `transformers.integrations.build_mup_param_groups`.

closes https://github.com/huggingface/transformers/issues/16157

## What is it?

[Tensor Programs V](https://hf.co/papers/2203.03466) shows that under a specific (init, forward, optimizer) triple, per-coordinate activation, logit, and weight-update magnitudes stay O(1) as the hidden size `n` grows. The practical consequence is *μTransfer*: the optimal learning rate at a small "base" width is also the optimal LR at any wider model. Tune once on a small proxy, reuse the LR on the full-size model.

## Design

Three rewrites, all generic, all driven by `config.mup`:

- **`PreTrainedModel.post_init`**: walk modules. Replace `get_output_embeddings()` with a `MuReadout` of the same shape; on every attention module that exposes `scaling` and `head_dim`, divide `scaling` by `√head_dim` (turns `1/√d_head` into `1/d_head`, preserves any layer-idx factor).
- **`PreTrainedModel._initialize_weights`**: after the model's own init, divide `Linear`/`Conv1D` weights by `√(d/d0)`. Guarded by `weight._is_hf_initialized` so weights loaded from a state dict are not re-rescaled.
- **`MuReadout.forward`**: `super().forward(x / m)` pre-multiplies the input, matching the Microsoft `mup` library so the bias is correctly unscaled.

The whole feature is ~40 LoC in `modeling_utils.py` plus a small `integrations/mup.py`. It depends only on conventions every modern HF model already follows (`module.scaling`/`module.head_dim`, `get/set_output_embeddings`), which is why it works on Llama and GPT-2 with zero model-file changes (and should work on any similar model).

Because the marker lives in the config, save/load round-trips: a checkpoint trained with `mup=True` rebuilds the μP architecture on `from_pretrained`, so generation doesn't go to garbage. There's a regression test for this (`test_save_load_round_trip`).

## Why this exact recipe

Table 8 of the paper admits several SGD-equivalent formulations that diverge under Adam. This PR matches the canonical `mup.MuAdam` split:

| Parameter category                       | Forward     | Init std       | AdamW LR  |
|------------------------------------------|-------------|----------------|-----------|
| matrix-like (hidden, both dims width)    | x1          | `std / √m`     | `lr / m`  |
| vector-like (readout, embed, bias, LN)   | x1          | unchanged      | `lr`      |
| readout (`MuReadout`)                    | `x / m`     | unchanged      | `lr`      |

Logit boundedness comes from the readout's input rescale, not from a readout LR adjustment. (I had the LR direction inverted in an earlier draft; the coord check below caught it immediately. That's the case for keeping these empirical tests.)

## Empirical validation

```$ python examples/pytorch/mup_demo.py --mode coord
========================================================================
COORDINATE CHECK — per-module mean(|activation|) after a few steps
========================================================================
What to look for:
  • Under μP each row should be roughly constant across w=64…256 (spread ≈ 1).
  • Under SP (standard parametrization, the default) values fan out with width.

=== μP (base_width=64) ===
layer                                          w=64      w=128     w=256       spread
-------------------------------------------------------------------------------------
lm_head                                          0.1484    0.1215    0.0975     1.52x
model.embed_tokens                               0.0159    0.0162    0.0162     1.02x
model.layers.0.mlp.down_proj                     0.0042    0.0042    0.0040     1.05x
model.layers.0.mlp.gate_proj                     0.1310    0.1316    0.1293     1.02x
model.layers.0.mlp.up_proj                       0.1317    0.1312    0.1289     1.02x
model.layers.0.self_attn.k_proj                  0.1322    0.1295    0.1286     1.03x
model.layers.0.self_attn.o_proj                  0.0084    0.0085    0.0084     1.02x
model.layers.0.self_attn.q_proj                  0.1295    0.1263    0.1303     1.03x
model.layers.0.self_attn.v_proj                  0.1291    0.1310    0.1292     1.01x
model.layers.1.mlp.down_proj                     0.0041    0.0041    0.0041     1.01x
model.layers.1.mlp.gate_proj                     0.1330    0.1316    0.1314     1.01x
model.layers.1.mlp.up_proj                       0.1302    0.1320    0.1313     1.01x
model.layers.1.self_attn.k_proj                  0.1326    0.1276    0.1310     1.04x
model.layers.1.self_attn.o_proj                  0.0142    0.0152    0.0155     1.09x
model.layers.1.self_attn.q_proj                  0.1324    0.1269    0.1306     1.04x
model.layers.1.self_attn.v_proj                  0.1291    0.1321    0.1313     1.02x

  → mean spread: 1.06x

=== SP ===
layer                                          w=64      w=128     w=256       spread
-------------------------------------------------------------------------------------
lm_head                                          0.1483    0.2377    0.3851     2.60x
model.embed_tokens                               0.0159    0.0163    0.0163     1.02x
model.layers.0.mlp.down_proj                     0.0042    0.0136    0.0451    10.66x
model.layers.0.mlp.gate_proj                     0.1310    0.1908    0.2769     2.11x
model.layers.0.mlp.up_proj                       0.1317    0.1895    0.2761     2.10x
model.layers.0.self_attn.k_proj                  0.1323    0.1842    0.2632     1.99x
model.layers.0.self_attn.o_proj                  0.0084    0.0158    0.0363     4.34x
model.layers.0.self_attn.q_proj                  0.1296    0.1804    0.2670     2.06x
model.layers.0.self_attn.v_proj                  0.1291    0.1851    0.2618     2.03x
model.layers.1.mlp.down_proj                     0.0041    0.0121    0.0403     9.87x
model.layers.1.mlp.gate_proj                     0.1329    0.1862    0.2718     2.05x
model.layers.1.mlp.up_proj                       0.1302    0.1883    0.2737     2.10x
model.layers.1.self_attn.k_proj                  0.1326    0.1816    0.2720     2.05x
model.layers.1.self_attn.o_proj                  0.0142    0.0306    0.0729     5.13x
model.layers.1.self_attn.q_proj                  0.1326    0.1782    0.2653     2.00x
model.layers.1.self_attn.v_proj                  0.1291    0.1875    0.2715     2.10x

  → mean spread: 3.39x

========================================================================
SUMMARY  μP mean spread = 1.06x   |   SP mean spread = 3.39x
μP keeps activations roughly width-invariant; SP fans out (≈ width range = 4x).
========================================================================
```

Hidden-layer activations are flat across the 4× width range under μP (mean spread 1.06×) and fan out under SP (3.39×). This is the desideratum the paper proves μP enforces.

```
$ python examples/pytorch/mup_demo.py --mode lr-transfer
========================================================================
LEARNING-RATE TRANSFER — final loss after a few steps
========================================================================
What to look for:
  • Under μP the optimal LR (← marker) should be the SAME row at every width.
  • Under SP (standard parametrization, the default) the optimal LR shifts as width grows.

=== μP (base_width=64) ===
        lr  w=64        w=128       w=256     
----------------------------------------------
   1.0e-04   4.782     4.793     4.783  
   3.0e-04   4.676     4.688     4.677  
   1.0e-03   4.437     4.443     4.435  
   3.0e-03   4.058     4.052     4.040  
   1.0e-02   3.865 ←   3.954 ←   3.661 ←
   1.0e-01   4.525     5.407     5.221  

=== SP ===
        lr  w=64        w=128       w=256     
----------------------------------------------
   1.0e-04   4.782     4.653     4.217  
   3.0e-04   4.675     4.361     3.623  
   1.0e-03   4.436     3.864     2.940 ←
   3.0e-03   4.056     3.503 ←   3.001  
   1.0e-02   3.846 ←   3.755     4.122  
   1.0e-01   4.690     6.336    14.042  

========================================================================
SUMMARY  optimal LR per width:
  μP : 1e-02, 1e-02, 1e-02   ✓ width-invariant
  SP : 1e-02, 3e-03, 1e-03   ✗ shifts
========================================================================
```

`←` marks the column-wise minimum. Under μP it stays in one row; under SP it walks, and `lr=1e-2` even diverges by `w=256`. That's μTransfer end-to-end.

## Coverage

Works on any causal LM exposing `module.scaling`/`module.head_dim` and `get/set_output_embeddings` — ~250 of the model files in `src/transformers/models/`. Tested on Llama (CI) and GPT-2 (smoke). Not covered:

- **Inlined-scale legacy LMs** (Falcon, GPT-J, MPT, BLOOM, CodeGen): the attention rewrite no-ops because there's no `module.scaling` to patch: silently leaves `1/√d_head` while the rest applies. One-line per-model fix; not done here.
- **State-space / RNN** (Mamba, RWKV, Falcon-Mamba, Recurrent-Gemma, xLSTM): μP recipe differs.
- **Non-standard attention** (T5/MT5/UMT5, Reformer, Perceiver, BigBird, LongT5, Transfo-XL): would need per-model adapters.
- **Encoder-only / vision / audio / multimodal wrappers**: `mup` exists on their config but is a no-op without a causal-LM head, by design.
- `models/afmoe`'s existing `mup_enabled` is unrelated and left alone.

Out of scope: SGD recipe, automatic base-width detection.

## Tests

`tests/integrations/test_mup.py` — 6 cases on Llama: readout swap, attention scale, hidden init rescale, save/load round-trip, MuAdam-style param groups, and the empirical coord check (μP spread strictly less than SP spread).

## AI assistance

Assisted by Claude Code (Opus 4.7).


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_45848). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=45848&sha=da8150