最大更新参数化 (μP) 和超参数迁移 (μTransfer)

论文链接 | 博客链接 | YouTube链接

在 Tensor Programs V: 通过零样本超参数迁移调优大型神经网络 中，我们表明，当以 最大更新参数化 (μP) 对模型进行参数化时，最优超参数会在不同规模的神经网络之间保持稳定。 这可用于调优极其庞大的神经网络，例如大型预训练 Transformer 模型，正如我们在工作中所做的那样。 更广泛地说，μP 能够降低从探索阶段过渡到大规模扩展时的脆弱性和不确定性，而这些往往在深度学习文献中并未被明确讨论。

上图：使用 Adam 优化器训练的不同 d_model 大小的 Transformer 的训练损失随学习率的变化。

事实证明，μP 是唯一一种具有这种跨宽度超参数稳定性特性的“自然”参数化方法，这一点在下方 GIF 动画中通过使用 SGD 训练的 MLP 得到了实证验证。在此过程中，我们随着时间推移，在 PyTorch 默认的学习率和初始化缩放与 μP 的缩放之间进行插值（右图），并利用这一插值后的缩放规则将宽度为 256 的模型（log2(width)=8）逐步扩展到宽度 2^13 = 8192（左图）。

本仓库包含 mup 软件包的源代码，这是我们用于在 PyTorch 模型中轻松且不易出错地实现 μP 的工具。

目录

• 安装
  • 从源码安装
• 基本用法
• mup 的工作原理
• 当前限制
• 检查参数化的正确性
  • 坐标检查
  • 自定义坐标检查图
  • 越宽越好
• 示例
• 运行测试
• 基础数学
• 贡献
• 商标

安装

pip install mup

从源码安装

克隆本仓库，进入其目录，并执行

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

基本用法

from mup import MuReadout, make_base_shapes, set_base_shapes, MuSGD, MuAdam

class MyModel(nn.Module):
    def __init__(self, width, ...):
        ...
        ### 在模型定义中，将输出层替换为 MuReadout
        # readout = nn.Linear(width, d_out)
        readout = MuReadout(width, d_out)
        ### 如果与输入的 nn.Embedding 层共享权重，则应使用
        # readout = MuSharedReadout(input_layer.weight)
        ...
    def forward(self, ...):
        ...
        ### 如果使用 Transformer，务必采用
        ###   1/d 而不是 1/sqrt(d) 的注意力缩放
        # attention_scores = query @ key.T / d**0.5
        attention_scores = query @ key.T * 8 / d
        ### 我们使用 8/d 而不是 1/d，以便向后兼容
        ###   当 d=64 时的 1/d**0.5，这是常见的头维度。
        ...

### 实例化一个基础模型
base_model = MyModel(width=1)
### 可选地，可以使用 `torchdistx.deferred_init.deferred_init` 来避免实例化参数
### 只需安装 `torchdistx` 并使用
# base_model = torchdistx.deferred_init.deferred_init(MyModel, width=1)
### 实例化一个“delta”模型，该模型在所有希望缩放的维度（即宽度）上都不同于基础模型。
### 这里很简单，但在 Transformer 中，你可能希望同时缩放 nhead 和 dhead，因此 delta 模型应在两者上都不同。
delta_model = MyModel(width=2) # 可选地使用 `torchdistx` 避免实例化

### 实例化目标模型（即你真正想要训练的模型）。
### 该模型应与基础模型相同，只是宽度可能有所不同。
### 特别要注意的是，基础模型和目标模型应具有相同的深度。
model = MyModel(width=100)

### 设置基础形状
### 当 `model` 的参数形状与 `base_model` 相同时，
###   `model` 的行为将完全等同于 `base_model`
###   （后者处于 PyTorch 的默认参数化方式下）。
###   这在当前模型尺寸下提供了向后兼容性。
###   否则，`model` 的初始化和学习率将按照 μP 进行缩放。
### 重要提示：应在重新初始化和定义优化器之前尽快调用此函数。
set_base_shapes(model, base_model, delta=delta_model)

### 或者，也可以将基础模型的形状保存到文件中
# make_base_shapes(base_model, delta_model, filename)
### 然后稍后直接从文件设置基础形状
# set_base_shapes(model, filename)
### 这在无法同时将 `base_model` 和 `model` 放入内存时非常有用

