fixes

src/pet/pet_model.py

@@ -396,6 +396,38 @@ class PETModelForCausalLM(PETModel):
             inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
             return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
 
+    def generate(self, **kwargs):
+        if self.pet_config.pet_type == PETType.LORA:
+            return self.base_model.generate(**kwargs)
+        else:
+            assert "input_ids" in kwargs, "input_ids must be provided for PET model generation"
+            if kwargs.get("attention_mask", None) is not None:
+                # concat prompt attention mask
+                prefix_attention_mask = torch.ones(
+                    kwargs["input_ids"].shape[0], self.pet_config.num_virtual_tokens
+                ).to(self.device)
+                kwargs["attention_mask"] = torch.cat((prefix_attention_mask, kwargs["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
+            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
+
+            if self.pet_config.pet_type == PETType.PREFIX_TUNING:
+                batch_size = kwargs["input_ids"].shape[0]
+                past_key_values = self.get_prompt(batch_size)
+                kwargs["past_key_values"] = past_key_values
+                return self.base_model.generate(**kwargs)
+            else:
+                prompts = self.get_prompt(batch_size=kwargs["input_ids"].shape[0])
+                kwargs["inputs_embeds"] = torch.cat((prompts, self.word_embeddings(kwargs["input_ids"])), dim=1)
+                kwargs["input_ids"] = None
+                return self.base_model.generate(**kwargs)
+
 
 class PETModelForSeq2SeqLM(PETModel):
     """
@@ -503,6 +535,43 @@ class PETModelForSeq2SeqLM(PETModel):
             )
             return self.base_model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs)
 
+    def generate(self, **kwargs):
+        if self.pet_config.pet_type == PETType.LORA:
+            return self.base_model.generate(**kwargs)
+        else:
+            assert "input_ids" in kwargs, "input_ids must be provided for PET model generation"
+
+            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
+
+            if self.pet_config.pet_type == PETType.PREFIX_TUNING:
+                batch_size = kwargs["input_ids"].shape[0]
+                past_key_values = self.get_prompt(batch_size)
+                kwargs["past_key_values"] = past_key_values
+                return self.base_model.generate(**kwargs)
+            else:
+                if kwargs.get("attention_mask", None) is not None:
+                    # concat prompt attention mask
+                    prefix_attention_mask = torch.ones(
+                        kwargs["input_ids"].shape[0], self.pet_config.num_virtual_tokens
+                    ).to(self.device)
+                    kwargs["attention_mask"] = torch.cat((prefix_attention_mask, kwargs["attention_mask"]), dim=1)
+                prompts = self.get_prompt(batch_size=kwargs["input_ids"].shape[0])
+                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
+                )
+                kwargs["inputs_embeds"] = inputs_embeds
+                kwargs["decoder_inputs_embeds"] = decoder_inputs_embeds
+                kwargs["input_ids"] = None
+                return self.base_model.generate(**kwargs)
+
 
 class PETModelForTokenClassification(PETModel):
     """

src/pet/tuners/lora.py

@@ -2,6 +2,7 @@
 import math
 from dataclasses import dataclass, field
 from typing import List, Optional
+import warnings
 
 import torch
 import torch.nn as nn
@@ -90,11 +91,20 @@ class LoRAModel(torch.nn.Module):
             if any(key.endswith(target_key) for target_key in self.config.target_modules):
                 parent, target, target_name = self._get_submodules(key)
                 bias = target.bias is not None
-                if isinstance(target, torch.nn.Linear):
+                if isinstance(target, torch.nn.Linear) and self.config.enable_lora is None:
                     new_module = Linear(target.in_features, target.out_features, bias=bias, **kwargs)
-                elif isinstance(target, Conv1D):
+                elif self.config.enable_lora is not None:
                     kwargs.update({"enable_lora": self.config.enable_lora})
-                    in_features, out_features = target.weight.shape
+                    if isinstance(target, Conv1D):
+                        in_features, out_features = target.weight.shape
+                    else:
+                        in_features, out_features = target.in_features, target.out_features
+                        if kwargs["fan_in_fan_out"]:
+                            warnings.warn(
+                                "fan_in_fan_out is set to True but the target module is not a Conv1D. "
+                                "Setting fan_in_fan_out to False."
+                            )
+                            kwargs["fan_in_fan_out"] = False
                     new_module = MergedLinear(in_features, out_features, bias=bias, **kwargs)
                 self._replace_module(parent, target_name, new_module, target)
 

src/pet/utils/save_and_load.py

@@ -3,14 +3,18 @@ from loralib import lora_state_dict
 from .config import PETType
 
 
-def get_pet_model_state_dict(model):
+def get_pet_model_state_dict(model, state_dict=None):
     """
     Get the state dict of the PET model.
 
     Args:
-        model (:obj:`PETModel`): The PET model.
+        model (:obj:`PETModel`): The PET model. When using torch.nn.DistributedDataParallel, DeepSpeed or FSDP,
+        the model should be teh underlying model/unwrapped model (i.e. model.module).
+        state_dict (:obj:`dict`, `optional`): The state dict of the model. If not provided, the state dict of the model
+        will be used.
     """
-    state_dict = model.state_dict()
+    if state_dict is None:
+        state_dict = model.state_dict()
     if model.pet_config.pet_type == PETType.LORA:
         to_return = lora_state_dict(model, bias=model.pet_config.bias)
     else: