Skip to content

Commit

Permalink
make sure to set alternate optimizer and set lr and eps from adam
Browse files Browse the repository at this point in the history
  • Loading branch information
winglian committed Oct 17, 2024
1 parent 446c685 commit c06731e
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion src/axolotl/core/trainer_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -487,7 +487,10 @@ def create_optimizer(self):
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)
Expand Down Expand Up @@ -1600,6 +1603,7 @@ def build(self, total_num_steps):
"ao_adamw_4bit",
"ao_adamw_8bit",
"ao_adamw_fp8",
"soap",
]:
# Set default so transformers doesn't throw
training_arguments_kwargs["optim"] = "adamw_hf"
Expand Down

0 comments on commit c06731e

Please sign in to comment.