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

add muon optimizer

Browse Source 
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
This commit is contained in:
Wing Lian
2025-02-27 23:15:41 -05:00
parent 0134093acc
commit 0542c7dd56
10 changed files with 296 additions and 158 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,