文章

LLM μP 超参数跨规模迁移生产实战:用 Base Shape 指纹、Coord Check 与 Proxy Sweep 降低大模型调参成本

本文系统拆解 μP 超参数跨规模迁移的原理与工程边界,说明如何固化 Base Shape、执行 Coord Check、组织代理模型扫参与回放门禁,并避免把宽度迁移误用于训练时长、批量大小或架构变化。

背景问题:小模型最优参数,为什么放大后经常失效

训练十亿级甚至更大模型时,直接在目标规模上搜索学习率、初始化尺度、输出层倍率和优化器参数,成本通常不可接受。工程团队常采用一个较小的 Proxy Model(代理模型)做扫参,再把结果复制到目标模型。

问题在于,标准参数化下,模型宽度变化会同时改变激活尺度、梯度尺度、参数更新幅度和输出层行为。小模型上的最优学习率在大模型上可能过激,也可能过于保守。此时继续增加代理实验数量并不能解决根因——代理模型与目标模型遵循的优化动力学并不一致。

Maximal Update Parametrization(μP) 的目标,是通过一组随宽度变化的初始化、学习率和输出层缩放规则,让不同宽度模型的特征学习行为尽可能保持一致。由此产生的 μTransfer 工作流,是在小规模模型上寻找超参数,再迁移到更宽的目标模型。

关键限定:μP 主要讨论宽度方向的迁移。它不是”任意规模、任意架构、任意训练预算都能零成本迁移”的通用定律。

核心原理:迁移的不是一个学习率,而是一套参数化契约

1. Base Shape 决定哪些维度是可扩展宽度

μP 实现通常需要三个模型形状:

模型作用
Base Model定义基准形状
Delta Model相对 Base 改变所有计划扩展的宽度维度
Target Model实际训练的目标宽度

框架根据 Base 和 Delta 推断每个参数维度是有限维度还是随规模变化的”无限维度”,再为初始化和优化器参数组计算缩放规则。

如果 Transformer 同时扩展 d_model、Attention Head 数量或 FFN 宽度,Delta 必须体现这些变化。漏掉一个维度,代码仍可能运行,但迁移行为已经不再符合预期。

2. 不同参数类型不能共用同一种缩放方式

隐藏层矩阵、输入 Embedding、输出 Readout、Bias 和归一化参数承担的角色不同。μP 实现通常需要:

  • 使用 μP 兼容初始化;
  • 对隐藏矩阵设置按宽度变化的学习率;
  • 使用 MuReadout 或等价输出层;
  • 对 Transformer Attention 缩放进行相应调整;
  • 用 μP 优化器构建细分参数组。

因此,“只把 AdamW 换成 MuAdam”不是完整改造。只修改部分模块会形成混合参数化,最难排查——Loss 可能正常下降,但超参数迁移曲线已经失真。

3. Coord Check 验证跨宽度坐标尺度

Coordinate Check 会在多个宽度模型上记录初始化和少量训练步后的激活坐标尺度。如果实现正确,关键层的统计量应随宽度保持大致稳定,而不是持续爆炸或趋近于零。

它的定位类似梯度检查:验证实现是否符合预期,而不是证明最终训练一定成功。Coord Check 通过后,仍需验证学习率峰值位置、训练稳定性和目标模型短跑。

工程落地:把 μP 做成可发布的训练配置制品

1. 建立不可变 Base Shape 指纹

不要只保存一个 .bsh 文件。建议为每次参数化发布生成 Manifest

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

Manifest 至少应记录:

  • 参数名、Shape 及有限/无限维度分类;
  • Base、Delta 和 Target 模型配置;
  • 初始化策略、Attention 缩放、输出层类型;
  • 优化器参数组及初始倍率;
  • Tokenizer、词表和权重绑定关系;
  • 深度、Token 预算、批量大小和 Scheduler。

这些字段决定了一次 Proxy Sweep 是否可以与目标训练建立可审计的对应关系。

2. 在构建优化器前设置 Base Shape

官方实现要求尽早调用 set_base_shapes,通常应在自定义重新初始化和优化器创建之前完成:

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)

⚠️ Checkpoint 恢复时要特别谨慎。部分实现不会把附着在参数上的 infshape 元数据完整写入普通 Checkpoint。加载已训练权重后重新设置 Base Shape 时,应避免再次缩放已有参数,并通过单元测试确认恢复前后参数逐项一致。

3. Proxy Sweep 必须保持非宽度条件可比

代理模型与目标模型至少应固定:

  • 数据分布和 Tokenizer;
  • 优化器类型及参数定义;
  • Scheduler 形状;
  • 训练目标和 Loss 归一化;
  • Global Batch Tokens;
  • Token Horizon 或明确的迁移修正规则;
  • 模型深度和模块拓扑;
  • 混合精度策略。

如果 Proxy 只训练很短的 Token 预算,而目标模型训练时间长得多,最优学习率可能发生系统性变化。已有研究表明,最优学习率与 Token Horizon 存在关系,学习率与 Batch Size、临界批量大小之间也存在耦合。因此,μP 的宽度迁移不能替代时间维度和批量维度的建模。

