最近，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的权衡上推进了帕累托前沿。
将Moonlight与类似规模的SOTA公开模型进行了比较：
LLAMA3-3B是一个使用9T个token训练的3B参数密集模型
Qwen2.5-3B是一个使用18T个token训练的3B参数密集模型
Deepseek-v2-Lite是一个使用5.7T个token训练的2.4B/16B参数MOE模型
推理代码
from modelscope import AutoModelForCausalLM, AutoTokenizermodel_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)
显存占用：
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.gitgit 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>'
粤ICP备17114055号