[question] Dose Muon Optimizer support deepspeed ?
Reminder
- [x] I have read the above rules and searched the existing issues.
System Info
NaN
Reproduction
我看 example 里的 muon 配置并没开 deepseed (https://github.com/hiyouga/LLaMA-Factory/blob/main/examples/extras/muon/qwen2_full_sft.yaml)
实际跑 deepseed + muon 崩了,self.tensor 会出现一个没有 use_muon value 。
简单绕过去,发现 loss 曲线没啥区别。 所以我怀疑现在 muon 还不支持 deepspeed ?
Others
No response
Can you try this PR? https://github.com/hiyouga/LLaMA-Factory/pull/7808
BTW, 代码仍需改一下 bug 才能执行。
Can you try this PR? #7808
看起来符合预期:
- 收敛更快
- 占显存更多,需要移植点别的进来
opencompass 精度在测,但应该问题不大。
哈喽,请问一下这两个曲线的batchsize设置的不一样吧?为啥epoch结束loss突降的点不在同一个step呢?
muon 要的显存多一些,batch 只能小一点,不然 OOM 了。
muon 要的显存多一些,batch 只能小一点,不然 OOM 了。
请问训练时间如何变化的 用了muon之后训练变慢显著吗
我已经跑起来了,确实要修改下才能跑 muon优化器 多占一倍显存,所以batch size需要变小 batch size变小的时候,loss是会更低的;作为对比实验我跑了一组使用adamW的batch size变小的实验 实验结果显示:muon优化器和adamW优化器loss曲线几乎重合,最后评估指标也无明显差异
怀疑这里muon的实现有问题,或者muon无用
我已经跑起来了,确实要修改下才能跑 muon优化器 多占一倍显存,所以batch size需要变小 batch size变小的时候,loss是会更低的;作为对比实验我跑了一组使用adamW的batch size变小的实验 实验结果显示:muon优化器和adamW优化器loss曲线几乎重合,最后评估指标也无明显差异
怀疑这里muon的实现有问题,或者muon无用 我也是发现了这个问题:muon优化器和adamW优化器loss曲线几乎重合,而且改了一下deepspeed zero3下的muon逻辑,发现即时有部分参数会传给muon,但是最后muon好像并没有对这些参数进行优化(或许adjust_lr_for_muon很小)
/
我已经跑起来了,确实要修改下才能跑 muon优化器 多占一倍显存,所以batch size需要变小 batch size变小的时候,loss是会更低的;作为对比实验我跑了一组使用adamW的batch size变小的实验 实验结果显示:muon优化器和adamW优化器loss曲线几乎重合,最后评估指标也无明显差异
怀疑这里muon的实现有问题,或者muon无用
你好,可以请问一下还需要修改哪里呢,现在还是会报KeyError: 'use_muon', 还有这个代码适配zero2吗
/
我已经跑起来了,确实要修改下才能跑 muon优化器 多占一倍显存,所以batch size需要变小 batch size变小的时候,loss是会更低的;作为对比实验我跑了一组使用adamW的batch size变小的实验 实验结果显示:muon优化器和adamW优化器loss曲线几乎重合,最后评估指标也无明显差异 怀疑这里muon的实现有问题,或者muon无用
你好,可以请问一下还需要修改哪里呢,现在还是会报KeyError: 'use_muon', 还有这个代码适配zero2吗
好像不支持zero3 、2,这里可以将将use_muon作为自定义属性直接附加到参数张量上,但是我分类了一下参数,发现 146 个参数由 Muon 优化器内部的 AdamW 逻辑更新。252 个参数由 Muon 优化器内部的 Muon 逻辑更新。这里可以看到有一部分参数传到了 Muon 逻辑,但实际上loss曲线还是和AdamW一致,我怀疑Muon 的更新来自于这一步:p.data.add_(u, alpha=-adjusted_lr)这个更新的效果要为零,要么是 u 几乎为零,要么是 adjusted_lr 几乎为零
使用muon训练 日志里显示 muon 优化器参数是0,参数都在adamw上,这个是配置文件哪里设置不对吗?
[INFO|2025-08-20 21:28:42] llamafactory.model.model_utils.checkpointing:143 >> Gradient checkpointing enabled.
[INFO|2025-08-20 21:28:42] llamafactory.model.model_utils.attention:143 >> Using torch SDPA for faster training and inference.
[INFO|2025-08-20 21:28:42] llamafactory.model.adapter:143 >> DeepSpeed ZeRO3 detected, remaining trainable params in float32.
[INFO|2025-08-20 21:28:42] llamafactory.model.adapter:143 >> Fine-tuning method: Full
[INFO|2025-08-20 21:28:42] llamafactory.model.loader:143 >> trainable params: 30,532,122,624 || all params: 30,532,122,624 || trainable%: 100.0000
[INFO|trainer.py:751] 2025-08-20 21:28:42,768 >> Using auto half precision backend
[WARNING|2025-08-20 21:28:42] llamafactory.train.callbacks:154 >> Previous trainer log in this folder will be deleted.
[INFO|deepspeed.py:380] 2025-08-20 21:28:43,631 >> Detected ZeRO Offload and non-DeepSpeed optimizers: This combination should work as long as the custom optimizer has both CPU and GPU implem
entation (except LAMB)
[INFO|2025-08-20 21:28:43] llamafactory.train.trainer_utils:143 >> Using Muon optimizer with 0 Muon params and 18867 AdamW params.
Installed CUDA version 12.8 does not match the version torch was compiled with 12.4 but since the APIs are compatible, accepting this combination
Using /root/.cache/torch_extensions/py311_cu124 as PyTorch extensions root...
Creating extension directory /root/.cache/torch_extensions/py311_cu124/cpu_adam...
Installed CUDA version 12.8 does not match the version torch was compiled with 12.4 but since the APIs are compatible, accepting this combination
Using /root/.cache/torch_extensions/py311_cu124 as PyTorch extensions root...
Emitting ninja build file /root/.cache/torch_extensions/py311_cu124/cpu_adam/build.ninja...
配置文件如下:
### model
model_name_or_path: /tmp/Qwen3-30B-A3B-Instruct-2507
trust_remote_code: true
### method
stage: sft
do_train: true
finetuning_type: full
deepspeed: examples/deepspeed/ds_z3_offload_config.json
packing: true
enable_liger_kernel: true
use_muon: true
### hardware 相关配置 V100 只能fp16,attention使用auto
flash_attn: auto # v100 卡选 auto,其他卡可以选 fa2
bf16: true # 其他机器使用 bf16 训练
fp16: false #V100 机器使用 fp16 训练
### dataset
dataset_dir: /LLaMA-Factory/data/ # 这里需要有一个dataset_info.json文件 配置具体数据文件的路径
# dataset: novel_writer
dataset: longwriter-6k
template: qwen3
cutoff_len: 8192
max_samples: 1000000
overwrite_cache: true
preprocessing_num_workers: 64
dataloader_num_workers: 8
### output
output_dir: /model_checkpoints/novel_writer_summary_qwen3_30B_a3B_2507 # 保存模型的路径
logging_steps: 1
save_steps: 5000
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: wandb # choices: [none, wandb, tensorboard, swanlab, mlflow]
run_name: novel_writer_summary_qwen3_30B_a3B_2507 # 用于 wandb 的 run name
### train
per_device_train_batch_size: 2
gradient_accumulation_steps: 1
learning_rate: 1.0e-5
num_train_epochs: 3
lr_scheduler_type: cosine
warmup_ratio: 0.1
ddp_timeout: 180000000
resume_from_checkpoint: null
### eval
# eval_dataset: alpaca_en_demo
# val_size: 0.1
# per_device_eval_batch_size: 1
eval_strategy: 'no'
# eval_steps: 500
不支持zero3, key error.
我看着这个 issue 一直有人评论。要不还是整理回复一下:
简单来说:不用试。
- 如果是想找更好的开源 optimizer 做 sft,想白嫖到更好的 LLM 效果,可能会失望
- kimi 用 muon 做的 pretrain, 直接用 muon 做 sft 不会好(目前 20250916)
- 苏神做的是前沿探索,光拿到他的公式,离验证实用还隔着 infra + 产品两个 team。
- 想跑起 muon training 需要 merge #7808 并且手动 fix 一下 key 不存在导致崩溃的问题
- 前面提了,跑起来也会爆显存
你说的那个最多支持zero2,你试试2个节点zero3
How about https://github.com/deepspeedai/DeepSpeed/pull/7509 ?
@tpoisonooo hi,方便问下为什么muon用在sft不会好吗?目前有哪些实验支撑这一点的?
@tpoisonooo hi,方便问下为什么muon用在sft不会好吗?目前有哪些实验支撑这一点的?
× language model + adam pretrain + muon sft 不好 √ 完全体(language model + muon pretrain + muon sft)肯定是好的 √ 图像任务是好的