4. 用多宽度曲线,而不是单一 Proxy 点

建议至少准备三个宽度:

宽度用途
小型扫参模型主搜索空间
中型验证模型验证迁移趋势
目标模型或短跑模型最终确认

对每个候选学习率,比较训练 Loss、梯度范数、激活统计和首批 Token 效率。如果最优学习率在小型和中型模型间已经明显漂移,就不应直接外推到目标规模。

比”某个学习率是否最好”更重要的是观察:

  • 排名是否稳定;
  • 最优区间是否重叠;
  • 对学习率误差是否鲁棒;
  • 更宽模型在相同步数下是否出现反常劣化。

5. 把 Embedding 与 Weight Decay 作为独立审计项

近期研究提出,在 AdamW 设置中,μP 相对标准参数化的显著收益可能很大程度来自 Embedding 层学习率不再成为瓶颈。另有研究对 μP 的长期动力学假设提出挑战,并强调 Weight Decay 和 Warmup 对跨宽度学习率迁移的重要性。

生产系统不应把 μP 当作黑盒。至少单独记录和扫描:

  • Embedding Learning Rate
  • Hidden Weight Learning Rate
  • Readout 倍率
  • Weight Decay
  • Warmup Token 数
  • 梯度裁剪阈值

如果迁移失败,先检查模块级更新幅度,而不是简单判断”μP 无效”。

Transfer Matrix:哪些参数可以迁移,哪些必须重新验证

变化维度建议
仅增加 Transformer 宽度μP 主要目标场景,可进行 Proxy Sweep 后迁移
改变模型深度不默认保证,需要独立验证或使用明确支持深度迁移的方法
改变总 Token 预算需要 Token Horizon 缩放研究或短跑校准
改变 Global Batch Tokens需要重新评估学习率与临界批量关系
Dense 改为 MoE架构和每个 Expert 有效 Token 变化,不应直接迁移
更换 Optimizer重新扫参,不视为同一迁移契约
更换 Tokenizer 或词表Embedding/Readout 形状与数据分布变化,重新验证
BF16 改为 FP8 训练数值稳定边界变化,需独立精度与 Scale 门禁

适用场景

μP 适合以下场景:

  • 计划训练多个同深度、同拓扑、不同宽度的模型;
  • 大模型直接扫参成本很高,但允许训练多个小型 Proxy;
  • 训练平台能够冻结代码、数据、Tokenizer 和优化器契约;
  • 团队愿意维护 Base Shape 和 Coord Check 测试。

不适合把它直接当作跨架构万能迁移方案。例如从 Dense 迁移到 MoE、同时改变深度和 Token 预算,或把短程实验的最优学习率直接用于超长预训练。

常见误区

误区一:Proxy 越小越省钱,迁移效果就越好

Proxy 过小可能处于不同优化区间,激活和梯度统计噪声也更大。应寻找”足够便宜且已进入稳定迁移区间”的最小模型,而不是无限缩小。

误区二:Coord Check 平坦就代表迁移成功

平坦曲线可能是学习率过小,模型几乎没有发生函数变化。官方项目建议使用足够大的学习率,使潜在爆炸问题能在少量步骤内暴露。

误区三:Base Shape 只是辅助文件

Base Shape 实际定义参数化语义。代码改名、层结构变化、权重绑定变化或新增投影层,都可能让旧 Base Shape 失效。它应像数据库 Schema 一样版本化。

误区四:目标模型完全不需要短跑

即使 Proxy Sweep 结果稳定,也应在目标规模执行受控短跑,检查 NaN/Inf、梯度范数、显存、吞吐、首段 Loss 下降和模块级更新倍率。

上线检查清单

  • Base、Delta、Target 深度与模块拓扑一致
  • Delta 覆盖全部计划扩展的宽度维度
  • Base Shape Manifest 已生成哈希并进入制品库
  • set_base_shapes 在初始化和优化器创建前执行
  • Readout、Attention 缩放、初始化和优化器均符合统一契约
  • 多宽度 Coord Check 在初始化及若干更新步后通过
  • Proxy Sweep 固定数据、Token Horizon、Batch 和 Scheduler
  • 小型与中型 Proxy 的最优区间具有重叠
  • Embedding LR、Weight Decay 和 Warmup 单独记录
  • 目标规模短跑通过数值稳定与 Loss 门禁
  • Checkpoint 恢复不会二次缩放参数
  • 迁移失败时可回退到标准参数化或上一版训练 Recipe

参考资料

  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

常见问题

μP 是否意味着大模型完全不需要调参?
不是。μP 主要改善随模型宽度变化的部分超参数迁移,训练时长、Token 预算、批量大小、深度、架构和优化器变化仍需独立验证。
Coord Check 通过后,学习率一定可以直接迁移吗?
不能保证。Coord Check 主要验证参数化实现是否随宽度保持稳定,还需要代理扫参一致性、目标规模短跑和训练稳定性门禁。
已经有标准参数化训练代码,能否只替换优化器实现 μP?
通常不行。还需处理 Base Shape、初始化、输出层、Attention 缩放、参数组和 Checkpoint 恢复等环节,否则会形成部分 μP、部分标准参数化的混合系统。