Skip to content

Commit

Permalink
Add more save strategies for DPO training. (axolotl-ai-cloud#1255)
Browse files Browse the repository at this point in the history
* Set save_strategy and save_steps in HFDPOTrainerBuilder

* fix doublicate save_steps
  • Loading branch information
PhilipMay authored and maximegmd committed Feb 6, 2024
1 parent 946196e commit ea852b0
Showing 1 changed file with 10 additions and 2 deletions.
12 changes: 10 additions & 2 deletions src/axolotl/core/trainer_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -1096,13 +1096,21 @@ def build_training_arguments(self, total_num_steps):
"use_reentrant": False
}

# set save_strategy and save_steps
if self.cfg.save_steps:
training_args_kwargs["save_strategy"] = "steps"
training_args_kwargs["save_steps"] = self.cfg.save_steps
elif self.cfg.save_strategy:
training_args_kwargs["save_strategy"] = self.cfg.save_strategy
else:
# default to saving each epoch if not defined
training_args_kwargs["save_strategy"] = "epoch"

training_args = TrainingArguments(
per_device_train_batch_size=self.cfg.micro_batch_size,
max_steps=self.cfg.max_steps or total_num_steps,
gradient_accumulation_steps=self.cfg.gradient_accumulation_steps,
learning_rate=self.cfg.learning_rate,
save_strategy="steps",
save_steps=self.cfg.save_steps,
output_dir=self.cfg.output_dir,
warmup_steps=self.cfg.warmup_steps,
logging_first_step=True,
Expand Down

0 comments on commit ea852b0

Please sign in to comment.