tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Optimizer refactor and add Muon support (#2367)

Browse Source 
* add muon optimizer

optimizer_cls_and_kwargs is on trainer_kwargs
only add adamw_kwargs if they're non-null
fix mocks
better handling of override and check the optimizer
unwrap optimizer

* fix import
This commit is contained in:
Wing Lian
2025-03-06 11:49:19 -05:00
committed by GitHub
parent 5e21b1a9da
commit ae66374156
10 changed files with 295 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/cli/test_cli_train.py
View File

@@ -28,7 +28,7 @@ class TestTrainCommand(BaseCliTest):
config_path.write_text(valid_test_config)
with patch("axolotl.cli.train.train") as mock_train:
mock_train.return_value = (MagicMock(), MagicMock())
mock_train.return_value = (MagicMock(), MagicMock(), MagicMock())
result = cli_runner.invoke(
cli,
@@ -48,7 +48,7 @@ class TestTrainCommand(BaseCliTest):
config_path = self._test_cli_overrides(tmp_path, valid_test_config)
with patch("axolotl.cli.train.train") as mock_train:
mock_train.return_value = (MagicMock(), MagicMock())
mock_train.return_value = (MagicMock(), MagicMock(), MagicMock())
result = cli_runner.invoke(
cli,