Enable optimized Adam backend for MuonWithAuxAdam optimizer

[jinyouzhi]

Motivation

This pull request aim to extend the Muon optimizer auxiliary Adam(w) optimizer to support more built-in optimized implementation such as FusedAdam.

Changes

• Extracted Adam/AdamW backend selection logic into a new _select_adam_optimizer method, improving modularity and allowing consistent backend selection for both Adam and Muon optimizers.
• Keep original naive implementation as fallback path and extend to support weight-decay (AdamW)
• Support CPUAdam and Zenflow

Tests

pytest tests/unit/ops/muon/test_muon_partial_training.py

tested on RTX 5090: 128 Tensors 16,384 FP32 elements, warmup 30, steps 200.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[sfc-gh-truwase]

Why not this?

Suggested change

	optimizer_class, optimizer_kwargs = DeepSpeedEngine._select_adam_optimizer(engine, parameters, adam_w_mode)
	optimizer_class, optimizer_kwargs = engine._select_adam_optimizer(parameters, adam_w_mode)

[jinyouzhi]

Yes, it's strange. Thank you for your suggestions.

[jinyouzhi]

Fixed.

[jinyouzhi]

@sfc-gh-truwase Sorry, I need to revert this modification. engine here is an instance of the local _AdamSelectionEngine test double, not a fully constructed DeepSpeedEngine. Because of that, we need to call DeepSpeedEngine's method directly (DeepSpeedEngine.get_optimizer_configuration(engine, ...)) rather than as a bound method on engine.

Fully instantiating a real DeepSpeedEngine (via deepspeed.initialize()) would require a model, config, and distributed backend — overhead that's disproportionate to what this test is actually checking, and it would distract from the test's real purpose (verifying Adam-backend selection logic).

[sfc-gh-truwase]

@PKUWZP FYI

[pengdurice]

Thank you for the great PR. I think this changes warrants an end to end test with loss value and curves reported on a realistic training (e.g. a 128m model for 100 steps) between baseline and this change. LMK what you think.

[jinyouzhi]

Thank you for the great PR. I think this changes warrants an end to end test with loss value and curves reported on a realistic training (e.g. a 128m model for 100 steps) between baseline and this change. LMK what you think.

Thank you for your suggestion. I agree that an end-to-end training loss comparison would be necessary. I’ll set up a training-loss curve check focusing on the auxiliary Adam backend with Muon. Do you have a particular model or configuration you would recommend?

[pengdurice]

Thank you for the great PR. I think this changes warrants an end to end test with loss value and curves reported on a realistic training (e.g. a 128m model for 100 steps) between baseline and this change. LMK what you think.

Thank you for your suggestion. I agree that an end-to-end training loss comparison would be necessary. I’ll set up a training-loss curve check focusing on the auxiliary Adam backend with Muon. Do you have a particular model or configuration you would recommend?

Thank you!
Let's maybe use a small model e.g. 128mb ish? and make sure the naming of the parameters are recognizable by how we split between muon and Adam (

DeepSpeed/deepspeed/__init__.py

Line 84 in aa3914d

def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) -> None:

)
test baseline and the possible arms after the PR (fused Adam, cpu Adam etc)
Also, make sure the lr are set appropriately (not sure if the grad clipping issue is solved or not), the seed is set so for each arm same data example is used by different arms at the same step.
LMK what you think, thank you!

[tohtana]

Hi @jinyouzhi,
Thank you for submitting this PR! This is a significant improvement. The change looks okay to me, but I want to give you a heads-up.

Before this PR, non-Muon groups under the Muon optimizer use the inline Adam path. After this PR is merged, the same Muon configuration automatically selects FusedAdam. A pre-PR checkpoint therefore loads under the same user-facing optimizer and configuration, but its first resumed step fails because the new backend expects fields absent from the saved state (e.g. KeyError: 'bias_correction').
Is there a good way to avoid it?

[jinyouzhi]

Hi @jinyouzhi, Thank you for submitting this PR! This is a significant improvement. The change looks okay to me, but I want to give you a heads-up.

Before this PR, non-Muon groups under the Muon optimizer use the inline Adam path. After this PR is merged, the same Muon configuration automatically selects FusedAdam. A pre-PR checkpoint therefore loads under the same user-facing optimizer and configuration, but its first resumed step fails because the new backend expects fields absent from the saved state (e.g. KeyError: 'bias_correction'). Is there a good way to avoid it?

Nice catch, @tohtana! Thanks, I missed that. I’ve pushed a fix in commit cf0f900.

[jinyouzhi]

