背景问题:小模型最优参数,为什么放大后经常失效
训练十亿级甚至更大模型时,直接在目标规模上搜索学习率、初始化尺度、输出层倍率和优化器参数,成本通常不可接受。工程团队常采用一个较小的 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
参考资料
- Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer — arXiv:2203.03466
- Microsoft μP Official Repository — github.com/microsoft/mup
- Scaling Optimal LR Across Token Horizons — OpenReview
- Time Transfer: On Optimal Learning Rate and Batch Size in the Infinite Data Limit — OpenReview
- Understanding the Mechanisms of Fast Hyperparameter Transfer — OpenReview
- Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate — arXiv:2605.21486
- Weight Decay may matter more than μP for Learning Rate Transfer in Practice — arXiv:2510.19093