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

Merge pull request #2 from huggingface/smangrul/add-examples-fixes-docs

Browse Source 
adding detailed docs, refactor and fixes
This commit is contained in:
Sourab Mangrulkar
2022-12-02 16:21:50 +05:30
committed by GitHub
parent a21f24d9e1 ae609d680e
commit 9ac43b90a6
16 changed files with 7609 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+1 -1
View File
@@ -1,6 +1,6 @@
.PHONY: quality style test docs
check_dirs := src
check_dirs := src examples
# Check that source code meets quality standards
README.md
+84 -1
View File
@@ -1,5 +1,5 @@
# 🤗 PET
Parameter-Efficient Tuning. Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP.
Parameter-Efficient Tuning methods enable . Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP.
Supported methods:
@@ -8,6 +8,87 @@ Supported methods:
3. P-Tuning
4. Prompt Tuning
## Getting started
```python
from transformers import AutoModelForSeq2SeqLM
from pet import get_pet_config, get_pet_model
model_name_or_path = "bigscience/mt0-large"
tokenizer_name_or_path = "bigscience/mt0-large"
config = {
"pet_type":"LORA",
"task_type":"SEQ_2_SEQ_LM",
"r": 8,
"lora_alpha": 32,
"lora_dropout": 0.1
}
pet_config = get_pet_config(config)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
model = get_pet_model(model, pet_config)
model.print_trainable_parameters()
# output: trainable params: 2359296 || all params: 1231940608 || trainable%: 0.19151053100118282
```
## PET + 🤗 Accelerate
PET models work with 🤗 Accelerate out of the box.
For scaling to large models, you can leverage 🤗 Accelerate's PyTorch FSDP integration as shown below.
PyTorch FSDP shards parameters, gradients and optimizer states across data parallel workers which enables
large language models to fit on available hardware.
It also supports CPU offloading to further enable distributed training at scale.
```python
from pet.utils.other import fsdp_auto_wrap_policy
...
if os.environ.get("ACCELERATE_USE_FSDP", None) is not None:
accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model)
model = accelerator.prepare(model)
```
Example of parameter efficient tuning with `mt0-xxl` base model using 🤗 Accelerate is provided in `~examples/pet_lora_seq2seq_accelerate_fsdp.py`.
1. First run `accelerate config --config_file fsdp_config.yaml` and answer the questionaire.
Below are the contents of the config file.
```
command_file: null
commands: null
compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: FSDP
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config:
fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
fsdp_backward_prefetch_policy: BACKWARD_PRE
fsdp_offload_params: true
fsdp_sharding_strategy: 1
fsdp_state_dict_type: FULL_STATE_DICT
fsdp_transformer_layer_cls_to_wrap: T5Block
gpu_ids: null
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config: {}
mixed_precision: 'no'
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_name: null
tpu_zone: null
use_cpu: false
```
2. run the below command to launch example script
```
accelerate launch --config_file fsdp_config.yaml examples/pet_lora_seq2seq_accelerate_fsdp.py
```
## Models support matrix
### Sequence Classification
@@ -41,5 +122,7 @@ Supported methods:
## Caveats:
1. Doesn't work currently with DeeSpeed ZeRO Stage-3. Extending support with DeeSpeed ZeRO Stage-3 is in backlog.
2. When using `P_TUNING` or `PROMPT_TUNING` with `SEQ_2_SEQ` task, remember to remove the `num_virtual_token` virtual prompt predictions from the left side of the model outputs during evaluations.
examples/pet_lora_seq2seq.ipynb
+4815
View File
File diff suppressed because it is too large Load Diff
examples/pet_lora_seq2seq_accelerate_fsdp.py
+137
View File
@@ -0,0 +1,137 @@
import torch
from accelerate import Accelerator
from torch.utils.data import DataLoader
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, default_data_collator, get_linear_schedule_with_warmup
from datasets import load_dataset
from pet import get_pet_config, get_pet_model, get_pet_model_state_dict
from pet.utils.other import fsdp_auto_wrap_policy
from tqdm import tqdm
def main():
accelerator = Accelerator()
model_name_or_path = "bigscience/mt0-xxl"
batch_size = 16
text_column = "sentence"
label_column = "text_label"
max_length = 64
lr = 1e-3
num_epochs = 1
config = {"pet_type": "LORA", "task_type": "SEQ_2_SEQ_LM", "r": 8, "lora_alpha": 32, "lora_dropout": 0.1}
pet_config = get_pet_config(config)
checkpoint_name = "financial_sentiment_analysis_lora_fsdp_v1.pt"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
model = get_pet_model(model, pet_config)
accelerator.print(model.print_trainable_parameters())
dataset = load_dataset("financial_phrasebank", "sentences_allagree")
dataset = dataset["train"].train_test_split(test_size=0.1)
dataset["validation"] = dataset["test"]
del dataset["test"]
classes = dataset["train"].features["label"].names
dataset = dataset.map(
lambda x: {"text_label": [classes[label] for label in x["label"]]},
batched=True,
num_proc=1,
)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
def preprocess_function(examples):
inputs = examples[text_column]
targets = examples[label_column]
model_inputs = tokenizer(
inputs, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
)
labels = tokenizer(targets, max_length=3, padding="max_length", truncation=True, return_tensors="pt")
labels = labels["input_ids"]
labels[labels == tokenizer.pad_token_id] = -100
model_inputs["labels"] = labels
return model_inputs
processed_datasets = dataset.map(
preprocess_function,
batched=True,
num_proc=1,
remove_columns=dataset["train"].column_names,
load_from_cache_file=False,
desc="Running tokenizer on dataset",
)
train_dataset = processed_datasets["train"]
eval_dataset = processed_datasets["validation"]
train_dataloader = DataLoader(
train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True
)
eval_dataloader = DataLoader(
eval_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True
)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
lr_scheduler = get_linear_schedule_with_warmup(
optimizer=optimizer,
num_warmup_steps=0,
num_training_steps=(len(train_dataloader) * num_epochs),
)
if accelerator.state.fsdp_plugin is not None:
accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model)
model, train_dataloader, eval_dataloader, optimizer, lr_scheduler = accelerator.prepare(
model, train_dataloader, eval_dataloader, optimizer, lr_scheduler
)
accelerator.print(model)
for epoch in range(num_epochs):
model.train()
total_loss = 0
for step, batch in enumerate(tqdm(train_dataloader)):
outputs = model(**batch)
loss = outputs.loss
total_loss += loss.detach().float()
loss.backward()
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
model.eval()
eval_loss = 0
eval_preds = []
for step, batch in enumerate(tqdm(eval_dataloader)):
with torch.no_grad():
outputs = model(**batch)
loss = outputs.loss
eval_loss += loss.detach().float()
eval_preds.extend(
tokenizer.batch_decode(
accelerator.gather_for_metrics(torch.argmax(outputs.logits, -1)).detach().cpu().numpy(),
skip_special_tokens=True,
)
)
eval_epoch_loss = eval_loss / len(train_dataloader)
eval_ppl = torch.exp(eval_epoch_loss)
train_epoch_loss = total_loss / len(eval_dataloader)
train_ppl = torch.exp(train_epoch_loss)
accelerator.print(f"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}")
correct = 0
total = 0
for pred, true in zip(eval_preds, dataset["validation"][label_column]):
if pred.strip() == true.strip():
correct += 1
total += 1
accuracy = correct / total * 100
accelerator.print(f"{accuracy=}")
accelerator.print(f"{eval_preds[:10]=}")
accelerator.wait_for_everyone()
accelerator.save(get_pet_model_state_dict(model), checkpoint_name)
accelerator.wait_for_everyone()
if __name__ == "__main__":
main()
examples/pet_prefix_tuning_seq2seq.ipynb
+1963
View File
File diff suppressed because it is too large Load Diff
setup.py
+2 -1
View File
@@ -39,9 +39,10 @@ setup(
"packaging>=20.0",
"psutil",
"pyyaml",
"torch>=1.4.0",
"torch>=1.13.0",
"transformers",
"accelerate",
"loralib",
],
extras_require=extras,
classifiers=[
src/pet/mapping.py
+19 -2
View File
@@ -18,6 +18,7 @@ PET_TYPE_TO_CONFIG_MAPPING = {
TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
"t5": ["q", "v"],
"mt5": ["q", "v"],
"bart": ["q_proj", "v_proj"],
"gpt2": ["c_attn"],
"bloom": ["query_key_value"],
@@ -27,6 +28,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
"gpt_neo": ["q_proj", "v_proj"],
"bert": ["query", "value"],
"roberta": ["query", "value"],
"xlm-roberta": ["query", "value"],
"electra": ["query", "value"],
"deberta-v2": ["query_proj", "value_proj"],
"deberta": ["in_proj"],
@@ -34,6 +36,13 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
def get_pet_config(config_dict):
"""
Returns a PET config object from a dictionary.
Args:
config_dict (:obj:`Dict[str, Any]`):
"""
return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
@@ -73,8 +82,8 @@ def _prepare_prompt_learning_config(pet_config, model_config):
raise ValueError("Please specify `num_attention_heads` in `pet_config`")
pet_config.num_attention_heads = num_attention_heads
if pet_config.encoder_hidden_size is None:
pet_config.encoder_hidden_size = token_dim
if getattr(pet_config, "encoder_hidden_size", None) is None:
setattr(pet_config, "encoder_hidden_size", token_dim)
return pet_config
@@ -93,6 +102,14 @@ def _prepare_lora_config(pet_config, model_config):
def get_pet_model(model, pet_config):
"""
Returns a PET model object from a model and a config.
Args:
model (:obj:`transformers.PreTrainedModel`):
pet_config (:obj:`transformers.PETConfig`):
"""
model_config = model.config.to_dict()
if pet_config.pet_type != PETType.LORA:
pet_config = _prepare_prompt_learning_config(pet_config, model_config)
src/pet/pet_model.py
+123 -29
View File
@@ -7,28 +7,50 @@ from transformers import PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from .tuners import LoRAModel, PrefixEncoder, PromptEmbedding, PromptEncoder
from .utils import PETConfig, PETType, TaskType, shift_tokens_right
from .utils import PETConfig, PETType, TaskType, _set_trainable, shift_tokens_right
class PETModel(torch.nn.Module):
"""
Parameter Efficient Tuning Model. Base model encompassing various PET methods.
Args:
model (:obj:`PreTrainedModel`): The base transformer model used for PET.
pet_config (:obj:`PETConfig`): The configuration of the PET model.
Attributes:
base_model (:obj:`PreTrainedModel`): The base transformer model used for PET. pet_config (:obj:`PETConfig`):
The configuration of the PET model. modules_to_save (:obj:`list` of :obj:`str`): The list of sub-module names
to save when saving the model. prompt_encoder (:obj:`PromptEncoder`): The prompt encoder used for PET if
`pet_config.pet_type != PETType.LORA`. prompt_tokens (:obj:`torch.Tensor`): The virtual prompt tokens used for
PET if `pet_config.pet_type != PETType.LORA`. transformer_backbone_name (:obj:`str`): The name of the
transformer backbone in the base model
if `pet_config.pet_type != PETType.LORA`.
word_embeddings (:obj:`torch.nn.Embedding`): The word embeddings of the transformer backbone
in the base model if `pet_config.pet_type != PETType.LORA`.
"""
def __init__(self, model, pet_config: PETConfig):
super().__init__()
self.pet_config = pet_config
self.base_model = model
self.config = self.base_model.config
self.modules_to_save = None
if pet_config.pet_type != PETType.LORA:
self._setup_prompt_encoder()
else:
self.base_model = LoRAModel(pet_config, model)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def _setup_prompt_encoder(self):
num_transformer_submodules = 0
transformer_backbone = None
for name, module in self.base_model.named_children():
for param in module.parameters():
param.requires_grad = False
if isinstance(module, PreTrainedModel):
# Make sure to freeze Tranformers model
for param in module.parameters():
param.requires_grad = False
if transformer_backbone is None:
transformer_backbone = module
self.transformer_backbone_name = name
@@ -54,14 +76,21 @@ class PETModel(torch.nn.Module):
).long()
def get_prompt_embedding_to_save(self):
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.base_model.device)
"""
Returns the prompt embedding to save when saving the model. Only applocable when `pet_config.pet_type !=
PETType.LORA`.
"""
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.device)
if self.pet_config.pet_type == PETType.PREFIX_TUNING:
prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
prompt_embeddings = self.prompt_encoder(prompt_tokens)
return prompt_embeddings[0].detach().cpu()
def get_prompt(self, batch_size):
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.base_model.device)
"""
Returns the virtual prompts to use for PET. Only applocable when `pet_config.pet_type != PETType.LORA`.
"""
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.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:
@@ -92,6 +121,9 @@ class PETModel(torch.nn.Module):
return prompts
def print_trainable_parameters(self):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in self.named_parameters():
@@ -102,11 +134,42 @@ class PETModel(torch.nn.Module):
f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
)
def __getattr__(self, name: str):
"""Forward missing attributes to the wrapped module."""
try:
return super().__getattr__(name) # defer to nn.Module's logic
except AttributeError:
return getattr(self.base_model, name)
class PETModelForSequenceClassification(PETModel):
"""
PET model for sequence classification tasks.
Args:
model (:obj:`PreTrainedModel`): Base transformer model
pet_config (:obj:`PETConfig`): PET config.
Attributes:
config (:obj:`PretrainedConfig`): The configuration object of the base model. cls_layer_name (:obj:`str`): The
name of the classification layer.
Example::
>>> from transformers import AutoModelForSequenceClassification >>> from pet import
PETModelForSequenceClassification, get_pet_config >>> config = {
'pet_type': 'PREFIX_TUNING', 'task_type': 'SEQ_CLS', 'inference_mode': False, 'num_virtual_tokens': 20,
'token_dim': 768, 'num_transformer_submodules': 1, 'num_attention_heads': 12, 'num_layers': 12,
'encoder_hidden_size': 768, 'prefix_projection': False, 'postprocess_past_key_value_function': None
}
>>> pet_config = get_pet_config(config) >>> model =
AutoModelForSequenceClassification.from_pretrained("bert-base-cased") >>> pet_model =
PETModelForSequenceClassification(model, pet_config) >>> pet_model.print_trainable_parameters() trainable
params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117
"""
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.base_model.config
self.modules_to_save = ["classifier"]
for name, module in self.base_model.named_children():
@@ -114,6 +177,9 @@ class PETModelForSequenceClassification(PETModel):
self.cls_layer_name = name
break
# to make sure classifier layer is trainable
_set_trainable(self.base_model)
def forward(
self,
input_ids=None,
@@ -142,9 +208,7 @@ class PETModelForSequenceClassification(PETModel):
batch_size = input_ids.shape[0]
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
)
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.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.")
@@ -160,12 +224,12 @@ class PETModelForSequenceClassification(PETModel):
)
if self.pet_config.pet_type == PETType.PREFIX_TUNING:
return self.prefix_tuning_forward(input_ids=input_ids, **kwargs)
return self._prefix_tuning_forward(input_ids=input_ids, **kwargs)
else:
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.base_model.device),
torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.device),
kwargs["token_type_ids"],
),
dim=1,
@@ -176,7 +240,7 @@ class PETModelForSequenceClassification(PETModel):
inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
def prefix_tuning_forward(
def _prefix_tuning_forward(
self,
input_ids=None,
attention_mask=None,
@@ -249,9 +313,29 @@ class PETModelForSequenceClassification(PETModel):
class PETModelForCausalLM(PETModel):
"""
PET model for Causal LM
Args:
model (:obj:`PreTrainedModel`): Base transformer model
pet_config (:obj:`PETConfig`): PET config.
Example::
>>> from transformers import AutoModelForCausalLM >>> from pet import PETModelForCausalLM, get_pet_config >>>
config = {
'pet_type': 'PREFIX_TUNING', 'task_type': 'CAUSAL_LM', 'inference_mode': False, 'num_virtual_tokens':
20, 'token_dim': 1280, 'num_transformer_submodules': 1, 'num_attention_heads': 20, 'num_layers': 36,
'encoder_hidden_size': 1280, 'prefix_projection': False, 'postprocess_past_key_value_function': None
}
>>> pet_config = get_pet_config(config) >>> model = AutoModelForCausalLM.from_pretrained("gpt2-large") >>>
pet_model = PETModelForCausalLM(model, pet_config) >>> pet_model.print_trainable_parameters() trainable params:
1843200 || all params: 775873280 || trainable%: 0.23756456724479544
"""
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.base_model.config
def forward(
self,
@@ -279,9 +363,7 @@ class PETModelForCausalLM(PETModel):
batch_size = input_ids.shape[0]
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
)
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device)
attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
if kwargs.get("position_ids", None) is not None:
@@ -308,9 +390,7 @@ class PETModelForCausalLM(PETModel):
inputs_embeds = self.word_embeddings(input_ids)
# 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
)
prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(self.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)
@@ -318,9 +398,29 @@ class PETModelForCausalLM(PETModel):
class PETModelForSeq2SeqLM(PETModel):
"""
PET model for Seq2Seq LM
Args:
model (:obj:`PreTrainedModel`): Base transformer model
pet_config (:obj:`PETConfig`): PET config.
Example::
>>> from transformers import AutoModelForSeq2SeqLM >>> from pet import PETModelForSeq2SeqLM, get_pet_config >>>
config = {
'pet_type': 'LORA', 'task_type': 'SEQ_2_SEQ_LM', 'inference_mode': False, 'r': 8, 'target_modules':
['q', 'v'], 'lora_alpha': 32, 'lora_dropout': 0.1, 'merge_weights': False, 'fan_in_fan_out': False,
'enable_lora': None, 'bias': 'none'
}
>>> pet_config = get_pet_config(config) >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>>
pet_model = PETModelForSeq2SeqLM(model, pet_config) >>> pet_model.print_trainable_parameters() trainable
params: 884736 || all params: 223843584 || trainable%: 0.3952474242013566
"""
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.base_model.config
def forward(
self,
@@ -354,9 +454,7 @@ class PETModelForSeq2SeqLM(PETModel):
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
)
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device)
decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1)
if kwargs.get("position_ids", None) is not None:
@@ -392,15 +490,11 @@ class PETModelForSeq2SeqLM(PETModel):
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
)
prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.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
)
prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(self.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)
src/pet/tuners/lora.py
+264 -11
View File
@@ -1,11 +1,14 @@
# todo
import math
from dataclasses import dataclass, field
from typing import Optional
from typing import List, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.pytorch_utils import Conv1D
import loralib as lora
import loralib as lora # noqa: F401
from loralib import mark_only_lora_as_trainable
from ..utils import PETConfig
@@ -13,6 +16,21 @@ from ..utils import PETConfig
@dataclass
class LoRAConfig(PETConfig):
"""
This is the configuration class to store the configuration of a :class:`~pet.LoRA`.
Args:
r: (:obj:`init`): LoRA attention dimension
target_modules (:obj: list of :obj: str): The names of the modules to apply LoRA to.
lora_alpha (:obj: float): The alpha parameter for LoRA scaling.
lora_dropout (:obj: float): The dropout probability for LoRA layers.
merge_weights (:obj: bool):
Whether to merge the weights of the LoRA layers with the base transformer model in `eval` mode.
fan_in_fan_out (:obj: bool): Set this to True if the layer to replace stores weight like (fan_in, fan_out)
enable_lora (:obj: list of :obj: bool): Used with `lora.MergedLinear`.
bias (:obj: str): Bias type for LoRA. Can be 'none', 'all' or 'lora_only'
"""
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"})
@@ -29,14 +47,37 @@ class LoRAConfig(PETConfig):
class LoRAModel(torch.nn.Module):
"""
Creates Low Rank Adapter (LoRA) model from a pretrained transformers model.
Args:
model (:obj:`transformers.PreTrainedModel`): The model to be adapted.
config (:obj:`LoRAConfig`): The configuration of the LoRA model.
Returns:
:obj:`torch.nn.Module`: The LoRA model.
Example::
>>> from transformers import AutoModelForSeq2SeqLM, LoRAConfig >>> from pet import LoRAModel, LoRAConfig >>>
config = LoRAConfig(
pet_type="LORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"],
lora_dropout=0.01, )
>>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> lora_model = LoRAModel(config, model)
Attributes:
model (:obj:`transformers.PreTrainedModel`): The model to be adapted. config (:obj:`LoRAConfig`): The
configuration of the LoRA model.
"""
def __init__(self, config, model):
super().__init__()
self.config = config
self.model = model
self.find_and_replace()
self._find_and_replace()
mark_only_lora_as_trainable(self.model, self.config.bias)
def find_and_replace(self):
def _find_and_replace(self):
kwargs = {
"r": self.config.r,
"lora_alpha": self.config.lora_alpha,
@@ -47,23 +88,23 @@ class LoRAModel(torch.nn.Module):
key_list = [key for key, _ in self.model.named_modules()]
for key in key_list:
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)
parent, target, target_name = self._get_submodules(key)
bias = target.bias is not None
if isinstance(target, torch.nn.Linear):
new_module = lora.Linear(target.in_features, target.out_features, **kwargs)
new_module = Linear(target.in_features, target.out_features, bias=bias, **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, **kwargs)
self.replace_module(parent, target_name, new_module, target)
new_module = MergedLinear(in_features, out_features, bias=bias, **kwargs)
self._replace_module(parent, target_name, new_module, target)
def get_submodules(self, key):
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)
return parent, target, target_name
def replace_module(self, parent_module, child_name, new_module, old_module):
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
if old_module.bias is not None:
@@ -71,3 +112,215 @@ class LoRAModel(torch.nn.Module):
def forward(self, *args, **kwargs):
return self.model(*args, **kwargs)
def __getattr__(self, name: str):
"""Forward missing attributes to the wrapped module."""
try:
return super().__getattr__(name) # defer to nn.Module's logic
except AttributeError:
return getattr(self.model, name)
# Below code is copied from https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
# and modified to work with PyTorch FSDP
# ------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
# ------------------------------------------------------------------------------------------
class LoRALayer:
def __init__(
self,
r: int,
lora_alpha: int,
lora_dropout: float,
merge_weights: bool,
):
self.r = r
self.lora_alpha = lora_alpha
# Optional dropout
if lora_dropout > 0.0:
self.lora_dropout = nn.Dropout(p=lora_dropout)
else:
self.lora_dropout = lambda x: x
# Mark the weight as unmerged
self.merged = False
self.merge_weights = merge_weights
class Linear(nn.Linear, LoRALayer):
# LoRA implemented in a dense layer
def __init__(
self,
in_features: int,
out_features: int,
r: int = 0,
lora_alpha: int = 1,
lora_dropout: float = 0.0,
fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
merge_weights: bool = True,
**kwargs,
):
nn.Linear.__init__(self, in_features, out_features, **kwargs)
LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
self.fan_in_fan_out = fan_in_fan_out
# Actual trainable parameters
if r > 0:
self.lora_A = nn.Linear(in_features, r, bias=False)
self.lora_B = nn.Linear(r, out_features, bias=False)
self.scaling = self.lora_alpha / self.r
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
self.reset_parameters()
if fan_in_fan_out:
self.weight.data = self.weight.data.T
def reset_parameters(self):
nn.Linear.reset_parameters(self)
if hasattr(self, "lora_A"):
# initialize A the same way as the default for nn.Linear and B to zero
nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
nn.init.zeros_(self.lora_B.weight)
def train(self, mode: bool = True):
def T(w):
return w.T if self.fan_in_fan_out else w
nn.Linear.train(self, mode)
self.lora_A.train(mode)
self.lora_B.train(mode)
if self.merge_weights and self.merged:
# Make sure that the weights are not merged
if self.r > 0:
self.weight.data -= T(self.lora_B.weight @ self.lora_A.weight) * self.scaling
self.merged = False
def eval(self):
def T(w):
return w.T if self.fan_in_fan_out else w
nn.Linear.eval(self)
self.lora_A.eval()
self.lora_B.eval()
if self.merge_weights and not self.merged:
# Merge the weights and mark it
if self.r > 0:
self.weight.data += T(self.lora_B.weight @ self.lora_A.weight) * self.scaling
self.merged = True
def forward(self, x: torch.Tensor):
def T(w):
return w.T if self.fan_in_fan_out else w
if self.r > 0 and not self.merged:
result = F.linear(x, T(self.weight), bias=self.bias)
if self.r > 0:
result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling
return result
else:
return F.linear(x, T(self.weight), bias=self.bias)
class MergedLinear(nn.Linear, LoRALayer):
# LoRA implemented in a dense layer
def __init__(
self,
in_features: int,
out_features: int,
r: int = 0,
lora_alpha: int = 1,
lora_dropout: float = 0.0,
enable_lora: List[bool] = [False],
fan_in_fan_out: bool = False,
merge_weights: bool = True,
**kwargs,
):
nn.Linear.__init__(self, in_features, out_features, **kwargs)
LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
assert out_features % len(enable_lora) == 0, "The length of enable_lora must divide out_features"
self.enable_lora = enable_lora
self.fan_in_fan_out = fan_in_fan_out
# Actual trainable parameters
if r > 0 and any(enable_lora):
self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False)
self.lora_B = nn.Conv1d(
r * sum(enable_lora),
out_features // len(enable_lora) * sum(enable_lora),
kernel_size=1,
groups=2,
bias=False,
)
self.scaling = self.lora_alpha / self.r
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
# Compute the indices
self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1)
self.lora_ind[enable_lora, :] = True
self.lora_ind = self.lora_ind.view(-1)
self.reset_parameters()
if fan_in_fan_out:
self.weight.data = self.weight.data.T
def reset_parameters(self):
nn.Linear.reset_parameters(self)
if hasattr(self, "lora_A"):
# initialize A the same way as the default for nn.Linear and B to zero
nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
nn.init.zeros_(self.lora_B.weight)
def zero_pad(self, x):
result = x.new_zeros((*x.shape[:-1], self.out_features))
result = result.view(-1, self.out_features)
result[:, self.lora_ind] = x.reshape(-1, self.out_features // len(self.enable_lora) * sum(self.enable_lora))
return result.view((*x.shape[:-1], self.out_features))
def train(self, mode: bool = True):
def T(w):
return w.T if self.fan_in_fan_out else w
nn.Linear.train(self, mode)
self.lora_A.train(mode)
self.lora_B.train(mode)
if self.merge_weights and self.merged:
# Make sure that the weights are not merged
if self.r > 0 and any(self.enable_lora):
delta_w = F.conv1d(
self.lora_A.weight.data.unsqueeze(0),
self.lora_B.weight.data.unsqueeze(-1),
groups=sum(self.enable_lora),
).squeeze(0)
self.weight.data -= self.zero_pad(T(delta_w * self.scaling))
self.merged = False
def eval(self):
def T(w):
return w.T if self.fan_in_fan_out else w
nn.Linear.eval(self)
self.lora_A.eval()
self.lora_B.eval()
if self.merge_weights and not self.merged:
# Merge the weights and mark it
if self.r > 0 and any(self.enable_lora):
delta_w = F.conv1d(
self.lora_A.weight.data.unsqueeze(0),
self.lora_B.weight.data.unsqueeze(-1),
groups=sum(self.enable_lora),
).squeeze(0)
self.weight.data += self.zero_pad(T(delta_w * self.scaling))
self.merged = True
def forward(self, x: torch.Tensor):
def T(w):
return w.T if self.fan_in_fan_out else w
if self.merged:
return F.linear(x, T(self.weight), bias=self.bias)
else:
result = F.linear(x, T(self.weight), bias=self.bias)
if self.r > 0:
after_A = self.lora_A(self.lora_dropout(x))
after_B = self.lora_B(after_A.transpose(-2, -1))
result += self.zero_pad(after_B) * self.scaling
return result
src/pet/tuners/p_tuning.py
+44 -3
View File
@@ -14,21 +14,33 @@ class PromptEncoderReparameterizationType(str, enum.Enum):
@dataclass
class PromptEncoderConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~pet.PromptEncoder`.
Args:
encoder_reparameterization_type
(:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]): The type of reparameterization to
use.
encoder_hidden_size (:obj:`int`): The hidden size of the prompt encoder.
encoder_num_layers (:obj:`int`): The number of layers of the prompt encoder.
encoder_dropout (:obj:`float`): The dropout probability of the prompt encoder.
"""
encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field(
default=PromptEncoderReparameterizationType.MLP,
metadata={"help": "How to reparameterize the prompt encoder"},
)
encoder_hidden_size: int = field(
default=None,
metadata={"help": "The hidden size of the prompt encoder reparameterization"},
metadata={"help": "The hidden size of the prompt encoder"},
)
encoder_num_layers: int = field(
default=2,
metadata={"help": "The number of layers of the prompt encoder reparameterization"},
metadata={"help": "The number of layers of the prompt encoder"},
)
encoder_dropout: float = field(
default=0.0,
metadata={"help": "The dropout of the prompt encoder reparameterization"},
metadata={"help": "The dropout of the prompt encoder"},
)
@@ -37,6 +49,35 @@ class PromptEncoderConfig(PromptLearningConfig):
class PromptEncoder(torch.nn.Module):
"""
The prompt encoder network that is used to generate the virtual token embeddings for p-tuning.
Args:
config (:class:`PromptEncoderConfig`): The configuration of the prompt encoder.
Example::
>>> from pet import PromptEncoder, PromptEncoderConfig >>> config = PromptEncoderConfig(
pet_type="P_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768,
num_transformer_submodules=1, num_attention_heads=12, num_layers=12,
encoder_reparameterization_type="MLP", encoder_hidden_size=768
)
>>> prompt_encoder = PromptEncoder(config)
Attributes:
embedding (:class:`~torch.nn.Embedding`): The embedding layer of the prompt encoder. mlp_head
(:class:`~torch.nn.Sequential`): The MLP head of the prompt encoder if `inference_mode=False`. lstm_head
(:class:`~torch.nn.LSTM`):
The LSTM head of the prompt encoder if `inference_mode=False` and `encoder_reparameterization_type="LSTM"`.
token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model. input_size (:obj:`int`):
The input size of the prompt encoder. output_size (:obj:`int`): The output size of the prompt encoder.
hidden_size (:obj:`int`): The hidden size of the prompt encoder. total_virtual_tokens (:obj:`int`): The total
number of virtual tokens of the prompt encoder. encoder_type
(:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]):
The encoder type of the prompt encoder.
Input shape: (batch_size, total_virtual_tokens)
Output shape: (batch_size, total_virtual_tokens, token_dim)
"""
def __init__(self, config):
src/pet/tuners/prefix_tuning.py
+28
View File
@@ -8,6 +8,16 @@ from ..utils import PromptLearningConfig
@dataclass
class PrefixTuningConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~pet.PrefixEncoder`.
Args:
encoder_hidden_size (:obj: int): The hidden size of the prompt encoder.
prefix_projection (:obj: bool): Whether to project the prefix embeddings.
postprocess_past_key_value_function (:
obj: Optional[Callable]): The function to postprocess the past key value.
"""
encoder_hidden_size: int = field(
default=None,
metadata={"help": "The hidden size of the encoder"},
@@ -28,6 +38,24 @@ class PrefixEncoder(torch.nn.Module):
r"""
The torch.nn model to encode the prefix
Args:
config (:class:`PrefixTuningConfig`): The configuration of the prefix encoder.
Example::
>>> from pet import PrefixEncoder, PrefixTuningConfig >>> config = PrefixTuningConfig(
pet_type="PREFIX_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768,
num_transformer_submodules=1, num_attention_heads=12, num_layers=12, encoder_hidden_size=768
)
>>> prefix_encoder = PrefixEncoder(config)
Attributes:
embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prefix encoder. trans
(:obj:`torch.nn.Sequential`): The two-layer MLP to transform the prefix embeddings
if :obj:`prefix_projection` is :obj:`True`.
prefix_projection (:obj:`bool`): Whether to project the prefix embeddings.
Input shape: (batch_size, num_virtual_tokens)
Output shape: (batch_size, num_virtual_tokens, 2*layers*hidden)
src/pet/tuners/prompt_tuning.py
+43 -4
View File
@@ -15,6 +15,18 @@ class PromptTuningInit(str, enum.Enum):
@dataclass
class PromptTuningConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~pet.PromptEmbedding`.
Args:
prompt_tuning_init (:
obj:Union[:class:`PromptTuningInit`, :obj:`str`]): The initialization of the prompt embedding.
prompt_tuning_init_text (:obj: Optional[:obj:`str`]): The text to initialize the prompt embedding.
Only used if `prompt_tuning_init` is `TEXT`
tokenizer_name_or_path (:obj: Optional[:obj:`str`]): The name or path of the tokenizer.
Only used if `prompt_tuning_init` is `TEXT`
"""
prompt_tuning_init: Union[PromptTuningInit, str] = field(
default=PromptTuningInit.RANDOM,
metadata={"help": "How to initialize the prompt tuning parameters"},
@@ -34,17 +46,44 @@ class PromptTuningConfig(PromptLearningConfig):
class PromptEmbedding(torch.nn.Module):
"""
The model to encode virtual tokens into prompt embeddings.
Args:
config (:class:`PromptTuningConfig`): The configuration of the prompt embedding.
word_embeddings (:obj:`torch.nn.Module`): The word embeddings of the base transformer model.
Attributes:
embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prompt embedding.
Example::
>>> from pet import PromptEmbedding, PromptTuningConfig >>> config = PromptTuningConfig(
pet_type="PROMPT_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768,
num_transformer_submodules=1, num_attention_heads=12, num_layers=12, prompt_tuning_init="TEXT",
prompt_tuning_init_text="Predict if sentiment of this review is positive, negative or neutral",
tokenizer_name_or_path="t5-base",
)
>>> # t5_model.shared is the word embeddings of the base model >>> prompt_embedding = PromptEmbedding(config,
t5_model.shared)
Input Shape: (batch_size, total_virtual_tokens)
Output Shape: (batch_size, total_virtual_tokens, token_dim)
"""
def __init__(self, config, word_embeddings):
super().__init__()
total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules
self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"])
self.embedding = torch.nn.Embedding(total_virtual_tokens, config.token_dim)
if config.prompt_tuning_init == PromptTuningInit.TEXT:
from transformers import AutoTokenizer
self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
self.init_text = config.prompt_tuning_init_text
init_token_ids = self.tokenizer(self.init_text)["input_ids"]
tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
init_text = config.prompt_tuning_init_text
init_token_ids = tokenizer(init_text)["input_ids"]
# Trim or iterate until num_text_tokens matches total_virtual_tokens
num_text_tokens = len(init_token_ids)
if num_text_tokens > total_virtual_tokens:
src/pet/utils/__init__.py
+1 -1
View File
@@ -3,5 +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 .other import _set_trainable, 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
+21 -2
View File
@@ -19,7 +19,12 @@ class TaskType(str, enum.Enum):
@dataclass
class PETConfig:
"""
This is the configuration class to store the configuration of a :class:`~pet.PETModel`.
This is the base configuration class to store the configuration of a :class:`~pet.PETModel`.
Args:
pet_type (:obj:Union[:class:`~pet.utils.config.PETType`, :obj:`str`]): The type of PET method to use.
task_type (:obj:Union[:class:`~pet.utils.config.TaskType`, :obj:`str`]): The type of task to perform.
inference_mode (:obj:`bool`, defaults to :obj:`False`): Whether to use the PET model in inference mode.
"""
pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"})
@@ -29,8 +34,22 @@ class PETConfig:
@dataclass
class PromptLearningConfig(PETConfig):
"""
This is the base configuration class to store the configuration of a :obj:Union[:class:`~pet.PrefixTuning`,
:class:`~pet.PromptEncoder`, :class:`~pet.PromptTuning`].
Args:
num_virtual_tokens (:obj:`int`): The number of virtual tokens to use.
token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model.
num_transformer_submodules (:obj:`int`): The number of transformer submodules in the base transformer model.
num_attention_heads (:obj:`int`): The number of attention heads in the base transformer model.
num_layers (:obj:`int`): The number of layers in the base transformer model.
"""
num_virtual_tokens: int = field(default=None, metadata={"help": "Number of virtual tokens"})
token_dim: int = field(default=None, metadata={"help": "Dimension of virtual tokens"})
token_dim: int = field(
default=None, metadata={"help": "The hidden embedding dimension of the base transformer model"}
)
num_transformer_submodules: Optional[int] = field(default=1, metadata={"help": "Number of transformer submodules"})
num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"})
num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"})
src/pet/utils/other.py
+49
View File
@@ -19,6 +19,11 @@ def bloom_model_postprocess_past_key_value(past_key_values):
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.
Args:
input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`): input ids
pad_token_id (:obj:`int`): The id of the `padding` token.
decoder_start_token_id (:obj:`int`): The id of the `start` token.
"""
shifted_input_ids = input_ids.new_zeros(input_ids.shape)
shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
@@ -30,3 +35,47 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start
shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
return shifted_input_ids
def _set_trainable(model):
if model.modules_to_save is not None:
for name, param in model.named_parameters():
if any(module_name in name for module_name in model.modules_to_save):
param.requires_grad = True
else:
param.requires_grad = False
def fsdp_auto_wrap_policy(model):
import functools
import os
from accelerate import FullyShardedDataParallelPlugin
from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy, transformer_auto_wrap_policy
from ..tuners import PrefixEncoder, PromptEmbedding, PromptEncoder
def lambda_policy_fn(module):
if (
len(list(module.named_children())) == 0
and getattr(module, "weight", None) is not None
and module.weight.requires_grad
):
return True
return False
lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn)
transformer_wrap_policy = functools.partial(
transformer_auto_wrap_policy,
transformer_layer_cls=(
PrefixEncoder,
PromptEncoder,
PromptEmbedding,
FullyShardedDataParallelPlugin.get_module_class_from_name(
model, os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "")
),
),
)
auto_wrap_policy = functools.partial(_or_policy, policies=[lambda_policy, transformer_wrap_policy])
return auto_wrap_policy
src/pet/utils/save_and_load.py
+15
View File
@@ -4,6 +4,13 @@ from .config import PETType
def get_pet_model_state_dict(model):
"""
Get the state dict of the PET model.
Args:
model (:obj:`PETModel`): The PET model.
"""
if model.pet_config.pet_type == PETType.LORA:
return lora_state_dict(model)
else:
@@ -19,6 +26,14 @@ def get_pet_model_state_dict(model):
def set_pet_model_state_dict(model, pet_model_state_dict):
"""
Set the state dict of the PET model.
Args:
model (:obj:`PETModel`): The PET model.
pet_model_state_dict (:obj:`dict`): The state dict of the PET model.
"""
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(