Add AutoGPTQ LoRA support

modules/LoRA.py

@@ -1,10 +1,13 @@
 from pathlib import Path
 
 import torch
+from auto_gptq import get_gptq_peft_model
+from auto_gptq.utils.peft_utils import GPTQLoraConfig
 from peft import PeftModel
 
 import modules.shared as shared
 from modules.logging_colors import logger
+from modules.models import reload_model
 
 
 def add_lora_to_model(lora_names):
@@ -13,43 +16,68 @@ def add_lora_to_model(lora_names):
     removed_set = prior_set - set(lora_names)
     shared.lora_names = list(lora_names)
 
-    # If no LoRA needs to be added or removed, exit
+    is_autogptq = 'GPTQForCausalLM' in shared.model.__class__.__name__
-    if len(added_set) == 0 and len(removed_set) == 0:
-        return
 
-    # Add a LoRA when another LoRA is already present
+    # AutoGPTQ case. It doesn't use the peft functions.
-    if len(removed_set) == 0 and len(prior_set) > 0:
+    # Copied from https://github.com/Ph0rk0z/text-generation-webui-testing
-        logger.info(f"Adding the LoRA(s) named {added_set} to the model...")
+    if is_autogptq:
-        for lora in added_set:
+        if len(prior_set) > 0:
-            shared.model.load_adapter(Path(f"{shared.args.lora_dir}/{lora}"), lora)
+            reload_model()
 
-        return
+        if len(shared.lora_names) == 0:
+            return
+        else:
+            if len(shared.lora_names) > 1:
+                logger.warning('AutoGPTQ can only work with 1 LoRA at the moment. Only the first one in the list will be loaded')
 
-    # If any LoRA needs to be removed, start over
+            peft_config = GPTQLoraConfig(
-    if len(removed_set) > 0:
+                inference_mode=True,
-        shared.model.disable_adapter()
+            )
-        shared.model = shared.model.base_model.model
 
-    if len(lora_names) > 0:
+            lora_path = Path(f"{shared.args.lora_dir}/{shared.lora_names[0]}")
-        logger.info("Applying the following LoRAs to {}: {}".format(shared.model_name, ', '.join(lora_names)))
+            logger.info("Applying the following LoRAs to {}: {}".format(shared.model_name, ', '.join([lora_names[0]])))
-        params = {}
+            shared.model = get_gptq_peft_model(shared.model, peft_config, lora_path)
-        if not shared.args.cpu:
+            return
-            params['dtype'] = shared.model.dtype
-            if hasattr(shared.model, "hf_device_map"):
-                params['device_map'] = {"base_model.model." + k: v for k, v in shared.model.hf_device_map.items()}
-            elif shared.args.load_in_8bit:
-                params['device_map'] = {'': 0}
 
-        shared.model = PeftModel.from_pretrained(shared.model, Path(f"{shared.args.lora_dir}/{lora_names[0]}"), **params)
+    # Transformers case
+    else:
+        # If no LoRA needs to be added or removed, exit
+        if len(added_set) == 0 and len(removed_set) == 0:
+            return
 
-        for lora in lora_names[1:]:
+        # Add a LoRA when another LoRA is already present
-            shared.model.load_adapter(Path(f"{shared.args.lora_dir}/{lora}"), lora)
+        if len(removed_set) == 0 and len(prior_set) > 0:
+            logger.info(f"Adding the LoRA(s) named {added_set} to the model...")
+            for lora in added_set:
+                shared.model.load_adapter(Path(f"{shared.args.lora_dir}/{lora}"), lora)
 
-        if not shared.args.load_in_8bit and not shared.args.cpu:
+            return
-            shared.model.half()
+
-            if not hasattr(shared.model, "hf_device_map"):
+        # If any LoRA needs to be removed, start over
-                if torch.has_mps:
+        if len(removed_set) > 0:
-                    device = torch.device('mps')
+            shared.model.disable_adapter()
-                    shared.model = shared.model.to(device)
+            shared.model = shared.model.base_model.model
-                else:
+
-                    shared.model = shared.model.cuda()
+        if len(lora_names) > 0:
+            logger.info("Applying the following LoRAs to {}: {}".format(shared.model_name, ', '.join(lora_names)))
+            params = {}
+            if not shared.args.cpu:
+                params['dtype'] = shared.model.dtype
+                if hasattr(shared.model, "hf_device_map"):
+                    params['device_map'] = {"base_model.model." + k: v for k, v in shared.model.hf_device_map.items()}
+                elif shared.args.load_in_8bit:
+                    params['device_map'] = {'': 0}
+
+            shared.model = PeftModel.from_pretrained(shared.model, Path(f"{shared.args.lora_dir}/{lora_names[0]}"), **params)
+
+            for lora in lora_names[1:]:
+                shared.model.load_adapter(Path(f"{shared.args.lora_dir}/{lora}"), lora)
+
+            if not shared.args.load_in_8bit and not shared.args.cpu:
+                shared.model.half()
+                if not hasattr(shared.model, "hf_device_map"):
+                    if torch.has_mps:
+                        device = torch.device('mps')
+                        shared.model = shared.model.to(device)
+                    else:
+                        shared.model = shared.model.cuda()