wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-09 11:28:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

Causal LM generation fix for prefix tuning: GPT2 model (#222)

Browse Source 
* 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 <vineeku6@in.ibm.com>
Co-authored-by: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
This commit is contained in:
Vineet Kumar
2023-03-29 15:20:14 +02:00
committed by GitHub
co-authored by Younes Belkada Vineet Kumar
parent 51f49a5fe4
commit d8d1007732
1 changed files with 15 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/peft/peft_model.py
+15 -1
View File
@@ -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