Don't reload model for the merge step;

Context.py

@@ -1,9 +1,12 @@
 import configparser
 import torch
 
+from peft import PeftModel
 from sympy.printing import str
 
 class Context:
+    model: PeftModel
+    
     accel: bool
     device: str
     purgeTargetDirectories: bool

Merger.py

@@ -7,15 +7,17 @@ class Merger:
     def mergeAndStore(self, context: Context):
         if not context.mergeFull:
             return
-        if not context.storeAdapter:
-            raise Exception("Configuration problem: mergeFull requires storeAdapter")
         
-        print("Reloading model")
+        if not context.model: 
-        baseModel = ModelLoader().load(context.locBaseModel, False, True, context)
+            if not context.storeAdapter:
-        peftModel = PeftModel.from_pretrained(baseModel, context.locAdapter)
+                raise Exception("Configuration problem: mergeFull requires storeAdapter")
+            
+            print("Loading model")
+            baseModel = ModelLoader().load(context.locBaseModel, False, True, context)
+            context.model = PeftModel.from_pretrained(baseModel, context.locAdapter)
         
         print("Merging")
-        mergedModel = peftModel.merge_and_unload()
+        mergedModel = context.model.merge_and_unload()
         
         print("Dequantize")
         try:

Trainer.py

@@ -16,42 +16,36 @@ class Trainer:
         
         baseModel = ModelLoader().load(context.locBaseModel, False, True, context)
         sftTrainer = self.__createTrainer(baseModel, context)
-        try:
+        context.model = sftTrainer.model
-            vInplace = context.validate and context.vInplace
-            validator = None
-            if vInplace:
-                validator = Validator()
-            
-            try:
-                cnt = 0
-                continueTraining = True
-                while continueTraining:
-                    cnt += 1
-                    print(f"Training run {cnt}")
-                    
-                    self.__train(sftTrainer)
-                    
-                    if vInplace:
-                        continueTraining = not validator.validateInPlace(sftTrainer.tokenizer, sftTrainer.model, context)
-                    else:
-                        continueTraining = False
-                
-                print(f"Total training runs: {cnt} epochs {cnt * (context.trEpochs or 1)}")
-                
-            finally:
-                if vInplace:
-                    print(f"Validation statistics:\n{validator.statistics}")
-                    validator.unload(context)
-            
-            if context.storeAdapter:
-                self.__storeAdapter(sftTrainer, context)
         
+        vInplace = context.validate and context.vInplace
+        validator = None
+        if vInplace:
+            validator = Validator()
+        
+        try:
+            cnt = 0
+            continueTraining = True
+            while continueTraining:
+                cnt += 1
+                print(f"Training run {cnt}")
+                
+                self.__train(sftTrainer)
+                
+                if vInplace:
+                    continueTraining = not validator.validateInPlace(sftTrainer.tokenizer, sftTrainer.model, context)
+                else:
+                    continueTraining = False
+            
+            print(f"Total training runs: {cnt} epochs {cnt * (context.trEpochs or 1)}")
+            
         finally:
-            del sftTrainer
+            if vInplace:
-            del baseModel
+                print(f"Validation statistics:\n{validator.statistics}")
-            gc.collect()
+                validator.unload(context)
-            if context.accel:
+        
-                torch.cuda.empty_cache()
+        if context.storeAdapter:
+            self.__storeAdapter(sftTrainer, context)
         
     def __createTrainer(self, baseModel, context: Context):
         tokenizer = AutoTokenizer.from_pretrained(