wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-09 11:28:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

FSDP auto wrap, 🐛 fixes and style

Browse Source
This commit is contained in:
Sourab Mangrulkar
2022-12-02 09:15:02 +05:30
parent e63f47ca52
commit 6c21f3bf38
8 changed files with 317 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pet/pet_model.py
+37 -70
View File
@@ -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):
src/pet/tuners/lora.py
+219 -18
View File
@@ -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
src/pet/tuners/p_tuning.py
+14 -22
View File
@@ -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.
src/pet/tuners/prefix_tuning.py
+7 -13
View File
@@ -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.
src/pet/tuners/prompt_tuning.py
+7 -13
View File
@@ -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)
src/pet/utils/__init__.py
+1 -1
View File
@@ -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
src/pet/utils/config.py
+2 -1
View File
@@ -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.
src/pet/utils/other.py
+30 -7
View File
@@ -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