import logging from contextlib import contextmanager from dataclasses import dataclass, field from enum import Enum from functools import partial, wraps from math import log, pi from typing import Any, Iterable, Optional, Tuple, Union import torch import torch.nn.functional as F from configs import AggregatorArguments, HypernetArguments, CtxEncoderArguments from einops import rearrange, repeat, unpack, einsum from einops.layers.torch import EinMix as Mix from einops.layers.torch import Reduce from hooks import add_generated_lora_hook, remove_hook_handles from jaxtyping import Float, Integer from model_loading import get_lora_config, get_model, get_model_and_tokenizer from peft import ( get_peft_config, load_peft_weights, LoraConfig, PeftConfig, PeftModel, LoraRuntimeConfig, ) from peft.utils import PeftType, TaskType from peft.tuners._buffer_dict import BufferDict from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists from pooling import POOL_FN, get_pooling_fn from torch import Tensor, nn from transformers import PerceiverConfig, PerceiverModel, PreTrainedModel from transformers.models.perceiver.modeling_perceiver import ( PerceiverBasicDecoder, ) from transformers.modeling_outputs import ModelOutput from utils import ( get_lora_module_names, get_num_layers, get_peft_in_out_features, get_base_model, ) logger = logging.getLogger() class AGGREGATOR_TYPE(str, Enum): POOLER = "pooler" PERCEIVER = "perceiver" @dataclass class AggregatorConfig: aggregator_type: AGGREGATOR_TYPE # pooler pooling_type: POOL_FN feature_size: int num_layers: int num_modules: int output_size: int # perceiver attention_probs_dropout_prob: float = 0.0 num_blocks: int = 1 num_self_attends_per_block: int = 16 self_attention_widening_factor: int = 4 cross_attention_widening_factor: int = 1 num_latent_factor: int = 8 def get_aggregator_config( model: PreTrainedModel, ctx_encoder_model: PreTrainedModel, output_size: int, aggregator_args: AggregatorArguments, ): lora_config = model.peft_config["default"] return AggregatorConfig( feature_size=ctx_encoder_model.config.hidden_size, output_size=output_size, num_layers=get_num_layers(model), num_modules=len(lora_config.target_modules), **vars(aggregator_args), ) @dataclass class HypernetConfig: latent_size: int use_light_weight_lora: bool light_weight_latent_size: int lora_config: LoraConfig module_names: dict[str, list[str]] layer_indices: Iterable[int] feature_sizes: tuple[dict[str, int], dict[str, int]] aggregator_config: AggregatorConfig def get_hypernet_config( model: PreTrainedModel, ctx_encoder_model: PreTrainedModel, hypernet_args: HypernetArguments, aggregator_args: AggregatorArguments, ): lora_config = model.peft_config["default"] indices = torch.arange(get_num_layers(model), device=model.device) return HypernetConfig( **vars(hypernet_args), lora_config=lora_config, module_names=get_lora_module_names(model, lora_config.target_modules, indices), layer_indices=indices, feature_sizes=get_peft_in_out_features(model, peft_config=lora_config), aggregator_config=get_aggregator_config( model, ctx_encoder_model, hypernet_args.latent_size, aggregator_args, ), ) class Perceiver(nn.Module): """perceiver w/ bottleneck size = n_modules * n_layers""" def __init__( self, feature_size, output_size, num_layers, num_modules, num_latent_factor, *args, **kwargs, ): super().__init__() self.num_layers = num_layers self.num_modules = num_modules self.config = PerceiverConfig( d_model=feature_size, # + num_bands num_latents=num_layers * num_modules * num_latent_factor, d_latents=output_size, # attention_probs_dropout_prob=0.0, # num_blocks=8, # num_self_attends_per_block=6, # self_attention_widening_factor=4, **kwargs, ) decoder = PerceiverBasicDecoder( self.config, output_num_channels=output_size, output_index_dims=num_layers * num_modules, num_channels=output_size, final_project=False, trainable_position_encoding_kwargs=dict( num_channels=output_size, index_dims=num_layers * num_modules, ), ) self.perceiver = PerceiverModel(self.config, decoder=decoder) def forward( self, ctx_features: Float[Tensor, "bs seq_len feature_dim"], ctx_attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, ): x = self.perceiver(ctx_features, ctx_attn_mask).logits x = rearrange( x, "bs (n_layers n_modules) d -> bs n_layers n_modules d", n_modules=self.num_modules, n_layers=self.num_layers, ) return x class Mixer(nn.Module): def __init__( self, input_size: int, intermediate_emb_size: int, output_size: int, ): super().__init__() self.gate_proj = nn.Linear(input_size, intermediate_emb_size, bias=False) self.up_proj = nn.Linear(input_size, intermediate_emb_size, bias=False) self.down_proj = nn.Linear(intermediate_emb_size, output_size, bias=False) self.act_fn = nn.SiLU() def forward(self, x: torch.Tensor) -> torch.Tensor: return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class MLPResidualBlock(nn.Module): def __init__( self, input_size: int, hidden_size: int, output_size: int, pre_layer_norm: bool = True, post_dropout: bool = True, ): super().__init__() layers = [] if pre_layer_norm: layers.append(nn.LayerNorm(input_size)) layers += [ nn.Dropout(0.1), nn.Linear(input_size, hidden_size), nn.SiLU(), nn.Dropout(0.1), nn.Linear(hidden_size, output_size), nn.SiLU(), ] if post_dropout: layers.append(nn.Dropout(0.1)) self.mlp = nn.Sequential(*layers) def forward(self, x): return x + self.mlp(x) class Pooler(nn.Module): def __init__( self, feature_size: int, output_size: int, pooling_type: POOL_FN, num_layers: int, num_modules: int, *args, **kwargs, ): super().__init__() self.num_layers = num_layers self.num_modules = num_modules # NOTE: features will be projected to size = output_size // 2 # then cat with layer and module embeddings (each with size output_size // 4) # which are collectively form features with size = output_size self.pool_fn = get_pooling_fn(pooling_type) self.feature_proj = nn.Linear(feature_size, output_size // 2) self.ln = nn.LayerNorm(output_size // 2) self.layer_embs = nn.Sequential( nn.Embedding(num_layers, output_size // 4), nn.LayerNorm(output_size // 4), ) self.module_embs = nn.Sequential( nn.Embedding(num_modules, output_size // 4), nn.LayerNorm(output_size // 4), ) self.mixer = Mixer(output_size, output_size * 4, output_size) self.mlp = MLPResidualBlock(output_size, output_size * 4, output_size) self.register_buffer("layer_indices", torch.arange(num_layers)) self.register_buffer("module_indices", torch.arange(num_modules)) def forward( self, features: Float[Tensor, "bs seq_len feature_dim"], attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, ): bs = features.shape[0] # [bs, feature_dim] x = self.ln(self.feature_proj(self.pool_fn(features, attn_mask).float())) x = repeat( x, "bs d -> bs n_layers n_modules d", n_modules=self.num_modules, n_layers=self.num_layers, ) layer_embs = self.layer_embs(self.layer_indices) # [num_layers, d] layer_embs = repeat( layer_embs, "n_layers d -> bs n_layers n_modules d", bs=bs, n_modules=self.num_modules, ) module_embs = self.module_embs(self.module_indices) # [num_modules, d] module_embs = repeat( module_embs, "n_modules d -> bs n_layers n_modules d", bs=bs, n_layers=self.num_layers, ) emb = torch.cat([x, layer_embs, module_embs], dim=3) return self.mlp(self.mixer(emb)) AG = { AGGREGATOR_TYPE.POOLER: Pooler, AGGREGATOR_TYPE.PERCEIVER: Perceiver, } @contextmanager def early_exit(base_model: PreTrainedModel, exit_layer: int): try: layers = base_model.layers base_model.layers = layers[:exit_layer] yield base_model finally: base_model.layers = layers @contextmanager def maybe_add_batch_dim(kwargs): try: batched_input = False batched_attn_mask = False if len(kwargs["input_ids"].shape) == 1: kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0) batched_input = True if len(kwargs["attention_mask"].shape) == 1: kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0) batched_attn_mask = True yield batched_input, batched_attn_mask finally: if batched_input: kwargs["input_ids"] = kwargs["input_ids"].squeeze(0) if batched_attn_mask: kwargs["attention_mask"] = kwargs["attention_mask"].squeeze(0) class EarlyExit(nn.Module): def __init__(self, base_model: PreTrainedModel, exit_layer: int): super().__init__() self.base_model = base_model self.exit_layer = exit_layer @torch.no_grad() def forward(self, **kwargs): # if len(kwargs["input_ids"].shape) == 1: # kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0) # if len(kwargs["attention_mask"].shape) == 1: # kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0) with ( early_exit(self.base_model, self.exit_layer), maybe_add_batch_dim(kwargs) as (batched_input, batched_attn_mask), ): model_outputs = self.base_model(**kwargs) if batched_input: model_outputs.last_hidden_state = model_outputs.last_hidden_state.squeeze(0) return model_outputs.last_hidden_state def get_init_peft_weights(model: PeftModel, peft_config: PeftConfig = None): if peft_config is None: peft_config = model.peft_config["default"] peft_weights = {module_name: dict() for module_name in peft_config.target_modules} adapter_name = "default" for module_name, module in model.named_modules(): if not check_target_module_exists(peft_config, module_name): continue if not isinstance(module, BaseTunerLayer): continue # support just Linear layer for now # all modules should be a leave module that is Linear layer assert isinstance( module.base_layer, nn.Linear ), "all modules should be a leave module that is Linear layer" # this should always pass name = module_name.split(".")[-1] assert name in peft_config.target_modules for submodule_name, submodule in module.named_modules(): if not isinstance(submodule, (nn.ModuleDict, nn.ParameterDict, BufferDict)): continue if adapter_name not in submodule: continue if submodule_name not in peft_weights[name]: peft_weights[name][submodule_name] = submodule[adapter_name] else: smod1 = peft_weights[name][submodule_name] smod2 = submodule[adapter_name] assert type(smod1) == type(smod2) return peft_weights class HyperLoRA(nn.Module): def __init__(self, config: HypernetConfig): super().__init__() # aggregator output [bs, n_layers, n_modules, feature_dim] # by mixing the pooled features with layer embs and module embs (for pooling) # or via a perceiver w/ bottleneck size = n_modules * n_layers self.config = config self._init_model() def _init_model(self): self.agg_config = self.config.aggregator_config self.aggregator = AG[self.agg_config.aggregator_type](**vars(self.agg_config)) self.lora_config = self.config.lora_config self.target_modules = self.lora_config.target_modules self.layer_indices = self.config.layer_indices self.d_in, self.d_out = self.config.feature_sizes self.layers = MLPResidualBlock( input_size=self.config.latent_size, hidden_size=self.config.latent_size * 4, output_size=self.config.latent_size, ) if self.config.use_light_weight_lora: # light-weight lora projection (per module) self.pre_lora_projection = nn.ParameterDict( { k: nn.Parameter( torch.randn(self.d_in[k], self.config.light_weight_latent_size) ) for k in self.target_modules } ) for param in self.pre_lora_projection.values(): nn.init.orthogonal_(param) self.post_lora_projection = nn.ParameterDict( { k: nn.Parameter( torch.randn(self.d_out[k], self.config.light_weight_latent_size) ) for k in self.target_modules } ) for param in self.post_lora_projection.values(): nn.init.orthogonal_(param) d_lora = self.config.light_weight_latent_size * 2 self.d_in = {k: self.config.light_weight_latent_size for k in self.d_in} self.d_out = {k: self.config.light_weight_latent_size for k in self.d_out} logger.info(f"Using light-weight LoRA with d_lora = {d_lora // 2}") else: d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules) n_modules = len(self.target_modules) # have to do this otherwise doesnt work with adamw_torch_fused # has something to do with the bias shape (n_modules r d_lora) # when n_modules == 1, adamw_torch_fused complains about device/layout # but when n_modules > 1, it works fine if n_modules == 1: self.head = Mix( "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", weight_shape="d_latent r d_lora", # bias_shape=None, # no bias bias_shape="r d_lora", d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) else: # each module processes d -> r d_out independently self.head = Mix( "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_modules d_latent r d_lora", # bias_shape=None, # no bias bias_shape="n_modules r d_lora", n_modules=len(self.target_modules), d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) # print(self.head) # print(self.head.weight.shape) # print(self.head.bias.shape) # breakpoint() def _to_lora_dict( self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"] ) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]: # list of [bs, n_layers, r, in_d_outim] # and in_d_outim might vary across modules loras = unpack( flat_loras, [[] for _ in range(len(self.target_modules))], "bs n_layers * r max_io_dim", ) # dict of {module: # {A: [bs, n_layers, r, d_inim], # B: [bs, n_layers, r, d_outim]}} lora_dict = dict() for module, lora in zip(self.target_modules, loras): A, B = unpack( lora[..., : self.d_in[module] + self.d_out[module]], [[self.d_in[module]], [self.d_out[module]]], "bs n_layers r *", ) # transpose B B = rearrange(B, "bs n_layers r d_out -> bs n_layers d_out r") if self.config.use_light_weight_lora: A = einsum( self.pre_lora_projection[module], A, "d_in d_latent, bs n_layers r d_latent -> bs n_layers r d_in", ) B = einsum( self.post_lora_projection[module], B, "d_out d_latent, bs n_layers d_latent r -> bs n_layers d_out r", ) lora_dict[module] = dict(A=A, B=B) return lora_dict def forward( self, features: Float[Tensor, "bs seq_len feature_dim"], attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, ): # [bs, n_layers, n_modules, feature_dim] emb = self.aggregator(features.to(torch.float32), attn_mask) # [bs, n_layers, n_modules, r, max_in_d_outim] flat_loras = self.head(self.layers(emb)) return flat_loras def generate_loras( self, features: Float[Tensor, "bs seq_len feature_dim"], attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, ): flat_loras = self.forward(features, attn_mask) return self._to_lora_dict(flat_loras) class ModulatedPretrainedModel(nn.Module): def __init__( self, base_model: PreTrainedModel, hypernet_config: HypernetConfig, ctx_encoder_args: CtxEncoderArguments, # use_kl_loss is only used for training use_kl_loss: bool = False, ): super().__init__() self.device = base_model.device self.hypernet_config = hypernet_config self.ctx_encoder_args = ctx_encoder_args self.use_kl_loss = use_kl_loss self.register_module("base_model", base_model) self._init_model() self._bias_hyper_init() # self.register_module("hypernet", hypernet) # self.register_module("ctx_encoder", ctx_encoder) @classmethod def from_state_dict(cls, state_dict: dict, train: bool = True): lora_config = state_dict["hypernet_config"].lora_config model_name_or_path = lora_config.base_model_name_or_path base_model = get_model( model_name_or_path, train=train, requires_grad=False, peft_config=lora_config, ) hypernet_config = state_dict["hypernet_config"] ctx_encoder_args = state_dict["ctx_encoder_args"] model = cls(base_model, hypernet_config, ctx_encoder_args) model.load_state_dict(state_dict) return model def _init_model(self): self.hypernet = HyperLoRA(self.hypernet_config).to(self.device) # TODO: allow ctx_encoder to be other models if self.ctx_encoder_args.ctx_encoder_model_name_or_path is not None: encoder_model = get_model( self.ctx_encoder_args.ctx_encoder_model_name_or_path, train=True, requires_grad=False, ) else: encoder_model = self.base_model self.ctx_encoder = EarlyExit( get_base_model(encoder_model), self.ctx_encoder_args.layer_idx ) def to(self, *args, **kwargs): # workaround to avoid the hypernet being wrapped by DeepSpeed self.base_model = self.base_model.to(*args, **kwargs) self.ctx_encoder = self.ctx_encoder.to(*args, **kwargs) # self.hypernet = self.hypernet.to(*args, **kwargs) self.hypernet.to(torch.float32) return self # Delegate to base_model @property def config(self): return self.base_model.config def get_input_embeddings(self): return self.base_model.get_input_embeddings() @torch.no_grad() def _bias_hyper_init(self): peft_weights = get_init_peft_weights(self.base_model, self.hypernet.lora_config) logger.debug(f"peft_weights: {peft_weights}") self.hypernet.head.weight.data[:] = 0 self.hypernet.head.bias.data[:] = 0 # d_in = self.hypernet.d_in # d_out = self.hypernet.d_out # m = max(self.hypernet.target_modules, key=lambda m: d_in[m] + d_out[m]) # A = peft_weights[m]["lora_A"].weight.clone()[:, : d_in[m]] # [r, d_in] # B = peft_weights[m]["lora_B"].weight.clone()[: d_out[m]] # [d_out, r] # biases = [A, B.T] # # bias-hyperinit # # init weights to zeros and bias to the base weights # bias_cat = torch.cat(biases, dim=1) # self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat for i, m in enumerate(self.hypernet.target_modules): A = peft_weights[m]["lora_A"].weight.clone() # [r, in_d] B = peft_weights[m]["lora_B"].weight.clone() # [out_d, r] if self.hypernet.config.use_light_weight_lora: A = A[:, : self.hypernet.config.light_weight_latent_size] B = B[: self.hypernet.config.light_weight_latent_size] biases = [A, B.T] # bias-hyperinit # init weights to zeros and bias to the base weights bias_cat = torch.cat(biases, dim=1) self.hypernet.head.bias.data[..., i, :, : bias_cat.shape[1]] = bias_cat # if len(self.hypernet.target_modules) > 1: # self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat # else: # self.hypernet.head.bias.data[..., :, : bias_cat.shape[1]] = bias_cat # self.hypernet.head.bias.requires_grad = False def state_dict(self, *args, **kwargs): state_dict = self.hypernet.state_dict(*args, **kwargs) state_dict["hypernet_config"] = self.hypernet_config state_dict["ctx_encoder_args"] = self.ctx_encoder_args return state_dict def load_state_dict(self, state_dict: dict, *args, **kwargs): # NOTE: might have to set `strict=False` as we don't save all the params self.hypernet_config = state_dict.pop("hypernet_config") self.ctx_encoder_args = state_dict.pop("ctx_encoder_args") if ( self.hypernet_config.lora_config.base_model_name_or_path != self.base_model.name_or_path ): raise ValueError( f"Base model name or path mismatch. " f"The base model given is: {self.base_model.name_or_path}, " f"but the hypernet config is for: {self.hypernet_config.lora_config.base_model_name_or_path}" ) self._init_model() return self.hypernet.load_state_dict(state_dict, *args, **kwargs) # @torch.no_grad() # def get_ctx_features( # self, # examples: dict[str, Any], # ): # # TODO: truncate the ctx_ids to the max_ctx_len # out = dict() # input_ids = torch.tensor(examples.get("ctx_ids")).to(self.device) # # we don't really use ctx_attn_mask here but it'll be padded # # and used by hypernet.aggregator which has to handle ctx_attn_mask # attention_mask = torch.tensor(examples.get("ctx_attn_mask")).to(self.device) # # NOTE: this might not work for batched inputs # if isinstance(self.ctx_encoder, nn.Embedding): # features = self.ctx_encoder(input_ids) # else: # features = self.ctx_encoder( # input_ids=input_ids, attention_mask=attention_mask # ) # out["ctx_features"] = features # return out def forward( self, # ctx_features: Optional[Float[Tensor, "bs ctx_length feature_dim"]] = None, ctx_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_len"]] = None, chat_ids: Optional[Integer[Tensor, "bs chat_len"]] = None, chat_attn_mask: Optional[Integer[Tensor, "bs chat_len"]] = None, chat_labels: Optional[Integer[Tensor, "bs chat_len"]] = None, **model_inputs_kwargs: dict[str, Any], ) -> Union[tuple, ModelOutput]: """Forward pass of the modulated model.""" generated_loras = None if ctx_ids is None: logger.warning( ( "*" * 100, "\n\nNo ctx_features provided, using the base model for forward pass\n\n", "*" * 100, ) ) # model_outputs = self.base_model(**model_inputs_kwargs) # model_outputs.generated_loras = None # return model_outputs else: with torch.no_grad(): ctx_features = self.ctx_encoder( input_ids=ctx_ids, attention_mask=ctx_attn_mask ) generated_loras = self.hypernet.generate_loras(ctx_features, ctx_attn_mask) # compute kl loss # - compute logits from the base model from tokenized chat [bs, chat_len, vocab_size] # - compute logits from model w/ generated loras [bs, answer_len, vocab_size] # - select only the position of the label tokens (but we're actually not using the labels themselves) # - compute kl loss between the two logits if self.use_kl_loss and "labels" in model_inputs_kwargs: bs = chat_ids.shape[0] labels = model_inputs_kwargs.pop("labels") # [bs, prompt_res_len] first_label_pos = torch.argmax((labels != -100).float(), dim=-1) # just a place holder, we dont use labels as groundtruth anyway labels[torch.arange(bs), first_label_pos - 1] = 1 # same for chat_labels first_label_pos = torch.argmax((chat_labels != -100).float(), dim=-1) chat_labels[torch.arange(bs), first_label_pos - 1] = 1 with torch.no_grad(): chat_outputs = self.base_model( input_ids=chat_ids, attention_mask=chat_attn_mask ) with apply_generated_loras( self.base_model, generated_loras, self.hypernet.layer_indices, self.training, ): model_outputs = self.base_model(**model_inputs_kwargs) pred_logits = model_outputs.logits # [bs, prompt_res_len, vocab_size] # [bs, ctx_prompt_res_len, vocab_size] base_chat_logits = chat_outputs.logits base_chat_logits = base_chat_logits[torch.where(chat_labels != -100)] pred_logits = pred_logits[torch.where(labels != -100)] kl_loss = F.kl_div( F.log_softmax(pred_logits, dim=-1), F.softmax(base_chat_logits, dim=-1), reduction="none", ) # only compute kl loss on tokens where labels != -100 kl_loss = kl_loss.sum(dim=-1).mean() model_outputs = ModelOutput(loss=kl_loss, **model_outputs) else: # input_ids in model_inputs_kwargs contains only # prompt + response (for hypernet training) with apply_generated_loras( self.base_model, generated_loras, self.hypernet.layer_indices, self.training, ): model_outputs = self.base_model(**model_inputs_kwargs) return model_outputs @torch.no_grad() def generate( 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], ): generated_loras = None if ctx_ids is None: logger.warning( ( "*" * 100, "\n\nNo ctx_ids provided, using the base model for generation\n\n", "*" * 100, ) ) # model_outputs = self.base_model(**model_inputs_kwargs) # model_outputs.generated_loras = None # return model_outputs else: ctx_features = self.ctx_encoder( input_ids=ctx_ids, attention_mask=ctx_attn_mask ) generated_loras = self.hypernet.generate_loras(ctx_features, ctx_attn_mask) # apply lora hook to the base model # self.apply_generated_loras(generated_loras) with apply_generated_loras( self.base_model, generated_loras, self.hypernet.layer_indices, self.training, ): model_outputs = self.base_model.generate(**model_inputs_kwargs) return model_outputs @contextmanager def apply_generated_loras( base_model: nn.Module, generated_loras: Optional[dict] = None, layer_indices: Optional[Iterable[int]] = None, training: bool = False, ): if generated_loras is None: yield base_model return try: hooks = [] for module_name in generated_loras: for layer_idx in layer_indices: hooks += add_generated_lora_hook( base_model, module_name, layer_idx, A=generated_loras[module_name]["A"][:, layer_idx], B=generated_loras[module_name]["B"][:, layer_idx], scaling=base_model.peft_config["default"].lora_alpha, input_dropout=base_model.peft_config["default"].lora_dropout, training=training, ) yield base_model finally: remove_hook_handles(hooks) # needed for loading model from checkpoint # see https://github.com/huggingface/transformers/pull/34632 torch.serialization.add_safe_globals( [ AggregatorConfig, LoraConfig, HypernetConfig, PeftType, TaskType, LoraRuntimeConfig, set, # for real? ] ) if __name__ == "__main__": # set torch randomness seed torch.manual_seed(42) model_name = "meta-llama/Llama-3.1-8B-Instruct" base_model, tokenizer = get_model_and_tokenizer( model_name, train=True, requires_grad=False, peft_config=get_lora_config(model_name), ) # TODO: base model shoul be init with target lora config print(base_model) ctx_model = base_model device = base_model.device # lora_config = base_model.peft_config["default"] # d_in, d_out = get_peft_in_out_features(base_model, peft_config=lora_config) hypernet_args = HypernetArguments(latent_size=512) aggregator_args = AggregatorArguments(aggregator_type=AGGREGATOR_TYPE.PERCEIVER) hypernet_config = get_hypernet_config( base_model, ctx_model, hypernet_args, aggregator_args ) ctx_encoder_args = CtxEncoderArguments(layer_idx=4) # ctx_encoder = EarlyExit(get_base_model(base_model), 4) # hypernet = HyperLoRA( # get_hypernet_config(base_model, hypernet_args, aggregator_args), # base_model, # ).to(device) model = ModulatedPretrainedModel(base_model, hypernet_config, ctx_encoder_args).to( device ) 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(device) ctx_ids = ctx_inputs["input_ids"] ctx_attn_mask = ctx_inputs["attention_mask"] ctx_features = model.ctx_encoder(input_ids=ctx_ids, attention_mask=ctx_attn_mask).to( torch.float32 ) print(ctx_ids.shape) model.eval() agg_features = model.hypernet.aggregator(ctx_features, ctx_attn_mask) print(agg_features) print(agg_features.shape) hnetout = model.hypernet(ctx_features, ctx_attn_mask) print(hnetout) print(hnetout.shape) prompt_msg = "hello" prompt_inputs = tokenizer(prompt_msg, return_tensors="pt").to(device) basemodelout = model.base_model(**prompt_inputs) print(basemodelout) modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs) print(modelout) # model.load_state_dict( # torch.load( # open( # "train_outputs/runs/Jan14_11-10-56_slurm0-a3nodeset-1_30c9e93d/checkpoint-1720/pytorch_model.bin", # "rb", # ) # ) # ) # modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs) # print(modelout) breakpoint()