Refactor emerging optimizer integration - #4113

Merged

skyw merged 11 commits into

NVIDIA:mainfrom

skyw:emerging-optimizer-refactor

Apr 3, 2026

Commits on Apr 2, 2026

[dev] refactor to support emerging optimizers beyond muon (#3618)
skyw
committed
5990608
Continue emerging optimizer refactoring (#3737)
skyw
committed
9fefb32
Fix emerging optimizer init_group for ckpt loading (#3897)
skyw
committed
710faf3
Fix missing imports in test_emerging_optimizers.py
skyw
committed
1f8ab3d
Fix TensorParallelMuon kwarg: mode -> tp_mode in test
skyw
committed
c65de58
remove test shouldn't exit yet
skyw
committed
5275694
Fix duplicate _EMERGING_OPTIMIZERS definition for autodoc2
skyw
committed
a4c8749
Fix test_layer_wise_optimizer: use get_megatron_optimizer's built-in LayerWise wrapping
skyw
committed
c74dd1a
Fix test overlap_param_gather: add use_layer_wise_distributed_optimizer=True
skyw
committed
762ca99
remove out of date assert
skyw
committed
90c5747
Pass overlap_param_gather config to LayerWiseDistributedOptimizer as async_allgather
skyw
committed
16f2e53