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