Add test coverage for Muon muon_lr/adam_lr overrides

[sowndappan5]

Summary

Add coverage for separate learning rate overrides in the Muon optimizer path and fix the related Muon blog documentation.

Background

Muon parameters and non-Muon parameters are automatically split into separate optimizer groups. The intended behavior is:

• muon_lr applies to Muon parameter groups
• adam_lr applies to Adam parameter groups
• lr remains the fallback for both groups when overrides are not provided

Changes

• add a parameterized test covering:
  • legacy lr fallback behavior
  • separate muon_lr / adam_lr override behavior
• fix the Muon blog table header to label muon_lr and adam_lr correctly

Validation

Ran:
python -m pytest DeepSpeed/tests/unit/ops/muon/test_muon_partial_training.py -k learning_rate_overrides -q -rs

Result:

• test collected successfully
• skipped locally because this distributed test requires 2 GPUs, while the local environment has 1 GPU

[delock]

It looks like change the title this way does not consistent with table content down below, i.e. AdamW learning rate is not muon learning rate. Does this file have to be modified?

[sowndappan5]

I removed the README table change and updated the PR to keep it focused on the test coverage for the existing muon_lr / adam_lr behavior.

[delock]

Hi @sowndappan5 thanks for your new test cases. Can you address the comments in README.md and also fix the DCO tests by sign-off your commits? Thanks!

[sowndappan5]

Thanks, I’ve addressed the README feedback and pushed the update. I’m fixing the DCO issue now by signing off the commits.

[sowndappan5]

I’ve addressed the README feedback and force-pushed signed-off commits to fix the DCO issue. The PR is now updated and pending review/CI.