From 702a669cadfa4c8247699f1663972479d9ec4dc7 Mon Sep 17 00:00:00 2001 From: Jeffrey Quesnelle Date: Thu, 16 May 2024 21:23:18 -0700 Subject: [PATCH] add save_only_model option (#1634) --- src/axolotl/core/trainer_builder.py | 2 ++ src/axolotl/utils/config/models/input/v0_4_1/__init__.py | 1 + 2 files changed, 3 insertions(+) diff --git a/src/axolotl/core/trainer_builder.py b/src/axolotl/core/trainer_builder.py index 2f38b12dc..7330a78ef 100644 --- a/src/axolotl/core/trainer_builder.py +++ b/src/axolotl/core/trainer_builder.py @@ -1123,6 +1123,8 @@ class HFCausalTrainerBuilder(TrainerBuilderBase): # default to saving each epoch if not defined training_arguments_kwargs["save_strategy"] = "epoch" + training_arguments_kwargs["save_only_model"] = self.cfg.save_only_model + if self.cfg.do_bench_eval: training_arguments_kwargs["do_bench_eval"] = self.cfg.do_bench_eval if self.cfg.bench_dataset: diff --git a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py index f1c12b2ba..9be0f6949 100644 --- a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py +++ b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py @@ -574,6 +574,7 @@ class AxolotlInputConfig( logging_steps: Optional[int] = None early_stopping_patience: Optional[int] = None load_best_model_at_end: Optional[bool] = False + save_only_model: Optional[bool] = False neftune_noise_alpha: Optional[float] = None