Allow user to specify huggingface link or local path to pretrained lora weights

ludwig/constants.py

@@ -281,6 +281,7 @@
 GENERATION = "generation"
 PROMPT = "prompt"
 ADAPTER = "adapter"
+PRETRAINED_WEIGHTS = "pretrained_weights"
 
 # CrossEntropyLoss for LLMs
 IGNORE_INDEX_TOKEN_ID = -100

ludwig/models/llm.py

@@ -221,12 +221,36 @@ def initialize_adapter(self):
                     "`finetune` or remove the adapter config."
                 )
 
-            from peft import get_peft_model, TaskType
+            from peft import get_peft_model
+
+            pretrained = False
+            if self.config_obj.adapter.pretrained_weights:
+                print(f"PRETRAINED_WEIGHTS: {self.config_obj.adapter.pretrained_weights}")
+                # If pretrained adapter weights are provided, we want to load them into the model
+                from peft import MODEL_TYPE_TO_PEFT_MODEL_MAPPING, PeftConfig
+
+                pretrained = True
+                peft_config = PeftConfig.from_pretrained(self.config_obj.adapter.pretrained_weights)
+                peft_dict = peft_config.to_dict()
+                for param_name, param_value in self.config_obj.adapter.to_config().to_dict().items():
+                    if param_name is None:
+                        continue
+
+                    if param_name not in peft_dict:
+                        setattr(peft_config, param_name, param_value)
+
+                self.model = MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type].from_pretrained(
+                    self.model, self.config_obj.adapter.pretrained_weights
+                )
+            else:
+                # If no pretrained adapter is provided, we want to load untrained weights into the model
+                from peft import TaskType
 
-            peft_config = self.config_obj.adapter.to_config(
-                task_type=TaskType.CAUSAL_LM, tokenizer_name_or_path=self.model_name
-            )
-            self.model = get_peft_model(self.model, peft_config)
+                peft_config = self.config_obj.adapter.to_config(
+                    task_type=TaskType.CAUSAL_LM, tokenizer_name_or_path=self.model_name
+                )
+
+                self.model = get_peft_model(self.model, peft_config, pretrained=pretrained)
 
             logger.info("==================================================")
             logger.info("Trainable Parameter Summary For Fine-Tuning")

ludwig/schema/llms/peft.py

@@ -69,6 +69,18 @@ class LoraConfig(BaseAdapterConfig):
         description="Bias type for Lora.",
     )
 
+    pretrained_weights: Optional[str] = schema_utils.String(
+        default="none",
+        description="Path to pretrained weights for Lora.",
+    )
+
+    target_modules: Optional[list] = schema_utils.List(
+        str,
+        default=None,
+        allow_none=True,
+        description="List of modules to apply Lora to. If None, apply to all modules.",
+    )
+
     def to_config(self, task_type: str = None, **kwargs) -> "PeftConfig":
         from peft import LoraConfig as _LoraConfig
 

ludwig/trainers/trainer.py

@@ -266,6 +266,7 @@ def closure():
                     targets, model_outputs, self.regularization_type, self.regularization_lambda
                 )
                 loss = loss / self.gradient_accumulation_steps
+                loss.requires_grad = True
 
         # Begin the backward pass
         variables = self.dist_model.parameters()

tests/integration_tests/test_llm.py

@@ -18,6 +18,7 @@
     MODEL_TYPE,
     OUTPUT_FEATURES,
     PREPROCESSING,
+    PRETRAINED_WEIGHTS,
     PROMPT,
     TRAINER,
     TYPE,
@@ -481,6 +482,36 @@ def test_llama_rope_scaling():
     assert model.model.config.rope_scaling["factor"] == 2.0
 
 
+def test_load_pretrained_adapter_weights():
+    from peft import PeftModel
+    from transformers import PreTrainedModel
+
+    config = {
+        MODEL_TYPE: MODEL_LLM,
+        BASE_MODEL: TEST_MODEL_NAME,
+        INPUT_FEATURES: [text_feature(name="input", encoder={"type": "passthrough"})],
+        OUTPUT_FEATURES: [text_feature(name="output")],
+        TRAINER: {
+            TYPE: "finetune",
+            BATCH_SIZE: 8,
+            EPOCHS: 2,
+        },
+        ADAPTER: {TYPE: "lora", PRETRAINED_WEIGHTS: "Infernaught/test_adapter_weights"},
+        BACKEND: {TYPE: "local"},
+    }
+
+    print(ModelConfig)
+    config_obj = ModelConfig.from_dict(config)
+    model = LLM(config_obj)
+
+    assert model.config_obj.adapter.pretrained_weights
+    assert model.config_obj.adapter.pretrained_weights == "Infernaught/test_adapter_weights"
+
+    model.prepare_for_training()
+    assert not isinstance(model.model, PreTrainedModel)
+    assert isinstance(model.model, PeftModel)
+
+
 def _compare_models(model_1: torch.nn.Module, model_2: torch.nn.Module) -> bool:
     # Source: https://discuss.pytorch.org/t/check-if-models-have-same-weights/4351/6
     for key_item_1, key_item_2 in zip(model_1.state_dict().items(), model_2.state_dict().items()):