### 替换自定义的初始化方法（如有）
for param in model.parameters():
    ### 如果手动使用固定的标准差或边界进行初始化，
    ### 则用 mup.init 中的相应函数替代
    # torch.nn.init.uniform_(param, -0.1, 0.1)
    mup.init.uniform_(param, -0.1, 0.1)
    ### 同样，如果使用
    ###   `xavier_uniform_, xavier_normal_, kaiming_uniform_, kaiming_normal_`
    ### 来自 `torch.nn.init` 的函数，也应替换为 `mup.init` 中的对应函数

### 使用 `mup.optim` 中的优化器，而非 `torch.optim`
# optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer = MuSGD(model.parameters(), lr=0.1)

### 之后即可正常训练

请注意，基础模型和 delta 模型 无需训练 —— 我们只是从中提取参数形状信息。 因此，可选地，我们可以通过使用 torchdistx 中的 deferred_init 函数来避免实例化这些潜在的大模型。 在安装 torchdistx 后，可用 torchdistx.deferred_init.deferred_init(MyModel, **args) 代替 MyModel(**args)。更多详情请参阅 此页面。 在我们提供的 MLP 和 Transformer 示例中（非 mutransformers），可通过传递 --deferred_init 来启用此功能。

mup 的内部工作机制

通过调用 set_base_shapes(model, ...), 模型的每个参数张量 p 都会获得一个 p.infshape 属性，该属性为每个维度存储对应的基维以及该维度是应被视为“无限”（即会被放大或缩小，例如 Transformer 中的 d_model）还是“有限”（即会被固定，例如词汇表大小）。这些信息会在初始化器和优化器中使用，以自动调整参数或学习率，使其符合 μP 规范。例如，隐藏权重 p 的 Adam 学习率计算公式为 globalLR / p.infshape.width_mult()，其中 p.infshape.width_mult() 实际上计算的是 fan_in / base_fan_in。

当前限制

• set_base_shapes(model, ...) 假设 model 刚刚以标准方式随机初始化，并根据基形状信息重新缩放其参数，从而使模型处于 μP 状态。
• 如果需要数据并行，建议使用 torch.nn.parallel.DistributedDataParallel 而不是 torch.nn.DataParallel。这是因为后者会移除 mup 包为模型的每个参数张量添加的属性。此外，出于性能考虑，PyTorch 也推荐使用前者。
• 我们通过从传递给 mup 优化器的参数中创建细化的参数组，并操作这些组中的 lr 属性，显式地按照 μP 规范缩放学习率。这种方式与 PyTorch 的学习率调度器兼容。然而，如果你自定义调度器，请确保它设置的学习率是相对于细化参数组中当前值的比例关系。以下是一个错误做法和正确做法的示例：

optimizer = mup.MuAdam(model.parameters(), lr=1e-3)
for pg in optimizer.param_groups:
  # 错误做法：绝对设置学习率
  # pg['lr'] = 1e-3 * 2
  # 正确做法：相对调整学习率
  pg['lr'] *= 2

• 默认情况下，任何具有两个“无限”维度（即不同于基维度的维度）的参数矩阵都被 mup 视为具有 (fan_out, fan_in) 形状，即在前向传播中，该矩阵会将输入右乘。这适用于 PyTorch 中的所有 nn.Linear 权重。如果你有一个自定义参数，比如 W，它违反了这一约定，你可以手动设置 W.infshape.main_idx = 0; W.infshape.main = W.infshape[0]，以告知 mup 其形状对应于 (fan_in, fan_out)。类似的情况也适用于具有多个维度但恰好有两个“无限”维度的参数张量，其中第一个维度是 fan_in，第二个维度是 fan_out。
• 目前，torch.save 不会保存附加到每个参数张量的 infshape 对象。在问题修复之前，你必须在加载模型检查点后手动设置基形状，如下所示：

model = torch.load('my/model/path.pt')
# 重要提示：注意 `rescale_params=False` 标志！
set_base_shapes(model, 'my/base/shape/path.bsh', rescale_params=False)

（set_base_shapes 默认会重新缩放由 PyTorch 刚刚初始化的 model 参数，使其与 μP 保持一致。rescale_params=False 标志会关闭此行为。）

检查参数化是否正确

