diff --git a/src/axolotl/core/trainer_builder/base.py b/src/axolotl/core/trainer_builder/base.py index 3293fd12c..4684c30ee 100644 --- a/src/axolotl/core/trainer_builder/base.py +++ b/src/axolotl/core/trainer_builder/base.py @@ -227,10 +227,6 @@ class TrainerBuilderBase(abc.ABC): training_args_kwargs["lr_scheduler_kwargs"] = ( self.cfg.lr_scheduler_kwargs if self.cfg.lr_scheduler_kwargs else {} ) - training_args_kwargs["cosine_min_lr_ratio"] = self.cfg.cosine_min_lr_ratio - training_args_kwargs["cosine_constant_lr_ratio"] = ( - self.cfg.cosine_constant_lr_ratio - ) # Handle custom optimizer custom_supported_optimizers = [opt.value for opt in CustomSupportedOptimizers] @@ -444,11 +440,15 @@ class TrainerBuilderBase(abc.ABC): # set arg into trainer_args_kwargs with same name if value not None for arg in [ + # optim/scheduler "adam_beta1", "adam_beta2", "adam_beta3", "adam_epsilon", "adam_epsilon2", + "cosine_min_lr_ratio", + "cosine_constant_lr_ratio", + # trainer "max_grad_norm", "dataloader_num_workers", "dataloader_pin_memory", diff --git a/tests/core/test_trainer_builder.py b/tests/core/test_trainer_builder.py index 1f79d8cc4..9c7753474 100644 --- a/tests/core/test_trainer_builder.py +++ b/tests/core/test_trainer_builder.py @@ -216,6 +216,7 @@ class TestHFRLTrainerBuilder: assert training_arguments.lr_scheduler_type == "cosine" assert training_arguments.warmup_steps == 10 assert training_arguments.cosine_min_lr_ratio == 0.1 + assert training_arguments.cosine_constant_lr_ratio == 0.2 # Other settings assert training_arguments.dataloader_num_workers == 1