Article

LLM μP Hyperparameter Cross-Scale Transfer in Production: Reducing Large Model Tuning Costs with Base Shape Fingerprints, Coord Checks, and Proxy Sweeps

A systematic breakdown of μP hyperparameter cross-scale transfer principles and engineering boundaries, covering how to solidify Base Shapes, run Coord Checks, organize proxy model sweeps with regression gates, and avoid misapplying width transfer to training duration, batch size, or architectural changes.

Background: Why Optimal Parameters for Small Models Often Fail at Scale

When training billion-parameter or larger models, directly searching for learning rates, initialization scales, output layer multipliers, and optimizer parameters at the target scale is usually prohibitively expensive. Engineering teams often use a smaller Proxy Model for hyperparameter sweeps and then replicate the results on the target model.

The problem is that under standard parameterization, changing model width simultaneously alters activation scales, gradient scales, parameter update magnitudes, and output layer behavior. The optimal learning rate on a small model may be too aggressive or too conservative on a large model. Simply increasing the number of proxy experiments doesn’t fix the root cause—the proxy and target models follow different optimization dynamics.

Maximal Update Parametrization (μP) aims to make feature learning behavior consistent across different model widths through a set of width-dependent initialization, learning rate, and output layer scaling rules. The resulting μTransfer workflow involves finding hyperparameters on a small-scale model and transferring them to a wider target model.

Key Limitation: μP primarily addresses width transfer. It is not a universal law for “zero-cost transfer across any scale, architecture, or training budget.”

Core Principles: Transferring Not a Learning Rate, But a Parameterization Contract

1. Base Shape Determines Which Dimensions Are Extensible Widths

A μP implementation typically requires three model shapes:

ModelRole
Base ModelDefines the reference shape
Delta ModelChanges all width dimensions planned for scaling relative to Base
Target ModelThe actual target width for training

The framework uses Base and Delta to infer which parameter dimensions are finite and which are “infinite” dimensions that scale with model size, then computes scaling rules for initialization and optimizer parameter groups.

If a Transformer simultaneously scales d_model, the number of Attention Heads, or FFN width, the Delta must reflect these changes. Missing a dimension means the code may still run, but the transfer behavior will no longer be as expected.

2. Different Parameter Types Cannot Share the Same Scaling

Hidden layer matrices, input Embeddings, output Readouts, Biases, and normalization parameters play different roles. A μP implementation typically requires:

  • Using μP-compatible initialization;
  • Setting width-dependent learning rates for hidden matrices;
  • Using MuReadout or an equivalent output layer;
  • Adjusting Transformer Attention scaling accordingly;
  • Building fine-grained parameter groups with a μP optimizer.

Therefore, “just swapping AdamW for MuAdam” is not a complete conversion. Modifying only some modules creates a hybrid parameterization that is hardest to debug—loss may decrease normally, but the hyperparameter transfer curves are already distorted.

3. Coord Check Validates Coordinate Scales Across Widths

Coordinate Check records activation coordinate scales after initialization and a few training steps across multiple width models. If the implementation is correct, key layer statistics should remain roughly stable as width increases, rather than exploding or trending toward zero.

Its role is analogous to a gradient check: verifying that the implementation behaves as expected, not proving that final training will succeed. After Coord Check passes, you still need to verify the learning rate peak location, training stability, and target model short runs.

Engineering Implementation: Making μP a Deployable Training Configuration Artifact

1. Establish an Immutable Base Shape Fingerprint

Don’t just save a .bsh file. It’s recommended to generate a Manifest for each parameterization release:

schema_version: 1
model_family: decoder-transformer
code_commit: "9f83c12"
base_shape_sha256: "..."
base_width:
  d_model: 512
  n_heads: 8
  d_ff: 2048
delta_width:
  d_model: 768
  n_heads: 12
  d_ff: 3072
fixed_dimensions:
  vocab_size: 128000
  n_layers: 24
mup_contract:
  readout: MuReadout
  attention_scale: "8 / head_dim"
  optimizer: MuAdamW-compatible
  init_policy: mup_init_v2
