[ISEEKYAN]

Summary

Add Muon optimizer support to verl's native Megatron backend.

verl's native Megatron backend previously supported only AdamW-type optimizers. Megatron-Core ships TensorParallelMuon (via the emerging_optimizers package), but verl did not expose it. This PR wires it through so verl users on the Megatron engine can select Muon.

What it does

• McoreOptimizerConfig (verl/workers/config/optimizer.py): expose optimizer=muon plus the muon_* hyperparameters, each mirroring the like-named field on Megatron-Core OptimizerConfig.
• verl/utils/megatron/optimizer.py: forward the muon_* fields to Megatron-Core and route parameters — Muon on 2D weight matrices, AdamW fallback on the rest (embeddings, norms, biases).
• LayerWise-DDP-aware wiring + master-weight buffer integration: avoids the double-buffer / extra fp32 master clone that would otherwise inflate optimizer memory.
• Fail-closed: requesting a muon_* field that the installed Megatron-Core build does not declare raises, instead of silently ignoring it.
• muon_match_adamw_update_rms: opt-in flag deriving muon_extra_scale_factor from sqrt((1 - beta1) / (1 + beta1)), the closed form that analytically matches AdamW's update RMS norm (emerging_optimizers 0.3.0 get_muon_scale_factor docstring; https://kexue.fm/archives/11267; https://arxiv.org/abs/2502.16982). Muon's effective step is lr * muon_extra_scale_factor and the Megatron-Core default 1.0 is not AdamW-comparable, so reusing an AdamW lr unchanged gives a ~4.4x larger effective step. The resolved value is logged on rank 0; the default stays 1.0 so no existing config changes behaviour.

Results

Qwen3-30B-A3B, 2 nodes × 8 GPUs (TP=2, PP=1, EP=8, DP=1), matched two-arm comparison (same data / seed / schedule; only the optimizer changes):

Metric	Muon	AdamW
Parameter update vs reference path	bitwise-identical	—
Peak optimizer memory (30B SFT, layer-wise)	39.66 GiB	46.63 GiB
Peak optimizer memory (30B GSM8K RL)	30.52 GiB	37.37 GiB
Actor update time (GSM8K RL)	7.50 s	15.30 s
GSM8K RL reward	level within run-to-run spread — see below

How these are judged. Two 30-step GSM8K A/Bs were run whose AdamW arms use a
byte-identical configuration, so the pair also measures this workload's run-to-run spread.
Every claim is reported against that spread rather than as a raw difference:

Metric	AdamW vs AdamW spread (n = 2)	Muon vs AdamW gap	ratio	verdict
Peak optimizer memory	0.14 GiB (0.4%)	6.85 GiB (18.3%)	49.5x	holds
Actor update time	1.89 s (11.6%)	7.80 s (51%)	4.1x	holds
Reward, last-3 mean	0.0104	0.0208 (Muon higher)	1.0x	within spread
Reward, last-10 mean	0.0234	0.0250 (AdamW higher)	1.1x	within spread
End-to-end step time	5.13 s (15.8%)	5.16 s	1.0x	within spread
Throughput	35.03 tok/s/GPU (86%)	—	—	not usable

Supported: Muon reduces peak optimizer memory by 18.3%, a gap 49.5x the run-to-run
spread of that metric, while preserving bitwise-correct parameter updates. Its
optimizer-update time is roughly half AdamW's, ~4x the spread — the portion of the step Muon
actually governs.

Not supported: on reward, Muon and AdamW are level within run-to-run spread. The gaps
match the AdamW-vs-AdamW spread in size and the sign flips between metrics (Muon ahead on
the last-3 mean, AdamW ahead on the last-10). An earlier total-gain difference of +0.032 is
likewise inside the 0.023 spread, so the previous "at least matching AdamW" wording is
withdrawn. No hyperparameter was tuned to produce a difference. End-to-end throughput and step
time are dominated by ~2x run-to-run variation in tokens per step and are not usable here.

Scoped to two 30-step short runs; the n = 2 spread estimate is itself crude.

Scale-factor caveat. The RL numbers were obtained with
muon_extra_scale_factor = 0.4227, a constant that came from a mis-calibration and is
not the recommended setting: upstream's often-quoted 0.2 is the value of the
factor, not a target for the measured update RMS. The recommended value is the closed
form sqrt((1 - beta1) / (1 + beta1)), which this PR exposes as the
muon_match_adamw_update_rms switch: it derives the factor from the configured
betas[0] at run time. Configure the switch, not a number -- a hand-entered constant
is the original problem regardless of how it was obtained. (For orientation only, the
expression evaluates to about 0.229 at beta1 = 0.9.) No A/B exists between the old
constant and the derived value; re-running the comparison is a pre-merge follow-up. The
memory and bitwise results above are independent of this factor.

Provenance of that constant: in this A/B the factor was injected as a literal value,
because the frozen experiment snapshot predates the muon_match_adamw_update_rms switch.
The injected value is bit-identical to what the switch computes: verified independently
with no factor present anywhere in the configuration, the switch resolves to the same
value on both the verl and Megatron-Lite sides, and tracks beta1 rather than being
fixed. The recommended usage is the switch, not a constant.

Compatibility

The optimizer wiring uses Megatron-Core's native get_megatron_optimizer and is independent of the model-build path. The LayerWise-DDP path is validated on both the vanilla Megatron-bridge and the full Megatron-Bridge (non-vanilla) model builds. On the full Megatron-Bridge path (Qwen3-30B-A3B, 12-step check): Muon peak 39.92 GiB vs AdamW 45.40 GiB (Muon saves ~12%), consistent with the vanilla path (39.64 GiB), confirming the layer-wise wiring holds under Megatron-Bridge.

Testing

• CPU: optimizer-config + wiring unit tests (tests/utils/megatron/, tests/workers/config/).
• GPU: muon_optim_gpu_smoke.py smoke test.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[CLAassistant]

All committers have signed the CLA.