[janEbert]

Add new Muon arguments from MCore #2241 and #3813 to distributed_muon_with_cosine_annealing.

Summary by CodeRabbit

• New Features
  • Added new configuration options for the Muon optimizer, enabling more granular control over optimization parameters including scale factors, optimizer selection, and beta coefficients.

[coderabbitai]

📝 Walkthrough

Walkthrough

The function distributed_muon_with_cosine_annealing() was extended with four new optional Muon optimizer parameters—muon_extra_scale_factor, muon_scalar_optimizer, lion_beta1, and lion_beta2—with documented defaults. These parameters are forwarded to the OptimizerConfig constructor without modifications to control flow or scheduling logic.

Changes

Cohort / File(s)	Summary
Muon Optimizer Parameters src/megatron/bridge/recipes/utils/optimizer_utils.py	Added four new parameters to distributed_muon_with_cosine_annealing(): muon_extra_scale_factor (default 1.0), muon_scalar_optimizer (default 'adam'), lion_beta1 (default 0.95), and lion_beta2 (default 0.98). Parameters are documented in the function docstring and passed through to OptimizerConfig.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Muon config update' clearly and concisely summarizes the main change: adding Muon configuration parameters to the distributed_muon_with_cosine_annealing function.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	Minor configuration parameter additions to optimizer utility function with sensible defaults; inherently low-risk change requiring no formal testing.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@src/megatron/bridge/recipes/utils/optimizer_utils.py`:
- Line 63: Fix the docstring typo for the lion_beta2 parameter in
optimizer_utils.py: change the description text that currently reads "momuntem
EMA" to "momentum EMA" (the docstring entry for lion_beta2 is the target).
- Line 37: The default value for muon_scalar_optimizer is using single quotes
which breaks ruff-format rules; update the declaration of muon_scalar_optimizer
in optimizer_utils.py (the variable named muon_scalar_optimizer) to use double
quotes for the string literal (i.e., change 'adam' to "adam") so it matches the
project's quoting convention and fixes CI formatting churn.
- Around line 99-102: The OptimizerConfig constructor is being given
muon_extra_scale_factor, muon_scalar_optimizer, lion_beta1, and lion_beta2
unconditionally which can raise TypeError on older megatron-core versions;
follow the existing pattern used for muon_nesterov and check
_OPTIMIZER_CONFIG_FIELDS before including each of these four keys into the
kwargs passed to OptimizerConfig—build a kwargs dict, conditionally add
muon_extra_scale_factor, muon_scalar_optimizer, lion_beta1, and lion_beta2 only
if they are present in _OPTIMIZER_CONFIG_FIELDS, then call
OptimizerConfig(**kwargs).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: ed485335-4649-4346-bc0a-32e63c158518

📥 Commits

Reviewing files that changed from the base of the PR and between f386fd6 and 5e95bb4.

📒 Files selected for processing (1)

• src/megatron/bridge/recipes/utils/optimizer_utils.py

[yaoyu-33]

/claude review

[claude]

The existing _MUON_NESTEROV_KWARG shim (line 25) guards against OptimizerConfig not having the muon_nesterov field. Should these four new fields (muon_extra_scale_factor, muon_scalar_optimizer, lion_beta1, lion_beta2) have the same guard? If Bridge needs to support MCore versions that predate these fields, passing them directly will raise a TypeError at runtime.

[janEbert]

I can add this, but I don't think I should?

[claude]

Missing test coverage: distributed_muon_with_cosine_annealing has no unit tests at all in tests/unit_tests/recipes/utils/test_optimizer_utils.py (only the Adam helpers are covered). Would be good to add a basic smoke test exercising the new muon_scalar_optimizer / lion_beta1 / lion_beta2 / muon_extra_scale_factor parameters.

[janEbert]

Added tests.

[yaoyu-33]

@janEbert please check contribute.md and fix lint.

[janEbert]

Fixed lints.

[suiyoubi]

LGTM