微软的Maximal Update Parametrization (μP) 工具使用教程

微软的Maximal Update Parametrization (μP) 工具使用教程

mup maximal update parametrization (µP) 项目地址: https://gitcode.com/gh_mirrors/mu/mup

1. 项目介绍

Maximal Update Parametrization (μP) 是微软提出的一种参数化方法，用于优化大型神经网络的训练过程。μP 通过稳定的超参数设置，使得神经网络在不同规模下都能保持性能的稳定性和可扩展性。本项目提供了 μP 的实现，帮助用户在 PyTorch 模型中轻松应用这一方法。

2. 项目快速启动

安装

首先，您需要安装 μP 包。可以使用 pip 命令从 PyPI 安装：

pip install mup

或者，如果您希望从源代码安装，可以克隆仓库并执行以下命令：

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

基本使用

以下是使用 μP 的一些基本步骤：

1. 定义您的模型，将输出层替换为 MuReadout。

import torch.nn as nn
from mup import MuReadout

class MyModel(nn.Module):
    def __init__(self, width, ...):
        ...
        # readout = nn.Linear(width, d_out)
        self.readout = MuReadout(width, d_out)
        ...

    def forward(self, ...):
        ...

2. 如果您使用的是变压器模型，确保使用正确的注意力分数计算方式。

attention_scores = query @ key.T * 8 / d

3. 实例化基础模型和目标模型。

base_model = MyModel(width=1)
model = MyModel(width=100)

4. 设置基础形状。

from mup import set_base_shapes

set_base_shapes(model, base_model, delta_model)

5. 使用 μP 提供的初始化器和优化器。

from mup import MuSGD

optimizer = MuSGD(model.parameters(), lr=0.1)

6. 按照正常流程训练模型。

3. 应用案例和最佳实践

• 案例一： 使用 μP 来训练不同规模的 MLP 和 Transformer 模型。
• 最佳实践： 在训练大规模模型前，确保已经调用了 set_base_shapes 函数，并且使用了 μP 提供的优化器。

4. 典型生态项目

• PyTorch： μP 是为 PyTorch 模型设计的，可以与 PyTorch 的生态系统无缝集成。
• torchdistx： 可以使用 torchdistx.deferred_init 来延迟初始化大型模型，减少内存使用。

以上就是 μP 工具的使用教程，希望对您的项目有所帮助！

mup maximal update parametrization (µP) 项目地址: https://gitcode.com/gh_mirrors/mu/mup