Add --muon-coefficient-type argument for Muon optimizer - #3927

Merged

mchrzanowski merged 25 commits into

NVIDIA:mainfrom

mchrzanowski:add-muon-coefficient-type

Mar 26, 2026

Commits on Mar 18, 2026

Add --muon-coefficient-type argument for Muon optimizer

root
and
claude
committed
444f17e
Use NSCoeffT Literal instead of private _COEFFICIENT_SETS

root
and
claude
committed
e357a41
Remove fallback coefficient types; require emerging_optimizers

root
and
claude
committed
bcac481
Run autoformatter on test file

root
and
claude
committed
17597ac
Derive test coefficient types dynamically from emerging_optimizers

root
and
claude
committed
bf834c1
Merge branch 'main' into add-muon-coefficient-type
mchrzanowski
authored
8a025d3

Commits on Mar 19, 2026

Merge branch 'main' into add-muon-coefficient-type
mchrzanowski
authored
678c832
Add emerging_optimizers >= 0.2 version gate for Lion import

root
and
claude
committed
7b8b042
Narrow broad Exception catch to PackageNotFoundError

root
and
claude
committed
02bad44

Commits on Mar 23, 2026

Merge branch 'main' into add-muon-coefficient-type
mchrzanowski
authored
c4d9c4a

Commits on Mar 24, 2026

Add emerging_optimizers >= 0.2 version gate for NSCoeffT and fix v0.2 API compat

root
and
claude
committed
07ee6fa
Consolidate HAVE_EO_V02 version check into __init__.py

root
and
claude
committed
4e929d2
Skip redundant HAVE_EMERGING_OPTIMIZERS check when HAVE_EO_V02 is true

root
and
claude
committed
6628801
Consolidate emerging_optimizers version detection into __init__.py

root
and
claude
committed
3de021e

Commits on Mar 25, 2026

Run black and isort autoformatting on PR files

root
and
claude
committed
420ebbd
Merge branch 'main' into add-muon-coefficient-type
mchrzanowski
authored
b945ada
Fix pylint E0606 for conditional imports of emerging_optimizers

root
and
claude
committed
99e8ad4
Fix black formatting for flash_attn imports in attention.py

root
and
claude
committed
8d9f7c5
Skip muon optimizer tests when emerging_optimizers is not installed

root
and
claude
committed
8d39d23
Run black autoformatting on test_muon_optimizer.py

root
and
claude
committed
a1dd63d
Merge branch 'main' into add-muon-coefficient-type
mchrzanowski
authored
b2936b2
Fix validate_coefficient_type crash when emerging_optimizers < 0.2

root
and
claude
committed
a8e82d0
Fall back to ("quintic",) for coefficient type validation without eo >= 0.2

root
and
claude
committed
0bc8f48
Narrow HAVE_EO_V02 skip to only tests that call get_supported_coefficient_types

root
and
claude
committed
0d613ec

Commits on Mar 26, 2026

Always pass coefficient_type to newton_schulz_tp regardless of EO version

root
and
claude
committed
10bf5eb