tokenizer_fingerprint: "..."
training_horizon_tokens: 3000000000
global_batch_tokens: 1048576

The Manifest should at least record:

  • Parameter names, shapes, and finite/infinite dimension classification;
  • Base, Delta, and Target model configurations;
  • Initialization strategy, Attention scaling, output layer type;
  • Optimizer parameter groups and initial multipliers;
  • Tokenizer, vocabulary, and weight tying relationships;
  • Depth, token budget, batch size, and scheduler.

These fields determine whether a Proxy Sweep can establish an auditable correspondence with target training.

2. Set Base Shape Before Building the Optimizer

The official implementation requires calling set_base_shapes early, typically before custom reinitialization and optimizer creation:

from mup import MuReadout, MuAdam, make_base_shapes, set_base_shapes

base = TransformerLM(d_model=256, n_heads=4, d_ff=1024)
delta = TransformerLM(d_model=512, n_heads=8, d_ff=2048)
make_base_shapes(base, delta, "transformer-v1.bsh")

model = TransformerLM(d_model=2048, n_heads=32, d_ff=8192)
set_base_shapes(model, "transformer-v1.bsh")
model.reset_parameters_with_mup_init()
optimizer = MuAdam(model.parameters(), lr=3e-4)

⚠️ Be especially careful when restoring checkpoints. Some implementations do not fully write the infshape metadata attached to parameters into regular checkpoints. When re-setting Base Shape after loading trained weights, avoid rescaling existing parameters, and use unit tests to confirm parameter-by-parameter consistency before and after restoration.

3. Proxy Sweep Must Keep Non-Width Conditions Comparable

The proxy and target models should at least fix:

  • Data distribution and Tokenizer;
  • Optimizer type and parameter definitions;
  • Scheduler shape;
  • Training objective and loss normalization;
  • Global Batch Tokens;
  • Token Horizon or explicit transfer correction rules;
  • Model depth and module topology;
  • Mixed precision strategy.

If the Proxy trains for a very short token budget while the target model trains much longer, the optimal learning rate may shift systematically. Research has shown a relationship between optimal learning rate and Token Horizon, as well as coupling between learning rate, Batch Size, and critical batch size. Therefore, μP width transfer cannot replace time-dimension and batch-dimension modeling.

4. Use Multi-Width Curves, Not a Single Proxy Point

It’s recommended to prepare at least three widths:

WidthPurpose
Small sweep modelMain search space
Medium validation modelVerify transfer trends
Target model or short-run modelFinal confirmation

For each candidate learning rate, compare training loss, gradient norms, activation statistics, and first-token efficiency. If the optimal learning rate drifts noticeably between small and medium models, do not extrapolate directly to the target scale.

More important than “which learning rate is best” is observing:

  • Whether rankings are stable;
  • Whether optimal intervals overlap;
  • Whether the model is robust to learning rate errors;
  • Whether wider models show anomalous degradation at the same step count.

5. Treat Embedding and Weight Decay as Independent Audit Items

Recent research suggests that in AdamW settings, the significant gains of μP over standard parameterization may largely come from removing the embedding layer learning rate bottleneck. Other studies challenge μP’s long-term dynamics assumptions and emphasize the importance of Weight Decay and Warmup for cross-width learning rate transfer.

Production systems should not treat μP as a black box. At a minimum, log and sweep independently:

  • Embedding Learning Rate
  • Hidden Weight Learning Rate
  • Readout Multiplier
  • Weight Decay
  • Warmup Token Count
  • Gradient Clipping Threshold

If transfer fails, first check module-level update magnitudes, rather than simply concluding “μP doesn’t work.”

Transfer Matrix: Which Parameters Can Be Transferred, Which Must Be Re-Validated