坐标检查

就像梯度检查是验证自动微分实现是否正确的一种简单方法一样，坐标检查也是一种验证你是否正确实现了 μP 的简便方法：在训练的几个步骤中，针对几种不同的宽度，计算模型输入和输出中每个激活向量坐标的平均大小（我们在下图的 y 轴上用 l1 表示）。如果实现正确，那么我们会看到这个 l1 在多种宽度下保持稳定；否则，随着宽度增加，l1 可能会爆炸式增长，也可能收缩至 0。 （我们实际上是在检验下面描述的第 1 个期望特性。） （l1 计算每个激活向量 x 的 x.abs().mean()，只是衡量 x 元素“平均大小”的一种指标；也可以使用类似定义的 l2、l4 等，尽管它们可能因随机种子的不同而出现更大的波动。）

例如，在下面的图表中，我们绘制了训练 2 个步骤时的 width 与 l1 曲线，其中 t=1 表示初始化时，尚未进行任何梯度更新。每条曲线对应于某一层的（预）激活向量或网络的输出。第一组 3 张图显示的是采用标准参数化（SP）的 MLP，由 adam 优化器训练。我们可以看到，经过 1 步更新后，激活/输出的 l1 随着宽度增加而迅速膨胀。这意味着 SP 是“不正确的”。 现在我们对采用最大更新参数化（μP）的 MLP 进行同样的操作（包括使用 mup.optim.MuAdam 而代之于 torch.optim.Adam）。与上述情况相反，所有曲线都保持水平，表明 μP 已被正确实现。 我们将这种检查实现正确性的方法称为“坐标检查”，简称“coord check”。

自己制作坐标检查图表

我们提供了一种简单的方法来实现这一检查，即使用 mup.coord_check 模块中的函数。典型的工作流程如下所示。

from mup.coord_check import get_coord_data, plot_coord_data
# 构建一个包含不同宽度懒加载 μP 模型的字典
def lazy_model(width):
    # `set_base_shapes` 返回模型
    return lambda: set_base_shapes(MyMuModel(width), 'my/base/shape/path.bsh')
    # 注意：任何使用 `mup.init` 的自定义初始化也需要在 lambda 内部完成
models = {64: lazy_model(64), ..., 1024: lazy_model(1024)}
# 创建一个批次大小/序列长度较小的数据加载器
#   仅用于测试
dataloader = ...
# 记录模型在训练的几个步骤中各层激活的数据
# 这将返回一个 pandas 数据框
df = get_coord_data(models, dataloader)
# 这会将坐标检查图表保存到指定文件名。
plot_coord_data(df, save_to=filename)
# 如果你在 Jupyter Notebook 中，也可以直接使用
#   `plt.show()`

# 用于展示图表

例如，mup.coord_check.example_plot_coord_check 函数就是以这种方式为玩具级的 MLP 和 CNN 模型实现的。

如果你在训练几轮后看到曲线随着宽度的增加而发散或收缩至零，那么你的 μP 实现中可能存在 bug（你是否忘记在 delta 模型中调整某些维度，比如 d_ffn？）。 相反，如果你看到曲线向右侧收敛，那么很可能你的实现是正确的。 不过，这里也有两种典型例外； 以下内容在 μP 的初始化阶段可能会收缩至零（以 1/√(宽度) 的速率）：

• 网络输出
• Transformer 中的注意力 logits

这些现象都是暂时的，经过几轮训练后，它们的曲线应该大致趋于平稳。 尽管如此，为了消除初始化时的差异，我们建议：

• 将输出层权重（应为 MuReadout 实例）通过 readout_zero_init=True 选项初始化为 0，
• 将 Transformer 中的查询矩阵手动初始化为 0。 如果希望在初始化时打破对称性以产生非零的注意力 logits，则可以使用非零方差来初始化（相对）位置偏置。

坐标检查技巧

• 使用较大的学习率（比实际训练时使用的还要大）。这样可以突出任何潜在的坐标爆炸问题，因为如果学习率过小，这些问题可能会被初始化掩盖。
• 如果你在前向传播中多次复用同一个模块，那么 mup.get_coord_data 只会记录最后一次使用的统计信息。在这种情况下，为了测试目的，你可以将不同的使用情况分别包裹在不同名称的 nn.Identity 模块中，以便区分它们。

