[Feature] Muon optimizer support for the DTensor backend
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
Muon is currently a Megatron-only optimizer in NeMo-RL (`docs/guides/muon-optimizer.md:11`). I'd like to make it available on the DTensor backend too.
## Why
Optimizer choice and backend choice are separable concerns everywhere else. AdamW, schedulers, and gradient clipping all work on both backends. Muon is the only optimizer that gates on backend.
It matters because DTensor is the on-ramp the project itself recommends. All three exemplar configs (`sft.yaml`, `grpo_math_1B.yaml`, `dpo.yaml`) default to `dtensor_cfg.enabled: true`, the README quickstart runs them as-is, and Megatron is the path the README explicitly flags as needing extra setup care (cuDNN env vars, submodule init, TE built from source). For a user trying NeMo-RL for the first time on a small model, the recommended path silently removes Muon as an option even though Muon is exactly the kind of thing a small-scale experimenter would want to try. Its documented wins in `docs/guides/muon-optimizer.md` (Qwen3-235B SFT, Qwen2.5-7B DAPO) are about sample efficiency, which is precisely what matters at small scale.
Contributor guide
Assessment
This issue has not been assessed yet.