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

refactor, lora support and utils for saving and loading

Browse Source
This commit is contained in:
Sourab Mangrulkar
2022-11-30 18:22:12 +05:30
parent 23aecc4f69
commit 751baf8aa7
9 changed files with 259 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+25 -25
View File
@@ -3,40 +3,40 @@ Parameter-Efficient Tuning. Intergrated with 🤗 Accelerate to scale seamlessly
Supported methods:
1. Prefix Tuning
2. P-Tuning
3. Prompt Tuning
4. LoRA [in backlog]
1. LoRA
2. Prefix Tuning
3. P-Tuning
4. Prompt Tuning
## Models support matrix
### Sequence Classification
| | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA |
| --------- | ---- | ---- | ---- | ---- |
| BERT | ✅ | ✅ | ✅ | |
| RoBERTa | ✅ | ✅ | ✅ | |
| GPT-2 | ✅ | ✅ | ✅ | |
| Bloom | ✅ | ✅ | ✅ | |
| OPT | ✅ | ✅ | ✅ | |
| GPT-Neo | ✅ | ✅ | ✅ | |
| GPT-J | ✅ | ✅ | ✅ | |
| Deberta | | | | |
| Deberta-v2 | | | | |
| Model | LoRA | Prefix Tuning | P-Tuning | Prompt Tuning |
| --------- | ---- | ---- | ---- | ---- |
| BERT | ✅ | ✅ | ✅ | ✅ |
| RoBERTa | ✅ | ✅ | ✅ | ✅ |
| GPT-2 | ✅ | ✅ | ✅ | ✅ |
| Bloom | ✅ | ✅ | ✅ | ✅ |
| OPT | ✅ | ✅ | ✅ | ✅ |
| GPT-Neo | ✅ | ✅ | ✅ | ✅ |
| GPT-J | ✅ | ✅ | ✅ | ✅ |
| Deberta | ✅ | | | |
| Deberta-v2 | ✅ | | | |
### Causal Language Modeling
| | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA |
| --------- | ---- | ---- | ---- | ---- |
| GPT-2 | ✅ | ✅ | ✅ | |
| Bloom | ✅ | ✅ | ✅ | |
| OPT | ✅ | ✅ | ✅ | |
| GPT-Neo | ✅ | ✅ | ✅ | |
| GPT-J | ✅ | ✅ | ✅ | |
| Model | LoRA | Prefix Tuning | P-Tuning | Prompt Tuning |
| --------- | ---- | ---- | ---- | ---- |
| GPT-2 | ✅ | ✅ | ✅ | ✅ |
| Bloom | ✅ | ✅ | ✅ | ✅ |
| OPT | ✅ | ✅ | ✅ | ✅ |
| GPT-Neo | ✅ | ✅ | ✅ | ✅ |
| GPT-J | ✅ | ✅ | ✅ | ✅ |
### Conditional Generation
| | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA |
| Model | LoRA | Prefix Tuning | P-Tuning | Prompt Tuning |
| --------- | ---- | ---- | ---- | ---- |
| T5 | ✅ | ✅ | ✅ | |
| BART | ✅ | ✅ | ✅ | |
| T5 | ✅ | ✅ | ✅ | ✅ |
| BART | ✅ | ✅ | ✅ | ✅ |
## Caveats:
src/pet/__init__.py
+10 -1
View File
@@ -18,4 +18,13 @@ from .tuners import (
PromptTuningConfig,
PromptTuningInit,
)
from .utils import PETConfig, PETType, PromptLearningConfig, TaskType
from .utils import (
PETConfig,
PETType,
PromptLearningConfig,
TaskType,
bloom_model_postprocess_past_key_value,
get_pet_model_state_dict,
set_pet_model_state_dict,
shift_tokens_right,
)
src/pet/mapping.py
+42 -3
View File
@@ -1,5 +1,6 @@
from .pet_model import PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification
from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig, LoRAConfig
from .tuners import LoRAConfig, PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
from .utils import PETType
MODEL_TYPE_TO_PET_MODEL_MAPPING = {
@@ -15,13 +16,28 @@ PET_TYPE_TO_CONFIG_MAPPING = {
"LORA": LoRAConfig,
}
TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
"t5": ["q", "v"],
"bart": ["q_proj", "v_proj"],
"gpt2": ["c_attn"],
"bloom": ["query_key_value"],
"opt": ["q_proj", "v_proj"],
"gptj": ["q_proj", "v_proj"],
"gpt_neox": ["query_key_value"],
"gpt_neo": ["q_proj", "v_proj"],
"bert": ["query", "value"],
"roberta": ["query", "value"],
"electra": ["query", "value"],
"deberta-v2": ["query_proj", "value_proj"],
"deberta": ["in_proj"],
}
def get_pet_config(config_dict):
return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
def get_pet_model(model, pet_config):
config = model.config.to_dict()
def _prepare_prompt_learning_config(pet_config, config):
if pet_config.num_layers is None:
if "num_hidden_layers" in config:
num_layers = config["num_hidden_layers"]
@@ -60,4 +76,27 @@ def get_pet_model(model, pet_config):
if pet_config.encoder_hidden_size is None:
pet_config.encoder_hidden_size = token_dim
return pet_config
def _prepare_lora_config(pet_config, config):
if pet_config.target_modules is None:
if config.model_type not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
raise ValueError("Please specify `target_modules` in `pet_config`")
pet_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[config.model_type]
if len(pet_config.target_modules) == 1:
pet_config.fan_in_fan_out = True
pet_config.enable_lora = [True, False, True]
if pet_config.inference_mode:
pet_config.merge_weights = True
return pet_config
def get_pet_model(model, pet_config):
config = model.config.to_dict()
if pet_config.pet_type != PETType.LORA:
pet_config = _prepare_prompt_learning_config(pet_config, config)
else:
pet_config = _prepare_lora_config(pet_config, config)
return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config)
src/pet/pet_model.py
+104 -65
View File
@@ -6,19 +6,24 @@ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from transformers import PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from .tuners import PrefixEncoder, PromptEmbedding, PromptEncoder
from .utils import PETConfig, PETType, TaskType
from .tuners import LoRAModel, PrefixEncoder, PromptEmbedding, PromptEncoder
from .utils import PETConfig, PETType, TaskType, shift_tokens_right
class PETModel(torch.nn.Module):
def __init__(self, model, pet_config: PETConfig):
super().__init__()
self.model = model
self.pet_config = pet_config
self.base_model = model
if pet_config.pet_type != PETType.LORA:
self._setup_prompt_encoder()
else:
self.base_model = LoRAModel(pet_config, model)
def _setup_prompt_encoder(self):
num_transformer_submodules = 0
transformer_backbone = None
for name, module in self.model.named_children():
for name, module in self.base_model.named_children():
if isinstance(module, PreTrainedModel):
# Make sure to freeze Tranformers model
for param in module.parameters():
@@ -30,7 +35,7 @@ class PETModel(torch.nn.Module):
self.pet_config.num_transformer_submodules = 2 if self.pet_config.task_type == TaskType.SEQ_2_SEQ_LM else 1
for named_param, value in list(transformer_backbone.named_parameters()):
if value.shape[0] == model.config.vocab_size:
if value.shape[0] == self.base_model.config.vocab_size:
self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", ""))
break
@@ -48,7 +53,7 @@ class PETModel(torch.nn.Module):
).long()
def get_prompt(self, batch_size):
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.model.device)
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.base_model.device)
if self.pet_config.pet_type == PETType.PREFIX_TUNING:
prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
if self.pet_config.inference_mode:
@@ -93,9 +98,9 @@ class PETModel(torch.nn.Module):
class PETModelForSequenceClassification(PETModel):
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.model.config
self.config = self.base_model.config
for name, module in self.model.named_children():
for name, module in self.base_model.named_children():
if isinstance(module, torch.nn.Linear):
self.cls_layer_name = name
break
@@ -113,10 +118,24 @@ class PETModelForSequenceClassification(PETModel):
):
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if self.pet_config.pet_type == PETType.LORA:
return self.base_model(
input_ids=input_ids,
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
labels=labels,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
**kwargs,
)
batch_size = input_ids.shape[0]
if attention_mask is not None and self.pet_config.pet_type != PETType.LORA:
if attention_mask is not None:
# concat prompt attention mask
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.model.device)
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
self.base_model.device
)
attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
if kwargs.get("position_ids", None) is not None:
warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
@@ -137,7 +156,7 @@ class PETModelForSequenceClassification(PETModel):
if kwargs.get("token_type_ids", None) is not None:
kwargs["token_type_ids"] = torch.cat(
(
torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.model.device),
torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.base_model.device),
kwargs["token_type_ids"],
),
dim=1,
@@ -146,7 +165,7 @@ class PETModelForSequenceClassification(PETModel):
inputs_embeds = self.word_embeddings(input_ids)
prompts = self.get_prompt(batch_size=batch_size)
inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
return self.model(inputs_embeds=inputs_embeds, **kwargs)
return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
def prefix_tuning_forward(
self,
@@ -161,7 +180,7 @@ class PETModelForSequenceClassification(PETModel):
):
batch_size = input_ids.shape[0]
past_key_values = self.get_prompt(batch_size)
fwd_params = list(inspect.signature(self.model.forward).parameters.keys())
fwd_params = list(inspect.signature(self.base_model.forward).parameters.keys())
kwargs.update(
{
"input_ids": input_ids,
@@ -174,37 +193,37 @@ class PETModelForSequenceClassification(PETModel):
}
)
if "past_key_values" in fwd_params:
return self.model(labels=labels, **kwargs)
return self.base_model(labels=labels, **kwargs)
else:
transformer_backbone_name = self.model.get_submodule(self.transformer_backbone_name)
transformer_backbone_name = self.base_model.get_submodule(self.transformer_backbone_name)
fwd_params = list(inspect.signature(transformer_backbone_name.forward).parameters.keys())
if "past_key_values" not in fwd_params:
raise ValueError("Model does not support past key values which are required for prefix tuning.")
outputs = transformer_backbone_name(**kwargs)
pooled_output = outputs[1] if len(outputs) > 1 else outputs[0]
if "dropout" in [name for name, _ in list(self.model.named_children())]:
pooled_output = self.model.dropout(pooled_output)
logits = self.model.get_submodule(self.cls_layer_name)(pooled_output)
if "dropout" in [name for name, _ in list(self.base_model.named_children())]:
pooled_output = self.base_model.dropout(pooled_output)
logits = self.base_model.get_submodule(self.cls_layer_name)(pooled_output)
loss = None
if labels is not None:
if self.config.problem_type is None:
if self.model.num_labels == 1:
if self.base_model.num_labels == 1:
self.config.problem_type = "regression"
elif self.model.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
elif self.base_model.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
self.config.problem_type = "single_label_classification"
else:
self.config.problem_type = "multi_label_classification"
if self.config.problem_type == "regression":
loss_fct = MSELoss()
if self.model.num_labels == 1:
if self.base_model.num_labels == 1:
loss = loss_fct(logits.squeeze(), labels.squeeze())
else:
loss = loss_fct(logits, labels)
elif self.config.problem_type == "single_label_classification":
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits.view(-1, self.model.num_labels), labels.view(-1))
loss = loss_fct(logits.view(-1, self.base_model.num_labels), labels.view(-1))
elif self.config.problem_type == "multi_label_classification":
loss_fct = BCEWithLogitsLoss()
loss = loss_fct(logits, labels)
@@ -223,7 +242,7 @@ class PETModelForSequenceClassification(PETModel):
class PETModelForCausalLM(PETModel):
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.model.config
self.config = self.base_model.config
def forward(
self,
@@ -236,14 +255,25 @@ class PETModelForCausalLM(PETModel):
return_dict=None,
**kwargs,
):
if self.pet_config.pet_type == PETType.LORA:
return self.base_model(
input_ids=input_ids,
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
labels=labels,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
**kwargs,
)
batch_size = input_ids.shape[0]
if self.pet_config.pet_type != PETType.LORA:
if attention_mask is not None:
# concat prompt attention mask
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
self.model.device
)
attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
if attention_mask is not None:
# concat prompt attention mask
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
self.base_model.device
)
attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
if kwargs.get("position_ids", None) is not None:
warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
@@ -263,26 +293,25 @@ class PETModelForCausalLM(PETModel):
if self.pet_config.pet_type == PETType.PREFIX_TUNING:
past_key_values = self.get_prompt(batch_size)
return self.model(input_ids=input_ids, past_key_values=past_key_values, **kwargs)
return self.base_model(input_ids=input_ids, past_key_values=past_key_values, **kwargs)
else:
if inputs_embeds is None:
inputs_embeds = self.word_embeddings(input_ids)
if self.pet_config.pet_type != PETType.LORA:
# concat prompt labels
if labels is not None:
prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(
self.model.device
)
kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
# concat prompt labels
if labels is not None:
prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(
self.base_model.device
)
kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
prompts = self.get_prompt(batch_size=batch_size)
inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
return self.model(inputs_embeds=inputs_embeds, **kwargs)
return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
class PETModelForSeq2SeqLM(PETModel):
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.model.config
self.config = self.base_model.config
def forward(
self,
@@ -298,15 +327,28 @@ class PETModelForSeq2SeqLM(PETModel):
return_dict=None,
**kwargs,
):
batch_size = input_ids.shape[0]
if self.pet_config.pet_type == PETType.LORA:
return self.base_model(
input_ids=input_ids,
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
decoder_input_ids=decoder_input_ids,
decoder_attention_mask=decoder_attention_mask,
decoder_inputs_embeds=decoder_inputs_embeds,
labels=labels,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
**kwargs,
)
if self.pet_config.pet_type != PETType.LORA:
if decoder_attention_mask is not None:
# concat prompt attention mask
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
self.model.device
)
decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1)
batch_size = input_ids.shape[0]
if decoder_attention_mask is not None:
# concat prompt attention mask
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
self.base_model.device
)
decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1)
if kwargs.get("position_ids", None) is not None:
warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
@@ -327,36 +369,33 @@ class PETModelForSeq2SeqLM(PETModel):
if self.pet_config.pet_type == PETType.PREFIX_TUNING:
past_key_values = self.get_prompt(batch_size)
return self.model(
return self.base_model(
input_ids=input_ids, decoder_input_ids=decoder_input_ids, past_key_values=past_key_values, **kwargs
)
else:
if inputs_embeds is None:
inputs_embeds = self.word_embeddings(input_ids)
if decoder_inputs_embeds is None and decoder_input_ids is None:
from transformers.models.bart.modeling_bart import shift_tokens_right
decoder_input_ids = shift_tokens_right(
labels, self.config.pad_token_id, self.config.decoder_start_token_id
)
decoder_inputs_embeds = self.word_embeddings(decoder_input_ids)
if self.pet_config.pet_type != PETType.LORA:
if attention_mask is not None:
# concat prompt attention mask
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
self.model.device
)
kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1)
# concat prompt labels
if labels is not None:
prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(
self.model.device
)
kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
if attention_mask is not None:
# concat prompt attention mask
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
self.base_model.device
)
kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1)
# concat prompt labels
if labels is not None:
prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(
self.base_model.device
)
kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
prompts = self.get_prompt(batch_size=batch_size)
inputs_embeds = torch.cat((prompts[:, : self.pet_config.num_virtual_tokens], inputs_embeds), dim=1)
decoder_inputs_embeds = torch.cat(
(prompts[:, self.pet_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1
)
return self.model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs)
return self.base_model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs)
src/pet/tuners/lora.py
+29 -16
View File
@@ -1,19 +1,20 @@
# todo
from dataclasses import asdict, dataclass, field
from dataclasses import dataclass, field
from typing import Optional
import torch
from transformers.pytorch_utils import Conv1D
import loralib as lora
from loralib import lora_state_dict, mark_only_lora_as_trainable # noqa: F401
from loralib import mark_only_lora_as_trainable
from ..utils import PETConfig
@dataclass
class LoRAConfig(PETConfig):
r: int = field(default=None, metadata={"help": "LoRA attention dimension"})
r: int = field(default=8, metadata={"help": "LoRA attention dimension"})
target_modules: Optional[list] = field(default=None, metadata={"help": "List of modules to replace with LoRA"})
lora_alpha: int = field(default=None, metadata={"help": "LoRA alpha"})
lora_dropout: float = field(default=None, metadata={"help": "LoRA dropout"})
merge_weights: bool = field(
@@ -23,7 +24,7 @@ class LoRAConfig(PETConfig):
default=False,
metadata={"help": "Set this to True if the layer to replace stores weight like (fan_in, fan_out)"},
)
target_modules: Optional[list] = field(default=None, metadata={"help": "List of modules to replace with LoRA"})
enable_lora: Optional[list[bool]] = field(default=None, metadata={"help": "Used with `lora.MergedLinear`."})
bias: str = field(default="none", metadata={"help": "Bias type for LoRA. Can be 'none', 'all' or 'lora_only'"})
@@ -31,30 +32,42 @@ class LoRAModel(torch.nn.Module):
def __init__(self, config, model):
super().__init__()
self.config = config
self.model = model
self.lora_model = model
self.find_and_replace()
mark_only_lora_as_trainable(self.model, self.config.bias)
mark_only_lora_as_trainable(self.lora_model, self.config.bias)
def find_and_replace(self):
key_list = [key for key, _ in self.model.named_modules()]
kwargs = {
"r": self.config.r,
"lora_alpha": self.config.lora_alpha,
"lora_dropout": self.config.lora_dropout,
"fan_in_fan_out": self.config.fan_in_fan_out,
"merge_weights": self.config.merge_weights,
}
key_list = [key for key, _ in self.lora_model.named_modules()]
for key in key_list:
if any(key.endswith(target_key) for target_key in self.config.target_module_keys):
if any(key.endswith(target_key) for target_key in self.config.target_modules):
parent, target, target_name = self.get_submodules(key)
# print(parent, target, target_name)
if isinstance(target, torch.nn.Linear):
new_module = lora.Linear(target.in_features, target.out_features, **asdict(self.config))
new_module = lora.Linear(target.in_features, target.out_features, **kwargs)
elif isinstance(target, Conv1D):
kwargs.update({"enable_lora": self.config.enable_lora})
in_features, out_features = target.weight.shape
new_module = lora.MergedLinear(in_features, out_features, **asdict(self.config))
self.replace_module(parent, target_name, new_module)
new_module = lora.MergedLinear(in_features, out_features, **kwargs)
self.replace_module(parent, target_name, new_module, target)
def get_submodules(self, key):
parent = self.model.get_submodule(".".join(key.split(".")[:-1]))
target_name = key.split(".")[:-1]
target = self.model.get_submodule(key)
parent = self.lora_model.get_submodule(".".join(key.split(".")[:-1]))
target_name = key.split(".")[-1]
target = self.lora_model.get_submodule(key)
return parent, target, target_name
def replace_module(self, parent_module, child_name, new_module, old_module):
setattr(parent_module, child_name, new_module)
new_module.weight = old_module.weight.clone()
new_module.weight = old_module.weight
if old_module.bias is not None:
new_module.bias = old_module.bias.clone()
new_module.bias = old_module.bias
def forward(self, *args, **kwargs):
return self.lora_model(*args, **kwargs)
src/pet/utils/__init__.py
+2
View File
@@ -3,3 +3,5 @@
# module, but to preserve other warnings. So, don't check this module at all
from .config import PETConfig, PETType, PromptLearningConfig, TaskType
from .other import bloom_model_postprocess_past_key_value, shift_tokens_right
from .save_and_load import get_pet_model_state_dict, set_pet_model_state_dict
src/pet/utils/config.py
+1 -1
View File
@@ -19,7 +19,7 @@ class TaskType(str, enum.Enum):
@dataclass
class PETConfig:
"""
This is the configuration class to store the configuration of a :class:`~transform
This is the configuration class to store the configuration of a :class:`~pet.PETModel`.
"""
pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"})
src/pet/utils/other.py
+18
View File
@@ -1,6 +1,7 @@
import torch
# needed for prefix-tuning of bloom model
def bloom_model_postprocess_past_key_value(past_key_values):
past_key_values = torch.cat(past_key_values)
total_layers, batch_size, num_attention_heads, num_virtual_tokens, head_dim = past_key_values.shape
@@ -12,3 +13,20 @@ def bloom_model_postprocess_past_key_value(past_key_values):
values = values.reshape(total_layers // 2, batch_size * num_attention_heads, num_virtual_tokens, head_dim)
return tuple(zip(keys, values))
# copied from transformers.models.bart.modeling_bart
def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
"""
Shift input ids one token to the right.
"""
shifted_input_ids = input_ids.new_zeros(input_ids.shape)
shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
shifted_input_ids[:, 0] = decoder_start_token_id
if pad_token_id is None:
raise ValueError("self.model.config.pad_token_id has to be defined.")
# replace possible -100 values in labels by `pad_token_id`
shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
return shifted_input_ids
src/pet/utils/save_and_load.py
+28
View File
@@ -0,0 +1,28 @@
from loralib import lora_state_dict
from .config import PETType
def get_pet_model_state_dict(model):
if model.pet_config.pet_type == PETType.LORA:
return lora_state_dict(model)
else:
to_return = {}
state_dict = model.state_dict()
prompt_tokens = model.prompt_tokens.unsqueeze(0).expand(1, -1).to(model.base_model.device)
prompt_embeddings = model.prompt_encoder(prompt_tokens).detach().cpu()
to_return["prompt_embeddings"] = prompt_embeddings
if model.modules_to_save is not None:
for key, value in state_dict.items():
if any(module_name in key for module_name in model.modules_to_save):
to_return[key] = value
return to_return
def set_pet_model_state_dict(model, pet_model_state_dict):
model.load_state_dict(pet_model_state_dict, strict=False)
if model.pet_config.pet_type != PETType.LORA:
model.prompt_encoder.embedding.load_state_dict(
{"weight": pet_model_state_dict["prompt_embeddings"]}, strict=True
)
return model