Change DimensionRecommendation
Increase Transformer width onlyμP’s primary target scenario; can transfer after Proxy Sweep
Change model depthNot guaranteed by default; needs independent validation or a method that explicitly supports depth transfer
Change total token budgetRequires Token Horizon scaling study or short-run calibration
Change Global Batch TokensRequires re-evaluation of learning rate vs. critical batch size relationship
Dense to MoEArchitecture and effective tokens per Expert change; do not transfer directly
Change OptimizerRe-sweep; not considered the same transfer contract
Change Tokenizer or vocabularyEmbedding/Readout shape and data distribution change; re-validate
BF16 to FP8 trainingNumerical stability boundaries change; requires independent precision and scale gates

Applicable Scenarios

μP is suitable for:

  • Planning to train multiple models with the same depth, topology, but different widths;
  • Direct hyperparameter sweeps on large models are very expensive, but training multiple small Proxies is feasible;
  • The training platform can freeze code, data, Tokenizer, and optimizer contracts;
  • The team is willing to maintain Base Shape and Coord Check tests.

Not suitable as a universal cross-architecture transfer solution. For example, transferring from Dense to MoE, simultaneously changing depth and token budget, or directly applying the optimal learning rate from a short experiment to very long pretraining.

Common Misconceptions

Misconception 1: The Smaller the Proxy, the Cheaper and Better the Transfer

A proxy that is too small may operate in a different optimization regime, with noisier activation and gradient statistics. Look for the smallest model that is “cheap enough and has entered a stable transfer regime,” rather than minimizing indefinitely.

Misconception 2: A Flat Coord Check Means Transfer Success

A flat curve could mean the learning rate is too small, and the model has barely changed functionally. The official project recommends using a sufficiently large learning rate so that potential explosion issues can be exposed within a few steps.

Misconception 3: Base Shape Is Just an Auxiliary File

Base Shape actually defines the parameterization semantics. Code renaming, layer structure changes, weight tying changes, or new projection layers can invalidate an old Base Shape. It should be versioned like a database schema.

Misconception 4: The Target Model Doesn’t Need Short Runs at All

Even if Proxy Sweep results are stable, perform controlled short runs at the target scale to check for NaN/Inf, gradient norms, memory, throughput, initial loss drop, and module-level update multipliers.

Go-Live Checklist

  • Base, Delta, and Target have consistent depth and module topology
  • Delta covers all width dimensions planned for scaling
  • Base Shape Manifest has been hashed and stored in the artifact repository
  • set_base_shapes is called before initialization and optimizer creation
  • Readout, Attention scaling, initialization, and optimizer all conform to a unified contract
  • Multi-width Coord Check passes at initialization and after several update steps
  • Proxy Sweep fixes data, Token Horizon, Batch, and Scheduler
  • Optimal intervals for small and medium Proxies overlap
  • Embedding LR, Weight Decay, and Warmup are logged separately
  • Target scale short run passes numerical stability and loss gates
  • Checkpoint restoration does not rescale parameters a second time
  • On transfer failure, can fall back to standard parameterization or previous training recipe

References

  1. Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter TransferarXiv:2203.03466
  2. Microsoft μP Official Repositorygithub.com/microsoft/mup
  3. Scaling Optimal LR Across Token HorizonsOpenReview
  4. Time Transfer: On Optimal Learning Rate and Batch Size in the Infinite Data LimitOpenReview
  5. Understanding the Mechanisms of Fast Hyperparameter TransferOpenReview
  6. Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning RatearXiv:2605.21486
  7. Weight Decay may matter more than μP for Learning Rate Transfer in PracticearXiv:2510.19093

FAQ

Does μP mean large models require no hyperparameter tuning at all?
No. μP primarily improves transfer of hyperparameters that vary with model width. Training duration, token budget, batch size, depth, architecture, and optimizer changes still require independent validation.
If Coord Check passes, can the learning rate be directly transferred?
Not guaranteed. Coord Check mainly verifies that the parameterization implementation stays stable across widths. Proxy sweep consistency, target model short runs, and training stability gates are still needed.
If I already have standard parameterized training code, can I just replace the optimizer to implement μP?
Usually not. You also need to handle Base Shape, initialization, output layer, attention scaling, parameter groups, and checkpoint restoration. Otherwise, you end up with a hybrid system that is part μP and part standard parameterization.