01

前言

最近，Muon优化器在训练小规模语言模型方面展示了强大的效果，但其在大规模模型上的可扩展性尚未得到验证。Kimi确定了两个扩展Muon的关键技术：

权重衰减：对扩展到更大模型至关重要

一致的RMS更新：在模型更新中保持一致的均方根

这些技术使Muon能够在大规模训练中开箱即用，无需超参数调优。扩展定律实验表明，在计算最优训练中，Muon相较通常默认使用的AdamW优化器，能提供高约2倍的样本效率。

基于这些改进，Kimi基于Muon训练了Moonlight-16B-A3B系列模型。这是一个具备16B参数（激活参数为3B）的专家混合(MoE)模型，使用5.7T个token数据训练得到。该模型模型改进了当前的帕累托前沿，与之前的模型相比，使用更少的训练FLOPs就能实现更好的性能。

同时Kimi开源了内存优化和通信效率高的Muon实现，还发布了预训练、指令微调和中间检查点，以支持未来的研究。

所有代码可在MoonshotAI/Moonlight获取。

代码链接：

https://github.com/MoonshotAI/Moonlight

模型链接：

Moonlight-16B-A3B

    https://modelscope.cn/models/moonshotai/Moonlight-16B-A3B

Moonlight-16B-A3B-Instruct

    https://modelscope.cn/models/moonshotai/Moonlight-16B-A3B-Instruct

体验链接：

https://www.modelscope.cn/studios/moonshotai/Moonlight-16B-Demo/summary

技术贡献包括：

Muon有效扩展分析：通过广泛分析，研究团队发现权重衰减在Muon的可扩展性中起着关键作用。此外，研究团队提出通过参数级更新尺度调整，在不同的矩阵和非矩阵参数之间保持一致的更新均方根(RMS)。这些调整显著提高了训练稳定性。

高效分布式实现：研究团队开发了具有ZeRO-1风格优化的Muon分布式版本，实现了最优的内存效率和减少的通信开销，同时保持了算法的数学特性。

扩展定律验证：研究团队进行了扩展定律研究，将Muon与强大的AdamW基线进行比较，展示了Muon的卓越性能（见图1）。基于扩展定律结果，Muon只需要约52%的训练FLOPs就能达到与AdamW训练的对应模型相当的性能。

Muon的扩展

(a) 比较Muon和Adam的扩展定律实验，Muon的样本效率是Adam的2倍；

(b) Moonlight模型（使用Muon优化）与其他可比较模型在MMLU上的表现。

Moonlight在性能与训练FLOPs的权衡上推进了帕累托前沿。

02

性能

将Moonlight与类似规模的SOTA公开模型进行了比较：

LLAMA3-3B是一个使用9T个token训练的3B参数密集模型

Qwen2.5-3B是一个使用18T个token训练的3B参数密集模型

Deepseek-v2-Lite是一个使用5.7T个token训练的2.4B/16B参数MOE模型

03

模型推理

推理代码

from modelscope import AutoModelForCausalLM, AutoTokenizer
model_name = "moonshotai/Moonlight-16B-A3B-Instruct"model = AutoModelForCausalLM.from_pretrained(    model_name,    torch_dtype="auto",    device_map="auto",    trust_remote_code=True,)tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
prompt = "1+1=2, 1+2="inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(model.device)generated_ids = model.generate(**inputs, max_new_tokens=100)response = tokenizer.batch_decode(generated_ids)[0]print(response)

显存占用：

04

Muon优化器微调

ms-swift第一时间提供了Muon优化器的对接。ms-swift是魔搭社区提供的，大模型训练部署框架，其开源地址为：https://github.com/modelscope/ms-swift

鉴于目前moonshotai/Moonlight-16B-A3B系列MoE模型不再支持进一步的微调（由于topk_method='noaux_tc'），我们这里选择使用由Moonshot改进的muon优化器，来实现对于dense模型的微调。具体而言，在以下范例中，我们采用了Qwen2.5-7B-Instruct，来验证通过swift进行基于Muon优化器的微调训练。

在开始微调之前，请确保您的环境已准备妥当。

# pip install git+https://github.com/modelscope/ms-swift.git
git clone https://github.com/modelscope/ms-swift.gitcd ms-swiftpip install -e .

微调脚本如下：

# 17GB# ref: https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.pyCUDA_VISIBLE_DEVICES=0 \swift sft \    --model Qwen/Qwen2.5-7B-Instruct \    --train_type lora \    --dataset 'AI-ModelScope/alpaca-gpt4-data-zh#500' \              'AI-ModelScope/alpaca-gpt4-data-en#500' \              'swift/self-cognition#500' \    --optimizer muon \    --torch_dtype bfloat16 \    --num_train_epochs 1 \    --per_device_train_batch_size 1 \    --per_device_eval_batch_size 1 \    --learning_rate 1e-4 \    --lora_rank 8 \    --lora_alpha 32 \    --target_modules all-linear \    --gradient_accumulation_steps 16 \    --eval_steps 50 \    --save_steps 50 \    --save_total_limit 5 \    --logging_steps 5 \    --max_length 2048 \    --output_dir output \    --system 'You are a helpful assistant.' \    --warmup_ratio 0.05 \    --dataloader_num_workers 4 \    --model_author swift \    --model_name swift-robot

训练显存占用：

如果要使用自定义数据集进行训练，你可以参考以下格式，并指定`--dataset <dataset_path>`。

{"messages": [{"role": "user", "content": "浙江的省会在哪？"}, {"role": "assistant", "content": "浙江的省会在杭州。"}]}{"messages": [{"role": "system", "content": "你是个有用无害的数学计算器"}, {"role": "user", "content": "1+1等于几"}, {"role": "assistant", "content": "等于2"}, {"role": "user", "content": "再加1呢"}, {"role": "assistant", "content": "等于3"}]}

训练完成后，使用以下命令对训练后的权重进行推理：

提示：这里的`--adapters`需要替换成训练生成的last checkpoint文件夹。由于adapters文件夹中包含了训练的参数文件`args.json`，因此不需要额外指定`--model`，swift会自动读取这些参数。如果要关闭此行为，可以设置`--load_args false`。

CUDA_VISIBLE_DEVICES=0 \swift infer \    --adapters output/vx-xxx/checkpoint-xxx \    --stream true \    --temperature 0

训练效果：

推送模型到ModelScope：

CUDA_VISIBLE_DEVICES=0 \swift export \    --adapters output/vx-xxx/checkpoint-xxx \    --push_to_hub true \    --hub_model_id '<your-model-id>' \    --hub_token '<your-sdk-token>'

点击阅读原文，直达体验~