Optimizer refactor and add Muon support (#2367)

* 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
2025-03-06 11:49:19 -05:00
parent 5e21b1a9da
commit ae66374156
10 changed files with 295 additions and 157 deletions
--- a/requirements.txt
+++ b/requirements.txt
@@ -63,3 +63,4 @@ torchao==0.7.0
 schedulefree==1.3.0

 axolotl-contribs-lgpl==0.0.3
+axolotl-contribs-mit==0.0.3