Basic PEFT support

aiserver.py

@@ -1345,6 +1345,7 @@ def general_startup(override_args=None):
     parser.add_argument("--summarizer_model", action='store', default="philschmid/bart-large-cnn-samsum", help="Huggingface model to use for summarization. Defaults to sshleifer/distilbart-cnn-12-6")
     parser.add_argument("--max_summary_length", action='store', default=75, help="Maximum size for summary to send to image generation")
     parser.add_argument("--multi_story", action='store_true', default=False, help="Allow multi-story mode (experimental)")
+    parser.add_argument("--peft", type=str, help="Specify the path or HuggingFace ID of a Peft to load it. Not supported on TPU. (Experimental)")
     
     parser.add_argument('-f', action='store', help="option for compatability with colab memory profiles")
     parser.add_argument('-v', '--verbosity', action='count', default=0, help="The default logging level is ERROR or higher. This value increases the amount of logging seen in your screen")

environments/huggingface.yml

@@ -45,3 +45,4 @@ dependencies:
     - ftfy
     - pydub
     - diffusers
+    - peft==0.3.0

environments/rocm.yml

@@ -41,3 +41,4 @@ dependencies:
     - ftfy
     - pydub
     - diffusers
+    - peft==0.3.0

modeling/inference_models/hf_torch.py

@@ -22,6 +22,7 @@ from transformers import (
     AutoModelForCausalLM,
     LogitsProcessorList,
 )
+from peft import PeftModel, PeftConfig
 
 import utils
 import modeling.lazy_loader as lazy_loader
@@ -211,6 +212,31 @@ class HFTorchInferenceModel(HFInferenceModel):
         new_sample.old_sample = transformers.GenerationMixin.sample
         use_core_manipulations.sample = new_sample
 
+        # PEFT Loading. This MUST be done after all save_pretrained calls are
+        # finished on the main model.
+        if utils.args.peft:
+            peft_local_path = os.path.join("models/peft", utils.args.peft.replace("/", "_"))
+            logger.debug(f"Loading PEFT '{utils.args.peft}', possible local path is '{peft_local_path}'.")
+
+            peft_installed_locally = True
+            possible_peft_locations = [peft_local_path, utils.args.peft]
+
+            for i, location in enumerate(possible_peft_locations):
+                try:
+                    m_self.model = PeftModel.from_pretrained(m_self.model, location)
+                    logger.debug(f"Loaded PEFT at '{location}'")
+                    break
+                except ValueError:
+                    peft_installed_locally = False
+                    if i == len(possible_peft_locations) - 1:
+                        raise RuntimeError(f"Unable to load PeftModel for given name '{utils.args.peft}'. Does it exist?")
+                except RuntimeError:
+                    raise RuntimeError("Error while loading PeftModel. Are you using the correct model?")
+
+            if not peft_installed_locally:
+                logger.debug(f"PEFT not saved to models folder; saving to '{peft_local_path}'")
+                m_self.model.save_pretrained(peft_local_path)
+
         return super()._post_load()
 
     def _raw_generate(
@@ -238,8 +264,13 @@ class HFTorchInferenceModel(HFInferenceModel):
 
         with torch.no_grad():
             start_time = time.time()
+
+            # HEED & BEWARE: All arguments passed to self.model.generate MUST be
+            # kwargs; see https://github.com/huggingface/peft/issues/232. If they
+            # aren't, PeftModel will EXPLODE!!!! But nothing will happen without
+            # a PEFT loaded so it's sneaky.
             genout = self.model.generate(
-                gen_in,
+                input_ids=gen_in,
                 do_sample=True,
                 max_length=min(
                     len(prompt_tokens) + max_new, utils.koboldai_vars.max_length

models/peft/README.txt

@@ -0,0 +1,2 @@
+PEFT models will be stored in this directory when downloaded.
+Please don't be too mean to this directory.

requirements.txt

@@ -37,3 +37,4 @@ pytest-html==3.2.0
 pytest-metadata==2.0.4
 requests-mock==1.10.0
 safetensors==0.3.1
+peft==0.3.0