emozilla commited on
Commit
702a669
·
unverified ·
1 Parent(s): 891ae8a

add save_only_model option (#1634)

Browse files
src/axolotl/core/trainer_builder.py CHANGED
@@ -1123,6 +1123,8 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
1123
  # default to saving each epoch if not defined
1124
  training_arguments_kwargs["save_strategy"] = "epoch"
1125
 
 
 
1126
  if self.cfg.do_bench_eval:
1127
  training_arguments_kwargs["do_bench_eval"] = self.cfg.do_bench_eval
1128
  if self.cfg.bench_dataset:
 
1123
  # default to saving each epoch if not defined
1124
  training_arguments_kwargs["save_strategy"] = "epoch"
1125
 
1126
+ training_arguments_kwargs["save_only_model"] = self.cfg.save_only_model
1127
+
1128
  if self.cfg.do_bench_eval:
1129
  training_arguments_kwargs["do_bench_eval"] = self.cfg.do_bench_eval
1130
  if self.cfg.bench_dataset:
src/axolotl/utils/config/models/input/v0_4_1/__init__.py CHANGED
@@ -574,6 +574,7 @@ class AxolotlInputConfig(
574
  logging_steps: Optional[int] = None
575
  early_stopping_patience: Optional[int] = None
576
  load_best_model_at_end: Optional[bool] = False
 
577
 
578
  neftune_noise_alpha: Optional[float] = None
579
 
 
574
  logging_steps: Optional[int] = None
575
  early_stopping_patience: Optional[int] = None
576
  load_best_model_at_end: Optional[bool] = False
577
+ save_only_model: Optional[bool] = False
578
 
579
  neftune_noise_alpha: Optional[float] = None
580