[v0ropaev]

AdaptiveMuon fails on the first step after a param group is added with add_param_group. For example, building AdaptiveMuon([w1], lr=0.01, momentum=0.9, weight_decay=0.0), then calling opt.add_param_group({"params": [w2]}) and opt.step() with gradients on both parameters raises KeyError: 'beta2' from the group["beta2"] lookup in step(). Any training loop that adds layers to the optimizer partway through, such as progressive unfreezing, runs into this.

The cause is that beta2 and eps were set with setdefault on the groups that existed at construction time, but they never went into self.defaults. torch's add_param_group fills missing hyperparameters only from self.defaults, so a later group ended up with lr, momentum and weight_decay and nothing else.

The fix passes beta2 and eps to OrthogonalizedOptimizer through its existing **kwargs, which already go into the defaults dict, and removes the setdefault loop. Per-group values still win, so {"params": [w], "beta2": 0.5} keeps 0.5.

I added test_step_with_param_group_added_after_init to tests/test_adaptive_muon.py, parameterized over adamuon, normuon and namo. On main it fails for all three with KeyError: 'beta2', and with the change it passes. I ran test_adaptive_muon.py, test_orthogonalized_optimizer.py, test_registry.py, test_weight_update_hooks.py and test_weight_decay_mixin.py with --device=cpu, once with no seed, once with --seed=42 and once with --seed=3735928559, all passing on torch 2.8.0 CPU on macOS. I did not run the CUDA suite. ruff 0.9.9 check and format, the end-of-file and trailing-whitespace hooks, and mypy 1.19.1 on adaptive_muon.py are clean.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge.

Summary

Fixes AdaptiveMuon for parameter groups added after construction.

• AdaptiveMuon gives added parameter groups the right moment settings.

_{Reviews (1) · Last reviewed commit: "Fix AdaptiveMuon step for param groups a..." · Reviewed by Greptile}

[skyw]

LGTM.

Note that add_param_group is not something regularly tested across the repo, there may be other holes.

[skyw]

/ok to test bc4b7ed