更宽总是更好

另一个表明 μP 未正确实现的迹象是：在训练过程中，当网络宽度超过某个值后，更宽的模型反而会导致训练损失变差。 上图展示了一系列训练曲线：（左）正确的实现应在训练的任何阶段都表现出随着宽度增加性能提升；（中）如果你使用的是标准参数化（SP），有时可能会观察到性能随宽度增加到一定程度后突然下降；（右）或者你甚至会发现即使对于较窄的模型，性能也会立即开始恶化。

示例

请参阅 examples/ 文件夹中的 MLP、Transformer 和 ResNet 子文件夹，以及 mup/test 中的测试用例作为示例。 熟悉 Huggingface Transformers 的用户也可以参考 examples/mutransformers 子模块（可通过 git submodule update --init 获取），该子模块也可单独在 https://github.com/microsoft/mutransformers 上找到。

与 Huggingface 的原生集成

是否曾因你的 Huggingface Transformer 在扩展规模时出现问题而感到沮丧？或者想在单个 GPU 上直接调试大型多 GPU Huggingface Transformer 的超参数吗？如果是的话，请为 这个 GitHub 问题 点赞吧！

运行测试

要运行测试，执行以下命令：

python -m mup.test

基本数学原理

μP 的设计旨在满足以下要求：

在训练过程中的任何时候

1. 网络中的每个（预）激活向量都应具有 Θ(1) 大小的坐标
2. 神经网络的输出应为 O(1)
3. 所有参数都应在不导致发散的前提下，尽可能地根据网络宽度进行缩放更新

事实证明，这些要求唯一地确定了 μP。 为了从这些要求推导出 μP，我们需要仔细考虑当矩阵 A 和向量 v“相关”时，由 A 乘以 v 所得到的向量 Av 的坐标大小如何依赖于 A 和 v 的坐标大小。 在这里，可以把 A 看作权重，v 看作激活向量。 而这又取决于 A 是什么类型的矩阵，v 又是什么类型的向量。 在训练宽神经网络的背景下，我们只需要考虑那些具有近似独立同分布坐标的向量，以及两类矩阵：1) 类似于这类向量外积的矩阵，2) 随机的独立同分布矩阵。 第一类矩阵涵盖了诸如权重梯度之类的内容；第二类则对应于权重初始化等场景。 因此，如果 A 和 v 的元素大小均为 Θ(1)，且它们在训练过程中以自然方式相关联，那么就会出现下表所示的情况。

	外积矩阵 A（类型 1）	独立同分布矩阵 A（类型 2）
Av 的元素大小	Θ(n)	Θ(sqrt(n))

基于这张表，我们可以直接追踪网络的前向和反向计算过程，从而推导出 μP。

有关更详细的介绍，请参阅我们的博客文章和论文。

贡献说明

本项目欢迎各种贡献和建议。大多数贡献都需要你签署一份贡献者许可协议（CLA），声明你有权并将你的贡献权利授予我们使用。详情请访问 https://cla.opensource.microsoft.com。

当你提交拉取请求时，CLA 机器人会自动判断你是否需要提供 CLA，并相应地标记 PR（例如状态检查、评论）。只需按照机器人提供的指示操作即可。对于所有使用我们 CLA 的仓库，你只需完成一次即可。

本项目已采纳 微软开源行为准则。 如需更多信息，请参阅行为准则常见问题解答，或发送邮件至 opencode@microsoft.com 提出任何其他问题或意见。

商标说明

本项目可能包含其他项目、产品或服务的商标或标识。未经授权使用微软商标或标识的行为必须遵守并遵循 微软商标与品牌指南。 在本项目的修改版本中使用微软商标或标识时，不得造成混淆或暗示微软的赞助关系。 任何第三方商标或标识的使用均须遵守其各自的政策。

mup 快速上手指南

mup (Maximal Update Parametrization) 是一个用于 PyTorch 的工具包，旨在实现“最大更新参数化”（μP）。它能让神经网络的最优超参数（如学习率）在不同模型宽度下保持稳定。这意味着你可以在小规模模型上调优超参数，然后直接将其迁移到超大模型上进行训练，无需重新搜索。

环境准备

