Muon优化器使用教程

Muon优化器使用教程

Muon Muon optimizer: +~30% sample efficiency with <3% wallclock overhead 项目地址: https://gitcode.com/gh_mirrors/muon4/Muon

1. 项目介绍

Muon是一个为神经网络隐藏层设计的优化器。它旨在提高样本效率，同时在运行时几乎不增加额外开销。Muon优化器特别适用于大规模神经网络模型的训练，能够有效提升训练速度，已经在多个场景中展示了其优越性。

2. 项目快速启动

首先，确保您的环境中已安装Python。接着，您可以通过以下命令安装Muon：

pip install git+https://github.com/KellerJordan/Muon

安装完成后，您可以按照以下步骤使用Muon优化器：

from muon import Muon
import torch
import torch.nn as nn

# 假设您已经定义了模型结构
model = YourModel()

# 查找模型中隐藏层的参数
muon_params = [
    p for p in model.body.parameters() if p.ndim >= 2
]

# 查找模型中其他参数
adamw_params = [
    p for p in model.body.parameters() if p.ndim < 2
] + list(model.head.parameters()) + list(model.embed.parameters())

# 创建Muon优化器和AdamW优化器
optimizers = [
    Muon(muon_params, lr=0.02, momentum=0.95, rank=0, world_size=1),
    torch.optim.AdamW(adamw_params, lr=3e-4, betas=(0.90, 0.95), weight_decay=0.01)
]

# 在训练循环中使用优化器
for opt in optimizers:
    opt.step()

请根据您的模型结构，适当替换YourModel，model.body，model.head，和model.embed。

3. 应用案例和最佳实践

• 案例1： 在NanoGPT速度运行中使用Muon优化器，显著提高了训练效率。
• 案例2： 在CIFAR-10数据集上，使用Muon优化器将训练时间从3.3 A100秒降低到2.6 A100秒。

最佳实践：

• 对于大多数情况，默认的动量（0.95）和步数（5）可以提供良好的性能。
• 学习率应该有恒定的muP缩放，即随着模型规模的增加，无需重新调整学习率。

4. 典型生态项目

目前，Muon优化器已经被一些典型的生态项目采用，例如Kimi.ai的规模化LLM训练等，它们利用Muon优化器来提升模型训练的效率。

以上是Muon优化器的使用教程，希望对您的项目有所帮助。

Muon Muon optimizer: +~30% sample efficiency with <3% wallclock overhead 项目地址: https://gitcode.com/gh_mirrors/muon4/Muon