From d8d1007732c464e4d86171f4741f8f2d1920d276 Mon Sep 17 00:00:00 2001 From: Vineet Kumar Date: Wed, 29 Mar 2023 18:50:14 +0530 Subject: [PATCH] Causal LM generation fix for prefix tuning: GPT2 model (#222) * expand attention mask after preparing generation inputs for prefix tuning * reformat * Update src/peft/peft_model.py Co-authored-by: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> * reformat as per black --------- Co-authored-by: Vineet Kumar Co-authored-by: Younes Belkada <49240599+younesbelkada@users.noreply.github.com> --- src/peft/peft_model.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index f73a66a..7491342 100644 --- a/src/peft/peft_model.py +++ b/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