make sure to save the lora adapter at the end of RL/dpo training (#1573)

axolotl-ai-cloud · May 8, 2024 · 796a085 · 796a085
1 parent cb78a36
commit 796a085
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/axolotl/train.py b/src/axolotl/train.py
@@ -212,6 +212,10 @@ def terminate_handler(_, __, model_weakref):
         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: