diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index fff3dd3..4ed76a1 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -7,7 +7,7 @@ from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput from .tuners import LoRAModel, PrefixEncoder, PromptEmbedding, PromptEncoder -from .utils import PETConfig, PETType, TaskType, shift_tokens_right, _set_trainable +from .utils import PETConfig, PETType, TaskType, _set_trainable, shift_tokens_right class PETModel(torch.nn.Module): @@ -20,12 +20,12 @@ class PETModel(torch.nn.Module): Attributes: - base_model (:obj:`PreTrainedModel`): The base transformer model used for PET. - pet_config (:obj:`PETConfig`): The configuration of the PET model. - modules_to_save (:obj:`list` of :obj:`str`): The list of sub-module names to save when saving the model. - prompt_encoder (:obj:`PromptEncoder`): The prompt encoder used for PET if `pet_config.pet_type != PETType.LORA`. - prompt_tokens (:obj:`torch.Tensor`): The virtual prompt tokens used for PET if `pet_config.pet_type != PETType.LORA`. - transformer_backbone_name (:obj:`str`): The name of the transformer backbone in the base model + base_model (:obj:`PreTrainedModel`): The base transformer model used for PET. pet_config (:obj:`PETConfig`): + The configuration of the PET model. modules_to_save (:obj:`list` of :obj:`str`): The list of sub-module names + to save when saving the model. prompt_encoder (:obj:`PromptEncoder`): The prompt encoder used for PET if + `pet_config.pet_type != PETType.LORA`. prompt_tokens (:obj:`torch.Tensor`): The virtual prompt tokens used for + PET if `pet_config.pet_type != PETType.LORA`. transformer_backbone_name (:obj:`str`): The name of the + transformer backbone in the base model if `pet_config.pet_type != PETType.LORA`. word_embeddings (:obj:`torch.nn.Embedding`): The word embeddings of the transformer backbone in the base model if `pet_config.pet_type != PETType.LORA`. @@ -76,8 +76,8 @@ class PETModel(torch.nn.Module): def get_prompt_embedding_to_save(self): """ - Returns the prompt embedding to save when saving the model. - Only applocable when `pet_config.pet_type != PETType.LORA`. + Returns the prompt embedding to save when saving the model. Only applocable when `pet_config.pet_type != + PETType.LORA`. """ prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.base_model.device) if self.pet_config.pet_type == PETType.PREFIX_TUNING: @@ -87,8 +87,7 @@ class PETModel(torch.nn.Module): def get_prompt(self, batch_size): """ - Returns the virtual prompts to use for PET. - Only applocable when `pet_config.pet_type != PETType.LORA`. + Returns the virtual prompts to use for PET. Only applocable when `pet_config.pet_type != PETType.LORA`. """ 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: @@ -151,31 +150,21 @@ class PETModelForSequenceClassification(PETModel): pet_config (:obj:`PETConfig`): PET config. Attributes: - config (:obj:`PretrainedConfig`): The configuration object of the base model. - cls_layer_name (:obj:`str`): The name of the classification layer. + config (:obj:`PretrainedConfig`): The configuration object of the base model. cls_layer_name (:obj:`str`): The + name of the classification layer. Example:: - >>> from transformers import AutoModelForSequenceClassification - >>> from pet import PETModelForSequenceClassification, get_pet_config - >>> config = { - 'pet_type': 'PREFIX_TUNING', - 'task_type': 'SEQ_CLS', - 'inference_mode': False, - 'num_virtual_tokens': 20, - 'token_dim': 768, - 'num_transformer_submodules': 1, - 'num_attention_heads': 12, - 'num_layers': 12, - 'encoder_hidden_size': 768, - 'prefix_projection': False, - 'postprocess_past_key_value_function': None + >>> from transformers import AutoModelForSequenceClassification >>> from pet import + PETModelForSequenceClassification, get_pet_config >>> config = { + 'pet_type': 'PREFIX_TUNING', 'task_type': 'SEQ_CLS', 'inference_mode': False, 'num_virtual_tokens': 20, + 'token_dim': 768, 'num_transformer_submodules': 1, 'num_attention_heads': 12, 'num_layers': 12, + 'encoder_hidden_size': 768, 'prefix_projection': False, 'postprocess_past_key_value_function': None } - >>> pet_config = get_pet_config(config) - >>> model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased") - >>> pet_model = PETModelForSequenceClassification(model, pet_config) - >>> pet_model.print_trainable_parameters() - trainable params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117 + >>> pet_config = get_pet_config(config) >>> model = + AutoModelForSequenceClassification.from_pretrained("bert-base-cased") >>> pet_model = + PETModelForSequenceClassification(model, pet_config) >>> pet_model.print_trainable_parameters() trainable + params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117 """ def __init__(self, model, pet_config: PETConfig): @@ -335,26 +324,15 @@ class PETModelForCausalLM(PETModel): Example:: - >>> from transformers import AutoModelForCausalLM - >>> from pet import PETModelForCausalLM, get_pet_config - >>> config = { - 'pet_type': 'PREFIX_TUNING', - 'task_type': 'CAUSAL_LM', - 'inference_mode': False, - 'num_virtual_tokens': 20, - 'token_dim': 1280, - 'num_transformer_submodules': 1, - 'num_attention_heads': 20, - 'num_layers': 36, - 'encoder_hidden_size': 1280, - 'prefix_projection': False, - 'postprocess_past_key_value_function': None + >>> from transformers import AutoModelForCausalLM >>> from pet import PETModelForCausalLM, get_pet_config >>> + config = { + 'pet_type': 'PREFIX_TUNING', 'task_type': 'CAUSAL_LM', 'inference_mode': False, 'num_virtual_tokens': + 20, 'token_dim': 1280, 'num_transformer_submodules': 1, 'num_attention_heads': 20, 'num_layers': 36, + 'encoder_hidden_size': 1280, 'prefix_projection': False, 'postprocess_past_key_value_function': None } - >>> pet_config = get_pet_config(config) - >>> model = AutoModelForCausalLM.from_pretrained("gpt2-large") - >>> pet_model = PETModelForCausalLM(model, pet_config) - >>> pet_model.print_trainable_parameters() - trainable params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544 + >>> pet_config = get_pet_config(config) >>> model = AutoModelForCausalLM.from_pretrained("gpt2-large") >>> + pet_model = PETModelForCausalLM(model, pet_config) >>> pet_model.print_trainable_parameters() trainable params: + 1843200 || all params: 775873280 || trainable%: 0.23756456724479544 """ def __init__(self, model, pet_config: PETConfig): @@ -435,26 +413,15 @@ class PETModelForSeq2SeqLM(PETModel): Example:: - >>> from transformers import AutoModelForSeq2SeqLM - >>> from pet import PETModelForSeq2SeqLM, get_pet_config - >>> config = { - 'pet_type': 'LORA', - 'task_type': 'SEQ_2_SEQ_LM', - 'inference_mode': False, - 'r': 8, - 'target_modules': ['q', 'v'], - 'lora_alpha': 32, - 'lora_dropout': 0.1, - 'merge_weights': False, - 'fan_in_fan_out': False, - 'enable_lora': None, - 'bias': 'none' + >>> from transformers import AutoModelForSeq2SeqLM >>> from pet import PETModelForSeq2SeqLM, get_pet_config >>> + config = { + 'pet_type': 'LORA', 'task_type': 'SEQ_2_SEQ_LM', 'inference_mode': False, 'r': 8, 'target_modules': + ['q', 'v'], 'lora_alpha': 32, 'lora_dropout': 0.1, 'merge_weights': False, 'fan_in_fan_out': False, + 'enable_lora': None, 'bias': 'none' } - >>> pet_config = get_pet_config(config) - >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") - >>> pet_model = PETModelForSeq2SeqLM(model, pet_config) - >>> pet_model.print_trainable_parameters() - trainable params: 884736 || all params: 223843584 || trainable%: 0.3952474242013566 + >>> pet_config = get_pet_config(config) >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> + pet_model = PETModelForSeq2SeqLM(model, pet_config) >>> pet_model.print_trainable_parameters() trainable + params: 884736 || all params: 223843584 || trainable%: 0.3952474242013566 """ def __init__(self, model, pet_config: PETConfig): diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index d0e9c71..cf57404 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -1,11 +1,14 @@ # todo +import math from dataclasses import dataclass, field -from typing import Optional +from typing import List, Optional import torch +import torch.nn as nn +import torch.nn.functional as F from transformers.pytorch_utils import Conv1D -import loralib as lora +import loralib as lora # noqa: F401 from loralib import mark_only_lora_as_trainable from ..utils import PETConfig @@ -56,22 +59,15 @@ class LoRAModel(torch.nn.Module): Example:: - >>> from transformers import AutoModelForSeq2SeqLM, LoRAConfig - >>> from pet import LoRAModel, LoRAConfig - >>> config = LoRAConfig( - pet_type="LORA", - task_type="SEQ_2_SEQ_LM", - r=8, - lora_alpha=32, - target_modules=["q", "v"], - lora_dropout=0.01, - ) - >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") - >>> lora_model = LoRAModel(config, model) + >>> from transformers import AutoModelForSeq2SeqLM, LoRAConfig >>> from pet import LoRAModel, LoRAConfig >>> + config = LoRAConfig( + pet_type="LORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"], + lora_dropout=0.01, ) + >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> lora_model = LoRAModel(config, model) Attributes: - model (:obj:`transformers.PreTrainedModel`): The model to be adapted. - config (:obj:`LoRAConfig`): The configuration of the LoRA model. + model (:obj:`transformers.PreTrainedModel`): The model to be adapted. config (:obj:`LoRAConfig`): The + configuration of the LoRA model. """ def __init__(self, config, model): @@ -95,11 +91,11 @@ class LoRAModel(torch.nn.Module): parent, target, target_name = self._get_submodules(key) bias = target.bias is not None if isinstance(target, torch.nn.Linear): - new_module = lora.Linear(target.in_features, target.out_features, bias=bias, **kwargs) + new_module = Linear(target.in_features, target.out_features, bias=bias, **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, bias=bias, **kwargs) + new_module = MergedLinear(in_features, out_features, bias=bias, **kwargs) self._replace_module(parent, target_name, new_module, target) def _get_submodules(self, key): @@ -123,3 +119,208 @@ class LoRAModel(torch.nn.Module): return super().__getattr__(name) # defer to nn.Module's logic except AttributeError: return getattr(self.model, name) + + +# Below code is copied from https://github.com/microsoft/LoRA/blob/main/loralib/layers.py +# and modified to work with PyTorch FSDP + +# ------------------------------------------------------------------------------------------ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information. +# ------------------------------------------------------------------------------------------ +class LoRALayer: + def __init__( + self, + r: int, + lora_alpha: int, + lora_dropout: float, + merge_weights: bool, + ): + self.r = r + self.lora_alpha = lora_alpha + # Optional dropout + if lora_dropout > 0.0: + self.lora_dropout = nn.Dropout(p=lora_dropout) + else: + self.lora_dropout = lambda x: x + # Mark the weight as unmerged + self.merged = False + self.merge_weights = merge_weights + + +class Linear(nn.Linear, LoRALayer): + # LoRA implemented in a dense layer + def __init__( + self, + in_features: int, + out_features: int, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + merge_weights: bool = True, + **kwargs, + ): + nn.Linear.__init__(self, in_features, out_features, **kwargs) + LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) + + self.fan_in_fan_out = fan_in_fan_out + # Actual trainable parameters + if r > 0: + self.lora_A = nn.Linear(in_features, r, bias=False) + self.lora_B = nn.Linear(r, out_features, bias=False) + self.scaling = self.lora_alpha / self.r + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + self.reset_parameters() + if fan_in_fan_out: + self.weight.data = self.weight.data.T + + def reset_parameters(self): + nn.Linear.reset_parameters(self) + if hasattr(self, "lora_A"): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B.weight) + + def train(self, mode: bool = True): + def T(w): + return w.T if self.fan_in_fan_out else w + + nn.Linear.train(self, mode) + self.lora_A.train(mode) + self.lora_B.train(mode) + if self.merge_weights and self.merged: + # Make sure that the weights are not merged + if self.r > 0: + self.weight.data -= T(self.lora_B.weight @ self.lora_A.weight) * self.scaling + self.merged = False + + def eval(self): + def T(w): + return w.T if self.fan_in_fan_out else w + + nn.Linear.eval(self) + self.lora_A.eval() + self.lora_B.eval() + if self.merge_weights and not self.merged: + # Merge the weights and mark it + if self.r > 0: + self.weight.data += T(self.lora_B.weight @ self.lora_A.weight) * self.scaling + self.merged = True + + def forward(self, x: torch.Tensor): + def T(w): + return w.T if self.fan_in_fan_out else w + + if self.r > 0 and not self.merged: + result = F.linear(x, T(self.weight), bias=self.bias) + if self.r > 0: + result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling + return result + else: + return F.linear(x, T(self.weight), bias=self.bias) + + +class MergedLinear(nn.Linear, LoRALayer): + # LoRA implemented in a dense layer + def __init__( + self, + in_features: int, + out_features: int, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + enable_lora: List[bool] = [False], + fan_in_fan_out: bool = False, + merge_weights: bool = True, + **kwargs, + ): + nn.Linear.__init__(self, in_features, out_features, **kwargs) + LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) + assert out_features % len(enable_lora) == 0, "The length of enable_lora must divide out_features" + self.enable_lora = enable_lora + self.fan_in_fan_out = fan_in_fan_out + # Actual trainable parameters + if r > 0 and any(enable_lora): + self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False) + self.lora_B = nn.Conv1d( + r * sum(enable_lora), + out_features // len(enable_lora) * sum(enable_lora), + kernel_size=1, + groups=2, + bias=False, + ) + self.scaling = self.lora_alpha / self.r + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + # Compute the indices + self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1) + self.lora_ind[enable_lora, :] = True + self.lora_ind = self.lora_ind.view(-1) + self.reset_parameters() + if fan_in_fan_out: + self.weight.data = self.weight.data.T + + def reset_parameters(self): + nn.Linear.reset_parameters(self) + if hasattr(self, "lora_A"): + # initialize A the same way as the default for nn.Linear and B to zero + nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B.weight) + + def zero_pad(self, x): + result = x.new_zeros((*x.shape[:-1], self.out_features)) + result = result.view(-1, self.out_features) + result[:, self.lora_ind] = x.reshape(-1, self.out_features // len(self.enable_lora) * sum(self.enable_lora)) + return result.view((*x.shape[:-1], self.out_features)) + + def train(self, mode: bool = True): + def T(w): + return w.T if self.fan_in_fan_out else w + + nn.Linear.train(self, mode) + self.lora_A.train(mode) + self.lora_B.train(mode) + if self.merge_weights and self.merged: + # Make sure that the weights are not merged + if self.r > 0 and any(self.enable_lora): + delta_w = F.conv1d( + self.lora_A.weight.data.unsqueeze(0), + self.lora_B.weight.data.unsqueeze(-1), + groups=sum(self.enable_lora), + ).squeeze(0) + self.weight.data -= self.zero_pad(T(delta_w * self.scaling)) + self.merged = False + + def eval(self): + def T(w): + return w.T if self.fan_in_fan_out else w + + nn.Linear.eval(self) + self.lora_A.eval() + self.lora_B.eval() + if self.merge_weights and not self.merged: + # Merge the weights and mark it + if self.r > 0 and any(self.enable_lora): + delta_w = F.conv1d( + self.lora_A.weight.data.unsqueeze(0), + self.lora_B.weight.data.unsqueeze(-1), + groups=sum(self.enable_lora), + ).squeeze(0) + self.weight.data += self.zero_pad(T(delta_w * self.scaling)) + self.merged = True + + def forward(self, x: torch.Tensor): + def T(w): + return w.T if self.fan_in_fan_out else w + + if self.merged: + return F.linear(x, T(self.weight), bias=self.bias) + else: + result = F.linear(x, T(self.weight), bias=self.bias) + if self.r > 0: + after_A = self.lora_A(self.lora_dropout(x)) + after_B = self.lora_B(after_A.transpose(-2, -1)) + result += self.zero_pad(after_B) * self.scaling + return result diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py index 0700ba8..af28acf 100644 --- a/src/pet/tuners/p_tuning.py +++ b/src/pet/tuners/p_tuning.py @@ -19,8 +19,8 @@ class PromptEncoderConfig(PromptLearningConfig): Args: encoder_reparameterization_type - (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]): - The type of reparameterization to use. + (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]): The type of reparameterization to + use. encoder_hidden_size (:obj:`int`): The hidden size of the prompt encoder. encoder_num_layers (:obj:`int`): The number of layers of the prompt encoder. encoder_dropout (:obj:`float`): The dropout probability of the prompt encoder. @@ -55,31 +55,23 @@ class PromptEncoder(torch.nn.Module): Example:: - >>> from pet import PromptEncoder, PromptEncoderConfig - >>> config = PromptEncoderConfig( - pet_type="P_TUNING", - task_type="SEQ_2_SEQ_LM", - num_virtual_tokens=20, - token_dim=768, - num_transformer_submodules=1, - num_attention_heads=12, - num_layers=12, - encoder_reparameterization_type="MLP", - encoder_hidden_size=768 + >>> from pet import PromptEncoder, PromptEncoderConfig >>> config = PromptEncoderConfig( + pet_type="P_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768, + num_transformer_submodules=1, num_attention_heads=12, num_layers=12, + encoder_reparameterization_type="MLP", encoder_hidden_size=768 ) >>> prompt_encoder = PromptEncoder(config) Attributes: - embedding (:class:`~torch.nn.Embedding`): The embedding layer of the prompt encoder. - mlp_head (:class:`~torch.nn.Sequential`): The MLP head of the prompt encoder if `inference_mode=False`. - lstm_head (:class:`~torch.nn.LSTM`): + embedding (:class:`~torch.nn.Embedding`): The embedding layer of the prompt encoder. mlp_head + (:class:`~torch.nn.Sequential`): The MLP head of the prompt encoder if `inference_mode=False`. lstm_head + (:class:`~torch.nn.LSTM`): The LSTM head of the prompt encoder if `inference_mode=False` and `encoder_reparameterization_type="LSTM"`. - token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model. - input_size (:obj:`int`): The input size of the prompt encoder. - output_size (:obj:`int`): The output size of the prompt encoder. - hidden_size (:obj:`int`): The hidden size of the prompt encoder. - total_virtual_tokens (:obj:`int`): The total number of virtual tokens of the prompt encoder. - encoder_type (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]): + token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model. input_size (:obj:`int`): + The input size of the prompt encoder. output_size (:obj:`int`): The output size of the prompt encoder. + hidden_size (:obj:`int`): The hidden size of the prompt encoder. total_virtual_tokens (:obj:`int`): The total + number of virtual tokens of the prompt encoder. encoder_type + (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]): The encoder type of the prompt encoder. diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py index cb4c575..39de89d 100644 --- a/src/pet/tuners/prefix_tuning.py +++ b/src/pet/tuners/prefix_tuning.py @@ -14,7 +14,8 @@ class PrefixTuningConfig(PromptLearningConfig): Args: encoder_hidden_size (:obj: int): The hidden size of the prompt encoder. prefix_projection (:obj: bool): Whether to project the prefix embeddings. - postprocess_past_key_value_function (:obj: Optional[Callable]): The function to postprocess the past key value. + postprocess_past_key_value_function (: + obj: Optional[Callable]): The function to postprocess the past key value. """ encoder_hidden_size: int = field( @@ -42,23 +43,16 @@ class PrefixEncoder(torch.nn.Module): Example:: - >>> from pet import PrefixEncoder, PrefixTuningConfig - >>> config = PrefixTuningConfig( - pet_type="PREFIX_TUNING", - task_type="SEQ_2_SEQ_LM", - num_virtual_tokens=20, - token_dim=768, - num_transformer_submodules=1, - num_attention_heads=12, - num_layers=12, - encoder_hidden_size=768 + >>> from pet import PrefixEncoder, PrefixTuningConfig >>> config = PrefixTuningConfig( + pet_type="PREFIX_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768, + num_transformer_submodules=1, num_attention_heads=12, num_layers=12, encoder_hidden_size=768 ) >>> prefix_encoder = PrefixEncoder(config) Attributes: - embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prefix encoder. - trans (:obj:`torch.nn.Sequential`): The two-layer MLP to transform the prefix embeddings + embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prefix encoder. trans + (:obj:`torch.nn.Sequential`): The two-layer MLP to transform the prefix embeddings if :obj:`prefix_projection` is :obj:`True`. prefix_projection (:obj:`bool`): Whether to project the prefix embeddings. diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py index 5ee6ef1..20e3a2e 100644 --- a/src/pet/tuners/prompt_tuning.py +++ b/src/pet/tuners/prompt_tuning.py @@ -19,7 +19,8 @@ class PromptTuningConfig(PromptLearningConfig): This is the configuration class to store the configuration of a :class:`~pet.PromptEmbedding`. Args: - prompt_tuning_init (:obj:Union[:class:`PromptTuningInit`, :obj:`str`]): The initialization of the prompt embedding. + prompt_tuning_init (: + obj:Union[:class:`PromptTuningInit`, :obj:`str`]): The initialization of the prompt embedding. prompt_tuning_init_text (:obj: Optional[:obj:`str`]): The text to initialize the prompt embedding. Only used if `prompt_tuning_init` is `TEXT` tokenizer_name_or_path (:obj: Optional[:obj:`str`]): The name or path of the tokenizer. @@ -57,21 +58,14 @@ class PromptEmbedding(torch.nn.Module): Example:: - >>> from pet import PromptEmbedding, PromptTuningConfig - >>> config = PromptTuningConfig( - pet_type="PROMPT_TUNING", - task_type="SEQ_2_SEQ_LM", - num_virtual_tokens=20, - token_dim=768, - num_transformer_submodules=1, - num_attention_heads=12, - num_layers=12, - prompt_tuning_init="TEXT", + >>> from pet import PromptEmbedding, PromptTuningConfig >>> config = PromptTuningConfig( + pet_type="PROMPT_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768, + num_transformer_submodules=1, num_attention_heads=12, num_layers=12, prompt_tuning_init="TEXT", prompt_tuning_init_text="Predict if sentiment of this review is positive, negative or neutral", tokenizer_name_or_path="t5-base", ) - >>> # t5_model.shared is the word embeddings of the base model - >>> prompt_embedding = PromptEmbedding(config, t5_model.shared) + >>> # t5_model.shared is the word embeddings of the base model >>> prompt_embedding = PromptEmbedding(config, + t5_model.shared) Input Shape: (batch_size, total_virtual_tokens) diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py index 582dad4..210d711 100644 --- a/src/pet/utils/__init__.py +++ b/src/pet/utils/__init__.py @@ -3,5 +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, _set_trainable +from .other import _set_trainable, 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 984016b..c6d4588 100644 --- a/src/pet/utils/config.py +++ b/src/pet/utils/config.py @@ -35,7 +35,8 @@ class PETConfig: @dataclass class PromptLearningConfig(PETConfig): """ - This is the base configuration class to store the configuration of a :obj:Union[:class:`~pet.PrefixTuning`, :class:`~pet.PromptEncoder`, :class:`~pet.PromptTuning`]. + This is the base configuration class to store the configuration of a :obj:Union[:class:`~pet.PrefixTuning`, + :class:`~pet.PromptEncoder`, :class:`~pet.PromptTuning`]. Args: num_virtual_tokens (:obj:`int`): The number of virtual tokens to use. diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py index a13b4f6..0b63248 100644 --- a/src/pet/utils/other.py +++ b/src/pet/utils/other.py @@ -46,13 +46,36 @@ def _set_trainable(model): param.requires_grad = False -# def fsdp_auto_wrap_policy(): -# from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy, lambda_auto_wrap_policy, _or_policy +def fsdp_auto_wrap_policy(model): + import functools + import os -# def lambda_policy(module): -# if len(module.named_children()) != 0 and -# return True -# return False + from accelerate import FullyShardedDataParallelPlugin + from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy, transformer_auto_wrap_policy + from ..tuners import PrefixEncoder, PromptEmbedding, PromptEncoder -# pass + def lambda_policy_fn(module): + if ( + len(module.named_children()) == 0 + and getattr(module, "weight", None) is not None + and module.weight.requires_grad + ): + return True + return False + + lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn) + transformer_wrap_policy = functools.partial( + transformer_auto_wrap_policy, + transformer_layer_cls=( + PrefixEncoder, + PromptEncoder, + PromptEmbedding, + FullyShardedDataParallelPlugin.get_module_class_from_name( + model, os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "") + ), + ), + ) + + auto_wrap_policy = functools.partial(_or_policy, policies=[lambda_policy, transformer_wrap_policy]) + return auto_wrap_policy