From 13e938149dd6a1eec317c4cdde74d3ce991a8b86 Mon Sep 17 00:00:00 2001 From: NanoCode012 Date: Sun, 17 Dec 2023 18:48:28 +0900 Subject: [PATCH] fix: add lr scheduler kwargs to Trainer (#972) --- src/axolotl/core/trainer_builder.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index ccd9d37c0..cc162d210 100644 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -692,6 +692,9 @@ class HFCausalTrainerBuilder(TrainerBuilderBase): and self.cfg.lr_scheduler not in ("one_cycle", "log_sweep") else "cosine" ) + training_arguments_kwargs["lr_scheduler_kwargs"] = ( + self.cfg.lr_scheduler_kwargs if self.cfg.lr_scheduler_kwargs else {} + ) training_arguments_kwargs["weight_decay"] = ( self.cfg.weight_decay if self.cfg.weight_decay is not None else 0.0 )