NVIDIA-NeMo / NVIDIA-NeMo/RL

[Feature] Muon optimizer support for the DTensor backend

Open
#2,510 2 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
community-request Feature waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.