DeepSpeed 正式集成 Muon 优化器：从原理到上手配置

Muon 优化器在过去几个月里迅速获得了前沿 AI 实验室的青睐——Moonshot AI 等团队已经在大规模训练中采用它。现在 DeepSpeed 完成了对 Muon 的原生支持，意味着用 DeepSpeed 做分布式训练的开发者可以直接在配置文件里切换到这个优化器，不再需要自己魔改训练循环。

Muon 做了什么不同的事

传统 AdamW 对每个参数维护一阶动量和二阶动量（逐元素平方），内存开销是参数量的 2 倍。Muon 的核心思路是：用矩阵级别的动量替代逐元素动量。对于权重矩阵 $W$，Muon 维护一个同形状的动量矩阵 $M$，然后对 $M$ 做正交化（通过 Newton-Schulz 迭代近似 SVD），再用正交化后的方向做更新。

这样做带来两个直接好处：

• 内存减半——不再需要逐元素的二阶统计量，只存一个动量矩阵。
• 更新方向更"干净"——正交化后的梯度方向彼此独立，减少了参数更新之间的干扰，在深层网络里收敛更稳定。

Moonshot AI 的实践表明，在相同 batch size 和模型规模下，Muon 达到目标 loss 所需的步数明显少于 AdamW，尤其在长上下文语言模型的训练中效果突出。

DeepSpeed 里的启用方式

DeepSpeed 把 Muon 作为 MuonOptimizer 类型接入，配置方式和 AdamW 几乎一致，只需要改 type 字段并调整对应的超参数。

最小可运行配置

下面是一个单机多卡的 DeepSpeed 配置文件，启用 Muon 优化器：

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "optimizer": {
    "type": "MuonOptimizer",
    "params": {
      "lr": 0.02,
      "momentum": 0.95,
      "nesterov": true,
      "weight_decay": 0.1,
      "ns_iters": 5,
      "scalable_lr": true
    }
  },
  "fp16": {
    "enabled": "auto"
  },
  "bf16": {
    "enabled": "auto"
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "none"
    },
    "overlap_comm": true,
    "contiguous_gradients": true
  }
}

几个关键参数说明：

注意：Muon 的学习率和 AdamW 不在同一量级。如果你从 AdamW 的 lr=2e-4 切换过来，直接用同样的 lr 会导致训练几乎不动。一般建议从 lr=0.02 开始试。

训练脚本集成

以 HuggingFace Transformers + Accelerate 为例，启动脚本只需要指定 DeepSpeed 配置：

# 单机 8 卡训练
accelerate launch \
    --config_file ds_muon_config.json \
    --num_processes 8 \
    train.py \
    --model_name_or_path meta-llama/Llama-3.1-8B \
    --dataset_path data/finetune.jsonl \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --max_seq_len 4096 \
    --logging_steps 10 \
    --save_steps 500 \
    --output_dir checkpoints/muon-run

如果你用的是原生 DeepSpeed 启动器：

deepspeed --num_gpus=8 train.py \
    --deepspeed ds_muon_config.json \
    --model_path meta-llama/Llama-3.1-8B \
    --batch_size 4

训练代码本身不需要改动优化器逻辑——DeepSpeed 会在引擎初始化时根据配置文件自动创建 MuonOptimizer。

ZeRO 和 Muon 的搭配

Muon 只维护一个动量矩阵，不存二阶状态，这让它天然适配 ZeRO 的分片策略。实际测试中：

• ZeRO Stage 1：梯度分片，Muon 的动量仍然完整保留在每个 rank 上，内存节省主要来自优化器状态减半。
• ZeRO Stage 2：优化器状态也分片，配合 Muon 后每个 rank 只存动量矩阵的 $1/N$，总优化器内存约为 AdamW + ZeRO-2 的 $1/4$。
• ZeRO Stage 3：参数+优化器+梯度全分片，Muon 的优势叠加，对超大模型（70B+）最明显。

一个 7B 模型在 8 卡 A100 上的粗略对比：

这意味着同样的硬件可以撑起更大的 batch size或更长的序列。

超参数调优的实操建议

从 AdamW 迁移到 Muon 不是简单改一行配置就能完事的。以下是几个容易踩的坑和对应的处理方式：

学习率缩放：开启 scalable_lr: true 后，DeepSpeed 会自动按权重矩阵的维度对 lr 做缩放（大致按 $1/\sqrt{n}$ 的比例），这样不同大小的层不会需要手动设不同 lr。如果你关掉这个选项，就需要自己按层维度调整，不建议。

warmup 策略：Muon 对 warmup 比敏感。建议用线性 warmup，步数占总训练步数的 5%–10%，和 AdamW 的 2%–3% 相比要更长一些：

# 在 Transformers Trainer 的自定义配置中
from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="checkpoints/muon-run",
    learning_rate=0.02,
    warmup_steps=200,       # 总步数 4000 的 5%
    lr_scheduler_type="linear",
    deepspeed="ds_muon_config.json",
)

权重衰减的范围：Muon 的 weight_decay 只作用于权重矩阵，不对 bias、LayerNorm 参数、embedding 矩阵做衰减。DeepSpeed 的实现已经内置了这个过滤逻辑，你不需要手动分组。

梯度裁剪：Muon 的更新方向经过正交化，梯度范数的行为和 AdamW 不同。建议把 gradient_clipping 从 AdamW 习惯的 1.0 放宽到 2.0–5.0，或者先不开裁剪观察几轮再决定：

{
  "gradient_clipping": 2.0
}

什么时候该考虑换 Muon

Muon 不是万能替代。根据目前公开的实验数据和社区反馈，大致规律是：

• 大模型（>7B）、长序列训练：收益最明确，收敛速度和内存都有显著改善。
• 小模型微调（<1B）：优势不明显，AdamW 的成熟调参经验可能更稳。
• 卷积网络、非 Transformer 架构：Muon 的正交化是为矩阵参数设计的，对卷积核等非矩阵参数目前退化为普通动量更新，效果不确定。
• 和 MoE 模型搭配：有初步正面反馈，但公开数据还不多，建议先小规模验证。

一个简单的决策检查清单：

1. 你的模型是否以 Transformer 为主、参数量 > 7B？→ 优先试 Muon。
2. 你是否在 ZeRO-2/3 下受优化器内存瓶颈限制？→ Muon 直接减半，值得试。
3. 你是否有成熟的 AdamW 调参方案且训练已经稳定？→ 可以先在单卡上跑对比实验再决定。
4. 你的模型里非矩阵参数占比高吗？→ 暂时留在 AdamW 更安全。

DeepSpeed 对 Muon 的集成降低了尝试门槛——改一行配置、调一下 lr 和 warmup 就能跑起来。如果你正在做大模型训练且受限于优化器内存或收敛速度，花半天时间跑一轮对比实验，大概率会有直观的感受。