From 336c5f9db95e57c35b7a3441ad6f405756f2c893 Mon Sep 17 00:00:00 2001 From: NanoCode012 Date: Fri, 16 May 2025 13:38:11 +0700 Subject: [PATCH] fix: move pop optimizer_cls_and_kwargs --- src/axolotl/core/trainer_builder/sft.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/axolotl/core/trainer_builder/sft.py b/src/axolotl/core/trainer_builder/sft.py index 311db6564..0303f5dae 100644 --- a/src/axolotl/core/trainer_builder/sft.py +++ b/src/axolotl/core/trainer_builder/sft.py @@ -343,6 +343,14 @@ class HFCausalTrainerBuilder(TrainerBuilderBase): self.cfg.kd_top_k_before_softmax ) + trainer_kwargs = {} + + # Pop optimizer_cls_and_kwargs to trainer_kwargs + if "optimizer_cls_and_kwargs" in training_arguments_kwargs: + trainer_kwargs["optimizer_cls_and_kwargs"] = training_arguments_kwargs.pop( + "optimizer_cls_and_kwargs" + ) + if self.cfg.reward_model: training_args_cls = AxolotlRewardConfig elif self.cfg.process_reward_model: @@ -374,13 +382,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase): data_collator_kwargs["pad_to_multiple_of"] = multiple trainer_cls = self._get_trainer_cls() - trainer_kwargs = {} - - # Pop optimizer_cls_and_kwargs to trainer_kwargs - if "optimizer_cls_and_kwargs" in training_arguments_kwargs: - trainer_kwargs["optimizer_cls_and_kwargs"] = training_arguments_kwargs.pop( - "optimizer_cls_and_kwargs" - ) trainer_kwargs, trainer_cls = self.hook_pre_create_trainer( trainer_kwargs, trainer_cls