clear memory before merging adapters

huggingface · Aug 26, 2024 · 1cb2f48 · 1cb2f48
1 parent 10defa9
commit 1cb2f48
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/configs/llm_finetuning/gpt2_sft.yml b/configs/llm_finetuning/gpt2_sft.yml
@@ -24,8 +24,9 @@ params:
   scheduler: linear
   gradient_accumulation: 4
   mixed_precision: fp16
+  merge_adapter: true
 
 hub:
   username: ${HF_USERNAME}
   token: ${HF_TOKEN}
-  push_to_hub: true
+  push_to_hub: false
diff --git a/src/autotrain/trainers/clm/utils.py b/src/autotrain/trainers/clm/utils.py
@@ -1,4 +1,5 @@
 import ast
+import gc
 import os
 from enum import Enum
 from itertools import chain
@@ -295,6 +296,9 @@ def post_training_steps(config, trainer):
         f.write(model_card)
 
     if config.peft and config.merge_adapter:
+        del trainer
+        gc.collect()
+        torch.cuda.empty_cache()
         logger.info("Merging adapter weights...")
         try:
             merge_adapter(