feat(muon): add Muon optimizer support to the DTensor backend - #2505
Open
bzantium wants to merge 8 commits into
Open
feat(muon): add Muon optimizer support to the DTensor backend#2505bzantium wants to merge 8 commits into
bzantium wants to merge 8 commits into
Conversation
Lay down the Muon adapter module that the DTensor backend will dispatch into. This commit is single-GPU only and bit-exact against the vanilla emerging_optimizers.Muon class so the math foundation is locked before the TP-aware (Newton-Schulz over partition_dim) and dispatch wiring commits land on top. * nemo_rl/algorithms/muon/dtensor_muon.py: DTensorMuon extends emerging_optimizers.OrthogonalizedOptimizer. The orthogonalize step routes through the same primitives that emerging_optimizers.Muon uses (newton_schulz, get_muon_scale_factor) so per-step updates are bit-identical for the same input gradients and state. The TP-aware override is stubbed as a single-GPU pass-through; a follow-up commit switches it to newton_schulz_tp once the placement translator lands. * nemo_rl/algorithms/muon/chained.py: ChainedTorchOptimizer wraps two torch.optim.Optimizer instances behind one step / state_dict / param_groups surface so the policy worker can keep treating self.optimizer as a single object. * nemo_rl/algorithms/muon/builder.py: build_dtensor_muon walks model.named_parameters(), routes 2D non-embedding/non-norm/non-lm_head weights to DTensorMuon and the rest to torch.optim.AdamW, matching the linear vs nonlinear split in Megatron's get_megatron_muon_optimizer at 3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/optimizer/muon.py:239-258. * tests/unit/algorithms/test_dtensor_muon.py: 22 cases covering bit-exact equivalence vs vanilla Muon (3 scale modes x 3 coefficient types), multi-step + nesterov + weight_decay_method matrix, extra_scale_factor linearity, state_dict resume, ChainedTorchOptimizer semantics, and the param-split heuristics. emerging_optimizers is imported lazily with a clear install-extras error so the import does not break environments that do not have the mcore / muon extras installed. Signed-off-by: Minho Ryu <ryumin93@gmail.com>
Replace the hard-coded `optimizer_cls(model.parameters(), **kwargs)` construction in dtensor_policy_worker.py and automodel/setup.py with a single helper, build_optimizer_from_cfg, that dispatches based on a marker attribute on the resolved callable. * nemo_rl/utils/optimizer_factory.py: new helper that resolves the configured optimizer name to either a class or a builder function. Builders advertise model-aware construction by setting `_builds_optimizer_from_model = True` and receive the whole module rather than `model.parameters()`. * nemo_rl/algorithms/muon/builder.py: tag build_dtensor_muon with the marker so the dispatch hands it the model and the linear vs nonlinear parameter split can run. * nemo_rl/models/policy/workers/dtensor_policy_worker.py and nemo_rl/models/automodel/setup.py: call build_optimizer_from_cfg instead of the inline cls(model.parameters(), ...) pattern. Existing AdamW recipes hit the default branch and are unchanged. * tests/unit/algorithms/test_dtensor_muon.py: two new tests cover the default AdamW dispatch and the muon builder dispatch end-to-end. The resolver uses importlib directly rather than hydra.utils.get_class because get_class rejects callables that are not classes, which would prevent dispatching to the build_dtensor_muon factory function. Signed-off-by: Minho Ryu <ryumin93@gmail.com>
…n_dim)
Override DTensorMuon.orthogonalize so a DTensor parameter is routed
through emerging_optimizers.muon_utils.newton_schulz_tp on its TP-mesh
process group. Plain tensors and DTensors with all-Replicate placements
keep the single-GPU pass-through that the previous commit established.
Mechanics
* Parse param.placements: take the first Shard placement as the TP axis
and use placement.dim as partition_dim. Multi-axis sharding falls back
to the single Shard path; multi-axis Muon TP is out of scope.
* Resolve the TP process group via param.device_mesh.get_group(axis) and
hand it to newton_schulz_tp together with the local shard.
* Translate muon_tp_mode through to newton_schulz_tp's tp_mode argument:
- blockwise and duplicated both call tp_mode="duplicated", which
all-gathers and runs Newton-Schulz on the full matrix (matches
Megatron's TensorParallelMuon at muon.py:87)
- distributed forwards as-is, running Newton-Schulz on the local shard
with cross-rank communication
* The scale factor is derived from the **global** matrix shape so the
per-shard call still divides by the same denominator the single-rank
computation uses.
* The orthogonalized local shard is wrapped back into a DTensor with the
original placements so the base class's `p.add_(orth_grad, alpha=-lr)`
stays type-consistent.
Tests (CPU + gloo backend via torch.multiprocessing.spawn)
* test_dtensor_muon_tp2_duplicated_matches_single_rank: TP=2 with
duplicated and blockwise modes is bit-exact (rtol=0, atol=0) against
the single-rank reference for both Shard(0) and Shard(1) layouts.
* test_dtensor_muon_tp2_distributed_matches_single_rank: distributed
mode matches the reference within 5e-4 (Newton-Schulz's per-step
all-reduces add expected FP slack).
Signed-off-by: Minho Ryu <ryumin93@gmail.com>
Add the per-head split-orthogonalize-concat path that Megatron's TensorParallelMuon.orthogonalize uses for fused linear_qkv.weight parameters, and detect those parameters automatically in build_dtensor_muon. HF-style separated q_proj / k_proj / v_proj weights are normal 2D matrices and stay on the regular Muon path with no special handling. * nemo_rl/algorithms/muon/dtensor_muon.py: add split_qkv, is_qkv_fn and qkv_split_shapes constructor knobs. orthogonalize() now branches into _qkv_split_orthogonalize when the param is reported as fused QKV; the split mirrors Megatron's logic at 3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/optimizer/muon.py:136-159 including the (num_query_groups, q+k+v, hidden) reshape and the per-slice orthogonalize. The DTensor branch unwraps to the local shard, runs the split there, and re-wraps with the original placements so the base class step stays type-consistent. * nemo_rl/algorithms/muon/builder.py: detect fused params via the ``*.linear_qkv.weight`` name pattern. Derive qkv_split_shapes from model.config (num_attention_heads / num_key_value_heads / head_dim, with HF and Megatron field names both supported), or accept an explicit override. Tests * test_build_dtensor_muon_separated_qkv_skips_split_path: HF-style q_proj/k_proj/v_proj go through Muon individually with split disabled. * test_build_dtensor_muon_fused_qkv_with_explicit_shapes: a synthetic ``self_attention.linear_qkv.weight`` enables the split path and the is_qkv_fn correctly identifies the fused param. * test_dtensor_muon_qkv_split_matches_per_slice_orthogonalize: the fused split-orth-concat is bit-exact against splitting the gradient manually and orthogonalizing each slice with scaled_orthogonalize_fn. Signed-off-by: Minho Ryu <ryumin93@gmail.com>
Wrap up the Muon-on-DTensor work with a user-visible delta: a runnable recipe, a nightly entry, the muon extras packaging, and the doc update. * docs/guides/muon-optimizer.md: replace the "Megatron only" Requirements section with separate Megatron and DTensor checklists. The DTensor checklist references the new uv `muon` extras, the `nemo_rl.algorithms.muon.build_dtensor_muon` builder, the supported knob set (subset of the Megatron knobs since use_distributed_optimizer and overlap_param_gather are Megatron-DDP specific), and the new recipe. * pyproject.toml: add a `muon` extras group that exposes `emerging-optimizers==0.2.0` to DTensor users without requiring the full mcore install. The mcore extras keep their existing pin so existing recipes are untouched. * examples/configs/recipes/llm/sft-llama3.2-1b-1n8g-fsdp2tp1-muon.v1.yaml: runnable DTensor SFT recipe that mirrors the existing sft-llama3.2-1b-1n8g-fsdp2tp1.v3 recipe with the optimizer block swapped over to the Muon builder. * tests/test_suites/llm/sft-llama3.2-1b-1n8g-fsdp2tp1-muon.v1.sh: nightly driver that runs the recipe and asserts a loosened loss budget vs the AdamW baseline (tighten in a follow-up after a full convergence sweep). * tests/test_suites/nightly.txt: append the new driver under the SFT section. Signed-off-by: Minho Ryu <ryumin93@gmail.com>
github-actions
Bot
added
Documentation
Earlier commits validate the optimizer with end-to-end bit-exact tests vs
vanilla emerging_optimizers.Muon and multi-rank torchrun comparisons. The
DTensorMuon-specific branches (placement parsing for the FSDP+TP
composition case, the QKV split / orth / concat reshape) were only
verified indirectly through those end-to-end paths. Add isolated tests
so regressions surface in the smallest possible failure unit.
Placement parsing (single-process gloo + single_rank_pg fixture):
* test_placement_parser_replicate_only_dtensor_matches_plain_tensor:
a Replicate-placement DTensor routes through the no-shard branch and
produces the same per-step update as a plain-tensor input.
* test_placement_parser_picks_shard_axis_on_2d_mesh: a 2D ("dp", "tp")
mesh with placements=[Replicate(), Shard(d)] (the FSDP+TP composition
shape) selects the TP axis correctly for both d=0 and d=1.
QKV reshape invariants (pure-tensor, no distributed):
* test_qkv_split_preserves_total_rows_and_hidden_dim: round-trip with an
identity orthogonalize must reconstruct the input shape exactly.
* test_qkv_split_preserves_slice_order_q_then_k_then_v: slices come back
in Q -> K -> V order, parameterized across num_query_groups ∈ {1, 2, 5}
(covers single-query-group and GQA cases).
* test_qkv_split_rejects_misaligned_grad_rows: row counts not divisible
by sum(qkv_split_shapes) raise ValueError instead of silently bending
per-head boundaries.
* test_qkv_split_mha_equal_q_k_v_shapes: per-head q == k == v (the MHA
case) goes through the split path without losing structure.
Also force-set MASTER_PORT in the existing multi-rank workers so the
new single_rank_pg fixture's parent env does not leak through and pin
the spawned children to a port already bound by the fixture.
Signed-off-by: Minho Ryu <ryumin93@gmail.com>
Post-validation cleanup discovered while running a Megatron-vs-DTensor Muon equivalence sweep on Llama-3.2-1B-Instruct (train-loss delta < 1e-3 over 100 steps): - Bundle ``emerging-optimizers`` into the ``automodel`` extras and drop the standalone ``muon`` extras group. DTensor users now get the Newton-Schulz primitives from the same install set that ships the rest of the DTensor backend, and we avoid asking them to combine two extras (``automodel`` is mutually exclusive with ``mcore``, so a separate ``muon`` group could not be combined with both backends cleanly). - Rename the DTensor builder kwarg ``muon_use_nesterov`` to ``muon_nesterov`` to match Megatron's ``OptimizerConfig`` field name. Mismatched naming made it easy to copy a Megatron recipe and silently drop nesterov; aligning the two backends removes the foot-gun. - Filter ``None``-valued kwargs in ``build_optimizer_from_cfg``. Recipes that inherit from a base optimizer config (e.g. ``sft.yaml``'s AdamW with ``betas``/``eps``/``foreach``/``fused``) and switch to an optimizer whose signature does not accept those keys need a way to neutralize them; OmegaConf deep-merge keeps the keys with ``null`` values, and the factory now drops nulls before ``**kwargs`` unpacking. Adds a unit test that builds the DTensor Muon optimizer through the factory with the four AdamW-only fields explicitly nulled. - Sync the docs guide (``automodel`` instead of ``muon`` extras, renamed knob) and the nightly recipe yaml. Signed-off-by: Minho Ryu <ryumin93@gmail.com>
Author
Megatron
|
Megatron dist_muon |
DTensor build_dtensor_muon |
|
|---|---|---|
| Model | meta-llama/Llama-3.2-1B-Instruct | meta-llama/Llama-3.2-1B-Instruct |
| Dataset | OpenMathInstruct-2 (split_validation_size 0.05) | (same) |
| Steps / GBS / MBS / seq | 100 / 32 / 2 / 2048 | (same) |
| Optimizer knobs | lr 2e-5, wd 0.1, momentum 0.9, scale=spectral, extra_scale_factor=0.2, split_qkv | (identical), tp_mode=duplicated |
| Train loss @ step 100 | 0.4355 | 0.4350 |
| Validation loss @ step 50 | 0.4961 | 0.4957 |
| max |Δ| over 100 steps | 0.0014 | |
| mean |Δ| over 100 steps | 0.0004 |
Δ stays within FP non-determinism for bf16 matmul and Newton-Schulz orthogonalization across the two backends, which confirms the DTensor builder reproduces the Megatron dist_muon math.
Top panel: train loss for both backends (mcore solid blue, dtensor dashed red). Bottom panel: per-step delta (dtensor − mcore).
Throughput is not a controlled comparison here (mcore ~144 TFLOPS, dtensor ~75 TFLOPS at this shape). The DTensor backend's per-step cost is dominated by HuggingFace model-fwd rather than Muon orthogonalization, so closing the gap is orthogonal to the equivalence check this PR is about and belongs in a separate follow-up.
Code-review polish: - Collapse the multi-line docstring on test_optimizer_factory_drops_none_valued_kwargs to a single line; the long form duplicated optimizer_factory.py's own WHY comment. - Mirror the sibling test by asserting the inner optimizer type (DTensorMuon) in addition to ChainedTorchOptimizer, so the test documents the dispatch result, not just "didn't TypeError". Signed-off-by: Minho Ryu <ryumin93@gmail.com>
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Brings the Muon optimizer (MomentUm Orthogonalized by Newton-Schulz, https://arxiv.org/abs/2502.16982) to the DTensor backend.
docs/guides/muon-optimizer.md:11previously documented the optimizer as Megatron-only; this PR lifts that restriction so users on FSDP2 / TP recipes can train with Muon without switching backends.The optimizer math is reused verbatim from NVIDIA's
emerging_optimizerspackage (the same one Megatron'sdist_muonshims over). The DTensor side is a thin adapter: a placement topartition_dimtranslator, an all-gather path fortp_mode="duplicated", andnewton_schulz_tpfor the distributed path. Equivalence is verified to bit-exact (atol=0, rtol=0) vs upstreamemerging_optimizers.Muonon single-GPU, bit-exact for TP=2duplicated/blockwise, within 5e-4 for TP=2distributed, and max |Δ| 1.4e-3 across a 100-step Llama-3.2-1B SFT smoke vs Megatrondist_muonon identical recipe. The Megatrondist_muonpath and existing AdamW recipes are untouched; the dispatch is name-based innemo_rl/utils/optimizer_factory.py.Issues
Closes #2510.
Usage
uv sync --extra automodelbrings inemerging-optimizersfor DTensor users; the existingmcoreextras still ship the same package for the Megatron path.Before your PR is "Ready for review"
Pre checks:
Additional Information
PytorchOptimizerConfigalready accepts an openkwargsdict); no hidden defaults in code, and the exemplar recipe carries the explicit knob set.importlibdirectly rather thanhydra.utils.get_classbecauseget_classrejects callables that are not classes, which would prevent dispatching to thebuild_dtensor_muonfactory function.