From 8febadd096deb206166c75894d748c9cd916a00a Mon Sep 17 00:00:00 2001 From: 51616 Date: Fri, 20 Dec 2024 15:09:32 +0000 Subject: [PATCH] embedding encoder + mean agg --- hyperlora/modeling_utils.py | 208 +++++++++++++++++++++++++++--------- 1 file changed, 159 insertions(+), 49 deletions(-) diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index 692f8f5..58e1749 100644 --- a/hyperlora/modeling_utils.py +++ b/hyperlora/modeling_utils.py @@ -1,75 +1,158 @@ +import logging from dataclasses import dataclass, field +from enum import Enum +from functools import partial from typing import Any, Optional, Tuple, Union import torch -from torch import nn +from jaxtyping import Float, Integer +from peft import LoraConfig +from torch import Tensor, nn from transformers import PreTrainedModel from transformers.modeling_outputs import ModelOutput -# @dataclass -# class ModelArguments: -# model_name_or_path: str = field(default="mistralai/Mistral-7B-v0.1") -# lora_r: int = field(default=128, metadata={"help": "lora rank"}) -# lora_dropout: float = field(default=0.05, metadata={"help": "lora dropout"}) -# train: bool = field( -# default=True, -# metadata={ -# "help": "if true, the model ckpt will be initialized for training; else, it's for inference" -# }, -# ) +from model_loading import get_lora_config, get_model_and_tokenizer + +logger = logging.getLogger(__name__) -class ModulatedPretrainedModel(nn.Module): - # def __init__(self, model_args, training_args, lora_config): - # super().__init__() - # self.model_args = model_args - # self.training_args = training_args - # self.model_name = model_args.model_name_or_path - # self.lora_config = lora_config - # self.model = AutoModelForCausalLM.from_pretrained( - # self.model_name, - # torch_dtype=torch.float16 if training_args.bf16 is False else torch.bfloat16, - # use_flash_attention_2=True, - # resume_download=True, - # ) - def __init__(self, base_model: PreTrainedModel): +def inv_bool_mask(m: Integer[Tensor, "bs seq_len"]) -> Integer[Tensor, "bs seq_len 1"]: + return (m - 1).bool().unsqueeze(-1) + + +# TODO: implement Perceiver +class Perceiver(nn.Module): ... + + +class MeanPool(nn.Module): + def forward( + self, + features: Float[Tensor, "bs seq_len feature_dim"], + attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, + ): + if attn_mask is not None: + features = features.masked_fill(inv_bool_mask(attn_mask), 0) + return features.sum(dim=1, keepdim=True) / attn_mask.sum( + dim=1, keepdim=True + ).unsqueeze(2) + + +class MaxPool(nn.Module): + def forward( + self, + features: Float[Tensor, "bs seq_len feature_dim"], + attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, + ): + if attn_mask is not None: + features = features.masked_fill(inv_bool_mask(attn_mask), -float("inf")) + return torch.max(features, dim=1, keepdim=True) + + +class LastTokenPool(nn.Module): + def forward( + self, + features: Float[Tensor, "bs seq_len feature_dim"], + attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, + ): + last_hidden_states = ( + features["hidden_states"][-1] + if "hidden_states" in features + else features["last_hidden_state"] + ) + + left_padding = attn_mask[:, -1].sum() == attn_mask.shape[0] + if left_padding: + return last_hidden_states[:, -1] + else: + sequence_lengths = attn_mask.sum(dim=1) - 1 + batch_size = last_hidden_states.shape[0] + return last_hidden_states[ + torch.arange(batch_size, device=last_hidden_states.device), + sequence_lengths, + ] + + +AGGREGATOR = Enum("AGGREGATOR", ["MEAN", "MAX", "PERCEIVER", "LAST_TOKEN"]) + +AGGREGATOR_CLS = { + AGGREGATOR.MEAN: MeanPool, + AGGREGATOR.MAX: MaxPool, + AGGREGATOR.LAST_TOKEN: LastTokenPool, + AGGREGATOR.PERCEIVER: Perceiver, +} + + +class HyperLoRA(nn.Module): + def __init__( + self, + lora_config: LoraConfig, + aggregator: AGGREGATOR, + aggregator_kwargs: Optional[dict[str, Any]] = None, + ): super().__init__() - self.base_model = base_model - # NOTE: we can reduce extractor_lm's depth by - # e.g., self.extractor_lm.model.layers = self.extractor_lm.model.layers[:num_extractor_layers] - # or just use the token embeddings extractor_lm.embed_tokens(input_ids) - self.extractor_lm = ... - self.hypernet = ... + self.lora_config = lora_config + self.aggregator = AGGREGATOR_CLS[aggregator](**(aggregator_kwargs or {})) def forward( self, - ctx_ids: Optional[torch.LongTensor] = None, - ctx_attention_mask: Optional[torch.LongTensor] = None, + features: Float[Tensor, "bs seq_len feature_dim"], + attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, + ) -> Float[Tensor, "bs num_modules num_layers lora_r in_out_dim"]: + + emb = self.aggregator(features, attn_mask) # [bs, n_features, feature_dim] + # loras = self.layers(emb) + # return loras + return emb + + +class ModulatedPretrainedModel(nn.Module): + def __init__(self, base_model: PreTrainedModel, hypernet: HyperLoRA): + super().__init__() + # self.base_model = base_model + self.device = base_model.device + # NOTE: we can reduce extractor_lm's depth by + # e.g., self.extractor_lm.model.layers = self.extractor_lm.model.layers[:num_extractor_layers] + # or just use the token embeddings extractor_lm.embed_tokens(input_ids) + + # HACK: hardcode the embedding layer for now + # TODO: add explicit encoder + self.encoder = base_model.get_input_embeddings() + + # register base_model as a submodule + self.register_module("base_model", base_model) + self.register_module("hypernet", hypernet) + + def get_ctx_features( + self, + input_ids: Integer[Tensor, "bs seq_len"], + attention_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, + ): + if isinstance(self.encoder, nn.Embedding): + features = self.encoder(input_ids) + else: + features = self.encoder(input_ids=input_ids, attention_mask=attention_mask) + return features + + def forward( + self, + ctx_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, + ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_length"]] = None, **model_inputs_kwargs: dict[str, Any], ) -> Union[tuple, ModelOutput]: - """Forward pass of the modulated model. + """Forward pass of the modulated model.""" - Args: - ctx_ids (torch.LongTensor, optional): Token IDs to be input to the hypernet - for generating LoRA parameters. Shape: (batch_size, ctx_length) - input_ids (torch.LongTensor, optional): Token IDs to be input to the base - language model. Shape: (batch_size, sequence_length) - labels (torch.LongTensor, optional): Labels for computing the language modeling - loss. Shape: (batch_size, sequence_length) - - Returns: - dict: Dictionary containing: - - loss (torch.Tensor): Language modeling loss if labels are provided - - logits (torch.Tensor): Output logits from the model - """ if ctx_ids is None: + logger.warning( + "No context ids provided, using the base model for the forward pass" + ) model_outputs = self.base_model(**model_inputs_kwargs) # model_outputs.generated_loras = None return model_outputs loss = ... logits = ... - generated_loras = ... + features = self.get_ctx_features(ctx_ids, ctx_attn_mask) + generated_loras = self.hypernet(features, ctx_attn_mask) # apply lora hook to the base model self.apply_lora_hook(generated_loras) @@ -80,3 +163,30 @@ class ModulatedPretrainedModel(nn.Module): def apply_lora_hook(self, generated_loras): pass + + +if __name__ == "__main__": + model_name = "meta-llama/Llama-3.2-1B-Instruct" + base_model, tokenizer = get_model_and_tokenizer( + model_name, + train=True, + requires_grad=False, + peft_config=get_lora_config(model_name), + ) + print(base_model) + hypernet = HyperLoRA(base_model.peft_config, aggregator=AGGREGATOR.MEAN) + model = ModulatedPretrainedModel(base_model, hypernet) + print(model) + + ctx_msg = "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum." + ctx_inputs = tokenizer(ctx_msg, return_tensors="pt").to(model.device) + ctx_features = model.get_ctx_features(**ctx_inputs) + ctx_attn_mask = ctx_inputs["attention_mask"] + print(ctx_features.shape) + + agg_features = hypernet.aggregator(ctx_features, ctx_attn_mask) + print(agg_features.shape) + + hnetout = hypernet(ctx_features, ctx_attn_mask) + print(hnetout.shape) + breakpoint()