From 67b9e31bbc0693d2698c0a5d93753e6f8cc9ce33 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 17 Oct 2024 17:21:19 -0400 Subject: [PATCH] make sure to set alternate optimizer and set lr and eps from adam --- src/axolotl/core/trainer_builder.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index cefbd09d4..7537f80fe 100755 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -487,7 +487,10 @@ class AxolotlTrainer(SchedulerMixin, Trainer): elif self.args.alternate_optimizer == "soap": from axolotl.utils.optimizers.soap import SOAP - optim_args = {} + optim_args = { + "lr": optimizer_kwargs.pop("lr"), + "eps": optimizer_kwargs.pop("eps"), + } if self.cfg.optim_args: optim_args.update(self.cfg.optim_args) @@ -1639,6 +1642,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase): "ao_adamw_4bit", "ao_adamw_8bit", "ao_adamw_fp8", + "soap", ]: # Set default so transformers doesn't throw training_arguments_kwargs["optim"] = "adamw_hf"