From 751baf8aa7ef39ea4e78d79a3b275783e6d7eccc Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 18:22:12 +0530 Subject: [PATCH] refactor, lora support and utils for saving and loading --- README.md | 50 +++++----- src/pet/__init__.py | 11 ++- src/pet/mapping.py | 45 ++++++++- src/pet/pet_model.py | 169 ++++++++++++++++++++------------- src/pet/tuners/lora.py | 45 +++++---- src/pet/utils/__init__.py | 2 + src/pet/utils/config.py | 2 +- src/pet/utils/other.py | 18 ++++ src/pet/utils/save_and_load.py | 28 ++++++ 9 files changed, 259 insertions(+), 111 deletions(-) create mode 100644 src/pet/utils/save_and_load.py diff --git a/README.md b/README.md index 9d6c7e7..0e1c270 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/src/pet/__init__.py b/src/pet/__init__.py index cce106b..734d7c2 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -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, +) diff --git a/src/pet/mapping.py b/src/pet/mapping.py index 4fdde10..f2a3969 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -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) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 45a6264..3f0f078 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -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) diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index 5fd7887..0036638 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -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) diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py index 2a32be1..aa359db 100644 --- a/src/pet/utils/__init__.py +++ b/src/pet/utils/__init__.py @@ -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 diff --git a/src/pet/utils/config.py b/src/pet/utils/config.py index 6f66966..d989601 100644 --- a/src/pet/utils/config.py +++ b/src/pet/utils/config.py @@ -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"}) diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py index 45d7da1..ac8b329 100644 --- a/src/pet/utils/other.py +++ b/src/pet/utils/other.py @@ -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 diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py new file mode 100644 index 0000000..39fa579 --- /dev/null +++ b/src/pet/utils/save_and_load.py @@ -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