merge

Lightning-AI · Mar 11, 2024 · 4f53463 · 4f53463
1 parent 2667fec
commit 4f53463
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/litgpt/finetune/lora.py b/litgpt/finetune/lora.py
@@ -20,6 +20,7 @@
 from litgpt.generate.base import generate
 from litgpt.lora import GPT, Block, Config, lora_filter, mark_only_lora_as_trainable
 from litgpt.prompts import save_prompt_style
+from litgpt.scripts.merge_lora import merge_lora
 from litgpt.tokenizer import Tokenizer
 from litgpt.utils import (
     CLI,
@@ -169,6 +170,8 @@ def main(fabric: L.Fabric, devices: int, seed: int, config: Config, data: LitDat
         copy_config_files(checkpoint_dir, save_path.parent)
         save_hyperparameters(setup, save_path.parent)
         save_prompt_style(data.prompt_style, save_path.parent)
+    fabric.barrier()
+    merge_lora(checkpoint_dir=save_path.parent)
 
 
 def fit(