Thank you for the great PR. I think this changes warrants an end to end test with loss value and curves reported on a realistic training (e.g. a 128m model for 100 steps) between baseline and this change. LMK what you think.

Thank you for your suggestion. I agree that an end-to-end training loss comparison would be necessary. I’ll set up a training-loss curve check focusing on the auxiliary Adam backend with Muon. Do you have a particular model or configuration you would recommend?

Thank you! Let's maybe use a small model e.g. 128mb ish? and make sure the naming of the parameters are recognizable by how we split between muon and Adam (

DeepSpeed/deepspeed/__init__.py

Line 84 in aa3914d

def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) -> None:

)
test baseline and the possible arms after the PR (fused Adam, cpu Adam etc)
Also, make sure the lr are set appropriately (not sure if the grad clipping issue is solved or not), the seed is set so for each arm same data example is used by different arms at the same step.
LMK what you think, thank you!

I am still working on it. Quick update, the losses of FusedAdam and PyTorch impl is close to base's, there's a gap about CPUAdam, still investigate into it. Thank you for your advice.

[jinyouzhi]

Dataset: tatsu-lab/alpaca
Tokenizer: /mnt/disk3/hf_models/Qwen3-8B-FP8
Seed: 1234
Steps: 100
Sequence length: 128
Vocab size:   151,669
Hidden size:   1,024
Layers:            7
Total params: 318,116,981
≈ 318M parameters
Weight decay: 0
GPU: RTX 5090 D

┌────────────────────┬──────────────────────┬─────────────┬─────────────┬─────────────────────┐
│ Configuration      │ Backend              │  First Loss │  Final Loss │ Max Diff vs. Pre-PR │
├────────────────────┼──────────────────────┼─────────────┼─────────────┼─────────────────────┤
│ Pre-PR (24402c3be) │ Original inline Adam │ 11.93690586 │ 11.51892757 │                   0 │
├────────────────────┼──────────────────────┼─────────────┼─────────────┼─────────────────────┤
│ New PR             │ FusedAdam            │ 11.93690586 │ 11.51893520 │             1.14e-5 │
├────────────────────┼──────────────────────┼─────────────┼─────────────┼─────────────────────┤
│ New PR             │ PyTorch AdamW        │ 11.93690586 │ 11.51893520 │             1.05e-5 │
├────────────────────┼──────────────────────┼─────────────┼─────────────┼─────────────────────┤
│ New PR             │ CPUAdam offload      │ 11.93690586 │ 11.46116829 │             9.37e-2 │
└────────────────────┴──────────────────────┴─────────────┴─────────────┴─────────────────────┘

CPUAdam itself is supported. However, Muon CPU offload is not currently correctness-safe because ZeRO’s flat CPU partitioning path bypasses Muon’s full-matrix update semantics. I plan to address the Muon CPU offload issues in a separate PR.

[pengdurice]

Dataset: tatsu-lab/alpaca
Tokenizer: /mnt/disk3/hf_models/Qwen3-8B-FP8
Seed: 1234
Steps: 100
Sequence length: 128
Vocab size:   151,669
Hidden size:   1,024
Layers:            7
Total params: 318,116,981
≈ 318M parameters
Weight decay: 0
GPU: RTX 5090 D

┌────────────────────┬──────────────────────┬─────────────┬─────────────┬─────────────────────┐
│ Configuration      │ Backend              │  First Loss │  Final Loss │ Max Diff vs. Pre-PR │
├────────────────────┼──────────────────────┼─────────────┼─────────────┼─────────────────────┤
│ Pre-PR (24402c3be) │ Original inline Adam │ 11.93690586 │ 11.51892757 │                   0 │
├────────────────────┼──────────────────────┼─────────────┼─────────────┼─────────────────────┤
│ New PR             │ FusedAdam            │ 11.93690586 │ 11.51893520 │             1.14e-5 │
├────────────────────┼──────────────────────┼─────────────┼─────────────┼─────────────────────┤
│ New PR             │ PyTorch AdamW        │ 11.93690586 │ 11.51893520 │             1.05e-5 │
├────────────────────┼──────────────────────┼─────────────┼─────────────┼─────────────────────┤
│ New PR             │ CPUAdam offload      │ 11.93690586 │ 11.46116829 │             9.37e-2 │
└────────────────────┴──────────────────────┴─────────────┴─────────────┴─────────────────────┘

CPUAdam itself is supported. However, Muon CPU offload is not currently correctness-safe because ZeRO’s flat CPU partitioning path bypasses Muon’s full-matrix update semantics. I plan to address the Muon CPU offload issues in a separate PR.

thank you for the great work! Look forward to the PR fixing Muon CPU offload issue.