• 操作系统: Linux, macOS, Windows
• Python: 3.8 或更高版本
• 核心依赖:
  • PyTorch: 需安装与你的 CUDA 版本匹配的 PyTorch。
  • torchdistx (可选但推荐): 用于延迟初始化（deferred initialization），在处理极大模型时可节省显存和内存。

国内加速建议： 推荐使用清华或阿里镜像源安装 PyTorch 和 pip 包，以提升下载速度。

# 设置 pip 使用清华镜像
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

方式一：通过 Pip 安装（推荐）

直接使用 pip 安装稳定版：

pip install mup

方式二：从源码安装

如果你需要最新功能或进行开发，可以克隆仓库并安装：

git clone https://github.com/microsoft/mup.git
cd mup
pip install -r requirements.txt
pip install -e .

(可选) 安装延迟初始化工具 torchdistx：

pip install torchdistx

基本使用

要在现有 PyTorch 项目中启用 μP，只需进行以下四步修改：

1. 修改模型定义

• 将输出层替换为 MuReadout。
• 如果是 Transformer 架构，将 Attention 的缩放因子从 1/sqrt(d) 改为 8/d（为了兼容常见头维度 64）。

import torch.nn as nn
from mup import MuReadout

class MyModel(nn.Module):
    def __init__(self, width, d_out, ...):
        super().__init__()
        # ... 其他层定义 ...
        
        # 【修改点 1】替换输出层
        # readout = nn.Linear(width, d_out)  # 原生写法
        self.readout = MuReadout(width, d_out) # mup 写法
        
        # 如果权重共享 (如 Transformer)，使用:
        # self.readout = MuSharedReadout(input_layer.weight)

    def forward(self, x, ...):
        # ... 前向传播逻辑 ...
        
        # 【修改点 2】调整 Attention 缩放 (仅针对 Transformer)
        # 原生: scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)
        # mup:  使用 8/d 代替 1/sqrt(d) 以保持向后兼容 (当 d=64 时等效)
        d = query.shape[-1]
        attention_scores = query @ key.transpose(-2, -1) * 8 / d
        
        # ... 后续逻辑 ...
        return self.readout(x)

2. 设置基础形状 (Base Shapes)

实例化一个极小的“基准模型”（base model）和一个稍大的“增量模型”（delta model），并将它们的形状信息应用到你要训练的目标模型上。注意：基准模型和增量模型不需要训练。

from mup import set_base_shapes

# 实例化基准模型 (宽度设为 1)
base_model = MyModel(width=1)

# 实例化增量模型 (宽度设为 2，用于识别哪些维度是可扩展的)
delta_model = MyModel(width=2)

# 实例化目标模型 (实际要训练的宽度，例如 1024)
model = MyModel(width=1024)

# 【关键步骤】应用 μP 缩放规则
# 必须在重新初始化参数或定义优化器之前调用
set_base_shapes(model, base_model, delta=delta_model)

# 可选：如果显存不足，可使用 torchdistx 延迟初始化上述模型
# from torchdistx.deferred_init import deferred_init
# base_model = deferred_init(MyModel, width=1)
# delta_model = deferred_init(MyModel, width=2)
# model = deferred_init(MyModel, width=1024)
# set_base_shapes(model, base_model, delta=delta_model)

3. 替换初始化函数 (如有自定义初始化)

如果你在代码中手动调用了 torch.nn.init 系列函数，请替换为 mup.init 中对应的函数，以确保缩放正确。

from mup import init

for param in model.parameters():
    # 原生写法: nn.init.uniform_(param, -0.1, 0.1)
    # mup 写法:
    init.uniform_(param, -0.1, 0.1)
    
    # 同样适用于 xavier_uniform_, kaiming_normal_ 等

注：如果没有自定义初始化逻辑，可跳过此步，set_base_shapes 会自动处理标准初始化的重缩放。

4. 使用专用优化器

将 torch.optim 替换为 mup 提供的优化器（如 MuSGD, MuAdam）。它们会自动根据参数形状调整学习率。

from mup import MuAdam

# 原生写法: optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# mup 写法:
optimizer = MuAdam(model.parameters(), lr=1e-3)

# 之后即可像往常一样进行训练循环
# for batch in dataloader:
#     ...

完成以上步骤后，你在小模型上调优得到的学习率等超参数，即可直接应用于大规模模型训练。