Kimi最近开源了Moonlight-16B-A3B，这个模型基于Muon优化器，在性能和训练效率上实现了双突破。这里有几个关键点：Muon优化器在大规模模型上的应用和扩展技术、Moonlight系列模型的参数规模与训练数据、以及Kimi开源的一系列资源，对后续研究和落地都很有价值。

01 前言

Muon优化器在训练小模型时表现很亮眼，但能否扩展到大规模模型，之前一直没得到验证。Kimi这次找到了两个扩展Muon的关键技术：

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

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

有了这两点，Muon在大规模训练中几乎可以开箱即用，不需要额外调超参。扩展定律实验也证实了这一点：在计算最优的训练配置下，Muon相比默认的AdamW优化器，样本效率能提高约2倍。

基于这些改进，Kimi用Muon训练了Moonlight-16B-A3B系列模型。这是一个16B参数的专家混合（MoE）模型，激活参数只有3B，训练数据用了5.7T个token。这个模型刷新了当前的帕累托前沿——用更少的训练FLOPs，就达到了比以往模型更好的性能。

同时，Kimi开源了内存优化和通信效率都很高的Muon实现，还发布了预训练、指令微调和中间检查点，给后续研究提供了很好的基础。

技术贡献主要包括：

• Muon有效扩展分析：通过大量实验，研究团队发现权重衰减在Muon的可扩展性中起到决定性作用。此外，通过参数级别的更新尺度调整，让不同的矩阵和非矩阵参数之间保持一致的更新均方根（RMS），大大提升了训练稳定性。

• 高效分布式实现：团队开发了带ZeRO-1风格优化的Muon分布式版本，既保证了内存效率最优，又减少了通信开销，同时还保留了算法的数学特性。

• 扩展定律验证：实验对比了Muon和强大的AdamW基线，结果非常清楚（见图1）。根据扩展定律，Muon只需要大约52%的训练FLOPs，就能达到与AdamW训练对应模型相当的性能。

【图1：扩展定律实验对比】
(a) Muon与AdamW的扩展定律实验对比，Muon样本效率是Adam的2倍；
(b) Moonlight模型（使用Muon优化）与其他可比较模型在MMLU上的表现。

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

02 性能

将Moonlight与类似规模的公开模型做对比，结果很有意思：

• LLAMA3-3B：3B参数密集模型，用了9T个token训练
• Qwen2.5-3B：3B参数密集模型，18T token
• Deepseek-v2-Lite：2.4B/16B参数的MoE模型，5.7T token

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优化器。这个框架是魔搭社区提供的大模型训练部署工具，开源在GitHub上（可自行搜索ms-swift获取）。

目前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.git
cd ms-swift
pip install -e .

微调脚本如下：

# 17GB
# ref: https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py
CUDA_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 \
  --sa ve_steps 50 \
  --sa ve_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 `即可：

{"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`为实际checkpoint路径，swift会自动从adapters文件夹读取参数）：

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

训练效果：

最后，如果需要推送模型到ModelScope，可以用export命令：

CUDA_VISIBLE_DEVICES=0 swift export \
  --adapters output/vx-xxx/checkpoint-xxx \
  --push_to_hub true \
  --hub_model_id '' \
  --hub_token ''