add code

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      |   |   |   |   |
+
+
+

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 (

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)

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

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):

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