Route Lion through DistributedOptimizer and support single-moment checkpointing - #5742
Conversation
|
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. |
467c807 to
74dd5e8
Compare
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Light review — one correctness concern (posted inline) and one test-coverage note.
Test coverage: This PR generalizes the DistributedOptimizer checkpoint save/load paths to support non-Adam optimizers, but there's no test exercising a distributed-checkpoint round-trip with a single-state-key optimizer (e.g. Lion). tests/unit_tests/test_lion_optimizer.py only covers the factory/config plumbing, and tests/unit_tests/dist_checkpointing/test_optimizer.py only uses Adam. Since setup_model_and_optimizer in tests/unit_tests/dist_checkpointing/utils.py already accepts an optimizer= argument, consider parametrizing an existing save/load test (e.g. TestDistributedOptimizer) over optimizer=['adam', 'lion'] so the new optimizer_state_keys path and the dp_zero gather/scatter are covered with a differing state-key count. This would also catch the step-key issue flagged inline.
Tested on a small MoE model, loss curve looks continuous |
latest test with routing scalar optimizer fix |
|
/ok to test 762efa8 |
|
/ok to test 2c97b67 |
5295186 to
3b21ad4
Compare
3b21ad4 to
2c5ed29
Compare
2c5ed29 to
7e70047
Compare
…kpointing The DistributedOptimizer checkpointing code hardcoded the assumption that every optimizer has exactly two moment states (exp_avg and exp_avg_sq), matching Adam. This broke optimizers like Lion that only use one moment (exp_avg). Add an optimizer_state_keys property that returns the correct state keys based on the optimizer type, and replace all hardcoded key tuples with this property. Also relax the optimizer-type assertion to allow any optimizer that provides an init_state_fn. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
Final dist ckpting check with muon + lion |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29629381676 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29630656179 |
7 participants
Summary
The
DistributedOptimizercheckpointing code hardcoded the assumption that every optimizer uses exactly two moment states (exp_avgandexp_avg_sq), matching Adam. Optimizers like Lion only use a single moment (exp_avg), causing failures during checkpoint save/load.Changes
megatron/core/optimizer/distrib_optimizer.pyoptimizer_state_keysproperty that returns the optimizer's tensor state keys via a hardcoded mapping (e.g.,("exp_avg",)for Lion, defaulting to("exp_avg", "exp_avg_sq")for Adam). When Muon is the top-level optimizer, resolves throughconfig.muon_scalar_optimizer._get_state_key_dtypehelper that maps state keys to their configured dtypes.("param", "exp_avg", "exp_avg_sq")tuples with("param",) + self.optimizer_state_keysinload_state_dict,get_parameter_state_dp_zero,load_parameter_state_from_dp_zero_legacy,load_parameter_state_from_dp_zero, andsplit_state_dict_if_needed.init_state_fn(not just Adam and HybridDeviceOptimizer).megatron/core/optimizer/__init__.pymuon_scalar_optimizer='lion'), rewrite the default param overrides so Lion scalar param groups route through the standard fallback path (creating aDistributedOptimizer) rather than the emerging-optimizer constructor._EMERGING_OPTIMIZERSmembership check to compare against the primary emerging optimizer name (eopt_name), so scalar optimizers that happen to also be registered (e.g., Lion) fall through correctly.Tests
tests/unit_tests/test_lion_optimizer.py: Unit tests foroptimizer_state_keys(Lion, Adam, Muon+Lion, Muon+Adam) and_get_state_key_dtype.tests/unit_tests/dist_checkpointing/test_optimizer.py: End-to-end checkpoint round-trip test using Muon + Lion (muon_scalar_optimizer='lion') withdp_reshardableandfully_reshardablesharding types. Verifies that Lion's single-moment state survives save/load throughDistributedOptimizer.tests/unit_tests/dist_checkpointing/utils.py: Plumbmuon_scalar_optimizerparameter and Lion learning-rate default throughsetup_model_and_optimizer.Test plan
tests/unit_tests/test_lion_optimizer.py— Lion config, factory, exactness, andDistributedOptimizerstate-key unit teststests/unit_tests/dist_checkpointing/test_optimizer.py::TestDistributedOptimizer::test_lion_optimizer_checkpoint_round_trip— checkpoint save/load round-trip with Muon+Lion🤖 Generated with Claude Code