From 4eaf613b6aac804982d38bd1bdf93c191fed5abf Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 25 Nov 2022 12:07:08 +0530 Subject: [PATCH 01/18] add code --- README.md | 26 +- requirements.txt | 4 + src/pet.py | 202 +++++++++++ src/prompt_learning_legacy.py | 639 ++++++++++++++++++++++++++++++++++ src/tuners/lora.py | 1 + src/tuners/p_tuning.py | 74 ++++ src/tuners/prefix_tuning.py | 36 ++ src/tuners/prompt_tuning.py | 38 ++ src/utils/constants.py | 1 + 9 files changed, 1019 insertions(+), 2 deletions(-) create mode 100644 requirements.txt create mode 100644 src/pet.py create mode 100644 src/prompt_learning_legacy.py create mode 100644 src/tuners/lora.py create mode 100644 src/tuners/p_tuning.py create mode 100644 src/tuners/prefix_tuning.py create mode 100644 src/tuners/prompt_tuning.py create mode 100644 src/utils/constants.py diff --git a/README.md b/README.md index 9d4e877..d01cdf3 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,24 @@ -# pets -Parameter-Efficient Tuning at Scale +# 🤗 pets +Parameter-Efficient Tuning at Scale with 🤗 Accelerate + +Supported moethods: +1. Prefix Tuning +2. P-Tuning +3. Prompt Tuning +4. LoRA [in progress] + +## Models support matrix + +### Sequence Classification +| | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA | +| --------- | ---- | ---- | ---- | ---- | +| RoBERTa | ✅ | ✅ | ✅ | | +| BERT | ✅ | ✅ | ✅ | | +| Deberta-v2 | | | | | +| BloomX | | | | | +| Bloom | | | | | +| mT-0 | | | | | +| T-0 | | | | | +| T5 | | | | | +| GPT-2 | | | | | +| BART | | | | | diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..10f15c0 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +transformers +accelerate +loralib +evaluate diff --git a/src/pet.py b/src/pet.py new file mode 100644 index 0000000..7ac208b --- /dev/null +++ b/src/pet.py @@ -0,0 +1,202 @@ +from collections import OrderedDict +import enum +import warnings +import torch +from transformers import PreTrainedModel +from tuners.p_tuning import PromptEncoder +from tuners.prefix_tuning import PrefixEncoder +from tuners.prompt_tuning import PromptEmbedding +from accelerate.state import AcceleratorState + + +class PromptEncoderType(str, enum.Enum): + PROMPT_TUNING = "PROMPT_TUNING" + P_TUNING = "P_TUNING" + PREFIX_TUNING = "PREFIX_TUNING" + LORA = "LORA" + + +class ParameterEfficientTuningModel(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + self.prompt_learning_config = model.config.prompt_learning_config + + modules = list(self.model._modules) + + for module in modules: + if isinstance(self.model.get_submodule(module), PreTrainedModel): + transformer_backbone = self.model.get_submodule(module) + break + + for named_param, value in list(transformer_backbone.named_parameters()): + if value.shape[0] == model.config.vocab_size: + self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", "")) + break + + # Make sure to freeze Tranformers model + for param in transformer_backbone.parameters(): + param.requires_grad = False + + if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PROMPT_TUNING: + prompt_encoder = PromptEmbedding(self.prompt_learning_config, self.word_embeddings) + elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING: + prompt_encoder = PromptEncoder(self.prompt_learning_config) + elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + prompt_encoder = PrefixEncoder(self.prompt_learning_config) + else: + raise ValueError("Not supported") + self.prompt_encoder = prompt_encoder + self.prompt_tokens = torch.arange(self.prompt_learning_config["num_virtual_tokens"]).long() + + def get_prompt(self, batch_size): + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.transformer_backbone.device) + if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + if self.prompt_learning_config.get("inference_mode", False): + past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + else: + past_key_values = self.prompt_encoder(prompt_tokens) + past_key_values = past_key_values.view( + batch_size, + self.prompt_learning_config["num_virtual_tokens"], + self.prompt_learning_config["num_layers"] * 2, + self.prompt_learning_config["num_attention_heads"], + self.prompt_learning_config["token_dim"] // self.prompt_learning_config["num_attention_heads"], + ) + past_key_values = self.dropout(past_key_values) + past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2) + return past_key_values + else: + if self.prompt_learning_config.get("inference_mode", False): + prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + else: + prompts = self.prompt_encoder(prompt_tokens) + return prompts + + def state_dict(self, destination=None, prefix=None, keep_vars=False): + """ + No frozen model parameters are stored in the state dict. + """ + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.model.device) + prompt_embeddings = self.prompt_encoder(prompt_tokens).detach().cpu() + if destination is None: + state_dict_ = OrderedDict() + else: + state_dict_ = destination + state_dict_["prompt_embeddings"] = prompt_embeddings[0] + return state_dict_ + + def load_state_dict(self, state_dict, strict: bool = True): + """ + Custom load state dict method that only loads prompt table and prompt encoder + parameters. Matching load method for this class' custom state dict method. + """ + self.prompt_encoder.embedding.load_state_dict({"weight": state_dict["prompt_embeddings"]}, strict) + + +class ParameterEfficientTuningModelForSequenceClassification(ParameterEfficientTuningModel): + def __init__(self, model): + super().__init__(model) + self.config = self.model.config + self.modules_to_save = ("prompt_encoder", "classifier") + + trainable_params = 0 + all_param = 0 + for _, param in self.named_parameters(): + all_param += param.numel() + if param.requires_grad: + trainable_params += param.numel() + print( + f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" + ) + + def forward( + self, + input_ids=None, + attention_mask=None, + inputs_embeds=None, + labels=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + batch_size = input_ids.shape[0] + # concat prompt attention mask + prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( + self.model.device + ) + attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) + if kwargs["token_type_ids"] is not None: + kwargs["token_type_ids"] = torch.cat( + ( + torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to(self.model.device), + kwargs["token_type_ids"], + ), + dim=1, + ) + + if kwargs["position_ids"] is not None: + warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") + kwargs["position_ids"] = None + + if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + past_key_values = self.get_prompt(batch_size=batch_size) + + return self.model( + input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + past_key_values=past_key_values, + **kwargs, + ) + else: + raw_embedding = self.word_embeddings(input_ids) + prompts = self.get_prompt(batch_size=batch_size) + inputs_embeds = torch.cat((prompts, raw_embedding), dim=1) + + return self.model( + # input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + # past_key_values=past_key_values, + ) + + def state_dict(self, destination=None, prefix=None, keep_vars=False): + """ + No frozen model parameters are stored in the state dict. + """ + if destination is None: + state_dict_ = OrderedDict() + else: + state_dict_ = destination + state_dict_["prompt_encoder"] = super().state_dict() + state_dict_["classifier"] = self.model.classifier.state_dict() + if AcceleratorState().fsdp_plugin is not None: + state_dict_["_flat_param"] = None + return state_dict_ + + def load_state_dict(self, state_dict, strict: bool = True): + """ + Custom load state dict method that only loads prompt table and prompt encoder + parameters. Matching load method for this class' custom state dict method. + """ + super().load_state_dict(state_dict["prompt_encoder"], strict) + self.model.classifier.load_state_dict(state_dict["classifier"], strict) + + def clean_state_dict(self, state_dict): + if AcceleratorState().fsdp_plugin is not None: + new_state_dict = OrderedDict() + for key in self.modules_to_save: + new_state_dict[key] = state_dict[key].copy() + state_dict = new_state_dict + return state_dict diff --git a/src/prompt_learning_legacy.py b/src/prompt_learning_legacy.py new file mode 100644 index 0000000..ffb6934 --- /dev/null +++ b/src/prompt_learning_legacy.py @@ -0,0 +1,639 @@ +import enum +import torch +import math +import os + +from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss +from transformers import PreTrainedModel +from transformers.modeling_outputs import SequenceClassifierOutput +from transformers import AutoModelForSequenceClassification +from datasets import load_dataset +import evaluate +import torch +from transformers import AutoTokenizer, get_linear_schedule_with_warmup, set_seed +from torch.utils.data import DataLoader +from accelerate import Accelerator +from accelerate.state import AcceleratorState +from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin +import functools +from torch.distributed.fsdp import ( + FullyShardedDataParallel, + CPUOffload, +) +from torch.distributed.fsdp.wrap import ( + enable_wrap, + wrap, + ModuleWrapPolicy, + transformer_auto_wrap_policy, + lambda_auto_wrap_policy, + _or_policy, +) +from collections import OrderedDict + + +class PromptEncoderReparameterizationType(str, enum.Enum): + MLP = "MLP" + LSTM = "LSTM" + + +class PromptEncoderType(str, enum.Enum): + PROMPT_TUNING = "PROMPT_TUNING" + P_TUNING_V1 = "P_TUNING_V1" + P_TUNING_V2 = "P_TUNING_V2" + + +class PromptTuningInit(str, enum.Enum): + TEXT = "TEXT" + RANDOM = "RANDOM" + + +class PromptEncoder(torch.nn.Module): + """ + The prompt encoder network that is used to generate the virtual + token embeddings for p-tuning. + """ + + def __init__(self, config): + super().__init__() + self.token_dim = config["token_dim"] + self.input_size = config["token_dim"] + self.output_size = config["token_dim"] + self.hidden_size = config["prompt_hidden_size"] + self.total_virtual_tokens = config["num_virtual_tokens"] + self.encoder_type = config["prompt_encoder_config"]["prompt_reparam_type"] + + # embedding + self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim) + if not config.get("inference_mode", False): + if self.encoder_type == PromptEncoderReparameterizationType.LSTM: + if "dropout" not in config["prompt_encoder_config"]: + lstm_dropout = 0.0 + else: + lstm_dropout = config["prompt_encoder_config"]["dropout"] + + if "num_layers" not in config["prompt_encoder_config"]: + num_layers = 2 + else: + num_layers = config["prompt_encoder_config"]["num_layers"] + # LSTM + self.lstm_head = torch.nn.LSTM( + input_size=self.input_size, + hidden_size=self.hidden_size, + num_layers=num_layers, + dropout=lstm_dropout, + bidirectional=True, + batch_first=True, + ) + + self.mlp_head = torch.nn.Sequential( + torch.nn.Linear(self.hidden_size * 2, self.hidden_size * 2), + torch.nn.ReLU(), + torch.nn.Linear(self.hidden_size * 2, self.output_size), + ) + + elif self.encoder_type == PromptEncoderReparameterizationType.MLP: + layers = [torch.nn.Linear(self.input_size, self.hidden_size), torch.nn.ReLU()] + layers.extend([torch.nn.Linear(self.hidden_size, self.hidden_size), torch.nn.ReLU()]) + layers.append(torch.nn.Linear(self.hidden_size, self.output_size)) + self.mlp_head = torch.nn.Sequential(*layers) + + else: + raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") + + def forward(self, indices): + input_embeds = self.embedding(indices) + if self.encoder_type == PromptEncoderReparameterizationType.LSTM: + output_embeds = self.mlp_head(self.lstm_head(input_embeds)[0]) + elif self.encoder_type == PromptEncoderReparameterizationType.MLP: + output_embeds = self.mlp_head(input_embeds) + else: + raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") + + return output_embeds + + +class PrefixEncoder(torch.nn.Module): + r""" + The torch.nn model to encode the prefix + + Input shape: (batch-size, prefix-length) + + Output shape: (batch-size, prefix-length, 2*layers*hidden) + """ + + def __init__(self, config): + super().__init__() + self.prefix_projection = config["prompt_encoder_config"]["prefix_projection"] + if self.prefix_projection and not config.get("inference_mode", False): + # Use a two-layer MLP to encode the prefix + self.embedding = torch.nn.Embedding(config["num_virtual_tokens"], config["token_dim"]) + self.trans = torch.nn.Sequential( + torch.nn.Linear(config["token_dim"], config["prompt_hidden_size"]), + torch.nn.Tanh(), + torch.nn.Linear(config["prompt_hidden_size"], config["num_layers"] * 2 * config["token_dim"]), + ) + else: + self.embedding = torch.nn.Embedding( + config["num_virtual_tokens"], config["num_layers"] * 2 * config["token_dim"] + ) + + def forward(self, prefix: torch.Tensor): + if self.prefix_projection: + prefix_tokens = self.embedding(prefix) + past_key_values = self.trans(prefix_tokens) + else: + past_key_values = self.embedding(prefix) + return past_key_values + + +class PromptEmbedding(torch.nn.Module): + def __init__(self, config, word_embeddings): + super().__init__() + + total_virtual_tokens = config["num_virtual_tokens"] + self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"]) + if config["prompt_encoder_config"]["prompt_tuning_init"] == PromptTuningInit.TEXT: + from transformers import AutoTokenizer + + self.tokenizer = AutoTokenizer.from_pretrained(config["prompt_encoder_config"]["tokenizer_name_or_path"]) + self.init_text = config["prompt_encoder_config"]["prompt_tuning_text"] + init_token_ids = self.tokenizer(self.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: + init_token_ids = init_token_ids[:total_virtual_tokens] + elif num_text_tokens < total_virtual_tokens: + num_reps = math.ceil(total_virtual_tokens / num_text_tokens) + init_token_ids = init_token_ids * num_reps + init_token_ids = init_token_ids[:total_virtual_tokens] + + word_embedding_weights = word_embeddings(torch.LongTensor(init_token_ids)).detach().clone() + self.embedding.weight = torch.nn.Parameter(word_embedding_weights) + + def forward(self, indices): + # Just get embeddings and dropout + prompt_embeddings = self.embedding(indices) + return prompt_embeddings + + +class PromptModel(torch.nn.Module): + def __init__(self, model): + super().__init__() + self.prompt_learning_config = model.config.prompt_learning_config + + modules = list(model._modules) + + for module in modules: + if isinstance(model.get_submodule(module), PreTrainedModel): + self.transformer_backbone = model.get_submodule(module) + break + + for named_param, value in list(self.transformer_backbone.named_parameters()): + if value.shape[0] == model.config.vocab_size: + self.word_embeddings = self.transformer_backbone.get_submodule(named_param.replace(".weight", "")) + break + + # Make sure to freeze Tranformers model + for param in self.transformer_backbone.parameters(): + param.requires_grad = False + + if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PROMPT_TUNING: + prompt_encoder = PromptEmbedding(self.prompt_learning_config, self.word_embeddings) + elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V1: + prompt_encoder = PromptEncoder(self.prompt_learning_config) + elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: + prompt_encoder = PrefixEncoder(self.prompt_learning_config) + else: + raise ValueError("Not supported") + self.prompt_encoder = prompt_encoder + self.prompt_tokens = torch.arange(self.prompt_learning_config["num_virtual_tokens"]).long() + + def get_prompt(self, batch_size): + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.transformer_backbone.device) + if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: + if self.prompt_learning_config.get("inference_mode", False): + past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + else: + past_key_values = self.prompt_encoder(prompt_tokens) + past_key_values = past_key_values.view( + batch_size, + self.prompt_learning_config["num_virtual_tokens"], + self.prompt_learning_config["num_layers"] * 2, + self.prompt_learning_config["num_attention_heads"], + self.prompt_learning_config["token_dim"] // self.prompt_learning_config["num_attention_heads"], + ) + past_key_values = self.dropout(past_key_values) + past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2) + return past_key_values + else: + if self.prompt_learning_config.get("inference_mode", False): + prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) + else: + prompts = self.prompt_encoder(prompt_tokens) + return prompts + + def state_dict(self, destination=None, prefix=None, keep_vars=False): + """ + No frozen model parameters are stored in the state dict. + """ + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.transformer_backbone.device) + prompt_embeddings = self.prompt_encoder(prompt_tokens).detach().cpu() + if destination is None: + state_dict_ = OrderedDict() + else: + state_dict_ = destination + state_dict_["prompt_embeddings"] = prompt_embeddings[0] + return state_dict_ + + def load_state_dict(self, state_dict, strict: bool = True): + """ + Custom load state dict method that only loads prompt table and prompt encoder + parameters. Matching load method for this class' custom state dict method. + """ + self.prompt_encoder.embedding.load_state_dict({"weight": state_dict["prompt_embeddings"]}, strict) + + +class PromptModelForSequenceClassification(PromptModel): + def __init__(self, model): + super().__init__(model) + if "dropout" in [name for name, _ in model.named_children()]: + self.dropout = model.dropout + else: + self.dropout = torch.nn.Dropout(model.config.hidden_dropout_prob) + self.classifier = model.classifier + self.num_labels = model.num_labels + self.config = model.config + self.modules_to_save = ("prompt_encoder", "classifier") + + trainable_params = 0 + all_param = 0 + for _, param in self.named_parameters(): + all_param += param.numel() + if param.requires_grad: + trainable_params += param.numel() + print( + f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" + ) + + def forward( + self, + input_ids=None, + attention_mask=None, + inputs_embeds=None, + labels=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + batch_size = input_ids.shape[0] + # concat prompt attention mask + prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( + self.transformer_backbone.device + ) + attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) + + if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: + past_key_values = self.get_prompt(batch_size=batch_size) + + outputs = self.transformer_backbone( + input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + past_key_values=past_key_values, + **kwargs, + ) + + pooled_output = outputs[1] if len(outputs) > 1 else outputs[0] + else: + raw_embedding = self.word_embeddings(input_ids) + prompts = self.get_prompt(batch_size=batch_size) + inputs_embeds = torch.cat((prompts, raw_embedding), dim=1) + + outputs = self.transformer_backbone( + # input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + # **kwargs, + # past_key_values=past_key_values, + ) + + sequence_output = outputs[0] + sequence_output = sequence_output[:, self.prompt_learning_config["num_virtual_tokens"] :, :].contiguous() + pooled_output = sequence_output[:, 0] + + if ( + "pooler" in [name for name, _ in self.transformer_backbone.named_children()] + and self.transformer_backbone.pooler is not None + ): + pooled_output = self.transformer_backbone.pooler.dense(pooled_output) + pooled_output = self.transformer_backbone.pooler.activation(pooled_output) + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.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.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.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def state_dict(self, destination=None, prefix=None, keep_vars=False): + """ + No frozen model parameters are stored in the state dict. + """ + if destination is None: + state_dict_ = OrderedDict() + else: + state_dict_ = destination + state_dict_["prompt_encoder"] = super().state_dict() + state_dict_["classifier"] = self.classifier.state_dict() + if AcceleratorState().fsdp_plugin is not None: + state_dict_["_flat_param"] = None + return state_dict_ + + def load_state_dict(self, state_dict, strict: bool = True): + """ + Custom load state dict method that only loads prompt table and prompt encoder + parameters. Matching load method for this class' custom state dict method. + """ + super().load_state_dict(state_dict["prompt_encoder"], strict) + self.classifier.load_state_dict(state_dict["classifier"], strict) + + def clean_state_dict(self, state_dict): + if AcceleratorState().fsdp_plugin is not None: + new_state_dict = OrderedDict() + for key in self.modules_to_save: + new_state_dict[key] = state_dict[key].copy() + state_dict = new_state_dict + return state_dict + + +model_type_to_prompt_model_mapping = {"SequenceClassification": PromptModelForSequenceClassification} +num_virtual_tokens = 30 +model_name_or_path = "roberta-large" +tokenizer_name_or_path = "roberta-large" + +prompt_tuning_config = { + "num_virtual_tokens": num_virtual_tokens, + "prompt_encoder_type": "PROMPT_TUNING", + "prompt_encoder_config": { + "prompt_tuning_init": "TEXT", + "tokenizer_name_or_path": tokenizer_name_or_path, + "prompt_tuning_text": "Output is true or false. Task requires to recognize" + " whether the meaning of one text is entailed (can be inferred) from the other text.", + }, +} + + +p_tuning_v1_mlp_config = { + "num_virtual_tokens": num_virtual_tokens, + "prompt_encoder_type": "P_TUNING_V1", + "prompt_encoder_config": {"prompt_reparam_type": "MLP"}, +} + +p_tuning_v1_lstm_config = { + "num_virtual_tokens": num_virtual_tokens, + "prompt_encoder_type": "P_TUNING_V1", + "prompt_encoder_config": {"prompt_reparam_type": "LSTM"}, +} + +p_tuning_v2_no_proj_config = { + "num_virtual_tokens": num_virtual_tokens, + "prompt_encoder_type": "P_TUNING_V2", + "prompt_encoder_config": {"prefix_projection": False}, +} + +p_tuning_v2_proj_config = { + "num_virtual_tokens": num_virtual_tokens, + "prompt_encoder_type": "P_TUNING_V2", + "prompt_encoder_config": {"prefix_projection": True}, +} + + +def prepare_prompt_model(model, prompt_learning_config): + config = model.config.to_dict() + if "num_layers" not in prompt_learning_config: + if "num_hidden_layers" in config: + num_layers = config["num_hidden_layers"] + elif "num_layers" in config: + num_layers = config["num_layers"] + else: + raise ValueError("Please specify `num_layers` in `prompt_learning_config`") + prompt_learning_config["num_layers"] = num_layers + + if "token_dim" not in prompt_learning_config: + if "hidden_size" in config: + token_dim = config["hidden_size"] + elif "n_embd" in config: + token_dim = config["n_embd"] + elif "d_model" in config: + token_dim = config["d_model"] + else: + raise ValueError("Please specify `token_dim` in `prompt_learning_config`") + prompt_learning_config["token_dim"] = token_dim + + if "num_attention_heads" not in prompt_learning_config: + if "num_attention_heads" in config: + num_attention_heads = config["num_attention_heads"] + elif "n_head" in config: + num_attention_heads = config["n_head"] + elif "num_heads" in config: + num_attention_heads = config["num_heads"] + else: + raise ValueError("Please specify `num_attention_heads` in `prompt_learning_config`") + prompt_learning_config["num_attention_heads"] = num_attention_heads + + if "prompt_hidden_size" not in prompt_learning_config: + prompt_learning_config["prompt_hidden_size"] = token_dim + + model.config.prompt_learning_config = prompt_learning_config + model_type = model.__class__.__name__.split("For") + if len(model_type) < 2: + raise ValueError("Model Type not supported") + model_cls = model_type_to_prompt_model_mapping[model_type[1]] + prompt_model = model_cls(model) + return prompt_model + + +def fsdp_auto_wrap_policy(model): + def wrap_layers_with_required_grads(module): + if ( + len(list(module.children())) == 0 + and len(list(module.named_parameters())) > 0 + and module.weight.requires_grad + ): + return True + return False + + transformer_cls_to_wrap = { + PrefixEncoder, + PromptEmbedding, + PromptEncoder, + PromptModel, + FullyShardedDataParallelPlugin.get_module_class_from_name( + model, os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "") + ), + } + policy_1 = functools.partial( + transformer_auto_wrap_policy, + transformer_layer_cls=transformer_cls_to_wrap, + ) + policy_2 = functools.partial( + lambda_auto_wrap_policy, + lambda_fn=wrap_layers_with_required_grads, + ) + auto_wrap_policy = functools.partial(_or_policy, policies=[policy_1, policy_2]) + return auto_wrap_policy + + +def main(): + accelerator = Accelerator() + task = "rte" + batch_size = 16 + lr = 5e-3 + num_epochs = 100 + device = "cuda" + seed = 11 + set_seed(seed) + + model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path) + model = prepare_prompt_model( + model, p_tuning_v2_no_proj_config + ) # p_tuning_v2_proj_config)#p_tuning_v2_no_proj_config) + # model = model.to("cuda") + + tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) + datasets = load_dataset("glue", task) + metric = evaluate.load("glue", task) + + def tokenize_function(examples): + # max_length=None => use the model max length (it's actually the default) + outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + return outputs + + # Apply the method we just defined to all the examples in all the splits of the dataset + # starting with the main process first: + tokenized_datasets = datasets.map( + tokenize_function, + batched=True, + remove_columns=["idx", "sentence1", "sentence2"], + ) + + # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the + # transformers library + tokenized_datasets = tokenized_datasets.rename_column("label", "labels") + + def collate_fn(examples): + return tokenizer.pad(examples, padding="longest", return_tensors="pt") + + # Instantiate dataloaders. + train_dataloader = DataLoader( + tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size + ) + eval_dataloader = DataLoader( + tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=batch_size + ) + + # Instantiate optimizer + optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr) + + # Instantiate scheduler + lr_scheduler = get_linear_schedule_with_warmup( + optimizer=optimizer, + num_warmup_steps=0, + num_training_steps=(len(train_dataloader) * num_epochs), + ) + + 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(train_dataloader): + # batch.to(device) + outputs = model(**batch) + loss = outputs.loss + total_loss += loss.detach().float() + loss.backward() + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + model.eval() + for step, batch in enumerate(eval_dataloader): + # batch.to(device) + with torch.no_grad(): + outputs = model(**batch) + predictions = outputs.logits.argmax(dim=-1) + predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"])) + metric.add_batch( + predictions=predictions, + references=references, + ) + + eval_metric = metric.compute() + accelerator.print(f"epoch {epoch}:", eval_metric) + accelerator.print(f"epoch {epoch} train loss:", total_loss / len(train_dataloader)) + + from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP + from torch.distributed.fsdp.fully_sharded_data_parallel import ( + BackwardPrefetch, + CPUOffload, + FullStateDictConfig, + ShardingStrategy, + StateDictType, + ) + + FSDP.set_state_dict_type( + model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + ) + state_dict = model.state_dict() + state_dict = model.clean_state_dict(state_dict) + accelerator.print(state_dict) + + torch.save(state_dict, "p_tuning_v2.pt") + + +if __name__ == "__main__": + main() diff --git a/src/tuners/lora.py b/src/tuners/lora.py new file mode 100644 index 0000000..044a482 --- /dev/null +++ b/src/tuners/lora.py @@ -0,0 +1 @@ +# todo diff --git a/src/tuners/p_tuning.py b/src/tuners/p_tuning.py new file mode 100644 index 0000000..2c103c3 --- /dev/null +++ b/src/tuners/p_tuning.py @@ -0,0 +1,74 @@ +import torch +import enum + + +class PromptEncoderReparameterizationType(str, enum.Enum): + MLP = "MLP" + LSTM = "LSTM" + + +# Based on https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/prompt_encoder.py +# with some refactor +class PromptEncoder(torch.nn.Module): + """ + The prompt encoder network that is used to generate the virtual + token embeddings for p-tuning. + """ + + def __init__(self, config): + super().__init__() + self.token_dim = config["token_dim"] + self.input_size = config["token_dim"] + self.output_size = config["token_dim"] + self.hidden_size = config["prompt_hidden_size"] + self.total_virtual_tokens = config["num_virtual_tokens"] + self.encoder_type = config["prompt_encoder_config"]["prompt_reparam_type"] + + # embedding + self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim) + if not config.get("inference_mode", False): + if self.encoder_type == PromptEncoderReparameterizationType.LSTM: + if "dropout" not in config["prompt_encoder_config"]: + lstm_dropout = 0.0 + else: + lstm_dropout = config["prompt_encoder_config"]["dropout"] + + if "num_layers" not in config["prompt_encoder_config"]: + num_layers = 2 + else: + num_layers = config["prompt_encoder_config"]["num_layers"] + # LSTM + self.lstm_head = torch.nn.LSTM( + input_size=self.input_size, + hidden_size=self.hidden_size, + num_layers=num_layers, + dropout=lstm_dropout, + bidirectional=True, + batch_first=True, + ) + + self.mlp_head = torch.nn.Sequential( + torch.nn.Linear(self.hidden_size * 2, self.hidden_size * 2), + torch.nn.ReLU(), + torch.nn.Linear(self.hidden_size * 2, self.output_size), + ) + + elif self.encoder_type == PromptEncoderReparameterizationType.MLP: + layers = [torch.nn.Linear(self.input_size, self.hidden_size), torch.nn.ReLU()] + layers.extend([torch.nn.Linear(self.hidden_size, self.hidden_size), torch.nn.ReLU()]) + layers.append(torch.nn.Linear(self.hidden_size, self.output_size)) + self.mlp_head = torch.nn.Sequential(*layers) + + else: + raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") + + def forward(self, indices): + input_embeds = self.embedding(indices) + if self.encoder_type == PromptEncoderReparameterizationType.LSTM: + output_embeds = self.mlp_head(self.lstm_head(input_embeds)[0]) + elif self.encoder_type == PromptEncoderReparameterizationType.MLP: + output_embeds = self.mlp_head(input_embeds) + else: + raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") + + return output_embeds diff --git a/src/tuners/prefix_tuning.py b/src/tuners/prefix_tuning.py new file mode 100644 index 0000000..c593b1e --- /dev/null +++ b/src/tuners/prefix_tuning.py @@ -0,0 +1,36 @@ +import torch + +# Based on https://github.com/THUDM/P-tuning-v2/blob/main/model/prefix_encoder.py +# with some refactor +class PrefixEncoder(torch.nn.Module): + r""" + The torch.nn model to encode the prefix + + Input shape: (batch-size, prefix-length) + + Output shape: (batch-size, prefix-length, 2*layers*hidden) + """ + + def __init__(self, config): + super().__init__() + self.prefix_projection = config["prompt_encoder_config"]["prefix_projection"] + if self.prefix_projection and not config.get("inference_mode", False): + # Use a two-layer MLP to encode the prefix + self.embedding = torch.nn.Embedding(config["num_virtual_tokens"], config["token_dim"]) + self.trans = torch.nn.Sequential( + torch.nn.Linear(config["token_dim"], config["prompt_hidden_size"]), + torch.nn.Tanh(), + torch.nn.Linear(config["prompt_hidden_size"], config["num_layers"] * 2 * config["token_dim"]), + ) + else: + self.embedding = torch.nn.Embedding( + config["num_virtual_tokens"], config["num_layers"] * 2 * config["token_dim"] + ) + + def forward(self, prefix: torch.Tensor): + if self.prefix_projection: + prefix_tokens = self.embedding(prefix) + past_key_values = self.trans(prefix_tokens) + else: + past_key_values = self.embedding(prefix) + return past_key_values diff --git a/src/tuners/prompt_tuning.py b/src/tuners/prompt_tuning.py new file mode 100644 index 0000000..cb4fc8d --- /dev/null +++ b/src/tuners/prompt_tuning.py @@ -0,0 +1,38 @@ +import torch +import enum +import math + + +class PromptTuningInit(str, enum.Enum): + TEXT = "TEXT" + RANDOM = "RANDOM" + + +class PromptEmbedding(torch.nn.Module): + def __init__(self, config, word_embeddings): + super().__init__() + + total_virtual_tokens = config["num_virtual_tokens"] + self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"]) + if config["prompt_encoder_config"]["prompt_tuning_init"] == PromptTuningInit.TEXT: + from transformers import AutoTokenizer + + self.tokenizer = AutoTokenizer.from_pretrained(config["prompt_encoder_config"]["tokenizer_name_or_path"]) + self.init_text = config["prompt_encoder_config"]["prompt_tuning_text"] + init_token_ids = self.tokenizer(self.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: + init_token_ids = init_token_ids[:total_virtual_tokens] + elif num_text_tokens < total_virtual_tokens: + num_reps = math.ceil(total_virtual_tokens / num_text_tokens) + init_token_ids = init_token_ids * num_reps + init_token_ids = init_token_ids[:total_virtual_tokens] + + word_embedding_weights = word_embeddings(torch.LongTensor(init_token_ids)).detach().clone() + self.embedding.weight = torch.nn.Parameter(word_embedding_weights) + + def forward(self, indices): + # Just get embeddings + prompt_embeddings = self.embedding(indices) + return prompt_embeddings diff --git a/src/utils/constants.py b/src/utils/constants.py new file mode 100644 index 0000000..044a482 --- /dev/null +++ b/src/utils/constants.py @@ -0,0 +1 @@ +# todo From 61157eaea1db20bce0e599dce108bd8426acc1c6 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 25 Nov 2022 12:55:35 +0530 Subject: [PATCH 02/18] refactoring --- .gitignore | 141 ++++++ MANIFEST.in | 1 + Makefile | 19 + pyproject.toml | 3 + setup.cfg | 23 + setup.py | 78 ++++ src/pet/__init__.py | 18 + src/{pet.py => pet/pet_model.py} | 14 +- src/{ => pet}/prompt_learning_legacy.py | 114 +++-- src/pet/tuners/__init__.py | 7 + src/{ => pet}/tuners/lora.py | 0 src/{ => pet}/tuners/p_tuning.py | 18 +- src/{ => pet}/tuners/prefix_tuning.py | 9 +- src/{ => pet}/tuners/prompt_tuning.py | 3 +- src/{ => pet}/utils/constants.py | 0 utils/style_doc.py | 556 ++++++++++++++++++++++++ 16 files changed, 942 insertions(+), 62 deletions(-) create mode 100644 .gitignore create mode 100644 MANIFEST.in create mode 100644 Makefile create mode 100644 pyproject.toml create mode 100644 setup.cfg create mode 100644 setup.py create mode 100644 src/pet/__init__.py rename src/{pet.py => pet/pet_model.py} (97%) rename src/{ => pet}/prompt_learning_legacy.py (91%) create mode 100644 src/pet/tuners/__init__.py rename src/{ => pet}/tuners/lora.py (100%) rename src/{ => pet}/tuners/p_tuning.py (87%) rename src/{ => pet}/tuners/prefix_tuning.py (81%) rename src/{ => pet}/tuners/prompt_tuning.py (99%) rename src/{ => pet}/utils/constants.py (100%) create mode 100644 utils/style_doc.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..da99824 --- /dev/null +++ b/.gitignore @@ -0,0 +1,141 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# VSCode +.vscode + +# IntelliJ +.idea + +# Mac .DS_Store +.DS_Store + +# More test things +wandb \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..1aba38f --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +include LICENSE diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..e1c15c5 --- /dev/null +++ b/Makefile @@ -0,0 +1,19 @@ +.PHONY: quality style test docs + +check_dirs := src + +# Check that source code meets quality standards + +# this target runs checks on all files +quality: + black --check $(check_dirs) + isort --check-only $(check_dirs) + flake8 $(check_dirs) + python utils/style_doc.py src --max_len 119 --check_only + +# Format source code automatically and check is there are any problems left that need manual fixing +style: + black $(check_dirs) + isort $(check_dirs) + python utils/style_doc.py src --max_len 119 + \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..b7465bb --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,3 @@ +[tool.black] +line-length = 119 +target-version = ['py36'] diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..6b26312 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,23 @@ +[isort] +default_section = FIRSTPARTY +ensure_newline_before_comments = True +force_grid_wrap = 0 +include_trailing_comma = True +known_first_party = pet +known_third_party = + numpy + torch + accelerate + transformers + +line_length = 119 +lines_after_imports = 2 +multi_line_output = 3 +use_parentheses = True + +[flake8] +ignore = E203, E722, E501, E741, W503, W605 +max-line-length = 119 + +[tool:pytest] +doctest_optionflags=NUMBER NORMALIZE_WHITESPACE ELLIPSIS \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..9c8edd9 --- /dev/null +++ b/setup.py @@ -0,0 +1,78 @@ +# Copyright 2021 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from setuptools import setup +from setuptools import find_packages + +extras = {} +extras["quality"] = ["black ~= 22.0", "isort >= 5.5.4", "flake8 >= 3.8.3"] +extras["dev"] = extras["quality"] + +setup( + name="pets", + version="0.1.0.dev0", + description="Parameter-Efficient Tuning at Scale (PETS)", + long_description=open("README.md", "r", encoding="utf-8").read(), + long_description_content_type="text/markdown", + keywords="deep learning", + license="Apache", + author="The HuggingFace team", + author_email="sourab@huggingface.co", + url="https://github.com/huggingface/pets", + package_dir={"": "src"}, + packages=find_packages("src"), + entry_points={}, + python_requires=">=3.7.0", + install_requires=[ + "numpy>=1.17", + "packaging>=20.0", + "psutil", + "pyyaml", + "torch>=1.4.0", + "transformers", + "accelerate", + ], + extras_require=extras, + classifiers=[ + "Development Status :: 5 - Production/Stable", + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: Apache Software License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + ], +) + +# Release checklist +# 1. Change the version in __init__.py and setup.py. +# 2. Commit these changes with the message: "Release: VERSION" +# 3. Add a tag in git to mark the release: "git tag VERSION -m 'Adds tag VERSION for pypi' " +# Push the tag to git: git push --tags origin main +# 4. Run the following commands in the top-level directory: +# python setup.py bdist_wheel +# python setup.py sdist +# 5. Upload the package to the pypi test server first: +# twine upload dist/* -r pypitest +# twine upload dist/* -r pypitest --repository-url=https://test.pypi.org/legacy/ +# 6. Check that you can install it in a virtualenv by running: +# pip install -i https://testpypi.python.org/pypi accelerate +# accelerate env +# accelerate test +# 7. Upload the final version to actual pypi: +# twine upload dist/* -r pypi +# 8. Add release notes to the tag in github once everything is looking hunky-dory. +# 9. Update the version in __init__.py, setup.py to the new version "-dev" and push to master diff --git a/src/pet/__init__.py b/src/pet/__init__.py new file mode 100644 index 0000000..2436302 --- /dev/null +++ b/src/pet/__init__.py @@ -0,0 +1,18 @@ +# flake8: noqa +# There's no way to ignore "F401 '...' imported but unused" warnings in this +# module, but to preserve other warnings. So, don't check this module at all. + +__version__ = "0.1.0.dev0" + +from .pet_model import ( + ParameterEfficientTuningModel, + ParameterEfficientTuningModelForSequenceClassification, + PromptEncoderType, +) +from .tuners import ( + PrefixEncoder, + PromptEmbedding, + PromptEncoder, + PromptEncoderReparameterizationType, + PromptTuningInit, +) diff --git a/src/pet.py b/src/pet/pet_model.py similarity index 97% rename from src/pet.py rename to src/pet/pet_model.py index 7ac208b..9345539 100644 --- a/src/pet.py +++ b/src/pet/pet_model.py @@ -1,12 +1,14 @@ -from collections import OrderedDict import enum import warnings +from collections import OrderedDict + import torch +from accelerate.state import AcceleratorState from transformers import PreTrainedModel + from tuners.p_tuning import PromptEncoder from tuners.prefix_tuning import PrefixEncoder from tuners.prompt_tuning import PromptEmbedding -from accelerate.state import AcceleratorState class PromptEncoderType(str, enum.Enum): @@ -88,8 +90,8 @@ class ParameterEfficientTuningModel(torch.nn.Module): def load_state_dict(self, state_dict, strict: bool = True): """ - Custom load state dict method that only loads prompt table and prompt encoder - parameters. Matching load method for this class' custom state dict method. + Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method + for this class' custom state dict method. """ self.prompt_encoder.embedding.load_state_dict({"weight": state_dict["prompt_embeddings"]}, strict) @@ -187,8 +189,8 @@ class ParameterEfficientTuningModelForSequenceClassification(ParameterEfficientT def load_state_dict(self, state_dict, strict: bool = True): """ - Custom load state dict method that only loads prompt table and prompt encoder - parameters. Matching load method for this class' custom state dict method. + Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method + for this class' custom state dict method. """ super().load_state_dict(state_dict["prompt_encoder"], strict) self.model.classifier.load_state_dict(state_dict["classifier"], strict) diff --git a/src/prompt_learning_legacy.py b/src/pet/prompt_learning_legacy.py similarity index 91% rename from src/prompt_learning_legacy.py rename to src/pet/prompt_learning_legacy.py index ffb6934..b08d018 100644 --- a/src/prompt_learning_legacy.py +++ b/src/pet/prompt_learning_legacy.py @@ -1,34 +1,27 @@ import enum -import torch +import functools import math import os +from collections import OrderedDict -from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss -from transformers import PreTrainedModel -from transformers.modeling_outputs import SequenceClassifierOutput -from transformers import AutoModelForSequenceClassification -from datasets import load_dataset -import evaluate import torch -from transformers import AutoTokenizer, get_linear_schedule_with_warmup, set_seed -from torch.utils.data import DataLoader from accelerate import Accelerator from accelerate.state import AcceleratorState from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin -import functools -from torch.distributed.fsdp import ( - FullyShardedDataParallel, - CPUOffload, +from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy, transformer_auto_wrap_policy +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from torch.utils.data import DataLoader +from transformers import ( + AutoModelForSequenceClassification, + AutoTokenizer, + PreTrainedModel, + get_linear_schedule_with_warmup, + set_seed, ) -from torch.distributed.fsdp.wrap import ( - enable_wrap, - wrap, - ModuleWrapPolicy, - transformer_auto_wrap_policy, - lambda_auto_wrap_policy, - _or_policy, -) -from collections import OrderedDict +from transformers.modeling_outputs import SequenceClassifierOutput + +import evaluate +from datasets import load_dataset class PromptEncoderReparameterizationType(str, enum.Enum): @@ -49,8 +42,7 @@ class PromptTuningInit(str, enum.Enum): class PromptEncoder(torch.nn.Module): """ - The prompt encoder network that is used to generate the virtual - token embeddings for p-tuning. + The prompt encoder network that is used to generate the virtual token embeddings for p-tuning. """ def __init__(self, config): @@ -92,13 +84,23 @@ class PromptEncoder(torch.nn.Module): ) elif self.encoder_type == PromptEncoderReparameterizationType.MLP: - layers = [torch.nn.Linear(self.input_size, self.hidden_size), torch.nn.ReLU()] - layers.extend([torch.nn.Linear(self.hidden_size, self.hidden_size), torch.nn.ReLU()]) + layers = [ + torch.nn.Linear(self.input_size, self.hidden_size), + torch.nn.ReLU(), + ] + layers.extend( + [ + torch.nn.Linear(self.hidden_size, self.hidden_size), + torch.nn.ReLU(), + ] + ) layers.append(torch.nn.Linear(self.hidden_size, self.output_size)) self.mlp_head = torch.nn.Sequential(*layers) else: - raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") + raise ValueError( + "Prompt encoder type not recognized. " " Please use one of MLP (recommended) or LSTM." + ) def forward(self, indices): input_embeds = self.embedding(indices) @@ -130,11 +132,15 @@ class PrefixEncoder(torch.nn.Module): self.trans = torch.nn.Sequential( torch.nn.Linear(config["token_dim"], config["prompt_hidden_size"]), torch.nn.Tanh(), - torch.nn.Linear(config["prompt_hidden_size"], config["num_layers"] * 2 * config["token_dim"]), + torch.nn.Linear( + config["prompt_hidden_size"], + config["num_layers"] * 2 * config["token_dim"], + ), ) else: self.embedding = torch.nn.Embedding( - config["num_virtual_tokens"], config["num_layers"] * 2 * config["token_dim"] + config["num_virtual_tokens"], + config["num_layers"] * 2 * config["token_dim"], ) def forward(self, prefix: torch.Tensor): @@ -247,8 +253,8 @@ class PromptModel(torch.nn.Module): def load_state_dict(self, state_dict, strict: bool = True): """ - Custom load state dict method that only loads prompt table and prompt encoder - parameters. Matching load method for this class' custom state dict method. + Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method + for this class' custom state dict method. """ self.prompt_encoder.embedding.load_state_dict({"weight": state_dict["prompt_embeddings"]}, strict) @@ -389,8 +395,8 @@ class PromptModelForSequenceClassification(PromptModel): def load_state_dict(self, state_dict, strict: bool = True): """ - Custom load state dict method that only loads prompt table and prompt encoder - parameters. Matching load method for this class' custom state dict method. + Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method + for this class' custom state dict method. """ super().load_state_dict(state_dict["prompt_encoder"], strict) self.classifier.load_state_dict(state_dict["classifier"], strict) @@ -528,7 +534,7 @@ def main(): batch_size = 16 lr = 5e-3 num_epochs = 100 - device = "cuda" + # device = "cuda" seed = 11 set_seed(seed) @@ -544,7 +550,12 @@ def main(): def tokenize_function(examples): # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None) + outputs = tokenizer( + examples["sentence1"], + examples["sentence2"], + truncation=True, + max_length=None, + ) return outputs # Apply the method we just defined to all the examples in all the splits of the dataset @@ -564,10 +575,16 @@ def main(): # Instantiate dataloaders. train_dataloader = DataLoader( - tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size + tokenized_datasets["train"], + shuffle=True, + collate_fn=collate_fn, + batch_size=batch_size, ) eval_dataloader = DataLoader( - tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=batch_size + tokenized_datasets["validation"], + shuffle=False, + collate_fn=collate_fn, + batch_size=batch_size, ) # Instantiate optimizer @@ -582,9 +599,13 @@ def main(): 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 - ) + ( + 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): @@ -616,17 +637,14 @@ def main(): accelerator.print(f"epoch {epoch}:", eval_metric) accelerator.print(f"epoch {epoch} train loss:", total_loss / len(train_dataloader)) + from torch.distributed.fsdp.fully_sharded_data_parallel import FullStateDictConfig from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - from torch.distributed.fsdp.fully_sharded_data_parallel import ( - BackwardPrefetch, - CPUOffload, - FullStateDictConfig, - ShardingStrategy, - StateDictType, - ) + from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType FSDP.set_state_dict_type( - model, StateDictType.FULL_STATE_DICT, FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + model, + StateDictType.FULL_STATE_DICT, + FullStateDictConfig(offload_to_cpu=True, rank0_only=True), ) state_dict = model.state_dict() state_dict = model.clean_state_dict(state_dict) diff --git a/src/pet/tuners/__init__.py b/src/pet/tuners/__init__.py new file mode 100644 index 0000000..41ae36c --- /dev/null +++ b/src/pet/tuners/__init__.py @@ -0,0 +1,7 @@ +# flake8: noqa +# There's no way to ignore "F401 '...' imported but unused" warnings in this +# module, but to preserve other warnings. So, don't check this module at all + +from .p_tuning import PromptEncoder, PromptEncoderReparameterizationType +from .prefix_tuning import PrefixEncoder +from .prompt_tuning import PromptEmbedding, PromptTuningInit diff --git a/src/tuners/lora.py b/src/pet/tuners/lora.py similarity index 100% rename from src/tuners/lora.py rename to src/pet/tuners/lora.py diff --git a/src/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py similarity index 87% rename from src/tuners/p_tuning.py rename to src/pet/tuners/p_tuning.py index 2c103c3..5b161c3 100644 --- a/src/tuners/p_tuning.py +++ b/src/pet/tuners/p_tuning.py @@ -1,6 +1,7 @@ -import torch import enum +import torch + class PromptEncoderReparameterizationType(str, enum.Enum): MLP = "MLP" @@ -11,8 +12,7 @@ class PromptEncoderReparameterizationType(str, enum.Enum): # with some refactor class PromptEncoder(torch.nn.Module): """ - The prompt encoder network that is used to generate the virtual - token embeddings for p-tuning. + The prompt encoder network that is used to generate the virtual token embeddings for p-tuning. """ def __init__(self, config): @@ -54,8 +54,16 @@ class PromptEncoder(torch.nn.Module): ) elif self.encoder_type == PromptEncoderReparameterizationType.MLP: - layers = [torch.nn.Linear(self.input_size, self.hidden_size), torch.nn.ReLU()] - layers.extend([torch.nn.Linear(self.hidden_size, self.hidden_size), torch.nn.ReLU()]) + layers = [ + torch.nn.Linear(self.input_size, self.hidden_size), + torch.nn.ReLU(), + ] + layers.extend( + [ + torch.nn.Linear(self.hidden_size, self.hidden_size), + torch.nn.ReLU(), + ] + ) layers.append(torch.nn.Linear(self.hidden_size, self.output_size)) self.mlp_head = torch.nn.Sequential(*layers) diff --git a/src/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py similarity index 81% rename from src/tuners/prefix_tuning.py rename to src/pet/tuners/prefix_tuning.py index c593b1e..761545f 100644 --- a/src/tuners/prefix_tuning.py +++ b/src/pet/tuners/prefix_tuning.py @@ -1,5 +1,6 @@ import torch + # Based on https://github.com/THUDM/P-tuning-v2/blob/main/model/prefix_encoder.py # with some refactor class PrefixEncoder(torch.nn.Module): @@ -20,11 +21,15 @@ class PrefixEncoder(torch.nn.Module): self.trans = torch.nn.Sequential( torch.nn.Linear(config["token_dim"], config["prompt_hidden_size"]), torch.nn.Tanh(), - torch.nn.Linear(config["prompt_hidden_size"], config["num_layers"] * 2 * config["token_dim"]), + torch.nn.Linear( + config["prompt_hidden_size"], + config["num_layers"] * 2 * config["token_dim"], + ), ) else: self.embedding = torch.nn.Embedding( - config["num_virtual_tokens"], config["num_layers"] * 2 * config["token_dim"] + config["num_virtual_tokens"], + config["num_layers"] * 2 * config["token_dim"], ) def forward(self, prefix: torch.Tensor): diff --git a/src/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py similarity index 99% rename from src/tuners/prompt_tuning.py rename to src/pet/tuners/prompt_tuning.py index cb4fc8d..c3bfed7 100644 --- a/src/tuners/prompt_tuning.py +++ b/src/pet/tuners/prompt_tuning.py @@ -1,7 +1,8 @@ -import torch import enum import math +import torch + class PromptTuningInit(str, enum.Enum): TEXT = "TEXT" diff --git a/src/utils/constants.py b/src/pet/utils/constants.py similarity index 100% rename from src/utils/constants.py rename to src/pet/utils/constants.py diff --git a/utils/style_doc.py b/utils/style_doc.py new file mode 100644 index 0000000..0422ebe --- /dev/null +++ b/utils/style_doc.py @@ -0,0 +1,556 @@ +# coding=utf-8 +# Copyright 2020 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Style utils for the .rst and the docstrings.""" + +import argparse +import os +import re +import warnings + +import black + + +BLACK_AVOID_PATTERNS = {} + + +# Regexes +# Re pattern that catches list introduction (with potential indent) +_re_list = re.compile(r"^(\s*-\s+|\s*\*\s+|\s*\d+\.\s+)") +# Re pattern that catches code block introduction (with potential indent) +_re_code = re.compile(r"^(\s*)```(.*)$") +# Re pattern that catches rst args blocks of the form `Parameters:`. +_re_args = re.compile("^\s*(Args?|Arguments?|Params?|Parameters?):\s*$") +# Re pattern that catches return blocks of the form `Return:`. +_re_returns = re.compile("^\s*Returns?:\s*$") +# Matches the special tag to ignore some paragraphs. +_re_doc_ignore = re.compile(r"(\.\.|#)\s*docstyle-ignore") +# Re pattern that matches , and blocks. +_re_tip = re.compile("^\s*|\s+warning={true}>)\s*$") + +DOCTEST_PROMPTS = [">>>", "..."] + + +def is_empty_line(line): + return len(line) == 0 or line.isspace() + + +def find_indent(line): + """ + Returns the number of spaces that start a line indent. + """ + search = re.search("^(\s*)(?:\S|$)", line) + if search is None: + return 0 + return len(search.groups()[0]) + + +def parse_code_example(code_lines): + """ + Parses a code example + + Args: + code_lines (`List[str]`): The code lines to parse. + max_len (`int`): The maximum length per line. + + Returns: + (List[`str`], List[`str`]): The list of code samples and the list of outputs. + """ + has_doctest = code_lines[0][:3] in DOCTEST_PROMPTS + + code_samples = [] + outputs = [] + in_code = True + current_bit = [] + + for line in code_lines: + if in_code and has_doctest and not is_empty_line(line) and line[:3] not in DOCTEST_PROMPTS: + code_sample = "\n".join(current_bit) + code_samples.append(code_sample.strip()) + in_code = False + current_bit = [] + elif not in_code and line[:3] in DOCTEST_PROMPTS: + output = "\n".join(current_bit) + outputs.append(output.strip()) + in_code = True + current_bit = [] + + # Add the line without doctest prompt + if line[:3] in DOCTEST_PROMPTS: + line = line[4:] + current_bit.append(line) + + # Add last sample + if in_code: + code_sample = "\n".join(current_bit) + code_samples.append(code_sample.strip()) + else: + output = "\n".join(current_bit) + outputs.append(output.strip()) + + return code_samples, outputs + + +def format_code_example(code: str, max_len: int, in_docstring: bool = False): + """ + Format a code example using black. Will take into account the doctest syntax as well as any initial indentation in + the code provided. + + Args: + code (`str`): The code example to format. + max_len (`int`): The maximum length per line. + in_docstring (`bool`, *optional*, defaults to `False`): Whether or not the code example is inside a docstring. + + Returns: + `str`: The formatted code. + """ + code_lines = code.split("\n") + + # Find initial indent + idx = 0 + while idx < len(code_lines) and is_empty_line(code_lines[idx]): + idx += 1 + if idx >= len(code_lines): + return "", "" + indent = find_indent(code_lines[idx]) + + # Remove the initial indent for now, we will had it back after styling. + # Note that l[indent:] works for empty lines + code_lines = [l[indent:] for l in code_lines[idx:]] + has_doctest = code_lines[0][:3] in DOCTEST_PROMPTS + + code_samples, outputs = parse_code_example(code_lines) + + # Let's blackify the code! We put everything in one big text to go faster. + delimiter = "\n\n### New code sample ###\n" + full_code = delimiter.join(code_samples) + line_length = max_len - indent + if has_doctest: + line_length -= 4 + + for k, v in BLACK_AVOID_PATTERNS.items(): + full_code = full_code.replace(k, v) + try: + mode = black.Mode(target_versions={black.TargetVersion.PY37}, line_length=line_length) + formatted_code = black.format_str(full_code, mode=mode) + error = "" + except Exception as e: + formatted_code = full_code + error = f"Code sample:\n{full_code}\n\nError message:\n{e}" + + # Let's get back the formatted code samples + for k, v in BLACK_AVOID_PATTERNS.items(): + formatted_code = formatted_code.replace(v, k) + # Triple quotes will mess docstrings. + if in_docstring: + formatted_code = formatted_code.replace('"""', "'''") + + code_samples = formatted_code.split(delimiter) + # We can have one output less than code samples + if len(outputs) == len(code_samples) - 1: + outputs.append("") + + formatted_lines = [] + for code_sample, output in zip(code_samples, outputs): + # black may have added some new lines, we remove them + code_sample = code_sample.strip() + in_triple_quotes = False + in_decorator = False + for line in code_sample.strip().split("\n"): + if has_doctest and not is_empty_line(line): + prefix = ( + "... " + if line.startswith(" ") or line in [")", "]", "}"] or in_triple_quotes or in_decorator + else ">>> " + ) + else: + prefix = "" + indent_str = "" if is_empty_line(line) else (" " * indent) + formatted_lines.append(indent_str + prefix + line) + + if '"""' in line: + in_triple_quotes = not in_triple_quotes + if line.startswith(" "): + in_decorator = False + if line.startswith("@"): + in_decorator = True + + formatted_lines.extend([" " * indent + line for line in output.split("\n")]) + if not output.endswith("===PT-TF-SPLIT==="): + formatted_lines.append("") + + result = "\n".join(formatted_lines) + return result.rstrip(), error + + +def format_text(text, max_len, prefix="", min_indent=None): + """ + Format a text in the biggest lines possible with the constraint of a maximum length and an indentation. + + Args: + text (`str`): The text to format + max_len (`int`): The maximum length per line to use + prefix (`str`, *optional*, defaults to `""`): A prefix that will be added to the text. + The prefix doesn't count toward the indent (like a - introducing a list). + min_indent (`int`, *optional*): The minimum indent of the text. + If not set, will default to the length of the `prefix`. + + Returns: + `str`: The formatted text. + """ + text = re.sub(r"\s+", " ", text) + if min_indent is not None: + if len(prefix) < min_indent: + prefix = " " * (min_indent - len(prefix)) + prefix + + indent = " " * len(prefix) + new_lines = [] + words = text.split(" ") + current_line = f"{prefix}{words[0]}" + for word in words[1:]: + try_line = f"{current_line} {word}" + if len(try_line) > max_len: + new_lines.append(current_line) + current_line = f"{indent}{word}" + else: + current_line = try_line + new_lines.append(current_line) + return "\n".join(new_lines) + + +def split_line_on_first_colon(line): + splits = line.split(":") + return splits[0], ":".join(splits[1:]) + + +def style_docstring(docstring, max_len): + """ + Style a docstring by making sure there is no useless whitespace and the maximum horizontal space is used. + + Args: + docstring (`str`): The docstring to style. + max_len (`int`): The maximum length of each line. + + Returns: + `str`: The styled docstring + """ + lines = docstring.split("\n") + new_lines = [] + + # Initialization + current_paragraph = None + current_indent = -1 + in_code = False + param_indent = -1 + prefix = "" + black_errors = [] + + # Special case for docstrings that begin with continuation of Args with no Args block. + idx = 0 + while idx < len(lines) and is_empty_line(lines[idx]): + idx += 1 + if ( + len(lines[idx]) > 1 + and lines[idx].rstrip().endswith(":") + and find_indent(lines[idx + 1]) > find_indent(lines[idx]) + ): + param_indent = find_indent(lines[idx]) + + for idx, line in enumerate(lines): + # Doing all re searches once for the one we need to repeat. + list_search = _re_list.search(line) + code_search = _re_code.search(line) + + # Are we starting a new paragraph? + # New indentation or new line: + new_paragraph = find_indent(line) != current_indent or is_empty_line(line) + # List item + new_paragraph = new_paragraph or list_search is not None + # Code block beginning + new_paragraph = new_paragraph or code_search is not None + # Beginning/end of tip + new_paragraph = new_paragraph or _re_tip.search(line) + + # In this case, we treat the current paragraph + if not in_code and new_paragraph and current_paragraph is not None and len(current_paragraph) > 0: + paragraph = " ".join(current_paragraph) + new_lines.append(format_text(paragraph, max_len, prefix=prefix, min_indent=current_indent)) + current_paragraph = None + + if code_search is not None: + if not in_code: + current_paragraph = [] + current_indent = len(code_search.groups()[0]) + current_code = code_search.groups()[1] + prefix = "" + if current_indent < param_indent: + param_indent = -1 + else: + current_indent = -1 + code = "\n".join(current_paragraph) + if current_code in ["py", "python"]: + formatted_code, error = format_code_example(code, max_len, in_docstring=True) + new_lines.append(formatted_code) + if len(error) > 0: + black_errors.append(error) + else: + new_lines.append(code) + current_paragraph = None + new_lines.append(line) + in_code = not in_code + + elif in_code: + current_paragraph.append(line) + elif is_empty_line(line): + current_paragraph = None + current_indent = -1 + prefix = "" + new_lines.append(line) + elif list_search is not None: + prefix = list_search.groups()[0] + current_indent = len(prefix) + current_paragraph = [line[current_indent:]] + elif _re_args.search(line): + new_lines.append(line) + param_indent = find_indent(lines[idx + 1]) + elif _re_tip.search(line): + # Add a new line before if not present + if not is_empty_line(new_lines[-1]): + new_lines.append("") + new_lines.append(line) + # Add a new line after if not present + if idx < len(lines) - 1 and not is_empty_line(lines[idx + 1]): + new_lines.append("") + elif current_paragraph is None or find_indent(line) != current_indent: + indent = find_indent(line) + # Special behavior for parameters intros. + if indent == param_indent: + # Special rules for some docstring where the Returns blocks has the same indent as the parameters. + if _re_returns.search(line) is not None: + param_indent = -1 + new_lines.append(line) + elif len(line) < max_len: + new_lines.append(line) + else: + intro, description = split_line_on_first_colon(line) + new_lines.append(intro + ":") + if len(description) != 0: + if find_indent(lines[idx + 1]) > indent: + current_indent = find_indent(lines[idx + 1]) + else: + current_indent = indent + 4 + current_paragraph = [description.strip()] + prefix = "" + else: + # Check if we have exited the parameter block + if indent < param_indent: + param_indent = -1 + + current_paragraph = [line.strip()] + current_indent = find_indent(line) + prefix = "" + elif current_paragraph is not None: + current_paragraph.append(line.lstrip()) + + if current_paragraph is not None and len(current_paragraph) > 0: + paragraph = " ".join(current_paragraph) + new_lines.append(format_text(paragraph, max_len, prefix=prefix, min_indent=current_indent)) + + return "\n".join(new_lines), "\n\n".join(black_errors) + + +def style_docstrings_in_code(code, max_len=119): + """ + Style all docstrings in some code. + + Args: + code (`str`): The code in which we want to style the docstrings. + max_len (`int`): The maximum number of characters per line. + + Returns: + `Tuple[str, str]`: A tuple with the clean code and the black errors (if any) + """ + # fmt: off + splits = code.split('\"\"\"') + splits = [ + (s if i % 2 == 0 or _re_doc_ignore.search(splits[i - 1]) is not None else style_docstring(s, max_len=max_len)) + for i, s in enumerate(splits) + ] + black_errors = "\n\n".join([s[1] for s in splits if isinstance(s, tuple) and len(s[1]) > 0]) + splits = [s[0] if isinstance(s, tuple) else s for s in splits] + clean_code = '\"\"\"'.join(splits) + # fmt: on + + return clean_code, black_errors + + +def style_file_docstrings(code_file, max_len=119, check_only=False): + """ + Style all docstrings in a given file. + + Args: + code_file (`str` or `os.PathLike`): The file in which we want to style the docstring. + max_len (`int`): The maximum number of characters per line. + check_only (`bool`, *optional*, defaults to `False`): + Whether to restyle file or just check if they should be restyled. + + Returns: + `bool`: Whether or not the file was or should be restyled. + """ + with open(code_file, "r", encoding="utf-8", newline="\n") as f: + code = f.read() + + clean_code, black_errors = style_docstrings_in_code(code, max_len=max_len) + + diff = clean_code != code + if not check_only and diff: + print(f"Overwriting content of {code_file}.") + with open(code_file, "w", encoding="utf-8", newline="\n") as f: + f.write(clean_code) + + return diff, black_errors + + +def style_mdx_file(mdx_file, max_len=119, check_only=False): + """ + Style a MDX file by formatting all Python code samples. + + Args: + mdx_file (`str` or `os.PathLike`): The file in which we want to style the examples. + max_len (`int`): The maximum number of characters per line. + check_only (`bool`, *optional*, defaults to `False`): + Whether to restyle file or just check if they should be restyled. + + Returns: + `bool`: Whether or not the file was or should be restyled. + """ + with open(mdx_file, "r", encoding="utf-8", newline="\n") as f: + content = f.read() + + lines = content.split("\n") + current_code = [] + current_language = "" + in_code = False + new_lines = [] + black_errors = [] + + for line in lines: + if _re_code.search(line) is not None: + in_code = not in_code + if in_code: + current_language = _re_code.search(line).groups()[1] + current_code = [] + else: + code = "\n".join(current_code) + if current_language in ["py", "python"]: + code, error = format_code_example(code, max_len) + if len(error) > 0: + black_errors.append(error) + new_lines.append(code) + + new_lines.append(line) + elif in_code: + current_code.append(line) + else: + new_lines.append(line) + + if in_code: + raise ValueError(f"There was a problem when styling {mdx_file}. A code block is opened without being closed.") + + clean_content = "\n".join(new_lines) + diff = clean_content != content + if not check_only and diff: + print(f"Overwriting content of {mdx_file}.") + with open(mdx_file, "w", encoding="utf-8", newline="\n") as f: + f.write(clean_content) + + return diff, "\n\n".join(black_errors) + + +def style_doc_files(*files, max_len=119, check_only=False): + """ + Applies doc styling or checks everything is correct in a list of files. + + Args: + files (several `str` or `os.PathLike`): The files to treat. + max_len (`int`): The maximum number of characters per line. + check_only (`bool`, *optional*, defaults to `False`): + Whether to restyle file or just check if they should be restyled. + + Returns: + List[`str`]: The list of files changed or that should be restyled. + """ + changed = [] + black_errors = [] + for file in files: + # Treat folders + if os.path.isdir(file): + files = [os.path.join(file, f) for f in os.listdir(file)] + files = [f for f in files if os.path.isdir(f) or f.endswith(".mdx") or f.endswith(".py")] + changed += style_doc_files(*files, max_len=max_len, check_only=check_only) + # Treat mdx + elif file.endswith(".mdx"): + try: + diff, black_error = style_mdx_file(file, max_len=max_len, check_only=check_only) + if diff: + changed.append(file) + if len(black_error) > 0: + black_errors.append( + f"There was a problem while formatting an example in {file} with black:\m{black_error}" + ) + except Exception: + print(f"There is a problem in {file}.") + raise + # Treat python files + elif file.endswith(".py"): + try: + diff, black_error = style_file_docstrings(file, max_len=max_len, check_only=check_only) + if diff: + changed.append(file) + if len(black_error) > 0: + black_errors.append( + f"There was a problem while formatting an example in {file} with black:\m{black_error}" + ) + except Exception: + print(f"There is a problem in {file}.") + raise + else: + warnings.warn(f"Ignoring {file} because it's not a py or an mdx file or a folder.") + if len(black_errors) > 0: + black_message = "\n\n".join(black_errors) + raise ValueError( + "Some code examples can't be interpreted by black, which means they aren't regular python:\n\n" + + black_message + + "\n\nMake sure to fix the corresponding docstring or doc file, or remove the py/python after ``` if it " + + "was not supposed to be a Python code sample." + ) + return changed + + +def main(*files, max_len=119, check_only=False): + changed = style_doc_files(*files, max_len=max_len, check_only=check_only) + if check_only and len(changed) > 0: + raise ValueError(f"{len(changed)} files should be restyled!") + elif len(changed) > 0: + print(f"Cleaned {len(changed)} files!") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("files", nargs="+", help="The file(s) or folder(s) to restyle.") + parser.add_argument("--max_len", type=int, help="The maximum length of lines.") + parser.add_argument("--check_only", action="store_true", help="Whether to only check and not fix styling issues.") + args = parser.parse_args() + + main(*args.files, max_len=args.max_len, check_only=args.check_only) From 6013c83dbcfdc319bcbc7c6c4a6c56bc0032c4b6 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Sat, 26 Nov 2022 18:49:53 +0530 Subject: [PATCH 03/18] add code --- README.md | 36 ++- src/pet/__init__.py | 4 +- src/pet/pet_model.py | 385 +++++++++++++++++++++++--------- src/pet/tuners/p_tuning.py | 2 +- src/pet/tuners/prefix_tuning.py | 8 +- src/pet/tuners/prompt_tuning.py | 2 +- 6 files changed, 320 insertions(+), 117 deletions(-) diff --git a/README.md b/README.md index d01cdf3..fe2639a 100644 --- a/README.md +++ b/README.md @@ -11,14 +11,32 @@ Supported moethods: ### Sequence Classification | | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA | -| --------- | ---- | ---- | ---- | ---- | -| RoBERTa | ✅ | ✅ | ✅ | | +| --------- | ---- | ---- | ---- | ---- | | BERT | ✅ | ✅ | ✅ | | +| RoBERTa | ✅ | ✅ | ✅ | | +| GPT-2 | ✅ | ✅ | ✅ | | +| Bloom | ✅ | ✅ | ✅ | | +| OPT | ✅ | ✅ | ✅ | | +| GPT-Neo | ✅ | ✅ | ✅ | | +| GPT-J | ✅ | ✅ | ✅ | | +| Deberta | | | | | | Deberta-v2 | | | | | -| BloomX | | | | | -| Bloom | | | | | -| mT-0 | | | | | -| T-0 | | | | | -| T5 | | | | | -| GPT-2 | | | | | -| BART | | | | | + +### Causal Language Modeling +| | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA | +| --------- | ---- | ---- | ---- | ---- | +| GPT-2 | | | | | +| Bloom | | | | | +| OPT | | | | | +| GPT-Neo | | | | | +| GPT-J | | | | | +| BART | | | | | + +### Conditional Generation +| | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA | +| --------- | ---- | ---- | ---- | ---- | +| T5 | | | | | +| BART | | | | | + + + diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 2436302..3e55142 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -5,8 +5,8 @@ __version__ = "0.1.0.dev0" from .pet_model import ( - ParameterEfficientTuningModel, - ParameterEfficientTuningModelForSequenceClassification, + PETModel, + PETModelForSequenceClassification, PromptEncoderType, ) from .tuners import ( diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 9345539..94e761c 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -1,14 +1,17 @@ import enum import warnings +import inspect from collections import OrderedDict import torch +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from accelerate.state import AcceleratorState from transformers import PreTrainedModel +from transformers.modeling_outputs import SequenceClassifierOutput -from tuners.p_tuning import PromptEncoder -from tuners.prefix_tuning import PrefixEncoder -from tuners.prompt_tuning import PromptEmbedding +from .tuners import PromptEncoder +from .tuners import PrefixEncoder +from .tuners import PromptEmbedding class PromptEncoderType(str, enum.Enum): @@ -18,18 +21,21 @@ class PromptEncoderType(str, enum.Enum): LORA = "LORA" -class ParameterEfficientTuningModel(torch.nn.Module): +class PETModel(torch.nn.Module): def __init__(self, model): super().__init__() self.model = model self.prompt_learning_config = model.config.prompt_learning_config - modules = list(self.model._modules) - - for module in modules: - if isinstance(self.model.get_submodule(module), PreTrainedModel): - transformer_backbone = self.model.get_submodule(module) - break + num_transformer_submodules = 0 + transformer_backbone = None + for name, module in self.model.named_children(): + if isinstance(module, PreTrainedModel): + if transformer_backbone is None: + transformer_backbone = module + self.transformer_backbone_name = name + num_transformer_submodules += 1 + self.prompt_learning_config["num_transformer_submodules"] = num_transformer_submodules for named_param, value in list(transformer_backbone.named_parameters()): if value.shape[0] == model.config.vocab_size: @@ -49,11 +55,15 @@ class ParameterEfficientTuningModel(torch.nn.Module): else: raise ValueError("Not supported") self.prompt_encoder = prompt_encoder - self.prompt_tokens = torch.arange(self.prompt_learning_config["num_virtual_tokens"]).long() + self.prompt_tokens = torch.arange( + self.prompt_learning_config["num_virtual_tokens"] + * self.prompt_learning_config["num_transformer_submodules"] + ).long() def get_prompt(self, batch_size): - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.transformer_backbone.device) + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.model.device) if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + prompt_tokens = prompt_tokens[:, : self.prompt_learning_config["num_virtual_tokens"]] if self.prompt_learning_config.get("inference_mode", False): past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: @@ -61,12 +71,20 @@ class ParameterEfficientTuningModel(torch.nn.Module): past_key_values = past_key_values.view( batch_size, self.prompt_learning_config["num_virtual_tokens"], - self.prompt_learning_config["num_layers"] * 2, + self.prompt_learning_config["num_layers"] + * self.prompt_learning_config["num_transformer_submodules"] + * 2, self.prompt_learning_config["num_attention_heads"], self.prompt_learning_config["token_dim"] // self.prompt_learning_config["num_attention_heads"], ) - past_key_values = self.dropout(past_key_values) - past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2) + past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split( + self.prompt_learning_config["num_transformer_submodules"] * 2 + ) + if "postprocess_past_key_value_function" in self.prompt_learning_config["prompt_encoder_config"]: + post_process_fn = self.prompt_learning_config["prompt_encoder_config"][ + "postprocess_past_key_value_function" + ] + past_key_values = post_process_fn(past_key_values) return past_key_values else: if self.prompt_learning_config.get("inference_mode", False): @@ -75,32 +93,16 @@ class ParameterEfficientTuningModel(torch.nn.Module): prompts = self.prompt_encoder(prompt_tokens) return prompts - def state_dict(self, destination=None, prefix=None, keep_vars=False): - """ - No frozen model parameters are stored in the state dict. - """ - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.model.device) - prompt_embeddings = self.prompt_encoder(prompt_tokens).detach().cpu() - if destination is None: - state_dict_ = OrderedDict() - else: - state_dict_ = destination - state_dict_["prompt_embeddings"] = prompt_embeddings[0] - return state_dict_ - def load_state_dict(self, state_dict, strict: bool = True): - """ - Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method - for this class' custom state dict method. - """ - self.prompt_encoder.embedding.load_state_dict({"weight": state_dict["prompt_embeddings"]}, strict) - - -class ParameterEfficientTuningModelForSequenceClassification(ParameterEfficientTuningModel): +class PETModelForSequenceClassification(PETModel): def __init__(self, model): super().__init__(model) self.config = self.model.config - self.modules_to_save = ("prompt_encoder", "classifier") + + for name, module in self.model.named_children(): + if isinstance(module, torch.nn.Linear): + self.cls_layer_name = name + break trainable_params = 0 all_param = 0 @@ -126,79 +128,262 @@ class ParameterEfficientTuningModelForSequenceClassification(ParameterEfficientT return_dict = return_dict if return_dict is not None else self.config.use_return_dict batch_size = input_ids.shape[0] - # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( - self.model.device - ) - attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) - if kwargs["token_type_ids"] is not None: - kwargs["token_type_ids"] = torch.cat( - ( - torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to(self.model.device), - kwargs["token_type_ids"], - ), - dim=1, + if attention_mask is not None and self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: + # concat prompt attention mask + prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( + self.model.device ) - - if kwargs["position_ids"] is not None: + 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.") kwargs["position_ids"] = None + kwargs.update( + { + "attention_mask": attention_mask, + "labels": labels, + "output_attentions": output_attentions, + "output_hidden_states": output_hidden_states, + "return_dict": return_dict, + } + ) if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: - past_key_values = self.get_prompt(batch_size=batch_size) - - return self.model( - input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - past_key_values=past_key_values, - **kwargs, - ) + return self.prefix_tuning_forward(input_ids=input_ids, **kwargs) else: - raw_embedding = self.word_embeddings(input_ids) + if kwargs.get("token_type_ids", None) is not None: + kwargs["token_type_ids"] = torch.cat( + ( + torch.zeros(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( + self.model.device + ), + kwargs["token_type_ids"], + ), + dim=1, + ).long() + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) prompts = self.get_prompt(batch_size=batch_size) - inputs_embeds = torch.cat((prompts, raw_embedding), dim=1) + inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1) + return self.model(inputs_embeds=inputs_embeds, **kwargs) - return self.model( - # input_ids, - inputs_embeds=inputs_embeds, - attention_mask=attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - **kwargs, - # past_key_values=past_key_values, + def prefix_tuning_forward( + self, + input_ids=None, + attention_mask=None, + inputs_embeds=None, + labels=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + batch_size = input_ids.shape[0] + past_key_values = self.get_prompt(batch_size) + fwd_params = list(inspect.signature(self.model.forward).parameters.keys()) + kwargs.update( + { + "input_ids": input_ids, + "attention_mask": attention_mask, + "inputs_embeds": inputs_embeds, + "output_attentions": output_attentions, + "output_hidden_states": output_hidden_states, + "return_dict": return_dict, + "past_key_values": past_key_values, + } + ) + if "past_key_values" in fwd_params: + return self.model(labels=labels, **kwargs) + else: + transformer_backbone_name = self.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) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.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): + 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: + 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)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, ) - def state_dict(self, destination=None, prefix=None, keep_vars=False): - """ - No frozen model parameters are stored in the state dict. - """ - if destination is None: - state_dict_ = OrderedDict() + +class PETModelForCausalLM(PETModel): + def __init__(self, model): + super().__init__(model) + self.config = self.model.config + + trainable_params = 0 + all_param = 0 + for _, param in self.named_parameters(): + all_param += param.numel() + if param.requires_grad: + trainable_params += param.numel() + print( + f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" + ) + + def forward( + self, + input_ids=None, + attention_mask=None, + labels=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + batch_size = input_ids.shape[0] + if self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: + if attention_mask is not None: + # concat prompt attention mask + prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( + self.model.device + ) + 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.prompt_learning_config["num_virtual_tokens"]), -100).to( + self.device + ) + labels = torch.cat((prefix_labels, labels), 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.") + kwargs["position_ids"] = None + if kwargs.get("token_type_ids", None) is not None: + warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids") + kwargs["token_type_ids"] = None + kwargs.update( + { + "attention_mask": attention_mask, + "labels": labels, + "output_attentions": output_attentions, + "output_hidden_states": output_hidden_states, + "return_dict": return_dict, + } + ) + + if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + past_key_values = self.get_prompt(batch_size) + return self.model(input_ids=input_ids, past_key_values=past_key_values, **kwargs) else: - state_dict_ = destination - state_dict_["prompt_encoder"] = super().state_dict() - state_dict_["classifier"] = self.model.classifier.state_dict() - if AcceleratorState().fsdp_plugin is not None: - state_dict_["_flat_param"] = None - return state_dict_ + if inputs_embeds is None: + 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) - def load_state_dict(self, state_dict, strict: bool = True): - """ - Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method - for this class' custom state dict method. - """ - super().load_state_dict(state_dict["prompt_encoder"], strict) - self.model.classifier.load_state_dict(state_dict["classifier"], strict) - def clean_state_dict(self, state_dict): - if AcceleratorState().fsdp_plugin is not None: - new_state_dict = OrderedDict() - for key in self.modules_to_save: - new_state_dict[key] = state_dict[key].copy() - state_dict = new_state_dict - return state_dict +class PETModelForSeq2SeqLM(PETModel): + def __init__(self, model): + super().__init__(model) + self.config = self.model.config + + trainable_params = 0 + all_param = 0 + for _, param in self.named_parameters(): + all_param += param.numel() + if param.requires_grad: + trainable_params += param.numel() + print( + f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" + ) + + def forward( + self, + input_ids=None, + attention_mask=None, + decoder_input_ids=None, + decoder_attention_mask=None, + labels=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + batch_size = input_ids.shape[0] + + if self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: + if attention_mask is not None: + # concat prompt attention mask + prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( + self.model.device + ) + attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) + if decoder_attention_mask is not None: + decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) + + # concat prompt labels + if labels is not None: + prefix_labels = torch.full((batch_size, self.prompt_learning_config["num_virtual_tokens"]), -100).to( + self.device + ) + labels = torch.cat((prefix_labels, labels), 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.") + kwargs["position_ids"] = None + if kwargs.get("token_type_ids", None) is not None: + warnings.warn("Token type ids are not supported for parameter efficient tuning. Ignoring token type ids") + kwargs["token_type_ids"] = None + kwargs.update( + { + "attention_mask": attention_mask, + "decoder_attention_mask": decoder_attention_mask, + "labels": labels, + "output_attentions": output_attentions, + "output_hidden_states": output_hidden_states, + "return_dict": return_dict, + } + ) + + if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + past_key_values = self.get_prompt(batch_size) + return self.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: + decoder_inputs_embeds = self.word_embeddings(decoder_input_ids) + prompts = self.get_prompt(batch_size=batch_size) + inputs_embeds = torch.cat( + (prompts[:, : self.prompt_learning_config["num_virtual_tokens"]], inputs_embeds), dim=1 + ) + decoder_inputs_embeds = torch.cat( + (prompts[:, self.prompt_learning_config["num_virtual_tokens"] :], decoder_inputs_embeds), dim=1 + ) + return self.model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs) diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py index 5b161c3..25149b9 100644 --- a/src/pet/tuners/p_tuning.py +++ b/src/pet/tuners/p_tuning.py @@ -21,7 +21,7 @@ class PromptEncoder(torch.nn.Module): self.input_size = config["token_dim"] self.output_size = config["token_dim"] self.hidden_size = config["prompt_hidden_size"] - self.total_virtual_tokens = config["num_virtual_tokens"] + self.total_virtual_tokens = config["num_virtual_tokens"] * config["num_transformer_submodules"] self.encoder_type = config["prompt_encoder_config"]["prompt_reparam_type"] # embedding diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py index 761545f..1319fe3 100644 --- a/src/pet/tuners/prefix_tuning.py +++ b/src/pet/tuners/prefix_tuning.py @@ -7,9 +7,9 @@ class PrefixEncoder(torch.nn.Module): r""" The torch.nn model to encode the prefix - Input shape: (batch-size, prefix-length) + Input shape: (batch_size, num_virtual_tokens) - Output shape: (batch-size, prefix-length, 2*layers*hidden) + Output shape: (batch_size, num_virtual_tokens, 2*(num_transformer_submodules)*layers*hidden) """ def __init__(self, config): @@ -23,13 +23,13 @@ class PrefixEncoder(torch.nn.Module): torch.nn.Tanh(), torch.nn.Linear( config["prompt_hidden_size"], - config["num_layers"] * 2 * config["token_dim"], + config["num_layers"] * 2 * config["num_transformer_submodules"] * config["token_dim"], ), ) else: self.embedding = torch.nn.Embedding( config["num_virtual_tokens"], - config["num_layers"] * 2 * config["token_dim"], + config["num_layers"] * 2 * config["num_transformer_submodules"] * config["token_dim"], ) def forward(self, prefix: torch.Tensor): diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py index c3bfed7..dca37c9 100644 --- a/src/pet/tuners/prompt_tuning.py +++ b/src/pet/tuners/prompt_tuning.py @@ -13,7 +13,7 @@ class PromptEmbedding(torch.nn.Module): def __init__(self, config, word_embeddings): super().__init__() - total_virtual_tokens = config["num_virtual_tokens"] + total_virtual_tokens = config["num_virtual_tokens"] * config["num_transformer_submodules"] self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"]) if config["prompt_encoder_config"]["prompt_tuning_init"] == PromptTuningInit.TEXT: from transformers import AutoTokenizer From d5c0c7e32c7da135ab30271b78b6462aa7b21e8e Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Sat, 26 Nov 2022 18:59:50 +0530 Subject: [PATCH 04/18] Update __init__.py --- src/pet/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 3e55142..6f832dd 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -7,6 +7,8 @@ __version__ = "0.1.0.dev0" from .pet_model import ( PETModel, PETModelForSequenceClassification, + PETModelForCausalLM, + PETModelForSeq2SeqLM, PromptEncoderType, ) from .tuners import ( From dec1c5c2eb0ff4796d5590f495267159b15efd57 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Sun, 27 Nov 2022 00:25:40 +0530 Subject: [PATCH 05/18] bug fixes --- src/pet/__init__.py | 2 +- src/pet/pet_model.py | 26 ++++++++++++-------------- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 6f832dd..9598427 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -6,9 +6,9 @@ __version__ = "0.1.0.dev0" from .pet_model import ( PETModel, - PETModelForSequenceClassification, PETModelForCausalLM, PETModelForSeq2SeqLM, + PETModelForSequenceClassification, PromptEncoderType, ) from .tuners import ( diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 94e761c..814473c 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -1,17 +1,13 @@ import enum -import warnings import inspect -from collections import OrderedDict +import warnings import torch from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from accelerate.state import AcceleratorState from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput -from .tuners import PromptEncoder -from .tuners import PrefixEncoder -from .tuners import PromptEmbedding +from .tuners import PrefixEncoder, PromptEmbedding, PromptEncoder class PromptEncoderType(str, enum.Enum): @@ -257,6 +253,7 @@ class PETModelForCausalLM(PETModel): self, input_ids=None, attention_mask=None, + inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, @@ -272,13 +269,6 @@ class PETModelForCausalLM(PETModel): ) 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.prompt_learning_config["num_virtual_tokens"]), -100).to( - self.device - ) - labels = torch.cat((prefix_labels, labels), 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.") kwargs["position_ids"] = None @@ -301,6 +291,12 @@ class PETModelForCausalLM(PETModel): else: if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) + # concat prompt labels + if kwargs["labels"] is not None: + prefix_labels = torch.full((batch_size, self.prompt_learning_config["num_virtual_tokens"]), -100).to( + self.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) @@ -325,8 +321,10 @@ class PETModelForSeq2SeqLM(PETModel): self, input_ids=None, attention_mask=None, + inputs_embeds=None, decoder_input_ids=None, decoder_attention_mask=None, + decoder_inputs_embeds=None, labels=None, output_attentions=None, output_hidden_states=None, @@ -348,7 +346,7 @@ class PETModelForSeq2SeqLM(PETModel): # concat prompt labels if labels is not None: prefix_labels = torch.full((batch_size, self.prompt_learning_config["num_virtual_tokens"]), -100).to( - self.device + self.model.device ) labels = torch.cat((prefix_labels, labels), dim=1) From a92a7876e0ff421215e7239afa04e26a75ff408c Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 29 Nov 2022 18:12:34 +0530 Subject: [PATCH 06/18] lot of refactoring --- README.md | 19 +- src/pet/__init__.py | 8 +- src/pet/pet_model.py | 193 ++++----- src/pet/prompt_learning_legacy.py | 657 ------------------------------ src/pet/task_mapping.py | 17 + src/pet/tuners/__init__.py | 6 +- src/pet/tuners/lora.py | 38 ++ src/pet/tuners/p_tuning.py | 49 ++- src/pet/tuners/prefix_tuning.py | 45 +- src/pet/tuners/prompt_tuning.py | 32 +- src/pet/utils/__init__.py | 5 + src/pet/utils/config.py | 36 ++ src/pet/utils/constants.py | 2 +- src/pet/utils/other.py | 14 + 14 files changed, 307 insertions(+), 814 deletions(-) delete mode 100644 src/pet/prompt_learning_legacy.py create mode 100644 src/pet/task_mapping.py create mode 100644 src/pet/utils/__init__.py create mode 100644 src/pet/utils/config.py create mode 100644 src/pet/utils/other.py diff --git a/README.md b/README.md index fe2639a..dde65c9 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,11 @@ # 🤗 pets Parameter-Efficient Tuning at Scale with 🤗 Accelerate -Supported moethods: +Supported methods: 1. Prefix Tuning 2. P-Tuning 3. Prompt Tuning -4. LoRA [in progress] +4. LoRA [in backlog] ## Models support matrix @@ -25,18 +25,17 @@ Supported moethods: ### Causal Language Modeling | | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA | | --------- | ---- | ---- | ---- | ---- | -| GPT-2 | | | | | -| Bloom | | | | | -| OPT | | | | | -| GPT-Neo | | | | | -| GPT-J | | | | | -| BART | | | | | +| GPT-2 | ✅ | ✅ | ✅ | | +| Bloom | ✅ | ✅ | ✅ | | +| OPT | ✅ | ✅ | ✅ | | +| GPT-Neo | ✅ | ✅ | ✅ | | +| GPT-J | ✅ | ✅ | ✅ | | ### Conditional Generation | | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA | | --------- | ---- | ---- | ---- | ---- | -| T5 | | | | | -| BART | | | | | +| T5 | ✅ | ✅ | ✅ | | +| BART | ✅ | ✅ | ✅ | | diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 9598427..77a32cc 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -9,12 +9,18 @@ from .pet_model import ( PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification, - PromptEncoderType, + PETPluginBase, + PETType, ) +from .task_mapping import MODEL_TYPE_TO_PROMPT_MODEL_MAPPING from .tuners import ( PrefixEncoder, + PrefixTuningConfig, PromptEmbedding, PromptEncoder, + PromptEncoderConfig, PromptEncoderReparameterizationType, + PromptTuningConfig, PromptTuningInit, ) +from .utils import PETConfig, PETType, PromptLearningConfig, TaskType diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 814473c..45a6264 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -1,4 +1,3 @@ -import enum import inspect import warnings @@ -8,98 +7,78 @@ from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput from .tuners import PrefixEncoder, PromptEmbedding, PromptEncoder - - -class PromptEncoderType(str, enum.Enum): - PROMPT_TUNING = "PROMPT_TUNING" - P_TUNING = "P_TUNING" - PREFIX_TUNING = "PREFIX_TUNING" - LORA = "LORA" +from .utils import PETConfig, PETType, TaskType class PETModel(torch.nn.Module): - def __init__(self, model): + def __init__(self, model, pet_config: PETConfig): super().__init__() self.model = model - self.prompt_learning_config = model.config.prompt_learning_config + self.pet_config = pet_config num_transformer_submodules = 0 transformer_backbone = None for name, module in self.model.named_children(): 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 num_transformer_submodules += 1 - self.prompt_learning_config["num_transformer_submodules"] = num_transformer_submodules + 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: self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", "")) break - # Make sure to freeze Tranformers model - for param in transformer_backbone.parameters(): - param.requires_grad = False - - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PROMPT_TUNING: - prompt_encoder = PromptEmbedding(self.prompt_learning_config, self.word_embeddings) - elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING: - prompt_encoder = PromptEncoder(self.prompt_learning_config) - elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: - prompt_encoder = PrefixEncoder(self.prompt_learning_config) + if self.pet_config.pet_type == PETType.PROMPT_TUNING: + prompt_encoder = PromptEmbedding(self.pet_config, self.word_embeddings) + elif self.pet_config.pet_type == PETType.P_TUNING: + prompt_encoder = PromptEncoder(self.pet_config) + elif self.pet_config.pet_type == PETType.PREFIX_TUNING: + prompt_encoder = PrefixEncoder(self.pet_config) else: raise ValueError("Not supported") self.prompt_encoder = prompt_encoder self.prompt_tokens = torch.arange( - self.prompt_learning_config["num_virtual_tokens"] - * self.prompt_learning_config["num_transformer_submodules"] + self.pet_config.num_virtual_tokens * self.pet_config.num_transformer_submodules ).long() def get_prompt(self, batch_size): prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.model.device) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: - prompt_tokens = prompt_tokens[:, : self.prompt_learning_config["num_virtual_tokens"]] - if self.prompt_learning_config.get("inference_mode", False): + 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: past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: past_key_values = self.prompt_encoder(prompt_tokens) past_key_values = past_key_values.view( batch_size, - self.prompt_learning_config["num_virtual_tokens"], - self.prompt_learning_config["num_layers"] - * self.prompt_learning_config["num_transformer_submodules"] - * 2, - self.prompt_learning_config["num_attention_heads"], - self.prompt_learning_config["token_dim"] // self.prompt_learning_config["num_attention_heads"], + self.pet_config.num_virtual_tokens, + self.pet_config.num_layers * 2, + self.pet_config.num_attention_heads, + self.pet_config.token_dim // self.pet_config.num_attention_heads, ) + if self.pet_config.num_transformer_submodules == 2: + past_key_values = torch.cat([past_key_values, past_key_values], dim=2) past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split( - self.prompt_learning_config["num_transformer_submodules"] * 2 + self.pet_config.num_transformer_submodules * 2 ) - if "postprocess_past_key_value_function" in self.prompt_learning_config["prompt_encoder_config"]: - post_process_fn = self.prompt_learning_config["prompt_encoder_config"][ - "postprocess_past_key_value_function" - ] + if self.pet_config.postprocess_past_key_value_function is not None: + post_process_fn = self.pet_config.postprocess_past_key_value_function past_key_values = post_process_fn(past_key_values) return past_key_values else: - if self.prompt_learning_config.get("inference_mode", False): + if self.pet_config.inference_mode: prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: prompts = self.prompt_encoder(prompt_tokens) return prompts - -class PETModelForSequenceClassification(PETModel): - def __init__(self, model): - super().__init__(model) - self.config = self.model.config - - for name, module in self.model.named_children(): - if isinstance(module, torch.nn.Linear): - self.cls_layer_name = name - break - + def print_trainable_parameters(self): trainable_params = 0 all_param = 0 for _, param in self.named_parameters(): @@ -110,6 +89,17 @@ class PETModelForSequenceClassification(PETModel): f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" ) + +class PETModelForSequenceClassification(PETModel): + def __init__(self, model, pet_config: PETConfig): + super().__init__(model, pet_config) + self.config = self.model.config + + for name, module in self.model.named_children(): + if isinstance(module, torch.nn.Linear): + self.cls_layer_name = name + break + def forward( self, input_ids=None, @@ -124,11 +114,9 @@ class PETModelForSequenceClassification(PETModel): return_dict = return_dict if return_dict is not None else self.config.use_return_dict batch_size = input_ids.shape[0] - if attention_mask is not None and self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: + if attention_mask is not None and self.pet_config.pet_type != PETType.LORA: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( - self.model.device - ) + 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 kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") @@ -143,15 +131,13 @@ class PETModelForSequenceClassification(PETModel): } ) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + if self.pet_config.pet_type == PETType.PREFIX_TUNING: 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.prompt_learning_config["num_virtual_tokens"]).to( - self.model.device - ), + torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.model.device), kwargs["token_type_ids"], ), dim=1, @@ -235,20 +221,10 @@ class PETModelForSequenceClassification(PETModel): class PETModelForCausalLM(PETModel): - def __init__(self, model): - super().__init__(model) + def __init__(self, model, pet_config: PETConfig): + super().__init__(model, pet_config) self.config = self.model.config - trainable_params = 0 - all_param = 0 - for _, param in self.named_parameters(): - all_param += param.numel() - if param.requires_grad: - trainable_params += param.numel() - print( - f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" - ) - def forward( self, input_ids=None, @@ -261,10 +237,10 @@ class PETModelForCausalLM(PETModel): **kwargs, ): batch_size = input_ids.shape[0] - if self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: + 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.prompt_learning_config["num_virtual_tokens"]).to( + 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) @@ -285,38 +261,29 @@ class PETModelForCausalLM(PETModel): } ) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + 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) else: if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) - # concat prompt labels - if kwargs["labels"] is not None: - prefix_labels = torch.full((batch_size, self.prompt_learning_config["num_virtual_tokens"]), -100).to( - self.model.device - ) - kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) + 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) 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) class PETModelForSeq2SeqLM(PETModel): - def __init__(self, model): - super().__init__(model) + def __init__(self, model, pet_config: PETConfig): + super().__init__(model, pet_config) self.config = self.model.config - trainable_params = 0 - all_param = 0 - for _, param in self.named_parameters(): - all_param += param.numel() - if param.requires_grad: - trainable_params += param.numel() - print( - f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" - ) - def forward( self, input_ids=None, @@ -333,22 +300,13 @@ class PETModelForSeq2SeqLM(PETModel): ): batch_size = input_ids.shape[0] - if self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: - if attention_mask is not None: + 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.prompt_learning_config["num_virtual_tokens"]).to( + 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 decoder_attention_mask is not None: - decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) - - # concat prompt labels - if labels is not None: - prefix_labels = torch.full((batch_size, self.prompt_learning_config["num_virtual_tokens"]), -100).to( - self.model.device - ) - labels = torch.cat((prefix_labels, labels), dim=1) + 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.") @@ -367,7 +325,7 @@ class PETModelForSeq2SeqLM(PETModel): } ) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + if self.pet_config.pet_type == PETType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.model( input_ids=input_ids, decoder_input_ids=decoder_input_ids, past_key_values=past_key_values, **kwargs @@ -375,13 +333,30 @@ class PETModelForSeq2SeqLM(PETModel): else: if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) - if decoder_inputs_embeds is None: + 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) prompts = self.get_prompt(batch_size=batch_size) - inputs_embeds = torch.cat( - (prompts[:, : self.prompt_learning_config["num_virtual_tokens"]], inputs_embeds), dim=1 - ) + inputs_embeds = torch.cat((prompts[:, : self.pet_config.num_virtual_tokens], inputs_embeds), dim=1) decoder_inputs_embeds = torch.cat( - (prompts[:, self.prompt_learning_config["num_virtual_tokens"] :], decoder_inputs_embeds), dim=1 + (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) diff --git a/src/pet/prompt_learning_legacy.py b/src/pet/prompt_learning_legacy.py deleted file mode 100644 index b08d018..0000000 --- a/src/pet/prompt_learning_legacy.py +++ /dev/null @@ -1,657 +0,0 @@ -import enum -import functools -import math -import os -from collections import OrderedDict - -import torch -from accelerate import Accelerator -from accelerate.state import AcceleratorState -from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin -from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy, transformer_auto_wrap_policy -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from torch.utils.data import DataLoader -from transformers import ( - AutoModelForSequenceClassification, - AutoTokenizer, - PreTrainedModel, - get_linear_schedule_with_warmup, - set_seed, -) -from transformers.modeling_outputs import SequenceClassifierOutput - -import evaluate -from datasets import load_dataset - - -class PromptEncoderReparameterizationType(str, enum.Enum): - MLP = "MLP" - LSTM = "LSTM" - - -class PromptEncoderType(str, enum.Enum): - PROMPT_TUNING = "PROMPT_TUNING" - P_TUNING_V1 = "P_TUNING_V1" - P_TUNING_V2 = "P_TUNING_V2" - - -class PromptTuningInit(str, enum.Enum): - TEXT = "TEXT" - RANDOM = "RANDOM" - - -class PromptEncoder(torch.nn.Module): - """ - The prompt encoder network that is used to generate the virtual token embeddings for p-tuning. - """ - - def __init__(self, config): - super().__init__() - self.token_dim = config["token_dim"] - self.input_size = config["token_dim"] - self.output_size = config["token_dim"] - self.hidden_size = config["prompt_hidden_size"] - self.total_virtual_tokens = config["num_virtual_tokens"] - self.encoder_type = config["prompt_encoder_config"]["prompt_reparam_type"] - - # embedding - self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim) - if not config.get("inference_mode", False): - if self.encoder_type == PromptEncoderReparameterizationType.LSTM: - if "dropout" not in config["prompt_encoder_config"]: - lstm_dropout = 0.0 - else: - lstm_dropout = config["prompt_encoder_config"]["dropout"] - - if "num_layers" not in config["prompt_encoder_config"]: - num_layers = 2 - else: - num_layers = config["prompt_encoder_config"]["num_layers"] - # LSTM - self.lstm_head = torch.nn.LSTM( - input_size=self.input_size, - hidden_size=self.hidden_size, - num_layers=num_layers, - dropout=lstm_dropout, - bidirectional=True, - batch_first=True, - ) - - self.mlp_head = torch.nn.Sequential( - torch.nn.Linear(self.hidden_size * 2, self.hidden_size * 2), - torch.nn.ReLU(), - torch.nn.Linear(self.hidden_size * 2, self.output_size), - ) - - elif self.encoder_type == PromptEncoderReparameterizationType.MLP: - layers = [ - torch.nn.Linear(self.input_size, self.hidden_size), - torch.nn.ReLU(), - ] - layers.extend( - [ - torch.nn.Linear(self.hidden_size, self.hidden_size), - torch.nn.ReLU(), - ] - ) - layers.append(torch.nn.Linear(self.hidden_size, self.output_size)) - self.mlp_head = torch.nn.Sequential(*layers) - - else: - raise ValueError( - "Prompt encoder type not recognized. " " Please use one of MLP (recommended) or LSTM." - ) - - def forward(self, indices): - input_embeds = self.embedding(indices) - if self.encoder_type == PromptEncoderReparameterizationType.LSTM: - output_embeds = self.mlp_head(self.lstm_head(input_embeds)[0]) - elif self.encoder_type == PromptEncoderReparameterizationType.MLP: - output_embeds = self.mlp_head(input_embeds) - else: - raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") - - return output_embeds - - -class PrefixEncoder(torch.nn.Module): - r""" - The torch.nn model to encode the prefix - - Input shape: (batch-size, prefix-length) - - Output shape: (batch-size, prefix-length, 2*layers*hidden) - """ - - def __init__(self, config): - super().__init__() - self.prefix_projection = config["prompt_encoder_config"]["prefix_projection"] - if self.prefix_projection and not config.get("inference_mode", False): - # Use a two-layer MLP to encode the prefix - self.embedding = torch.nn.Embedding(config["num_virtual_tokens"], config["token_dim"]) - self.trans = torch.nn.Sequential( - torch.nn.Linear(config["token_dim"], config["prompt_hidden_size"]), - torch.nn.Tanh(), - torch.nn.Linear( - config["prompt_hidden_size"], - config["num_layers"] * 2 * config["token_dim"], - ), - ) - else: - self.embedding = torch.nn.Embedding( - config["num_virtual_tokens"], - config["num_layers"] * 2 * config["token_dim"], - ) - - def forward(self, prefix: torch.Tensor): - if self.prefix_projection: - prefix_tokens = self.embedding(prefix) - past_key_values = self.trans(prefix_tokens) - else: - past_key_values = self.embedding(prefix) - return past_key_values - - -class PromptEmbedding(torch.nn.Module): - def __init__(self, config, word_embeddings): - super().__init__() - - total_virtual_tokens = config["num_virtual_tokens"] - self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"]) - if config["prompt_encoder_config"]["prompt_tuning_init"] == PromptTuningInit.TEXT: - from transformers import AutoTokenizer - - self.tokenizer = AutoTokenizer.from_pretrained(config["prompt_encoder_config"]["tokenizer_name_or_path"]) - self.init_text = config["prompt_encoder_config"]["prompt_tuning_text"] - init_token_ids = self.tokenizer(self.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: - init_token_ids = init_token_ids[:total_virtual_tokens] - elif num_text_tokens < total_virtual_tokens: - num_reps = math.ceil(total_virtual_tokens / num_text_tokens) - init_token_ids = init_token_ids * num_reps - init_token_ids = init_token_ids[:total_virtual_tokens] - - word_embedding_weights = word_embeddings(torch.LongTensor(init_token_ids)).detach().clone() - self.embedding.weight = torch.nn.Parameter(word_embedding_weights) - - def forward(self, indices): - # Just get embeddings and dropout - prompt_embeddings = self.embedding(indices) - return prompt_embeddings - - -class PromptModel(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.prompt_learning_config = model.config.prompt_learning_config - - modules = list(model._modules) - - for module in modules: - if isinstance(model.get_submodule(module), PreTrainedModel): - self.transformer_backbone = model.get_submodule(module) - break - - for named_param, value in list(self.transformer_backbone.named_parameters()): - if value.shape[0] == model.config.vocab_size: - self.word_embeddings = self.transformer_backbone.get_submodule(named_param.replace(".weight", "")) - break - - # Make sure to freeze Tranformers model - for param in self.transformer_backbone.parameters(): - param.requires_grad = False - - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PROMPT_TUNING: - prompt_encoder = PromptEmbedding(self.prompt_learning_config, self.word_embeddings) - elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V1: - prompt_encoder = PromptEncoder(self.prompt_learning_config) - elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: - prompt_encoder = PrefixEncoder(self.prompt_learning_config) - else: - raise ValueError("Not supported") - self.prompt_encoder = prompt_encoder - self.prompt_tokens = torch.arange(self.prompt_learning_config["num_virtual_tokens"]).long() - - def get_prompt(self, batch_size): - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.transformer_backbone.device) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: - if self.prompt_learning_config.get("inference_mode", False): - past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) - else: - past_key_values = self.prompt_encoder(prompt_tokens) - past_key_values = past_key_values.view( - batch_size, - self.prompt_learning_config["num_virtual_tokens"], - self.prompt_learning_config["num_layers"] * 2, - self.prompt_learning_config["num_attention_heads"], - self.prompt_learning_config["token_dim"] // self.prompt_learning_config["num_attention_heads"], - ) - past_key_values = self.dropout(past_key_values) - past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2) - return past_key_values - else: - if self.prompt_learning_config.get("inference_mode", False): - prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) - else: - prompts = self.prompt_encoder(prompt_tokens) - return prompts - - def state_dict(self, destination=None, prefix=None, keep_vars=False): - """ - No frozen model parameters are stored in the state dict. - """ - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.transformer_backbone.device) - prompt_embeddings = self.prompt_encoder(prompt_tokens).detach().cpu() - if destination is None: - state_dict_ = OrderedDict() - else: - state_dict_ = destination - state_dict_["prompt_embeddings"] = prompt_embeddings[0] - return state_dict_ - - def load_state_dict(self, state_dict, strict: bool = True): - """ - Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method - for this class' custom state dict method. - """ - self.prompt_encoder.embedding.load_state_dict({"weight": state_dict["prompt_embeddings"]}, strict) - - -class PromptModelForSequenceClassification(PromptModel): - def __init__(self, model): - super().__init__(model) - if "dropout" in [name for name, _ in model.named_children()]: - self.dropout = model.dropout - else: - self.dropout = torch.nn.Dropout(model.config.hidden_dropout_prob) - self.classifier = model.classifier - self.num_labels = model.num_labels - self.config = model.config - self.modules_to_save = ("prompt_encoder", "classifier") - - trainable_params = 0 - all_param = 0 - for _, param in self.named_parameters(): - all_param += param.numel() - if param.requires_grad: - trainable_params += param.numel() - print( - f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" - ) - - def forward( - self, - input_ids=None, - attention_mask=None, - inputs_embeds=None, - labels=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - **kwargs, - ): - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - batch_size = input_ids.shape[0] - # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( - self.transformer_backbone.device - ) - attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) - - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: - past_key_values = self.get_prompt(batch_size=batch_size) - - outputs = self.transformer_backbone( - input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - past_key_values=past_key_values, - **kwargs, - ) - - pooled_output = outputs[1] if len(outputs) > 1 else outputs[0] - else: - raw_embedding = self.word_embeddings(input_ids) - prompts = self.get_prompt(batch_size=batch_size) - inputs_embeds = torch.cat((prompts, raw_embedding), dim=1) - - outputs = self.transformer_backbone( - # input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - # **kwargs, - # past_key_values=past_key_values, - ) - - sequence_output = outputs[0] - sequence_output = sequence_output[:, self.prompt_learning_config["num_virtual_tokens"] :, :].contiguous() - pooled_output = sequence_output[:, 0] - - if ( - "pooler" in [name for name, _ in self.transformer_backbone.named_children()] - and self.transformer_backbone.pooler is not None - ): - pooled_output = self.transformer_backbone.pooler.dense(pooled_output) - pooled_output = self.transformer_backbone.pooler.activation(pooled_output) - - pooled_output = self.dropout(pooled_output) - logits = self.classifier(pooled_output) - - loss = None - if labels is not None: - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.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.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.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - if not return_dict: - output = (logits,) + outputs[2:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutput( - loss=loss, - logits=logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - def state_dict(self, destination=None, prefix=None, keep_vars=False): - """ - No frozen model parameters are stored in the state dict. - """ - if destination is None: - state_dict_ = OrderedDict() - else: - state_dict_ = destination - state_dict_["prompt_encoder"] = super().state_dict() - state_dict_["classifier"] = self.classifier.state_dict() - if AcceleratorState().fsdp_plugin is not None: - state_dict_["_flat_param"] = None - return state_dict_ - - def load_state_dict(self, state_dict, strict: bool = True): - """ - Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method - for this class' custom state dict method. - """ - super().load_state_dict(state_dict["prompt_encoder"], strict) - self.classifier.load_state_dict(state_dict["classifier"], strict) - - def clean_state_dict(self, state_dict): - if AcceleratorState().fsdp_plugin is not None: - new_state_dict = OrderedDict() - for key in self.modules_to_save: - new_state_dict[key] = state_dict[key].copy() - state_dict = new_state_dict - return state_dict - - -model_type_to_prompt_model_mapping = {"SequenceClassification": PromptModelForSequenceClassification} -num_virtual_tokens = 30 -model_name_or_path = "roberta-large" -tokenizer_name_or_path = "roberta-large" - -prompt_tuning_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "PROMPT_TUNING", - "prompt_encoder_config": { - "prompt_tuning_init": "TEXT", - "tokenizer_name_or_path": tokenizer_name_or_path, - "prompt_tuning_text": "Output is true or false. Task requires to recognize" - " whether the meaning of one text is entailed (can be inferred) from the other text.", - }, -} - - -p_tuning_v1_mlp_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "P_TUNING_V1", - "prompt_encoder_config": {"prompt_reparam_type": "MLP"}, -} - -p_tuning_v1_lstm_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "P_TUNING_V1", - "prompt_encoder_config": {"prompt_reparam_type": "LSTM"}, -} - -p_tuning_v2_no_proj_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "P_TUNING_V2", - "prompt_encoder_config": {"prefix_projection": False}, -} - -p_tuning_v2_proj_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "P_TUNING_V2", - "prompt_encoder_config": {"prefix_projection": True}, -} - - -def prepare_prompt_model(model, prompt_learning_config): - config = model.config.to_dict() - if "num_layers" not in prompt_learning_config: - if "num_hidden_layers" in config: - num_layers = config["num_hidden_layers"] - elif "num_layers" in config: - num_layers = config["num_layers"] - else: - raise ValueError("Please specify `num_layers` in `prompt_learning_config`") - prompt_learning_config["num_layers"] = num_layers - - if "token_dim" not in prompt_learning_config: - if "hidden_size" in config: - token_dim = config["hidden_size"] - elif "n_embd" in config: - token_dim = config["n_embd"] - elif "d_model" in config: - token_dim = config["d_model"] - else: - raise ValueError("Please specify `token_dim` in `prompt_learning_config`") - prompt_learning_config["token_dim"] = token_dim - - if "num_attention_heads" not in prompt_learning_config: - if "num_attention_heads" in config: - num_attention_heads = config["num_attention_heads"] - elif "n_head" in config: - num_attention_heads = config["n_head"] - elif "num_heads" in config: - num_attention_heads = config["num_heads"] - else: - raise ValueError("Please specify `num_attention_heads` in `prompt_learning_config`") - prompt_learning_config["num_attention_heads"] = num_attention_heads - - if "prompt_hidden_size" not in prompt_learning_config: - prompt_learning_config["prompt_hidden_size"] = token_dim - - model.config.prompt_learning_config = prompt_learning_config - model_type = model.__class__.__name__.split("For") - if len(model_type) < 2: - raise ValueError("Model Type not supported") - model_cls = model_type_to_prompt_model_mapping[model_type[1]] - prompt_model = model_cls(model) - return prompt_model - - -def fsdp_auto_wrap_policy(model): - def wrap_layers_with_required_grads(module): - if ( - len(list(module.children())) == 0 - and len(list(module.named_parameters())) > 0 - and module.weight.requires_grad - ): - return True - return False - - transformer_cls_to_wrap = { - PrefixEncoder, - PromptEmbedding, - PromptEncoder, - PromptModel, - FullyShardedDataParallelPlugin.get_module_class_from_name( - model, os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "") - ), - } - policy_1 = functools.partial( - transformer_auto_wrap_policy, - transformer_layer_cls=transformer_cls_to_wrap, - ) - policy_2 = functools.partial( - lambda_auto_wrap_policy, - lambda_fn=wrap_layers_with_required_grads, - ) - auto_wrap_policy = functools.partial(_or_policy, policies=[policy_1, policy_2]) - return auto_wrap_policy - - -def main(): - accelerator = Accelerator() - task = "rte" - batch_size = 16 - lr = 5e-3 - num_epochs = 100 - # device = "cuda" - seed = 11 - set_seed(seed) - - model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path) - model = prepare_prompt_model( - model, p_tuning_v2_no_proj_config - ) # p_tuning_v2_proj_config)#p_tuning_v2_no_proj_config) - # model = model.to("cuda") - - tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) - datasets = load_dataset("glue", task) - metric = evaluate.load("glue", task) - - def tokenize_function(examples): - # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer( - examples["sentence1"], - examples["sentence2"], - truncation=True, - max_length=None, - ) - return outputs - - # Apply the method we just defined to all the examples in all the splits of the dataset - # starting with the main process first: - tokenized_datasets = datasets.map( - tokenize_function, - batched=True, - remove_columns=["idx", "sentence1", "sentence2"], - ) - - # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the - # transformers library - tokenized_datasets = tokenized_datasets.rename_column("label", "labels") - - def collate_fn(examples): - return tokenizer.pad(examples, padding="longest", return_tensors="pt") - - # Instantiate dataloaders. - train_dataloader = DataLoader( - tokenized_datasets["train"], - shuffle=True, - collate_fn=collate_fn, - batch_size=batch_size, - ) - eval_dataloader = DataLoader( - tokenized_datasets["validation"], - shuffle=False, - collate_fn=collate_fn, - batch_size=batch_size, - ) - - # Instantiate optimizer - optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr) - - # Instantiate scheduler - lr_scheduler = get_linear_schedule_with_warmup( - optimizer=optimizer, - num_warmup_steps=0, - num_training_steps=(len(train_dataloader) * num_epochs), - ) - - 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(train_dataloader): - # batch.to(device) - outputs = model(**batch) - loss = outputs.loss - total_loss += loss.detach().float() - loss.backward() - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - model.eval() - for step, batch in enumerate(eval_dataloader): - # batch.to(device) - with torch.no_grad(): - outputs = model(**batch) - predictions = outputs.logits.argmax(dim=-1) - predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"])) - metric.add_batch( - predictions=predictions, - references=references, - ) - - eval_metric = metric.compute() - accelerator.print(f"epoch {epoch}:", eval_metric) - accelerator.print(f"epoch {epoch} train loss:", total_loss / len(train_dataloader)) - - from torch.distributed.fsdp.fully_sharded_data_parallel import FullStateDictConfig - from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType - - FSDP.set_state_dict_type( - model, - StateDictType.FULL_STATE_DICT, - FullStateDictConfig(offload_to_cpu=True, rank0_only=True), - ) - state_dict = model.state_dict() - state_dict = model.clean_state_dict(state_dict) - accelerator.print(state_dict) - - torch.save(state_dict, "p_tuning_v2.pt") - - -if __name__ == "__main__": - main() diff --git a/src/pet/task_mapping.py b/src/pet/task_mapping.py new file mode 100644 index 0000000..00e7415 --- /dev/null +++ b/src/pet/task_mapping.py @@ -0,0 +1,17 @@ +from .pet_model import PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification +from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig +from .utils import PETConfig + + +MODEL_TYPE_TO_PROMPT_MODEL_MAPPING = { + "SEQ_CLS": PETModelForSequenceClassification, + "SEQ_2_SEQ_LM": PETModelForSeq2SeqLM, + "CAUSAL_LM": PETModelForCausalLM, +} + +PET_TYPE_TO_CONFIG_MAPPING = { + "PROMPT_TUNING": PromptTuningConfig, + "PREFIX_TUNING": PrefixTuningConfig, + "P_TUNING": PromptEncoderConfig, + "LORA": PETConfig, +} diff --git a/src/pet/tuners/__init__.py b/src/pet/tuners/__init__.py index 41ae36c..79fcc15 100644 --- a/src/pet/tuners/__init__.py +++ b/src/pet/tuners/__init__.py @@ -2,6 +2,6 @@ # There's no way to ignore "F401 '...' imported but unused" warnings in this # module, but to preserve other warnings. So, don't check this module at all -from .p_tuning import PromptEncoder, PromptEncoderReparameterizationType -from .prefix_tuning import PrefixEncoder -from .prompt_tuning import PromptEmbedding, PromptTuningInit +from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType +from .prefix_tuning import PrefixEncoder, PrefixTuningConfig +from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index 044a482..7fbb0b4 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -1 +1,39 @@ # todo +import torch +from transformers import Conv1D + +import loralib as lora + + +class LoRAModel(torch.nn.Module): + def __init__(self, config, model): + super().__init__() + self.config = config + self.model = model + + def find_and_replace(self): + 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_module_keys"]): + parent, target_name, target = self.get_submodules(key) + if isinstance(target, torch.nn.Linear): + new_module = lora.Linear( + target.in_features, target.out_features, **self.config["prompt_encoder_config"] + ) + elif isinstance(target, torch.nn.Conv1d, Conv1D): + new_module = lora.LoRAConv1d( + target.in_channels, target.out_channels, target.kernel_size, bias=target.bias is not None + ) + self.replace_module(parent, target_name, new_module) + + 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_name, target + + 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() + if old_module.bias is not None: + new_module.bias = old_module.bias.clone() diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py index 25149b9..8f8e030 100644 --- a/src/pet/tuners/p_tuning.py +++ b/src/pet/tuners/p_tuning.py @@ -1,13 +1,37 @@ import enum +from dataclasses import dataclass, field +from typing import Union import torch +from ..utils import PromptLearningConfig + class PromptEncoderReparameterizationType(str, enum.Enum): MLP = "MLP" LSTM = "LSTM" +@dataclass +class PromptEncoderConfig(PromptLearningConfig): + encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field( + default=PromptEncoderReparameterizationType.MLP, + metadata={"help": "How to reparameterize the prompt encoder"}, + ) + encoder_hidden_size: int = field( + default=256, + metadata={"help": "The hidden size of the prompt encoder reparameterization"}, + ) + encoder_num_layers: int = field( + default=2, + metadata={"help": "The number of layers of the prompt encoder reparameterization"}, + ) + encoder_dropout: float = field( + default=0.0, + metadata={"help": "The dropout of the prompt encoder reparameterization"}, + ) + + # Based on https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/prompt_encoder.py # with some refactor class PromptEncoder(torch.nn.Module): @@ -17,26 +41,19 @@ class PromptEncoder(torch.nn.Module): def __init__(self, config): super().__init__() - self.token_dim = config["token_dim"] - self.input_size = config["token_dim"] - self.output_size = config["token_dim"] - self.hidden_size = config["prompt_hidden_size"] - self.total_virtual_tokens = config["num_virtual_tokens"] * config["num_transformer_submodules"] - self.encoder_type = config["prompt_encoder_config"]["prompt_reparam_type"] + self.token_dim = config.token_dim + self.input_size = self.token_dim + self.output_size = self.token_dim + self.hidden_size = config.encoder_hidden_size + self.total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules + self.encoder_type = config.encoder_reparameterization_type # embedding self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim) - if not config.get("inference_mode", False): + if not config.inference_mode: if self.encoder_type == PromptEncoderReparameterizationType.LSTM: - if "dropout" not in config["prompt_encoder_config"]: - lstm_dropout = 0.0 - else: - lstm_dropout = config["prompt_encoder_config"]["dropout"] - - if "num_layers" not in config["prompt_encoder_config"]: - num_layers = 2 - else: - num_layers = config["prompt_encoder_config"]["num_layers"] + lstm_dropout = config.encoder_dropout + num_layers = config.encoder_num_layers # LSTM self.lstm_head = torch.nn.LSTM( input_size=self.input_size, diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py index 1319fe3..049138b 100644 --- a/src/pet/tuners/prefix_tuning.py +++ b/src/pet/tuners/prefix_tuning.py @@ -1,5 +1,26 @@ +from dataclasses import dataclass, field +from typing import Callable, Optional + import torch +from ..utils import PromptLearningConfig + + +@dataclass +class PrefixTuningConfig(PromptLearningConfig): + encoder_hidden_size: int = field( + default=256, + metadata={"help": "The hidden size of the encoder"}, + ) + prefix_projection: bool = field( + default=False, + metadata={"help": "Whether to project the prefix tokens"}, + ) + postprocess_past_key_value_function: Optional[Callable] = field( + default=None, + metadata={"help": "The function to postprocess the past key value"}, + ) + # Based on https://github.com/THUDM/P-tuning-v2/blob/main/model/prefix_encoder.py # with some refactor @@ -9,28 +30,26 @@ class PrefixEncoder(torch.nn.Module): Input shape: (batch_size, num_virtual_tokens) - Output shape: (batch_size, num_virtual_tokens, 2*(num_transformer_submodules)*layers*hidden) + Output shape: (batch_size, num_virtual_tokens, 2*layers*hidden) """ def __init__(self, config): super().__init__() - self.prefix_projection = config["prompt_encoder_config"]["prefix_projection"] - if self.prefix_projection and not config.get("inference_mode", False): + self.prefix_projection = config.prefix_projection + token_dim = config.token_dim + num_layers = config.num_layers + encoder_hidden_size = config.encoder_hidden_size + num_virtual_tokens = config.num_virtual_tokens + if self.prefix_projection and not config.inference_mode: # Use a two-layer MLP to encode the prefix - self.embedding = torch.nn.Embedding(config["num_virtual_tokens"], config["token_dim"]) + self.embedding = torch.nn.Embedding(num_virtual_tokens, token_dim) self.trans = torch.nn.Sequential( - torch.nn.Linear(config["token_dim"], config["prompt_hidden_size"]), + torch.nn.Linear(token_dim, encoder_hidden_size), torch.nn.Tanh(), - torch.nn.Linear( - config["prompt_hidden_size"], - config["num_layers"] * 2 * config["num_transformer_submodules"] * config["token_dim"], - ), + torch.nn.Linear(encoder_hidden_size, num_layers * 2 * token_dim), ) else: - self.embedding = torch.nn.Embedding( - config["num_virtual_tokens"], - config["num_layers"] * 2 * config["num_transformer_submodules"] * config["token_dim"], - ) + self.embedding = torch.nn.Embedding(num_virtual_tokens, num_layers * 2 * token_dim) def forward(self, prefix: torch.Tensor): if self.prefix_projection: diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py index dca37c9..fdad421 100644 --- a/src/pet/tuners/prompt_tuning.py +++ b/src/pet/tuners/prompt_tuning.py @@ -1,25 +1,49 @@ import enum import math +from dataclasses import dataclass, field +from typing import Optional, Union import torch +from ..utils import PromptLearningConfig + class PromptTuningInit(str, enum.Enum): TEXT = "TEXT" RANDOM = "RANDOM" +@dataclass +class PromptTuningConfig(PromptLearningConfig): + prompt_tuning_init: Union[PromptTuningInit, str] = field( + default=PromptTuningInit.RANDOM, + metadata={"help": "How to initialize the prompt tuning parameters"}, + ) + prompt_tuning_init_text: Optional[str] = field( + default=None, + metadata={ + "help": "The text to use for prompt tuning initialization. Only used if prompt_tuning_init is `TEXT`" + }, + ) + tokenizer_name_or_path: Optional[str] = field( + default=None, + metadata={ + "help": "The tokenizer to use for prompt tuning initialization. Only used if prompt_tuning_init is `TEXT`" + }, + ) + + class PromptEmbedding(torch.nn.Module): def __init__(self, config, word_embeddings): super().__init__() - total_virtual_tokens = config["num_virtual_tokens"] * config["num_transformer_submodules"] + total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"]) - if config["prompt_encoder_config"]["prompt_tuning_init"] == PromptTuningInit.TEXT: + if config.prompt_tuning_init == PromptTuningInit.TEXT: from transformers import AutoTokenizer - self.tokenizer = AutoTokenizer.from_pretrained(config["prompt_encoder_config"]["tokenizer_name_or_path"]) - self.init_text = config["prompt_encoder_config"]["prompt_tuning_text"] + 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"] # Trim or iterate until num_text_tokens matches total_virtual_tokens num_text_tokens = len(init_token_ids) diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py new file mode 100644 index 0000000..2a32be1 --- /dev/null +++ b/src/pet/utils/__init__.py @@ -0,0 +1,5 @@ +# flake8: noqa +# There's no way to ignore "F401 '...' imported but unused" warnings in this +# module, but to preserve other warnings. So, don't check this module at all + +from .config import PETConfig, PETType, PromptLearningConfig, TaskType diff --git a/src/pet/utils/config.py b/src/pet/utils/config.py new file mode 100644 index 0000000..6f66966 --- /dev/null +++ b/src/pet/utils/config.py @@ -0,0 +1,36 @@ +import enum +from dataclasses import dataclass, field +from typing import Optional, Union + + +class PETType(str, enum.Enum): + PROMPT_TUNING = "PROMPT_TUNING" + P_TUNING = "P_TUNING" + PREFIX_TUNING = "PREFIX_TUNING" + LORA = "LORA" + + +class TaskType(str, enum.Enum): + SEQ_CLS = "SEQ_CLS" + SEQ_2_SEQ_LM = "SEQ_2_SEQ_LM" + CAUSAL_LM = "CAUSAL_LM" + + +@dataclass +class PETConfig: + """ + This is the configuration class to store the configuration of a :class:`~transform + """ + + pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"}) + task_type: Union[str, TaskType] = field(default=None, metadata={"help": "Task type"}) + inference_mode: bool = field(default=False, metadata={"help": "Whether to use inference mode"}) + + +@dataclass +class PromptLearningConfig(PETConfig): + 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"}) + 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"}) diff --git a/src/pet/utils/constants.py b/src/pet/utils/constants.py index 044a482..e832fa3 100644 --- a/src/pet/utils/constants.py +++ b/src/pet/utils/constants.py @@ -1 +1 @@ -# todo +# ToDo diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py new file mode 100644 index 0000000..45d7da1 --- /dev/null +++ b/src/pet/utils/other.py @@ -0,0 +1,14 @@ +import torch + + +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 + keys = past_key_values[: total_layers // 2] + keys = keys.transpose(2, 3).reshape( + total_layers // 2, batch_size * num_attention_heads, head_dim, num_virtual_tokens + ) + values = past_key_values[total_layers // 2 :] + values = values.reshape(total_layers // 2, batch_size * num_attention_heads, num_virtual_tokens, head_dim) + + return tuple(zip(keys, values)) From 7c21cb456c6f3f5f9610a26620286ef67e115518 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 29 Nov 2022 18:20:23 +0530 Subject: [PATCH 07/18] Update __init__.py --- src/pet/__init__.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 77a32cc..51ab10c 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -4,14 +4,7 @@ __version__ = "0.1.0.dev0" -from .pet_model import ( - PETModel, - PETModelForCausalLM, - PETModelForSeq2SeqLM, - PETModelForSequenceClassification, - PETPluginBase, - PETType, -) +from .pet_model import PETModel, PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification from .task_mapping import MODEL_TYPE_TO_PROMPT_MODEL_MAPPING from .tuners import ( PrefixEncoder, From 91deee81e2f5b2cf6d74daccbf99b6b85fdf892a Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 29 Nov 2022 18:40:41 +0530 Subject: [PATCH 08/18] refactor --- src/pet/__init__.py | 2 +- src/pet/{task_mapping.py => mapping.py} | 10 +++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) rename src/pet/{task_mapping.py => mapping.py} (65%) diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 51ab10c..43b71b2 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -5,7 +5,7 @@ __version__ = "0.1.0.dev0" from .pet_model import PETModel, PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification -from .task_mapping import MODEL_TYPE_TO_PROMPT_MODEL_MAPPING +from .mapping import MODEL_TYPE_TO_PET_MODEL_MAPPING, PET_TYPE_TO_CONFIG_MAPPING, get_pet_config, get_pet_model from .tuners import ( PrefixEncoder, PrefixTuningConfig, diff --git a/src/pet/task_mapping.py b/src/pet/mapping.py similarity index 65% rename from src/pet/task_mapping.py rename to src/pet/mapping.py index 00e7415..39f7b89 100644 --- a/src/pet/task_mapping.py +++ b/src/pet/mapping.py @@ -3,7 +3,7 @@ from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig from .utils import PETConfig -MODEL_TYPE_TO_PROMPT_MODEL_MAPPING = { +MODEL_TYPE_TO_PET_MODEL_MAPPING = { "SEQ_CLS": PETModelForSequenceClassification, "SEQ_2_SEQ_LM": PETModelForSeq2SeqLM, "CAUSAL_LM": PETModelForCausalLM, @@ -15,3 +15,11 @@ PET_TYPE_TO_CONFIG_MAPPING = { "P_TUNING": PromptEncoderConfig, "LORA": PETConfig, } + + +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): + return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config) From 513630dbc7e9bd08127837b3a9aa6712095e4cd0 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Tue, 29 Nov 2022 18:52:23 +0530 Subject: [PATCH 09/18] `get_pet_model` fn --- src/pet/mapping.py | 39 +++++++++++++++++++++++++++++++++ src/pet/tuners/p_tuning.py | 2 +- src/pet/tuners/prefix_tuning.py | 2 +- 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/src/pet/mapping.py b/src/pet/mapping.py index 39f7b89..4c82bea 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -22,4 +22,43 @@ def get_pet_config(config_dict): def get_pet_model(model, pet_config): + config = model.config.to_dict() + if pet_config.num_layers is None: + if "num_hidden_layers" in config: + num_layers = config["num_hidden_layers"] + elif "num_layers" in config: + num_layers = config["num_layers"] + elif "n_layer" in config: + num_layers = config["n_layer"] + else: + raise ValueError("Please specify `num_layers` in `pet_config`") + pet_config.num_layers = num_layers + + if pet_config.token_dim is None: + if "hidden_size" in config: + token_dim = config["hidden_size"] + elif "n_embd" in config: + token_dim = config["n_embd"] + elif "d_model" in config: + token_dim = config["d_model"] + else: + raise ValueError("Please specify `token_dim` in `pet_config`") + pet_config.token_dim = token_dim + + if pet_config.num_attention_heads is None: + if "num_attention_heads" in config: + num_attention_heads = config["num_attention_heads"] + elif "n_head" in config: + num_attention_heads = config["n_head"] + elif "num_heads" in config: + num_attention_heads = config["num_heads"] + elif "encoder_attention_heads" in config: + num_attention_heads = config["encoder_attention_heads"] + else: + 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 + return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config) diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py index 8f8e030..bb5a5a8 100644 --- a/src/pet/tuners/p_tuning.py +++ b/src/pet/tuners/p_tuning.py @@ -19,7 +19,7 @@ class PromptEncoderConfig(PromptLearningConfig): metadata={"help": "How to reparameterize the prompt encoder"}, ) encoder_hidden_size: int = field( - default=256, + default=None, metadata={"help": "The hidden size of the prompt encoder reparameterization"}, ) encoder_num_layers: int = field( diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py index 049138b..d709baf 100644 --- a/src/pet/tuners/prefix_tuning.py +++ b/src/pet/tuners/prefix_tuning.py @@ -9,7 +9,7 @@ from ..utils import PromptLearningConfig @dataclass class PrefixTuningConfig(PromptLearningConfig): encoder_hidden_size: int = field( - default=256, + default=None, metadata={"help": "The hidden size of the encoder"}, ) prefix_projection: bool = field( From e8160370247b3b61f57e59eb3f49acf9e3618b4b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 14:51:26 +0530 Subject: [PATCH 10/18] add lora support --- README.md | 8 ++++++-- setup.py | 2 +- src/pet/__init__.py | 2 ++ src/pet/tuners/__init__.py | 1 + src/pet/tuners/lora.py | 42 ++++++++++++++++++++++++++++---------- 5 files changed, 41 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index dde65c9..9d6c7e7 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,8 @@ -# 🤗 pets -Parameter-Efficient Tuning at Scale with 🤗 Accelerate +# 🤗 PET +Parameter-Efficient Tuning. Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP. Supported methods: + 1. Prefix Tuning 2. P-Tuning 3. Prompt Tuning @@ -38,4 +39,7 @@ Supported methods: | BART | ✅ | ✅ | ✅ | | +## Caveats: +1. Doesn't work currently with DeeSpeed ZeRO Stage-3. Extending support with DeeSpeed ZeRO Stage-3 is in backlog. + diff --git a/setup.py b/setup.py index 9c8edd9..6f74e46 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,7 @@ extras["dev"] = extras["quality"] setup( name="pets", version="0.1.0.dev0", - description="Parameter-Efficient Tuning at Scale (PETS)", + description="Parameter-Efficient Tuning (PET)", long_description=open("README.md", "r", encoding="utf-8").read(), long_description_content_type="text/markdown", keywords="deep learning", diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 43b71b2..8843808 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -15,5 +15,7 @@ from .tuners import ( PromptEncoderReparameterizationType, PromptTuningConfig, PromptTuningInit, + LoRAModel, + LoRAConfig, ) from .utils import PETConfig, PETType, PromptLearningConfig, TaskType diff --git a/src/pet/tuners/__init__.py b/src/pet/tuners/__init__.py index 79fcc15..d066951 100644 --- a/src/pet/tuners/__init__.py +++ b/src/pet/tuners/__init__.py @@ -5,3 +5,4 @@ from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType from .prefix_tuning import PrefixEncoder, PrefixTuningConfig from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit +from .lora import LoRAModel, LoRAConfig diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index 7fbb0b4..5538605 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -1,8 +1,29 @@ # todo +from typing import Callable, Optional import torch -from transformers import Conv1D +from transformers.pytorch_utils import Conv1D +from dataclasses import dataclass, asdict, field import loralib as lora +from loralib import mark_only_lora_as_trainable, lora_state_dict # flake8: noqa + +from ..utils import PETConfig + + +@dataclass +class LoRAConfig(PETConfig): + r: int = field(default=None, metadata={"help": "LoRA attention dimension"}) + lora_alpha: int = field(default=None, metadata={"help": "LoRA alpha"}) + lora_dropout: float = field(default=None, metadata={"help": "LoRA dropout"}) + merge_weights: bool = field( + default=False, metadata={"help": "Merge weights of the original model and the LoRA model"} + ) + fan_in_fan_out: bool = field( + 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"}) + bias: str = field(default="none", metadata={"help": "Bias type for LoRA. Can be 'none', 'all' or 'lora_only'"}) class LoRAModel(torch.nn.Module): @@ -10,27 +31,26 @@ class LoRAModel(torch.nn.Module): super().__init__() self.config = config self.model = model + self.find_and_replace() + mark_only_lora_as_trainable(self.model, self.config.bias) def find_and_replace(self): 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_module_keys"]): - parent, target_name, target = self.get_submodules(key) + if any(key.endswith(target_key) for target_key in self.config.target_module_keys): + parent, target, target_name = self.get_submodules(key) if isinstance(target, torch.nn.Linear): - new_module = lora.Linear( - target.in_features, target.out_features, **self.config["prompt_encoder_config"] - ) - elif isinstance(target, torch.nn.Conv1d, Conv1D): - new_module = lora.LoRAConv1d( - target.in_channels, target.out_channels, target.kernel_size, bias=target.bias is not None - ) + new_module = lora.Linear(target.in_features, target.out_features, **asdict(self.config)) + elif isinstance(target, Conv1D): + 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) 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_name, target + return parent, target, target_name def replace_module(self, parent_module, child_name, new_module, old_module): setattr(parent_module, child_name, new_module) From 23aecc4f6903c39fff225d9f17c504f39d9504d5 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 14:57:23 +0530 Subject: [PATCH 11/18] fix --- src/pet/__init__.py | 6 +++--- src/pet/mapping.py | 5 ++--- src/pet/tuners/__init__.py | 2 +- src/pet/tuners/lora.py | 7 ++++--- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 8843808..cce106b 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -4,9 +4,11 @@ __version__ = "0.1.0.dev0" -from .pet_model import PETModel, PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification from .mapping import MODEL_TYPE_TO_PET_MODEL_MAPPING, PET_TYPE_TO_CONFIG_MAPPING, get_pet_config, get_pet_model +from .pet_model import PETModel, PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification from .tuners import ( + LoRAConfig, + LoRAModel, PrefixEncoder, PrefixTuningConfig, PromptEmbedding, @@ -15,7 +17,5 @@ from .tuners import ( PromptEncoderReparameterizationType, PromptTuningConfig, PromptTuningInit, - LoRAModel, - LoRAConfig, ) from .utils import PETConfig, PETType, PromptLearningConfig, TaskType diff --git a/src/pet/mapping.py b/src/pet/mapping.py index 4c82bea..4fdde10 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -1,6 +1,5 @@ from .pet_model import PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification -from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig -from .utils import PETConfig +from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig, LoRAConfig MODEL_TYPE_TO_PET_MODEL_MAPPING = { @@ -13,7 +12,7 @@ PET_TYPE_TO_CONFIG_MAPPING = { "PROMPT_TUNING": PromptTuningConfig, "PREFIX_TUNING": PrefixTuningConfig, "P_TUNING": PromptEncoderConfig, - "LORA": PETConfig, + "LORA": LoRAConfig, } diff --git a/src/pet/tuners/__init__.py b/src/pet/tuners/__init__.py index d066951..22fd5c8 100644 --- a/src/pet/tuners/__init__.py +++ b/src/pet/tuners/__init__.py @@ -2,7 +2,7 @@ # There's no way to ignore "F401 '...' imported but unused" warnings in this # module, but to preserve other warnings. So, don't check this module at all +from .lora import LoRAConfig, LoRAModel from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType from .prefix_tuning import PrefixEncoder, PrefixTuningConfig from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit -from .lora import LoRAModel, LoRAConfig diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index 5538605..5fd7887 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -1,11 +1,12 @@ # todo -from typing import Callable, Optional +from dataclasses import asdict, dataclass, field +from typing import Optional + import torch from transformers.pytorch_utils import Conv1D -from dataclasses import dataclass, asdict, field import loralib as lora -from loralib import mark_only_lora_as_trainable, lora_state_dict # flake8: noqa +from loralib import lora_state_dict, mark_only_lora_as_trainable # noqa: F401 from ..utils import PETConfig 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 12/18] 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 From 992422100fd86adf954bd438baaa5d72ca8d0171 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 18:27:10 +0530 Subject: [PATCH 13/18] fix --- src/pet/pet_model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 3f0f078..a3f6fa6 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -15,6 +15,7 @@ class PETModel(torch.nn.Module): super().__init__() self.pet_config = pet_config self.base_model = model + self.modules_to_save = None if pet_config.pet_type != PETType.LORA: self._setup_prompt_encoder() else: @@ -99,6 +100,7 @@ class PETModelForSequenceClassification(PETModel): 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(): if isinstance(module, torch.nn.Linear): From a8350a57fe2dd5c3f040f53cd66e741abb99a8cf Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 18:39:36 +0530 Subject: [PATCH 14/18] fixes --- src/pet/pet_model.py | 7 +++++++ src/pet/utils/save_and_load.py | 3 +-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index a3f6fa6..1398146 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -53,6 +53,13 @@ class PETModel(torch.nn.Module): self.pet_config.num_virtual_tokens * self.pet_config.num_transformer_submodules ).long() + def get_prompt_embedding_to_save(self): + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -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] + 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) if self.pet_config.pet_type == PETType.PREFIX_TUNING: diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py index 39fa579..e2eea05 100644 --- a/src/pet/utils/save_and_load.py +++ b/src/pet/utils/save_and_load.py @@ -9,8 +9,7 @@ def get_pet_model_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() + prompt_embeddings = model.get_prompt_embedding_to_save() to_return["prompt_embeddings"] = prompt_embeddings if model.modules_to_save is not None: for key, value in state_dict.items(): From 3626d4cf76b78e416d7f9712cda081d90cd2df1b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 18:50:18 +0530 Subject: [PATCH 15/18] Delete constants.py --- src/pet/utils/constants.py | 1 - 1 file changed, 1 deletion(-) delete mode 100644 src/pet/utils/constants.py diff --git a/src/pet/utils/constants.py b/src/pet/utils/constants.py deleted file mode 100644 index e832fa3..0000000 --- a/src/pet/utils/constants.py +++ /dev/null @@ -1 +0,0 @@ -# ToDo From 207eb07f03137da29e4b0e917f6330368766e7c2 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 18:58:28 +0530 Subject: [PATCH 16/18] refactor --- src/pet/mapping.py | 54 +++++++++++++++++++++++----------------------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/src/pet/mapping.py b/src/pet/mapping.py index f2a3969..e0b934a 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -37,38 +37,38 @@ def get_pet_config(config_dict): return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict) -def _prepare_prompt_learning_config(pet_config, config): +def _prepare_prompt_learning_config(pet_config, model_config): if pet_config.num_layers is None: - if "num_hidden_layers" in config: - num_layers = config["num_hidden_layers"] - elif "num_layers" in config: - num_layers = config["num_layers"] - elif "n_layer" in config: - num_layers = config["n_layer"] + if "num_hidden_layers" in model_config: + num_layers = model_config["num_hidden_layers"] + elif "num_layers" in model_config: + num_layers = model_config["num_layers"] + elif "n_layer" in model_config: + num_layers = model_config["n_layer"] else: raise ValueError("Please specify `num_layers` in `pet_config`") pet_config.num_layers = num_layers if pet_config.token_dim is None: - if "hidden_size" in config: - token_dim = config["hidden_size"] - elif "n_embd" in config: - token_dim = config["n_embd"] - elif "d_model" in config: - token_dim = config["d_model"] + if "hidden_size" in model_config: + token_dim = model_config["hidden_size"] + elif "n_embd" in model_config: + token_dim = model_config["n_embd"] + elif "d_model" in model_config: + token_dim = model_config["d_model"] else: raise ValueError("Please specify `token_dim` in `pet_config`") pet_config.token_dim = token_dim if pet_config.num_attention_heads is None: - if "num_attention_heads" in config: - num_attention_heads = config["num_attention_heads"] - elif "n_head" in config: - num_attention_heads = config["n_head"] - elif "num_heads" in config: - num_attention_heads = config["num_heads"] - elif "encoder_attention_heads" in config: - num_attention_heads = config["encoder_attention_heads"] + if "num_attention_heads" in model_config: + num_attention_heads = model_config["num_attention_heads"] + elif "n_head" in model_config: + num_attention_heads = model_config["n_head"] + elif "num_heads" in model_config: + num_attention_heads = model_config["num_heads"] + elif "encoder_attention_heads" in model_config: + num_attention_heads = model_config["encoder_attention_heads"] else: raise ValueError("Please specify `num_attention_heads` in `pet_config`") pet_config.num_attention_heads = num_attention_heads @@ -79,11 +79,11 @@ def _prepare_prompt_learning_config(pet_config, config): return pet_config -def _prepare_lora_config(pet_config, config): +def _prepare_lora_config(pet_config, model_config): if pet_config.target_modules is None: - if config.model_type not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: + if model_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] + pet_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config.model_type] if len(pet_config.target_modules) == 1: pet_config.fan_in_fan_out = True pet_config.enable_lora = [True, False, True] @@ -93,10 +93,10 @@ def _prepare_lora_config(pet_config, config): def get_pet_model(model, pet_config): - config = model.config.to_dict() + model_config = model.config.to_dict() if pet_config.pet_type != PETType.LORA: - pet_config = _prepare_prompt_learning_config(pet_config, config) + pet_config = _prepare_prompt_learning_config(pet_config, model_config) else: - pet_config = _prepare_lora_config(pet_config, config) + pet_config = _prepare_lora_config(pet_config, model_config) return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config) From 9aba69efb563d061297ddce1abfc64cad8a862ea Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 19:00:42 +0530 Subject: [PATCH 17/18] fix --- src/pet/mapping.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pet/mapping.py b/src/pet/mapping.py index e0b934a..0cb2210 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -81,9 +81,9 @@ def _prepare_prompt_learning_config(pet_config, model_config): def _prepare_lora_config(pet_config, model_config): if pet_config.target_modules is None: - if model_config.model_type not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING: + if model_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[model_config.model_type] + pet_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]] if len(pet_config.target_modules) == 1: pet_config.fan_in_fan_out = True pet_config.enable_lora = [True, False, True] From 2ab8b95867fe052df4e098b3f397ff7068cb1a90 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Wed, 30 Nov 2022 19:10:05 +0530 Subject: [PATCH 18/18] fixes --- src/pet/tuners/lora.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index 0036638..eb30bad 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -32,9 +32,9 @@ class LoRAModel(torch.nn.Module): def __init__(self, config, model): super().__init__() self.config = config - self.lora_model = model + self.model = model self.find_and_replace() - mark_only_lora_as_trainable(self.lora_model, self.config.bias) + mark_only_lora_as_trainable(self.model, self.config.bias) def find_and_replace(self): kwargs = { @@ -44,7 +44,7 @@ class LoRAModel(torch.nn.Module): "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()] + 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) @@ -58,9 +58,9 @@ class LoRAModel(torch.nn.Module): self.replace_module(parent, target_name, new_module, target) def get_submodules(self, key): - parent = self.lora_model.get_submodule(".".join(key.split(".")[:-1])) + parent = self.model.get_submodule(".".join(key.split(".")[:-1])) target_name = key.split(".")[-1] - target = self.lora_model.get_submodule(key) + target = self.model.get_submodule(key) return parent, target, target_name def replace_module(self, parent_module, child_name, new_module, old_module): @@ -70,4 +70,4 @@ class LoRAModel(torch.nn.Module): new_module.bias = old_module.bias def forward(self, *args, **kwargs): - return self.lora_model(*args, **kwargs) + return self.model(*args, **kwargs)