Add --muon-use-syrk option for Triton SYRK kernel in Newton-Schulz

[deepakn94]

Summary

Adds an opt-in --muon-use-syrk flag that routes the Gram-matrix computation inside Newton-Schulz through the Triton SYRK kernel in emerging_optimizers. The Gram matrix X @ X.T is symmetric, so SYRK computes roughly half the FLOPs of the equivalent GEMM.

Off by default, so behavior is unchanged unless the flag is set.

Wiring

• OptimizerConfig.muon_use_syrk (default False).
• --muon-use-syrk in _add_regularization_args.
• use_syrk parameter on TensorParallelMuon.__init__ and TensorParallelAdaptiveMuon.__init__.

No plumbing code is needed between the config and the optimizer: _kwargs_from_config already maps muon_<name> config fields onto matching __init__ parameters, so muon_use_syrk reaches use_syrk automatically.

Version compatibility

newton_schulz_tp — the only Newton-Schulz entry point Megatron calls — gained the use_syrk kwarg in emerging_optimizers 0.4.0:

emerging_optimizers	newton_schulz_tp accepts use_syrk
v0.2.0	no
v0.3.0	no
v0.3.1	no
v0.4.0	yes

Note that use_syrk appears on the non-TP newton_schulz as far back as v0.2.0, so the presence of the name in muon_utils.py is not a reliable signal; only the TP variant matters here.

The check is a version gate built on a new get_emerging_optimizers_version() / is_emerging_optimizers_min_version() pair in megatron/core/utils.py, following the existing is_te_min_version / is_mamba_min_version / is_causal_conv1d_min_version helpers. The minimum is spelled 0.4.0.dev0 so pre-release builds of that line are accepted, matching how the TransformerEngine minimums are written elsewhere in the tree.

Requesting --muon-use-syrk against an older release raises ValueError at optimizer construction rather than warning and silently falling back. A perf flag that appears to take effect but doesn't is worse than a hard failure, and the fallback made it impossible to tell from the logs whether SYRK was actually running.

Dependency pin

pyproject.toml currently pins emerging_optimizers at rev = "v0.2.0", so --muon-use-syrk will raise under the locked dependency until that pin moves to v0.4.0. This is intentional for now: the flag is opt-in and aimed at runs on newer builds, and bumping the pin pulls in everything else that changed across 0.2.0 → 0.4.0 plus a uv.lock regeneration, which belongs in its own PR.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[deepakn94]

Addressed the review feedback and rebased onto latest main.

@claude's two items:

1. TensorParallelAdaptiveMuon docstring — added the missing use_syrk entry after tp_mode.

2. Test coverage for the version gate — added two parametrized tests over both TensorParallelMuon and TensorParallelAdaptiveMuon:

   • test_muon_use_syrk_rejected_on_old_emerging_optimizers — asserts the ValueError fires when the version helper reports an older release. This is the only reachable branch under CI today given the v0.2.0 pin, so it was worth covering.
   • test_muon_use_syrk_default_off_ignores_version — asserts the gate does not fire on the default path, so an old emerging_optimizers still constructs fine when use_syrk is left off.

Correction to the version number. The gate is 0.4.0.dev0, not 0.3.0 as I said in the earlier thread. use_syrk exists on the non-TP newton_schulz back to v0.2.0, but newton_schulz_tp — the only Newton-Schulz entry point Megatron calls — only gained it at v0.4.0. A 0.3.0 minimum would have let v0.3.x through the check and then failed with TypeError: newton_schulz_tp() got an unexpected keyword argument 'use_syrk' at the first optimizer step, which is worse than the warning this replaces.

Note pyproject.toml still pins emerging_optimizers at rev = "v0.2.0", so --muon-use-syrk raises under the locked dependency until that pin moves to v0.4.0. Left for a separate PR since it needs a uv.lock regeneration and pulls in everything else that changed across 0.2.0 → 0.4.0.

[deepakn94]

/claude review

[claude]

Light review. Wiring looks correct — I traced muon_use_syrk -> use_syrk through _kwargs_from_config, and confirmed the second _kwargs_from_config(TensorParallelAdaptiveMuon, "adaptive_muon", config) in _adaptive_muon_config_to_kwargs does not clobber it (neither adaptive_muon_use_syrk nor use_syrk exists on the config, so the key is absent from the update dict). Two test-coverage gaps:

1. The positive path — the actual behavior change — is untested.

Both new tests only exercise the rejection branch. Nothing asserts that use_syrk=True actually reaches newton_schulz_tp, which is the entire point of the flag. Combined with the rev = "v0.2.0" pin, the ns_kwargs forwarding in scaled_orthogonalize_fn (megatron/core/optimizer/emerging_optimizers.py:221-227) will never execute in CI — a typo like {"syrk": True} would pass every test in this PR.

Suggest a test that patches the gate open and captures the kwargs:

@pytest.mark.parametrize("use_syrk,expected", [(True, True), (False, None)])
def test_muon_use_syrk_forwarded_to_newton_schulz(monkeypatch, use_syrk, expected):
    """use_syrk=True must reach newton_schulz_tp; use_syrk=False must not pass the kwarg."""
    import megatron.core.optimizer.emerging_optimizers as eo_module

    monkeypatch.setattr(eo_module, "is_emerging_optimizers_min_version", lambda _v: True)

    captured = {}

    def fake_newton_schulz_tp(grad, **kwargs):
        captured.update(kwargs)
        return grad

    monkeypatch.setattr(eo_module, "newton_schulz_tp", fake_newton_schulz_tp)

    model = torch.nn.Linear(60, 30, bias=False, dtype=torch.float32, device="cuda")
    optimizer = TensorParallelMuon(
        params=[model.weight], lr=0.01, pg_collection=None, tp_mode="duplicated",
        use_syrk=use_syrk,
    )
    model(torch.randn(8, 60, device="cuda")).sum().backward()
    optimizer.step()

    assert captured.get("use_syrk") is expected

The (False, None) case also pins the "do not forward the kwarg at all when disabled" contract that the new comment above ns_kwargs describes.

2. The new megatron/core/utils.py helpers have zero coverage.

get_emerging_optimizers_version() / is_emerging_optimizers_min_version() are never exercised — both new tests monkeypatch is_emerging_optimizers_min_version out of emerging_optimizers.py entirely, and that whole test file is pytest.mark.skipif-ed on LTS and when the package is missing. So the check_equality=False branch and the hyphenated-distribution-name fallback (version("emerging-optimizers")) are unverified.

tests/unit_tests/test_utils.py:57 already has the right shape for this (test_is_flashinfer_min_version), and putting it there gets it out from under the LTS skip:

@pytest.mark.skipif(not util.HAVE_PACKAGING, reason="packaging is not installed")
@pytest.mark.parametrize("check_equality", [True, False])
def test_is_emerging_optimizers_min_version(check_equality):
    from packaging.version import Version as PkgVersion

    with patch.object(util, "get_emerging_optimizers_version", return_value=PkgVersion("0.4.0")):
        assert util.is_emerging_optimizers_min_version("0.4.0.dev0", check_equality=check_equality) is True
        assert util.is_emerging_optimizers_min_version("0.5.0", check_equality=check_equality) is False
        assert (
            util.is_emerging_optimizers_min_version("0.4.0", check_equality=check_equality)
            is check_equality
        )

Nothing blocking beyond that — raising rather than silently falling back is the right call, and the get_*_version helper matches the existing is_mamba_min_version / is_causal_conv1d_min_version pattern (including letting the ImportError propagate when the package is absent).

[deepakn94]

/ok to test a5e91d4

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31429558540

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31438323017