From e3d6568a1978233902eecc921a7b0818854f5785 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 1 Dec 2022 14:34:09 +0530 Subject: [PATCH 01/13] adding detailed docs, refactor and fixes --- setup.py | 3 +- src/pet/mapping.py | 15 ++++ src/pet/pet_model.py | 141 ++++++++++++++++++++++++++++++-- src/pet/tuners/lora.py | 63 ++++++++++++-- src/pet/tuners/p_tuning.py | 55 ++++++++++++- src/pet/tuners/prefix_tuning.py | 34 ++++++++ src/pet/tuners/prompt_tuning.py | 51 +++++++++++- src/pet/utils/__init__.py | 2 +- src/pet/utils/config.py | 22 ++++- src/pet/utils/other.py | 14 ++++ src/pet/utils/save_and_load.py | 15 ++++ 11 files changed, 391 insertions(+), 24 deletions(-) diff --git a/setup.py b/setup.py index 6f74e46..fcd1b39 100644 --- a/setup.py +++ b/setup.py @@ -39,9 +39,10 @@ setup( "packaging>=20.0", "psutil", "pyyaml", - "torch>=1.4.0", + "torch>=1.13.0", "transformers", "accelerate", + "loralib", ], extras_require=extras, classifiers=[ diff --git a/src/pet/mapping.py b/src/pet/mapping.py index 0cb2210..1f2780f 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -34,6 +34,13 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { def get_pet_config(config_dict): + """ + Returns a PET config object from a dictionary. + + Args: + config_dict (:obj:`Dict[str, Any]`): + """ + return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict) @@ -93,6 +100,14 @@ def _prepare_lora_config(pet_config, model_config): def get_pet_model(model, pet_config): + """ + Returns a PET model object from a model and a config. + + Args: + model (:obj:`transformers.PreTrainedModel`): + pet_config (:obj:`transformers.PETConfig`): + """ + model_config = model.config.to_dict() if pet_config.pet_type != PETType.LORA: pet_config = _prepare_prompt_learning_config(pet_config, model_config) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 1398146..e1805fa 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -7,10 +7,30 @@ 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 +from .utils import PETConfig, PETType, TaskType, shift_tokens_right, _set_trainable class PETModel(torch.nn.Module): + """ + Parameter Efficient Tuning Model. Base model encompassing various PET methods. + + Args: + model (:obj:`PreTrainedModel`): The base transformer model used for PET. + pet_config (:obj:`PETConfig`): The configuration of the PET model. + + + 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 + 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`. + """ + def __init__(self, model, pet_config: PETConfig): super().__init__() self.pet_config = pet_config @@ -54,6 +74,10 @@ class PETModel(torch.nn.Module): ).long() 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`. + """ prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.base_model.device) if self.pet_config.pet_type == PETType.PREFIX_TUNING: prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens] @@ -61,6 +85,10 @@ class PETModel(torch.nn.Module): return prompt_embeddings[0].detach().cpu() def get_prompt(self, batch_size): + """ + 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: prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens] @@ -92,6 +120,9 @@ class PETModel(torch.nn.Module): return prompts def print_trainable_parameters(self): + """ + Prints the number of trainable parameters in the model. + """ trainable_params = 0 all_param = 0 for _, param in self.named_parameters(): @@ -104,6 +135,41 @@ class PETModel(torch.nn.Module): class PETModelForSequenceClassification(PETModel): + """ + PET model for sequence classification tasks. + + Args: + model (:obj:`PreTrainedModel`): Base transformer model + 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. + + 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 + } + >>> 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): super().__init__(model, pet_config) self.config = self.base_model.config @@ -114,6 +180,9 @@ class PETModelForSequenceClassification(PETModel): self.cls_layer_name = name break + # to make sure classifier layer is trainable + _set_trainable(self.base_model) + def forward( self, input_ids=None, @@ -160,7 +229,7 @@ class PETModelForSequenceClassification(PETModel): ) if self.pet_config.pet_type == PETType.PREFIX_TUNING: - return self.prefix_tuning_forward(input_ids=input_ids, **kwargs) + return self._prefix_tuning_forward(input_ids=input_ids, **kwargs) else: if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( @@ -176,7 +245,7 @@ class PETModelForSequenceClassification(PETModel): inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1) return self.base_model(inputs_embeds=inputs_embeds, **kwargs) - def prefix_tuning_forward( + def _prefix_tuning_forward( self, input_ids=None, attention_mask=None, @@ -249,9 +318,40 @@ class PETModelForSequenceClassification(PETModel): class PETModelForCausalLM(PETModel): + """ + PET model for Causal LM + + Args: + model (:obj:`PreTrainedModel`): Base transformer model + pet_config (:obj:`PETConfig`): PET config. + + + 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 + } + >>> 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): super().__init__(model, pet_config) - self.config = self.base_model.config def forward( self, @@ -318,9 +418,40 @@ class PETModelForCausalLM(PETModel): class PETModelForSeq2SeqLM(PETModel): + """ + PET model for Seq2Seq LM + + Args: + model (:obj:`PreTrainedModel`): Base transformer model + pet_config (:obj:`PETConfig`): PET config. + + + 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' + } + >>> 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): super().__init__(model, pet_config) - self.config = self.base_model.config def forward( self, diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index eb30bad..9c3899a 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -13,6 +13,21 @@ from ..utils import PETConfig @dataclass class LoRAConfig(PETConfig): + """ + This is the configuration class to store the configuration of a :class:`~pet.LoRA`. + + Args: + r: (:obj:`init`): LoRA attention dimension + target_modules (:obj: list of :obj: str): The names of the modules to apply LoRA to. + lora_alpha (:obj: float): The alpha parameter for LoRA scaling. + lora_dropout (:obj: float): The dropout probability for LoRA layers. + merge_weights (:obj: bool): + Whether to merge the weights of the LoRA layers with the base transformer model in `eval` mode. + fan_in_fan_out (:obj: bool): Set this to True if the layer to replace stores weight like (fan_in, fan_out) + enable_lora (:obj: list of :obj: bool): Used with `lora.MergedLinear`. + bias (:obj: str): Bias type for LoRA. Can be 'none', 'all' or 'lora_only' + """ + r: int = field(default=8, metadata={"help": "LoRA attention dimension"}) target_modules: Optional[list] = field(default=None, metadata={"help": "List of modules to replace with LoRA"}) lora_alpha: int = field(default=None, metadata={"help": "LoRA alpha"}) @@ -29,14 +44,44 @@ class LoRAConfig(PETConfig): class LoRAModel(torch.nn.Module): + """ + Creates Low Rank Adapter (LoRA) model from a pretrained transformers model. + + Args: + model (:obj:`transformers.PreTrainedModel`): The model to be adapted. + config (:obj:`LoRAConfig`): The configuration of the LoRA model. + + Returns: + :obj:`torch.nn.Module`: The LoRA model. + + 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) + + Attributes: + model (:obj:`transformers.PreTrainedModel`): The model to be adapted. + config (:obj:`LoRAConfig`): The configuration of the LoRA model. + """ + def __init__(self, config, model): super().__init__() self.config = config self.model = model - self.find_and_replace() + self._find_and_replace() mark_only_lora_as_trainable(self.model, self.config.bias) - def find_and_replace(self): + def _find_and_replace(self): kwargs = { "r": self.config.r, "lora_alpha": self.config.lora_alpha, @@ -47,23 +92,23 @@ class LoRAModel(torch.nn.Module): key_list = [key for key, _ in self.model.named_modules()] for key in key_list: if any(key.endswith(target_key) for target_key in self.config.target_modules): - parent, target, target_name = self.get_submodules(key) - # print(parent, target, target_name) + 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, **kwargs) + new_module = lora.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, **kwargs) - self.replace_module(parent, target_name, new_module, target) + new_module = lora.MergedLinear(in_features, out_features, bias=bias, **kwargs) + self._replace_module(parent, target_name, new_module, target) - def get_submodules(self, key): + def _get_submodules(self, key): parent = self.model.get_submodule(".".join(key.split(".")[:-1])) target_name = key.split(".")[-1] target = self.model.get_submodule(key) return parent, target, target_name - def replace_module(self, parent_module, child_name, new_module, old_module): + def _replace_module(self, parent_module, child_name, new_module, old_module): setattr(parent_module, child_name, new_module) new_module.weight = old_module.weight if old_module.bias is not None: diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py index bb5a5a8..0700ba8 100644 --- a/src/pet/tuners/p_tuning.py +++ b/src/pet/tuners/p_tuning.py @@ -14,21 +14,33 @@ class PromptEncoderReparameterizationType(str, enum.Enum): @dataclass class PromptEncoderConfig(PromptLearningConfig): + """ + This is the configuration class to store the configuration of a :class:`~pet.PromptEncoder`. + + Args: + encoder_reparameterization_type + (: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. + """ + encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field( default=PromptEncoderReparameterizationType.MLP, metadata={"help": "How to reparameterize the prompt encoder"}, ) encoder_hidden_size: int = field( default=None, - metadata={"help": "The hidden size of the prompt encoder reparameterization"}, + metadata={"help": "The hidden size of the prompt encoder"}, ) encoder_num_layers: int = field( default=2, - metadata={"help": "The number of layers of the prompt encoder reparameterization"}, + metadata={"help": "The number of layers of the prompt encoder"}, ) encoder_dropout: float = field( default=0.0, - metadata={"help": "The dropout of the prompt encoder reparameterization"}, + metadata={"help": "The dropout of the prompt encoder"}, ) @@ -37,6 +49,43 @@ class PromptEncoderConfig(PromptLearningConfig): class PromptEncoder(torch.nn.Module): """ The prompt encoder network that is used to generate the virtual token embeddings for p-tuning. + + Args: + config (:class:`PromptEncoderConfig`): The configuration of the prompt encoder. + + 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 + ) + >>> 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`): + 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`]): + The encoder type of the prompt encoder. + + + Input shape: (batch_size, total_virtual_tokens) + + Output shape: (batch_size, total_virtual_tokens, token_dim) """ def __init__(self, config): diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py index d709baf..cb4c575 100644 --- a/src/pet/tuners/prefix_tuning.py +++ b/src/pet/tuners/prefix_tuning.py @@ -8,6 +8,15 @@ from ..utils import PromptLearningConfig @dataclass class PrefixTuningConfig(PromptLearningConfig): + """ + This is the configuration class to store the configuration of a :class:`~pet.PrefixEncoder`. + + 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. + """ + encoder_hidden_size: int = field( default=None, metadata={"help": "The hidden size of the encoder"}, @@ -28,6 +37,31 @@ class PrefixEncoder(torch.nn.Module): r""" The torch.nn model to encode the prefix + Args: + config (:class:`PrefixTuningConfig`): The configuration of the prefix encoder. + + 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 + ) + >>> 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 + if :obj:`prefix_projection` is :obj:`True`. + prefix_projection (:obj:`bool`): Whether to project the prefix embeddings. + Input shape: (batch_size, num_virtual_tokens) Output shape: (batch_size, num_virtual_tokens, 2*layers*hidden) diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py index fdad421..5ee6ef1 100644 --- a/src/pet/tuners/prompt_tuning.py +++ b/src/pet/tuners/prompt_tuning.py @@ -15,6 +15,17 @@ class PromptTuningInit(str, enum.Enum): @dataclass 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_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. + Only used if `prompt_tuning_init` is `TEXT` + """ + prompt_tuning_init: Union[PromptTuningInit, str] = field( default=PromptTuningInit.RANDOM, metadata={"help": "How to initialize the prompt tuning parameters"}, @@ -34,6 +45,40 @@ class PromptTuningConfig(PromptLearningConfig): class PromptEmbedding(torch.nn.Module): + """ + The model to encode virtual tokens into prompt embeddings. + + Args: + config (:class:`PromptTuningConfig`): The configuration of the prompt embedding. + word_embeddings (:obj:`torch.nn.Module`): The word embeddings of the base transformer model. + + Attributes: + embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prompt embedding. + + 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", + 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) + + + Input Shape: (batch_size, total_virtual_tokens) + + Output Shape: (batch_size, total_virtual_tokens, token_dim) + """ + def __init__(self, config, word_embeddings): super().__init__() @@ -42,9 +87,9 @@ class PromptEmbedding(torch.nn.Module): if config.prompt_tuning_init == PromptTuningInit.TEXT: from transformers import AutoTokenizer - self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path) - self.init_text = config.prompt_tuning_init_text - init_token_ids = self.tokenizer(self.init_text)["input_ids"] + tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path) + init_text = config.prompt_tuning_init_text + init_token_ids = tokenizer(init_text)["input_ids"] # Trim or iterate until num_text_tokens matches total_virtual_tokens num_text_tokens = len(init_token_ids) if num_text_tokens > total_virtual_tokens: diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py index aa359db..582dad4 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 +from .other import bloom_model_postprocess_past_key_value, shift_tokens_right, _set_trainable 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 d989601..984016b 100644 --- a/src/pet/utils/config.py +++ b/src/pet/utils/config.py @@ -19,7 +19,12 @@ class TaskType(str, enum.Enum): @dataclass class PETConfig: """ - This is the configuration class to store the configuration of a :class:`~pet.PETModel`. + This is the base configuration class to store the configuration of a :class:`~pet.PETModel`. + + Args: + pet_type (:obj:Union[:class:`~pet.utils.config.PETType`, :obj:`str`]): The type of PET method to use. + task_type (:obj:Union[:class:`~pet.utils.config.TaskType`, :obj:`str`]): The type of task to perform. + inference_mode (:obj:`bool`, defaults to :obj:`False`): Whether to use the PET model in inference mode. """ pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"}) @@ -29,8 +34,21 @@ 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`]. + + Args: + num_virtual_tokens (:obj:`int`): The number of virtual tokens to use. + token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model. + num_transformer_submodules (:obj:`int`): The number of transformer submodules in the base transformer model. + num_attention_heads (:obj:`int`): The number of attention heads in the base transformer model. + num_layers (:obj:`int`): The number of layers in the base transformer model. + """ + num_virtual_tokens: int = field(default=None, metadata={"help": "Number of virtual tokens"}) - token_dim: int = field(default=None, metadata={"help": "Dimension of virtual tokens"}) + token_dim: int = field( + default=None, metadata={"help": "The hidden embedding dimension of the base transformer model"} + ) num_transformer_submodules: Optional[int] = field(default=1, metadata={"help": "Number of transformer submodules"}) num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"}) num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"}) diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py index ac8b329..c6274db 100644 --- a/src/pet/utils/other.py +++ b/src/pet/utils/other.py @@ -19,6 +19,11 @@ def bloom_model_postprocess_past_key_value(past_key_values): def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): """ Shift input ids one token to the right. + + Args: + input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`): input ids + pad_token_id (:obj:`int`): The id of the `padding` token. + decoder_start_token_id (:obj:`int`): The id of the `start` token. """ shifted_input_ids = input_ids.new_zeros(input_ids.shape) shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() @@ -30,3 +35,12 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) return shifted_input_ids + + +def _set_trainable(model): + if model.modules_to_save is not None: + for name, param in model.named_parameters(): + if any(module_name in name for module_name in model.modules_to_save): + param.requires_grad = True + else: + param.requires_grad = False diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py index e2eea05..1418dfa 100644 --- a/src/pet/utils/save_and_load.py +++ b/src/pet/utils/save_and_load.py @@ -4,6 +4,13 @@ from .config import PETType def get_pet_model_state_dict(model): + """ + Get the state dict of the PET model. + + Args: + model (:obj:`PETModel`): The PET model. + """ + if model.pet_config.pet_type == PETType.LORA: return lora_state_dict(model) else: @@ -19,6 +26,14 @@ def get_pet_model_state_dict(model): def set_pet_model_state_dict(model, pet_model_state_dict): + """ + Set the state dict of the PET model. + + Args: + model (:obj:`PETModel`): The PET model. + pet_model_state_dict (:obj:`dict`): The state dict of the PET model. + """ + model.load_state_dict(pet_model_state_dict, strict=False) if model.pet_config.pet_type != PETType.LORA: model.prompt_encoder.embedding.load_state_dict( From 870d36039074c5436cbe32e488244cb8048a923b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 1 Dec 2022 14:47:35 +0530 Subject: [PATCH 02/13] fix --- README.md | 22 ++++++++++++++++++++++ src/pet/mapping.py | 1 + 2 files changed, 23 insertions(+) diff --git a/README.md b/README.md index 0e1c270..0b4fe38 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,28 @@ Supported methods: 3. P-Tuning 4. Prompt Tuning +## Getting started + +```python +from transformers import AutoModelForSeq2SeqLM +from pet import get_pet_config,get_pet_model +model_name_or_path = "bigscience/mt0-large" +tokenizer_name_or_path = "bigscience/mt0-large" + +config = { + "pet_type":"LORA", + "task_type":"SEQ_2_SEQ_LM", + "r": 8, + "lora_alpha": 32, + "lora_dropout": 0.1 +} +pet_config = get_pet_config(config) + +model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) +model = get_pet_model(model, pet_config) +model.print_trainable_parameters() +``` + ## Models support matrix ### Sequence Classification diff --git a/src/pet/mapping.py b/src/pet/mapping.py index 1f2780f..daeb650 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -18,6 +18,7 @@ PET_TYPE_TO_CONFIG_MAPPING = { TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { "t5": ["q", "v"], + "mt5": ["q", "v"], "bart": ["q_proj", "v_proj"], "gpt2": ["c_attn"], "bloom": ["query_key_value"], From 67d980f13beafeb598237671ee93e70bc79892b4 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 1 Dec 2022 14:51:26 +0530 Subject: [PATCH 03/13] fix --- README.md | 1 + src/pet/pet_model.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 0b4fe38..8babeed 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ pet_config = get_pet_config(config) model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) model = get_pet_model(model, pet_config) model.print_trainable_parameters() +# output: ``` ## Models support matrix diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index e1805fa..a865623 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -35,6 +35,7 @@ class PETModel(torch.nn.Module): super().__init__() self.pet_config = pet_config self.base_model = model + self.config = self.base_model.config self.modules_to_save = None if pet_config.pet_type != PETType.LORA: self._setup_prompt_encoder() @@ -172,7 +173,6 @@ class PETModelForSequenceClassification(PETModel): def __init__(self, model, pet_config: PETConfig): super().__init__(model, pet_config) - self.config = self.base_model.config self.modules_to_save = ["classifier"] for name, module in self.base_model.named_children(): From 8920caeb1b3986d9c172a5e8cdbd518a5a80f465 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 1 Dec 2022 15:28:45 +0530 Subject: [PATCH 04/13] fix --- src/pet/pet_model.py | 7 +++++++ src/pet/tuners/lora.py | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index a865623..bcc97b3 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -134,6 +134,13 @@ class PETModel(torch.nn.Module): f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" ) + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + return getattr(self.base_model, name) + class PETModelForSequenceClassification(PETModel): """ diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index 9c3899a..d0e9c71 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -116,3 +116,10 @@ class LoRAModel(torch.nn.Module): def forward(self, *args, **kwargs): return self.model(*args, **kwargs) + + def __getattr__(self, name: str): + """Forward missing attributes to the wrapped module.""" + try: + return super().__getattr__(name) # defer to nn.Module's logic + except AttributeError: + return getattr(self.model, name) From 5f2062e90b1a28fd578ac2e80fba3592e5a70542 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 1 Dec 2022 18:51:05 +0530 Subject: [PATCH 05/13] Update mapping.py --- src/pet/mapping.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pet/mapping.py b/src/pet/mapping.py index daeb650..7d7d3c6 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -28,6 +28,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = { "gpt_neo": ["q_proj", "v_proj"], "bert": ["query", "value"], "roberta": ["query", "value"], + "xlm-roberta": ["query", "value"], "electra": ["query", "value"], "deberta-v2": ["query_proj", "value_proj"], "deberta": ["in_proj"], From e63f47ca520a95ddc9a91ec479efa6451472842c Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Thu, 1 Dec 2022 21:15:01 +0530 Subject: [PATCH 06/13] fix --- src/pet/pet_model.py | 4 ++-- src/pet/utils/other.py | 12 ++++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index bcc97b3..fff3dd3 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -46,10 +46,10 @@ class PETModel(torch.nn.Module): num_transformer_submodules = 0 transformer_backbone = None for name, module in self.base_model.named_children(): + for param in module.parameters(): + param.requires_grad = False if isinstance(module, PreTrainedModel): # Make sure to freeze Tranformers model - for param in module.parameters(): - param.requires_grad = False if transformer_backbone is None: transformer_backbone = module self.transformer_backbone_name = name diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py index c6274db..a13b4f6 100644 --- a/src/pet/utils/other.py +++ b/src/pet/utils/other.py @@ -44,3 +44,15 @@ def _set_trainable(model): param.requires_grad = True else: 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 lambda_policy(module): +# if len(module.named_children()) != 0 and +# return True +# return False + + +# pass From 6c21f3bf380cea3f420abbfdfb4ce44ea871e78a Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 2 Dec 2022 09:15:02 +0530 Subject: [PATCH 07/13] =?UTF-8?q?FSDP=20auto=20wrap,=20=F0=9F=90=9B=20fixe?= =?UTF-8?q?s=20and=20style?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/pet/pet_model.py | 107 +++++--------- src/pet/tuners/lora.py | 237 +++++++++++++++++++++++++++++--- src/pet/tuners/p_tuning.py | 36 ++--- src/pet/tuners/prefix_tuning.py | 20 +-- src/pet/tuners/prompt_tuning.py | 20 +-- src/pet/utils/__init__.py | 2 +- src/pet/utils/config.py | 3 +- src/pet/utils/other.py | 37 ++++- 8 files changed, 317 insertions(+), 145 deletions(-) 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 From 7a9c4a62875efd56d5c2fefa9e6c4eeacea2203b Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 2 Dec 2022 09:56:29 +0530 Subject: [PATCH 08/13] fix --- src/pet/utils/other.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py index 0b63248..b7a81d5 100644 --- a/src/pet/utils/other.py +++ b/src/pet/utils/other.py @@ -57,7 +57,7 @@ def fsdp_auto_wrap_policy(model): def lambda_policy_fn(module): if ( - len(module.named_children()) == 0 + len(list(module.named_children())) == 0 and getattr(module, "weight", None) is not None and module.weight.requires_grad ): From 16ba5fcf035c84f50c4cc92d201111507e642c5c Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 2 Dec 2022 11:05:30 +0530 Subject: [PATCH 09/13] FSDP cpu offloading fix --- src/pet/pet_model.py | 30 +++++++++--------------------- 1 file changed, 9 insertions(+), 21 deletions(-) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 4ed76a1..8e96c04 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -79,7 +79,7 @@ class PETModel(torch.nn.Module): 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) + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to("cuda") if self.pet_config.pet_type == PETType.PREFIX_TUNING: prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens] prompt_embeddings = self.prompt_encoder(prompt_tokens) @@ -89,7 +89,7 @@ class PETModel(torch.nn.Module): """ 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) + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to("cuda") if self.pet_config.pet_type == PETType.PREFIX_TUNING: prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens] if self.pet_config.inference_mode: @@ -207,9 +207,7 @@ class PETModelForSequenceClassification(PETModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( - self.base_model.device - ) + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to("cuda") 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.") @@ -230,7 +228,7 @@ class PETModelForSequenceClassification(PETModel): if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( ( - torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.base_model.device), + torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to("cuda"), kwargs["token_type_ids"], ), dim=1, @@ -364,9 +362,7 @@ class PETModelForCausalLM(PETModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( - self.base_model.device - ) + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to("cuda") attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: @@ -393,9 +389,7 @@ class PETModelForCausalLM(PETModel): inputs_embeds = self.word_embeddings(input_ids) # concat prompt labels if labels is not None: - prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to( - self.base_model.device - ) + prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to("cuda") kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1) @@ -459,9 +453,7 @@ class PETModelForSeq2SeqLM(PETModel): batch_size = input_ids.shape[0] if decoder_attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( - self.base_model.device - ) + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to("cuda") decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: @@ -497,15 +489,11 @@ class PETModelForSeq2SeqLM(PETModel): if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( - self.base_model.device - ) + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to("cuda") kwargs["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.pet_config.num_virtual_tokens), -100).to( - self.base_model.device - ) + prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to("cuda") kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) inputs_embeds = torch.cat((prompts[:, : self.pet_config.num_virtual_tokens], inputs_embeds), dim=1) From d8a25bf6e69009b82dbd6fda090e313ce5459bab Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 2 Dec 2022 11:09:38 +0530 Subject: [PATCH 10/13] refactor --- src/pet/pet_model.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 8e96c04..dbe2b6e 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -41,6 +41,7 @@ class PETModel(torch.nn.Module): self._setup_prompt_encoder() else: self.base_model = LoRAModel(pet_config, model) + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def _setup_prompt_encoder(self): num_transformer_submodules = 0 @@ -79,7 +80,7 @@ class PETModel(torch.nn.Module): 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("cuda") + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.device) if self.pet_config.pet_type == PETType.PREFIX_TUNING: prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens] prompt_embeddings = self.prompt_encoder(prompt_tokens) @@ -89,7 +90,7 @@ class PETModel(torch.nn.Module): """ 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("cuda") + prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.device) if self.pet_config.pet_type == PETType.PREFIX_TUNING: prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens] if self.pet_config.inference_mode: @@ -207,7 +208,7 @@ class PETModelForSequenceClassification(PETModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to("cuda") + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device) 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.") @@ -228,7 +229,7 @@ class PETModelForSequenceClassification(PETModel): if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( ( - torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to("cuda"), + torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.device), kwargs["token_type_ids"], ), dim=1, @@ -362,7 +363,7 @@ class PETModelForCausalLM(PETModel): batch_size = input_ids.shape[0] if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to("cuda") + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: @@ -389,7 +390,7 @@ class PETModelForCausalLM(PETModel): inputs_embeds = self.word_embeddings(input_ids) # concat prompt labels if labels is not None: - prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to("cuda") + prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(self.device) kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1) @@ -453,7 +454,7 @@ class PETModelForSeq2SeqLM(PETModel): batch_size = input_ids.shape[0] if decoder_attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to("cuda") + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device) decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: @@ -489,11 +490,11 @@ class PETModelForSeq2SeqLM(PETModel): if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to("cuda") + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device) kwargs["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.pet_config.num_virtual_tokens), -100).to("cuda") + prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(self.device) kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) inputs_embeds = torch.cat((prompts[:, : self.pet_config.num_virtual_tokens], inputs_embeds), dim=1) From 3862620467532bd224e7cfe852cfa92bf6d74c54 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 2 Dec 2022 11:23:00 +0530 Subject: [PATCH 11/13] fix --- README.md | 25 +++++++++++++++++++++++-- src/pet/mapping.py | 4 ++-- 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 8babeed..95eeccf 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ Supported methods: ```python from transformers import AutoModelForSeq2SeqLM -from pet import get_pet_config,get_pet_model +from pet import get_pet_config, get_pet_model model_name_or_path = "bigscience/mt0-large" tokenizer_name_or_path = "bigscience/mt0-large" @@ -28,9 +28,30 @@ pet_config = get_pet_config(config) model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path) model = get_pet_model(model, pet_config) model.print_trainable_parameters() -# output: +# output: trainable params: 2359296 || all params: 1231940608 || trainable%: 0.19151053100118282 ``` +## PET + 🤗 Accelerate + +PET models work with 🤗 Accelerate out of the box. +For scaling to large models, you can leverage 🤗 Accelerate's PyTorch FSDP integration as shown below. +PyTorch FSDP shards parameters, gradients and optimizer states across data parallel workers which enables +large language models to fit on available hardware. +It also supports CPU offloading to further enable distributed training at scale. +The support for DeepSpeed ZeRO Stage-3 is currently in backlog. + +```python +from pet.utils.other import fsdp_auto_wrap_policy + +... + +if accelerator.state.fsdp_plugin is not None: + accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) + +model = accelerator.prepare(model) +``` + + ## Models support matrix ### Sequence Classification diff --git a/src/pet/mapping.py b/src/pet/mapping.py index 7d7d3c6..115372e 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -82,8 +82,8 @@ def _prepare_prompt_learning_config(pet_config, model_config): raise ValueError("Please specify `num_attention_heads` in `pet_config`") pet_config.num_attention_heads = num_attention_heads - if pet_config.encoder_hidden_size is None: - pet_config.encoder_hidden_size = token_dim + if getattr(pet_config, "encoder_hidden_size", None) is None: + setattr(pet_config, "encoder_hidden_size", token_dim) return pet_config From d89882a55850b0c258dee0eef1950b191bb5e307 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 2 Dec 2022 11:24:26 +0530 Subject: [PATCH 12/13] fix --- src/pet/tuners/prompt_tuning.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py index 20e3a2e..5b6c313 100644 --- a/src/pet/tuners/prompt_tuning.py +++ b/src/pet/tuners/prompt_tuning.py @@ -77,7 +77,7 @@ class PromptEmbedding(torch.nn.Module): super().__init__() total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules - self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"]) + self.embedding = torch.nn.Embedding(total_virtual_tokens, config.token_dim) if config.prompt_tuning_init == PromptTuningInit.TEXT: from transformers import AutoTokenizer From ae609d680e967b6bc9999415a88bbd30abb2ff9a Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 2 Dec 2022 16:21:16 +0530 Subject: [PATCH 13/13] add examples and update README --- Makefile | 2 +- README.md | 45 +- examples/pet_lora_seq2seq.ipynb | 4815 ++++++++++++++++++ examples/pet_lora_seq2seq_accelerate_fsdp.py | 137 + examples/pet_prefix_tuning_seq2seq.ipynb | 1963 +++++++ 5 files changed, 6958 insertions(+), 4 deletions(-) create mode 100644 examples/pet_lora_seq2seq.ipynb create mode 100644 examples/pet_lora_seq2seq_accelerate_fsdp.py create mode 100644 examples/pet_prefix_tuning_seq2seq.ipynb diff --git a/Makefile b/Makefile index e1c15c5..888056d 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ .PHONY: quality style test docs -check_dirs := src +check_dirs := src examples # Check that source code meets quality standards diff --git a/README.md b/README.md index 95eeccf..1dfafca 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ # 🤗 PET -Parameter-Efficient Tuning. Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP. +Parameter-Efficient Tuning methods enable . Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP. Supported methods: @@ -38,19 +38,56 @@ For scaling to large models, you can leverage 🤗 Accelerate's PyTorch FSDP int PyTorch FSDP shards parameters, gradients and optimizer states across data parallel workers which enables large language models to fit on available hardware. It also supports CPU offloading to further enable distributed training at scale. -The support for DeepSpeed ZeRO Stage-3 is currently in backlog. ```python from pet.utils.other import fsdp_auto_wrap_policy ... -if accelerator.state.fsdp_plugin is not None: +if os.environ.get("ACCELERATE_USE_FSDP", None) is not None: accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) model = accelerator.prepare(model) ``` +Example of parameter efficient tuning with `mt0-xxl` base model using 🤗 Accelerate is provided in `~examples/pet_lora_seq2seq_accelerate_fsdp.py`. +1. First run `accelerate config --config_file fsdp_config.yaml` and answer the questionaire. +Below are the contents of the config file. +``` +command_file: null +commands: null +compute_environment: LOCAL_MACHINE +deepspeed_config: {} +distributed_type: FSDP +downcast_bf16: 'no' +dynamo_backend: 'NO' +fsdp_config: + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_backward_prefetch_policy: BACKWARD_PRE + fsdp_offload_params: true + fsdp_sharding_strategy: 1 + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_transformer_layer_cls_to_wrap: T5Block +gpu_ids: null +machine_rank: 0 +main_process_ip: null +main_process_port: null +main_training_function: main +megatron_lm_config: {} +mixed_precision: 'no' +num_machines: 1 +num_processes: 2 +rdzv_backend: static +same_network: true +tpu_name: null +tpu_zone: null +use_cpu: false +``` +2. run the below command to launch example script +``` +accelerate launch --config_file fsdp_config.yaml examples/pet_lora_seq2seq_accelerate_fsdp.py +``` + ## Models support matrix @@ -85,5 +122,7 @@ model = accelerator.prepare(model) ## Caveats: 1. Doesn't work currently with DeeSpeed ZeRO Stage-3. Extending support with DeeSpeed ZeRO Stage-3 is in backlog. +2. When using `P_TUNING` or `PROMPT_TUNING` with `SEQ_2_SEQ` task, remember to remove the `num_virtual_token` virtual prompt predictions from the left side of the model outputs during evaluations. + diff --git a/examples/pet_lora_seq2seq.ipynb b/examples/pet_lora_seq2seq.ipynb new file mode 100644 index 0000000..f8d86c6 --- /dev/null +++ b/examples/pet_lora_seq2seq.ipynb @@ -0,0 +1,4815 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 17, + "id": "5f93b7d1", + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoModelForSeq2SeqLM\n", + "from pet import get_pet_config,get_pet_model, get_pet_model_state_dict\n", + "import torch\n", + "from datasets import load_dataset\n", + "import os\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "from transformers import AutoTokenizer\n", + "from torch.utils.data import DataLoader\n", + "from transformers import default_data_collator,get_linear_schedule_with_warmup\n", + "from tqdm import tqdm\n", + "from datasets import load_dataset\n", + "\n", + "device = \"cuda\"\n", + "model_name_or_path = \"bigscience/mt0-large\"\n", + "tokenizer_name_or_path = \"bigscience/mt0-large\"\n", + "\n", + "config = {\n", + " \"pet_type\":\"LORA\",\n", + " \"task_type\":\"SEQ_2_SEQ_LM\",\n", + " \"r\":16,\n", + " \"lora_alpha\": 32,\n", + " \"lora_dropout\": 0.1\n", + "}\n", + "checkpoint_name = \"financial_sentiment_analysis_lora_v1.pt\"\n", + "text_column = \"sentence\"\n", + "label_column = \"text_label\"\n", + "max_length=128\n", + "lr = 1e-3\n", + "num_epochs = 3\n", + "batch_size=8\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "8d0850ac", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "trainable params: 4718592 || all params: 1234299904 || trainable%: 0.38228893842642636\n" + ] + }, + { + "data": { + "text/plain": [ + "PETModelForSeq2SeqLM(\n", + " (base_model): LoRAModel(\n", + " (model): MT5ForConditionalGeneration(\n", + " (shared): Embedding(250112, 1024)\n", + " (encoder): T5Stack(\n", + " (embed_tokens): Embedding(250112, 1024)\n", + " (block): ModuleList(\n", + " (0): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (relative_attention_bias): Embedding(32, 16)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (1): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (2): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (3): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (4): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (5): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (6): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (7): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (8): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (9): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (10): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (11): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (12): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (13): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (14): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (15): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (16): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (17): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (18): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (19): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (20): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (21): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (22): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (23): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " )\n", + " (final_layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (decoder): T5Stack(\n", + " (embed_tokens): Embedding(250112, 1024)\n", + " (block): ModuleList(\n", + " (0): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (relative_attention_bias): Embedding(32, 16)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (1): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (2): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (3): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (4): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (5): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (6): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (7): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (8): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (9): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (10): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (11): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (12): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (13): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (14): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (15): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (16): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (17): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (18): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (19): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (20): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (21): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (22): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " (23): T5Block(\n", + " (layer): ModuleList(\n", + " (0): T5LayerSelfAttention(\n", + " (SelfAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (1): T5LayerCrossAttention(\n", + " (EncDecAttention): T5Attention(\n", + " (q): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (k): Linear(in_features=1024, out_features=1024, bias=False)\n", + " (v): Linear(\n", + " in_features=1024, out_features=1024, bias=False\n", + " (lora_dropout): Dropout(p=0.1, inplace=False)\n", + " (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n", + " (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n", + " )\n", + " (o): Linear(in_features=1024, out_features=1024, bias=False)\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (2): T5LayerFF(\n", + " (DenseReluDense): T5DenseGatedActDense(\n", + " (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n", + " (wo): Linear(in_features=2816, out_features=1024, bias=False)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " (act): NewGELUActivation()\n", + " )\n", + " (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " )\n", + " )\n", + " )\n", + " (final_layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n", + " (dropout): Dropout(p=0.1, inplace=False)\n", + " )\n", + " (lm_head): Linear(in_features=1024, out_features=250112, bias=False)\n", + " )\n", + " )\n", + ")" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# creating model\n", + "pet_config = get_pet_config(config)\n", + "\n", + "model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)\n", + "model = get_pet_model(model, pet_config)\n", + "model.print_trainable_parameters()\n", + "model" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "4ee2babf", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/sourab/miniconda3/envs/ml/lib/python3.10/site-packages/huggingface_hub/utils/_deprecation.py:97: FutureWarning: Deprecated argument(s) used in 'dataset_info': token. Will not be supported from version '0.12'.\n", + " warnings.warn(message, FutureWarning)\n", + "Found cached dataset financial_phrasebank (/home/sourab/.cache/huggingface/datasets/financial_phrasebank/sentences_allagree/1.0.0/550bde12e6c30e2674da973a55f57edde5181d53f5a5a34c1531c53f93b7e141)\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "6de075f8208349108291ac5ab7f5c980", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + " 0%| | 0/1 [00:00