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

make sure to set alternate optimizer and set lr and eps from adam

Browse Source
This commit is contained in:
Wing Lian
2024-10-17 17:21:19 -04:00
parent ad60916323
commit 67b9e31bbc
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/axolotl/core/trainer_builder.py
View File

@@ -487,7 +487,10 @@ class AxolotlTrainer(SchedulerMixin, Trainer):
elif self.args.alternate_optimizer == "soap": elif self.args.alternate_optimizer == "soap":
from axolotl.utils.optimizers.soap import 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: if self.cfg.optim_args:
optim_args.update(self.cfg.optim_args) optim_args.update(self.cfg.optim_args)
@@ -1639,6 +1642,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
"ao_adamw_4bit", "ao_adamw_4bit",
"ao_adamw_8bit", "ao_adamw_8bit",
"ao_adamw_fp8", "ao_adamw_fp8",
"soap",
]: ]:
# Set default so transformers doesn't throw # Set default so transformers doesn't throw
training_arguments_kwargs["optim"] = "adamw_hf" training_arguments_kwargs["optim"] = "adamw_hf"