Merge branch 'main' into smangrul/add-docs

Makefile

@@ -17,4 +17,4 @@ style:
 	doc-builder style src/peft tests docs/source --max_len 119
 
 test:
-	pytest tests/
+	pytest -n 3 tests/

src/peft/peft_model.py

@@ -582,7 +582,13 @@ class PeftModelForCausalLM(PeftModel):
             else:
                 if "input_ids" not in kwargs:
                     raise ValueError("input_ids must be provided for Peft model generation")
-                if kwargs.get("attention_mask", None) is not None:
+                # For gpt2 models, we construct postion_ids on the fly by using attention mask, and position ids need to match input_shape.
+                # for prefix tuning, input shape is determined using `input_ids`. Thus we should not expand 'attention_mask' here
+                # for prompt tuning input_ids is not passed but a concatenated input_embeds is passed. Thus attention_mask needs to be of same size of num_virtual_tokens + input_ids
+                if kwargs.get("attention_mask", None) is not None and self.peft_config.peft_type in [
+                    PeftType.PROMPT_TUNING,
+                    PeftType.P_TUNING,
+                ]:
                     # concat prompt attention mask
                     prefix_attention_mask = torch.ones(
                         kwargs["input_ids"].shape[0], self.peft_config.num_virtual_tokens
@@ -611,6 +617,14 @@ class PeftModelForCausalLM(PeftModel):
     def prepare_inputs_for_generation(self, *args, **kwargs):
         model_kwargs = self.base_model_prepare_inputs_for_generation(*args, **kwargs)
         if isinstance(self.peft_config, PromptLearningConfig):
+            if self.peft_config.peft_type == PeftType.PREFIX_TUNING:
+                prefix_attention_mask = torch.ones(
+                    model_kwargs["input_ids"].shape[0], self.peft_config.num_virtual_tokens
+                ).to(model_kwargs["input_ids"].device)
+                model_kwargs["attention_mask"] = torch.cat(
+                    (prefix_attention_mask, model_kwargs["attention_mask"]), dim=1
+                )
+
             if model_kwargs["past_key_values"] is None and self.peft_config.peft_type == PeftType.PREFIX_TUNING:
                 past_key_values = self.get_prompt(batch_size=model_kwargs["input_ids"].shape[0])
                 model_kwargs["past_key_values"] = past_key_values

tests/test_peft_model.py

@@ -31,8 +31,14 @@ from .testing_common import PeftTestConfigManager
 
 
 # This has to be in the order: model_id, lora_kwargs, prefix_tuning_kwargs, prompt_encoder_kwargs, prompt_tuning_kwargs
-PEFT_MODELS_TO_TEST = [
-    ("hf-internal-testing/tiny-random-OPTForCausalLM", {"target_modules": ["q_proj", "v_proj"]}, {}, {}, {}),
+PEFT_DECODER_MODELS_TO_TEST = [
+    # ("HuggingFaceM4/tiny-random-LlamaForCausalLM", {}, {}, {}, {}), wait until the next `transformers` release
+    ("hf-internal-testing/tiny-random-OPTForCausalLM", {}, {}, {}, {}),
+    ("hf-internal-testing/tiny-random-GPTNeoXForCausalLM", {}, {}, {}, {}),
+    ("hf-internal-testing/tiny-random-GPT2LMHeadModel", {}, {}, {}, {}),
+    ("hf-internal-testing/tiny-random-BloomForCausalLM", {}, {}, {}, {}),
+    ("hf-internal-testing/tiny-random-gpt_neo", {}, {}, {}, {}),
+    ("hf-internal-testing/tiny-random-GPTJForCausalLM", {}, {}, {}, {}),
 ]
 
 
@@ -48,7 +54,7 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
     We use parametrized.expand for debugging purposes to test each model individually.
     """
 
-    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_MODELS_TO_TEST))
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_DECODER_MODELS_TO_TEST))
     def test_attributes_parametrized(self, test_name, model_id, config_cls, config_kwargs):
         self._test_model_attr(model_id, config_cls, config_kwargs)
 
@@ -105,7 +111,7 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
 
         self.assertTrue(dummy_output.requires_grad)
 
-    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_MODELS_TO_TEST))
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_DECODER_MODELS_TO_TEST))
     def test_prepare_for_training_parametrized(self, test_name, model_id, config_cls, config_kwargs):
         self._test_prepare_for_training(model_id, config_cls, config_kwargs)
 
@@ -151,6 +157,29 @@ class PeftModelTester(unittest.TestCase, PeftTestMixin):
             # check if `config.json` is not present
             self.assertFalse(os.path.exists(os.path.join(tmp_dirname, "config.json")))
 
-    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_MODELS_TO_TEST))
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_DECODER_MODELS_TO_TEST))
     def test_save_pretrained(self, test_name, model_id, config_cls, config_kwargs):
         self._test_save_pretrained(model_id, config_cls, config_kwargs)
+
+    def _test_generate(self, model_id, config_cls, config_kwargs):
+        model = AutoModelForCausalLM.from_pretrained(model_id)
+        config = config_cls(
+            base_model_name_or_path=model_id,
+            **config_kwargs,
+        )
+        model = get_peft_model(model, config)
+        model = model.to(self.torch_device)
+
+        input_ids = torch.LongTensor([[1, 1, 1], [2, 1, 2]]).to(self.torch_device)
+        attention_mask = torch.LongTensor([[1, 1, 1], [1, 0, 1]]).to(self.torch_device)
+
+        # check if `generate` works
+        _ = model.generate(input_ids=input_ids, attention_mask=attention_mask)
+
+        with self.assertRaises(TypeError):
+            # check if `generate` raises an error if no positional arguments are passed
+            _ = model.generate(input_ids, attention_mask=attention_mask)
+
+    @parameterized.expand(PeftTestConfigManager.get_grid_parameters(PEFT_DECODER_MODELS_TO_TEST))
+    def test_generate(self, test_name, model_id, config_cls, config_kwargs):
+        self._test_generate(model_id, config_cls, config_kwargs)

tests/testing_common.py

@@ -79,23 +79,24 @@ class ClassInstantier(OrderedDict):
         for model_tuple in model_list:
             model_id, lora_kwargs, prefix_tuning_kwargs, prompt_encoder_kwargs, prompt_tuning_kwargs = model_tuple
             for key, value in self.items():
+                peft_method = value[1].copy()
                 if key == "lora":
                     # update value[1] if necessary
                     if lora_kwargs is not None:
-                        value[1].update(lora_kwargs)
+                        peft_method.update(lora_kwargs)
                 elif key == "prefix_tuning":
                     # update value[1] if necessary
                     if prefix_tuning_kwargs is not None:
-                        value[1].update(prefix_tuning_kwargs)
+                        peft_method.update(prefix_tuning_kwargs)
                 elif key == "prompt_encoder":
                     # update value[1] if necessary
                     if prompt_encoder_kwargs is not None:
-                        value[1].update(prompt_encoder_kwargs)
+                        peft_method.update(prompt_encoder_kwargs)
                 else:
                     # update value[1] if necessary
                     if prompt_tuning_kwargs is not None:
-                        value[1].update(prompt_tuning_kwargs)
-                grid_parameters.append((f"test_{model_id}_{key}", model_id, value[0], value[1]))
+                        peft_method.update(prompt_tuning_kwargs)
+                grid_parameters.append((f"test_{model_id}_{key}", model_id, value[0], peft_method))
 
         return grid_parameters