From bebd92046253751e9d878b1136da2042f72afff8 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 22 Dec 2022 17:30:13 +0530 Subject: [PATCH] fixes --- src/pet/pet_model.py | 69 ++++++++++++++++++++++++++++++++++ src/pet/tuners/lora.py | 16 ++++++-- src/pet/utils/save_and_load.py | 10 +++-- 3 files changed, 89 insertions(+), 6 deletions(-) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index e9d353c..48de251 100644 --- a/src/pet/pet_model.py +++ b/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): """ diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index 62d1ad8..06a7449 100644 --- a/src/pet/tuners/lora.py +++ b/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) diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py index 9ac53db..87e3709 100644 --- a/src/pet/utils/save_and_load.py +++ b/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: