Skip to content

Commit

Permalink
make sure to save the lora adapter at the end of RL/dpo training
Browse files Browse the repository at this point in the history
  • Loading branch information
winglian committed Apr 29, 2024
1 parent 5294653 commit 98bee77
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions src/axolotl/train.py
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,10 @@ def terminate_handler(_, __, model):
if cfg.flash_optimum and BetterTransformer:
model = BetterTransformer.reverse(model)

if cfg.rl and cfg.adapter and not cfg.rl_adapter_ref_model:
trainer.model.save_pretrained(
cfg.output_dir, safe_serialization=safe_serialization
)
model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)

if not cfg.hub_model_id:
Expand Down

0 comments on commit 98bee77

Please sign in to comment.