From cd2395987ecc18f1fd1841736780c88dd4fda2d2 Mon Sep 17 00:00:00 2001
From: NanoCode012 <kevinvong@rocketmail.com>
Date: Wed, 10 May 2023 03:18:38 +0900
Subject: [PATCH] Don't save full model for lora

---
 scripts/finetune.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/scripts/finetune.py b/scripts/finetune.py
index a7c3d6c77..c05023f40 100644
--- a/scripts/finetune.py
+++ b/scripts/finetune.py
@@ -228,11 +228,12 @@ def train(
     logging.info(
         f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}"
     )
-    # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
-    trainer.save_model(cfg.output_dir)
     
     if cfg.adapter == 'lora':
         trainer.save_pretrained(cfg.output_dir)
+    else:
+        # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
+        trainer.save_model(cfg.output_dir)
 
 
 if __name__ == "__main__":