add SOAP optimizer

src/axolotl/core/trainer_builder.py

@@ -663,6 +663,11 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
 
                 optimizer_cls = MuonOptimizerFactory
                 optimizer_kwargs.update(adam_kwargs)
+            elif self.cfg.optimizer == "soap":
+                from axolotl.utils.optimizers.soap import SOAP
+
+                optimizer_cls = SOAP
+                optimizer_kwargs.update(adam_kwargs)
             elif self.cfg.optimizer == "optimi_adamw":
                 from optimi import AdamW
 

src/axolotl/utils/schemas/enums.py

@@ -52,3 +52,4 @@ class CustomSupportedOptimizers(str, Enum):
     ao_adamw_fp8 = "ao_adamw_fp8"  # pylint: disable=invalid-name
     adopt_adamw = "adopt_adamw"  # pylint: disable=invalid-name
     muon = "muon"  # pylint: disable=invalid-name
+    soap = "soap"  # pylint: disable=invalid-name