Skip to content

added muon optimizer component#449

Open
BlueCrescent wants to merge 1 commit into
mainfrom
muon
Open

added muon optimizer component#449
BlueCrescent wants to merge 1 commit into
mainfrom
muon

Conversation

@BlueCrescent
Copy link
Copy Markdown
Member

What does this PR do?

This PR adds muon optimizer component

General Changes

  • Adde component
  • Added some tests for optimizer factory

Checklist before submitting final PR

  • My PR is minimal and addresses one issue in isolation
  • I have merged the latest version of the target branch into this feature branch
  • I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
  • I have run a sample config for model training
  • I have checked that all tests run through (python tests/tests.py)
  • I have updated the internal changelog (CHANGELOG_DEV.md)

Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds support for a new Muon optimizer component to the Modalities component/registry system, including configuration wiring and unit tests to ensure the optimizer factory constructs the optimizer using the expected parameter groups.

Changes:

  • Added OptimizerFactory.get_muon(...) that builds a Muon optimizer from get_optimizer_groups(...).
  • Introduced MuonOptimizerConfig and registered the "optimizer" / "muon" component variant.
  • Added unit tests verifying optimizer factory methods construct Adam/AdamW/Muon using computed optimizer groups.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File Description
tests/test_optimizer_factory.py Adds unit tests for Adam/AdamW/Muon factory construction and checkpointed optimizer loading behavior.
src/modalities/registry/components.py Registers the new "muon" optimizer component variant and wires it to MuonOptimizerConfig.
src/modalities/optimizers/optimizer_factory.py Adds get_muon factory method and imports Muon from torch.optim.
src/modalities/config/config.py Adds MuonOptimizerConfig to represent Muon optimizer configuration in the component system.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants