import logging from copy import deepcopy 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 einops import rearrange, repeat, unpack, einsum from einops.layers.torch import EinMix as Mix from einops.layers.torch import Reduce from jaxtyping import Float, Integer from peft import ( get_peft_config, load_peft_weights, LoraConfig, PeftConfig, PeftModel, LoraRuntimeConfig, get_peft_model_state_dict, set_peft_model_state_dict, ) from peft.utils import PeftType, TaskType, ModulesToSaveWrapper from peft.tuners._buffer_dict import BufferDict from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists from torch import Tensor, nn from transformers import ( PerceiverConfig, PerceiverModel, PreTrainedModel, PretrainedConfig, PreTrainedTokenizerBase, ) from transformers.models.perceiver.modeling_perceiver import ( PerceiverBasicDecoder, ) from transformers.modeling_outputs import ModelOutput from transformers.models.modernbert.modeling_modernbert import ModernBertModel from ctx_to_lora.configs import ( AggregatorArguments, HypernetArguments, CtxEncoderArguments, ) from ctx_to_lora.hooks import ( add_generated_layernorm_hook, add_generated_lora_hook, remove_hook_handles, ) from ctx_to_lora.model_loading import get_lora_config, get_model, get_model_and_tokenizer from ctx_to_lora.pooling import POOL_FN, get_pooling_fn from ctx_to_lora.utils import ( get_lora_module_names, get_num_layers, get_peft_in_out_features, get_base_model, ) from ctx_to_lora.modeling_idefics2 import Idefics2PerceiverConfig, Idefics2Perceiver 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 num_extra_modules: int output_size: int # perceiver # attention_probs_dropout_prob: float = 0.0 # num_blocks: int = 1 num_self_attends_per_block: int # = 16 decoder_depth: int # = 1 # 1 = only cross-attention # self_attention_widening_factor: int = 4 # cross_attention_widening_factor: int = 1 num_latent_factor: int = 8 lora_r: int = 8 per_rank_gen: bool = False def get_aggregator_config( model: PreTrainedModel, ctx_encoder_model_config: PretrainedConfig, output_size: int, num_modules: int, num_extra_modules: int, lora_r: int, per_rank_gen: bool, aggregator_args: AggregatorArguments, ): return AggregatorConfig( feature_size=ctx_encoder_model_config.hidden_size, output_size=output_size, num_layers=get_num_layers(model), num_modules=num_modules, num_extra_modules=num_extra_modules, lora_r=lora_r, per_rank_gen=per_rank_gen, **vars(aggregator_args), ) @dataclass class HypernetConfig: latent_size: int use_light_weight_lora: bool light_weight_latent_size: int per_rank_gen: bool per_layer_processing: bool dropout_rate: float lora_config: LoraConfig # module_names: dict[str, list[str]] extra_modules: Optional[list[str]] base_hidden_size: int 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_config: PretrainedConfig, hypernet_args: HypernetArguments, aggregator_args: AggregatorArguments, ): num_modules = 0 lora_config = getattr(model, "peft_config", None) if lora_config is not None: lora_config = lora_config["default"] num_modules += len(lora_config.target_modules) num_extra_modules = len(hypernet_args.extra_modules or []) indices = torch.arange(get_num_layers(model), device=model.device) return HypernetConfig( **vars(hypernet_args), base_hidden_size=model.config.hidden_size, 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_config, hypernet_args.latent_size, num_modules, num_extra_modules, lora_config.r, hypernet_args.per_rank_gen, 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_extra_modules, per_rank_gen, lora_r, num_latent_factor, *args, **kwargs, ): super().__init__() self.num_layers = num_layers self.num_modules = num_modules self.num_extra_modules = num_extra_modules self.per_rank_gen = per_rank_gen self.r = lora_r if self.per_rank_gen else 1 # 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) n_output_queries = num_layers * (num_modules * self.r + num_extra_modules) self.config = Idefics2PerceiverConfig( input_size=feature_size, # the first layer is xattn resampler_depth=kwargs["num_self_attends_per_block"] + 1, resampler_n_latents=n_output_queries * num_latent_factor, intermediate_size_factor=4, hidden_size=output_size, attn_implementation="flash_attention_2", ) self.decoder_config = Idefics2PerceiverConfig( input_size=output_size, resampler_depth=kwargs["decoder_depth"], resampler_n_latents=n_output_queries, intermediate_size_factor=4, hidden_size=output_size, attn_implementation="flash_attention_2", ) self.perceiver = Idefics2Perceiver(self.config, self.decoder_config) def forward( self, ctx_features: Float[Tensor, "bs seq_len feature_dim"], ctx_attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, ctx_position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, ): x = self.perceiver(ctx_features, ctx_attn_mask, ctx_position_ids) lora_x, extra_x = unpack( x, [ [self.num_layers * self.num_modules * self.r], [self.num_layers * self.num_extra_modules], ], "bs * feature_dim", ) lora_x = rearrange( lora_x, "bs (n_layers n_modules r) d -> bs n_layers n_modules r d", n_modules=self.num_modules, n_layers=self.num_layers, r=self.r, ) if not self.per_rank_gen: lora_x = lora_x.squeeze(3) extra_x = rearrange( extra_x, "bs (n_layers n_extra_modules) d -> bs n_layers n_extra_modules d", n_extra_modules=self.num_extra_modules, n_layers=self.num_layers, ) # 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, # ) # lora_emb, extra_emb = unpack( # emb, # [[self.num_modules], [self.num_extra_modules]], # "bs n_layers * feature_dim", # ) return lora_x, extra_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.GELU(approximate="tanh") 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, dropout_rate: float = 0, ): super().__init__() layers = [] # if pre_layer_norm: # layers.append(nn.LayerNorm(input_size)) layers = [ Gemma3RMSNorm(input_size), nn.Dropout(dropout_rate), nn.Linear(input_size, hidden_size), nn.GELU(approximate="tanh"), nn.Dropout(dropout_rate), nn.Linear(hidden_size, output_size), Gemma3RMSNorm(output_size), # nn.GELU(approximate="tanh"), ] # if post_dropout: # layers.append(nn.Dropout(dropout_rate)) 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 ( "input_ids" in kwargs and kwargs["input_ids"] is not None and len(kwargs["input_ids"].shape) == 1 ): kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0) batched_input = True if ( "attention_mask" in kwargs and kwargs["attention_mask"] is not None and isinstance(kwargs["attention_mask"], torch.Tensor) and 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 if "gte" in base_model.config.name_or_path: self.base_model.encoder.layer = base_model.encoder.layer[:exit_layer] else: self.base_model.layers = base_model.layers[:exit_layer] # self.exit_layer = exit_layer @property def config(self): return self.base_model.config @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 Gemma3RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.zeros(dim)) def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, x): output = self._norm(x.float()) # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16) # See https://github.com/huggingface/transformers/pull/29402 output = output * (1.0 + self.weight.float()) return output.type_as(x) def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.eps}" class ResMLPBlock(nn.Module): ... class UpMixPerLayer(nn.Module): def __init__(self, n_layers: int, n_modules: int, r: int, d_in: int, d_up: int): super().__init__() self.linear = Mix( "bs n_layers n_modules r d_in -> bs n_layers n_modules r d_up", weight_shape="n_layers d_in d_up", bias_shape=None, # "n_layers n_modules d_up", n_layers=n_layers, # n_modules=n_modules, r=r, d_in=d_in, d_up=d_up, ) def forward(self, x): return self.linear(x) class DownMixPerLayer(nn.Module): def __init__(self, n_layers: int, n_modules: int, r: int, d_up: int, d_out: int): super().__init__() self.linear = Mix( "bs n_layers n_modules r d_up -> bs n_layers n_modules r d_out", weight_shape="n_layers d_up d_out", bias_shape=None, # "n_layers n_modules d_out", n_layers=n_layers, # n_modules=n_modules, r=r, d_up=d_up, d_out=d_out, ) def forward(self, x): return self.linear(x) class MixerPerLayer(nn.Module): def __init__(self, n_layers: int, n_modules: int, r: int, d_in: int): super().__init__() self.up_proj = UpMixPerLayer(n_layers, n_modules, r, d_in, d_in * 4) self.gate_proj = UpMixPerLayer(n_layers, n_modules, r, d_in, d_in * 4) self.down_proj = DownMixPerLayer(n_layers, n_modules, r, d_in * 4, d_in) self.act_fn = nn.GELU(approximate="tanh") def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) class ResMLPBlockPerLayer(nn.Module): def __init__(self, n_layers: int, n_modules: int, r: int, d_in: int): super().__init__() # input shape: [bs, n_layers, n_modules, feature_dim] self.pre_norm = Gemma3RMSNorm(d_in) self.post_norm = Gemma3RMSNorm(d_in) self.mixer = MixerPerLayer(n_layers, n_modules, r, d_in) def forward(self, x): inp = x x = self.pre_norm(x) x = self.mixer(x) x = self.post_norm(x) return x + inp 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 logger.debug(f"HyperLoRA config: {self.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 if self.lora_config else None ) self.num_modules = len(self.target_modules) if self.target_modules else 0 self.extra_modules = ( self.config.extra_modules if self.config.extra_modules else None ) self.num_extra_modules = len(self.extra_modules) if self.extra_modules else 0 self.layer_indices = self.config.layer_indices self.n_layers = len(self.layer_indices) self.d_in, self.d_out = self.config.feature_sizes self.d_latent = self.config.latent_size if self.target_modules: # self.layers = nn.Sequential( # *[ # MLPResidualBlock( # input_size=self.config.latent_size, # hidden_size=self.config.latent_size * 4, # output_size=self.config.latent_size, # dropout_rate=getattr(self.config, "dropout_rate", 0), # ) # for _ in range(4) # ] # ) if self.config.per_layer_processing: self.layers = nn.Sequential( nn.Linear(self.d_latent, self.d_latent), ResMLPBlockPerLayer( self.n_layers, self.num_modules, self.lora_config.r, self.d_latent, ), ) else: self.layers = nn.Sequential( nn.Linear(self.d_latent, self.d_latent), MLPResidualBlock( input_size=self.config.latent_size, hidden_size=self.config.latent_size * 4, output_size=self.config.latent_size, dropout_rate=getattr(self.config, "dropout_rate", 0), ), ) d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules) # 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}") 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 self.config.per_rank_gen: if n_modules == 1: if self.config.per_layer_processing: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers r d_latent d_lora", # bias_shape=None, # no bias bias_shape="n_layers r d_lora", n_layers=len(self.layer_indices), d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) else: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="r d_latent d_lora", # bias_shape=None, # no bias bias_shape="r d_lora", # n_layers=len(self.layer_indices), d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) else: if self.config.per_layer_processing: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers n_modules r d_latent d_lora", # bias_shape=None, # no bias bias_shape="n_layers n_modules r d_lora", n_layers=len(self.layer_indices), n_modules=n_modules, d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) else: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_modules r d_latent d_lora", # bias_shape=None, # no bias bias_shape="n_modules r d_lora", n_layers=len(self.layer_indices), n_modules=n_modules, d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) else: if n_modules == 1: if self.config.per_layer_processing: self.head = Mix( "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers d_latent r d_lora", # bias_shape=None, # no bias bias_shape="n_layers r d_lora", d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) 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 if self.config.per_layer_processing: self.head = Mix( "bs n_layers n_modules d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers n_modules d_latent r d_lora", # bias_shape=None, # no bias bias_shape="n_layers n_modules r d_lora", n_modules=n_modules, d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) else: 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=n_modules, d_latent=self.config.latent_size, r=self.config.lora_config.r, d_lora=d_lora, ) # if self.extra_modules: # # self.extra_layers = nn.Sequential( # # *[ # # MLPResidualBlock( # # input_size=self.config.latent_size, # # hidden_size=self.config.latent_size * 4, # # output_size=self.config.latent_size, # # dropout_rate=getattr(self.config, "dropout_rate", 0), # # ) # # for _ in range(4) # # ] # # ) # self.extra_layers = MLPResidualBlock( # input_size=self.config.latent_size, # hidden_size=self.config.latent_size * 4, # output_size=self.config.latent_size, # dropout_rate=getattr(self.config, "dropout_rate", 0), # ) # # self.extra_layers = nn.Identity() # # self.pre_extra_layers_norm = nn.LayerNorm(self.config.latent_size) # # self.extra_norm = nn.LayerNorm(self.config.latent_size) # # separate head for extra_modules, e.g., layernorm # n_extra_modules = len(self.extra_modules) if self.extra_modules else 0 # if n_extra_modules == 1: # self.extra_head = Mix( # "bs n_layers n_modules d_latent -> bs n_layers n_modules hidden_size", # weight_shape="d_latent hidden_size", # bias_shape="hidden_size", # d_latent=self.config.latent_size, # hidden_size=self.config.base_hidden_size, # ) # else: # self.extra_head = Mix( # "bs n_layers n_modules d_latent -> bs n_layers n_modules hidden_size", # weight_shape="n_modules d_latent hidden_size", # bias_shape="n_modules hidden_size", # n_modules=n_extra_modules, # d_latent=self.config.latent_size, # hidden_size=self.config.base_hidden_size, # ) 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 _"]]]: if self.target_modules is None: return None # 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 _to_layernorm_dict( self, flat_layernorms: Float[Tensor, "bs n_layers n_modules hidden_size"] ) -> dict[str, Float[Tensor, "bs n_layers hidden_size"]]: if self.extra_modules is None: return None layernorms = unpack( flat_layernorms, [[] for _ in range(len(self.extra_modules))], "bs n_layers * hidden_size", ) return {k: v for k, v in zip(self.extra_modules, layernorms)} def forward( self, features: Float[Tensor, "bs seq_len feature_dim"], attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, ): # [bs, n_layers, n_total_modules, feature_dim] lora_emb, extra_emb = self.aggregator(features, attn_mask, position_ids) # lora_emb, extra_emb = unpack( # emb, # [[self.num_modules], [self.num_extra_modules]], # "bs n_layers * feature_dim", # ) # [bs, n_layers, n_modules, r, max_in_d_outim] flat_loras = None if self.target_modules: lora_emb = self.layers(lora_emb) flat_loras = self.head(lora_emb) flat_layernorms = None if self.num_extra_modules: # [bs, n_layers, n_extra_modules, base_hidden_size] extra_emb = self.extra_layers(extra_emb) flat_layernorms = self.extra_head(extra_emb) return flat_loras, flat_layernorms def generate_weights( self, features: Float[Tensor, "bs seq_len feature_dim"], attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, ): flat_loras, flat_layernorms = self.forward(features, attn_mask, position_ids) # print(f"flat_loras: {flat_loras.shape}") # if flat_layernorms is not None: # print(f"flat_layernorms: {flat_layernorms.shape}") return self._to_lora_dict(flat_loras), self._to_layernorm_dict(flat_layernorms) 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, use_base_input_as_ctx: 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.use_base_input_as_ctx = use_base_input_as_ctx 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, use_flash_attn: bool = True, ): lora_config = state_dict["hypernet_config"].lora_config model_name_or_path = state_dict["base_model_name_or_path"] base_model = get_model( model_name_or_path, train=train, requires_grad=False, peft_config=lora_config, use_flash_attn=use_flash_attn, ) 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): # for name, module in self.base_model.named_modules(): # if isinstance(module, ModulesToSaveWrapper): # module.set_adapter("default") # trainable_base_modules = getattr( # self.hypernet_config, "trainable_base_modules", None # ) # if trainable_base_modules: # for name, param in self.base_model.named_parameters(): # if any(module_name in name for module_name in trainable_base_modules): # param.requires_grad = True self.hypernet = HyperLoRA(self.hypernet_config).to(self.device).to(torch.float32) # ctx_encoder_name = # 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 ctx_model_name = self.ctx_encoder_args.ctx_encoder_model_name_or_path if ctx_model_name is None: ctx_model_name = self.base_model.config.name_or_path # use an explicit copy of the base model # for using with "modules_to_save" base_model_attn_impl = self.base_model.config._attn_implementation logger.debug(f"ctx_model_name: {ctx_model_name}") logger.debug(f"base_model.config._attn_implementation: {base_model_attn_impl}") encoder_model = get_model( ctx_model_name, train=self.base_model.training, requires_grad=False, use_flash_attn=base_model_attn_impl == "flash_attention_2", ) 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): if self.hypernet.extra_modules: self.hypernet.extra_head.weight.data[:] = 0 self.hypernet.extra_head.bias.data[:] = 0 if self.hypernet.target_modules: 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] if self.hypernet.config.per_rank_gen: A = A[0:1] B = B[:, 0:1] 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): # we assume ctx_encoder and base model is frozen here if len([p for p in self.ctx_encoder.parameters() if p.requires_grad]): raise ValueError("ctx_encoder contains trainable parameters") if len([p for p in self.base_model.parameters() if p.requires_grad]): raise ValueError("base model contains trainable parameters") state_dict = self.hypernet.state_dict(*args, **kwargs) # base_state_dict = dict() # if self.hypernet_config.trainable_base_modules: # # for k, v in get_peft_model_state_dict(self.base_model).items(): # # if any(module in k for module in self.hypernet_config.trainable_base_modules): # # # save only "modules_to_save" # # base_state_dict[k] = v # for name, param in self.base_model.named_parameters(): # if param.requires_grad and any( # module in name # for module in self.hypernet_config.trainable_base_modules # ): # base_state_dict[name] = param.data # state_dict["base_model"] = base_state_dict state_dict["base_model_name_or_path"] = self.base_model.name_or_path 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.base_model_name_or_path = state_dict.pop("base_model_name_or_path") self.hypernet_config = state_dict.pop("hypernet_config") self.ctx_encoder_args = state_dict.pop("ctx_encoder_args") if self.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 loaded name is: {self.base_model_name_or_path}" ) self._init_model() state_dict.pop("base_model", None) # base_state_dict = state_dict.pop("base_model", None) # if base_state_dict: # print("Loading base model state dict") # print(base_state_dict.keys()) # # peft_load_result = set_peft_model_state_dict( # # self.base_model, # # base_state_dict, # # *args, # # **kwargs, # # ) # load_result = self.base_model.load_state_dict(base_state_dict, strict=False) # print(f"base_load_result: {load_result}") 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 generate_weights( self, ctx_ids: Integer[Tensor, "bs ctx_len"], ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_len"]] = None, ctx_position_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, **kwargs: Any, ): with torch.no_grad(): ctx_encoder_kwargs = dict( input_ids=ctx_ids, attention_mask=ctx_attn_mask, position_ids=ctx_position_ids, ) # TODO: for modernbert ctx_encoder pass # `cu_seq_len` and `max_seq_len` to the forward call if isinstance(self.ctx_encoder.base_model, ModernBertModel): position_ids = ctx_position_ids.flatten() indices = torch.arange( position_ids.size(0), device=position_ids.device, dtype=torch.int32 ) # [bsz + 1] cu_seqlens = torch.cat( ( indices[position_ids == 0], torch.tensor( position_ids.size(), device=position_ids.device, dtype=torch.int32, ), ) ) ctx_encoder_kwargs = dict( input_ids=ctx_ids.squeeze(0), cu_seqlens=cu_seqlens, max_seqlen=position_ids.max() + 1, attention_mask=-1, seq_len=-1, batch_size=-1, ) ctx_features = self.ctx_encoder(**ctx_encoder_kwargs, **kwargs) if isinstance(self.ctx_encoder.base_model, ModernBertModel): ctx_features = ctx_features.unsqueeze(0) # print(f"padded ctx_features: {ctx_features.shape}") # if ctx_attn_mask is not None: # print(f"padded ctx_attn_mask: {ctx_attn_mask.shape}") # if ctx_position_ids is not None: # print(f"padded ctx_position_ids: {ctx_position_ids.shape}") return self.hypernet.generate_weights( ctx_features, ctx_attn_mask, ctx_position_ids ) # @torch.inference_mode() # def _unpack(self, ctx_position_ids, ctx_features): # # [1, len, d] -> [len, d] # ctx_features = ctx_features.squeeze(0) # lengths = torch.where(ctx_position_ids == 0)[1] # total_len = torch.tensor([len(ctx_features)], device=self.device) # lengths = torch.cat([lengths, total_len]) # # compute the difference between the lengths # lens = torch.diff(lengths).unsqueeze(1) # ctx_features = unpack(ctx_features, lens, "* d") # ctx_attn_mask = [torch.ones(len(x), device=self.device) for x in ctx_features] # ctx_features = torch.nn.utils.rnn.pad_sequence( # ctx_features, # batch_first=True, # padding_value=0, # ) # ctx_attn_mask = torch.nn.utils.rnn.pad_sequence( # ctx_attn_mask, # batch_first=True, # padding_value=0, # ) # return ctx_features, ctx_attn_mask 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, ctx_position_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, return_generated_lora: Optional[bool] = False, # 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_args: Any, **model_inputs_kwargs: dict[str, Any], ) -> Union[tuple, ModelOutput]: """Forward pass of the modulated model.""" generated_loras = None generated_layernorms = None if ctx_ids is None and not self.use_base_input_as_ctx: 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: if self.use_base_input_as_ctx: ctx_ids = ( model_inputs_kwargs["input_ids"] if "input_ids" in model_inputs_kwargs else model_inputs_args[0] ) ctx_attn_mask = ( model_inputs_kwargs["attention_mask"] if "attention_mask" in model_inputs_kwargs else None ) ctx_position_ids = ( model_inputs_kwargs["position_ids"] if "position_ids" in model_inputs_kwargs else None ) # with torch.inference_mode(): # ctx_features = self.ctx_encoder( # input_ids=ctx_ids, # attention_mask=ctx_attn_mask, # position_ids=ctx_position_ids, # ) # generated_loras, generated_layernorms = self.hypernet.generate_weights( # ctx_features, ctx_attn_mask # ) generated_loras, generated_layernorms = self.generate_weights( ctx_ids, ctx_attn_mask, ctx_position_ids ) # input_ids in model_inputs_kwargs contains only # prompt + response (for hypernet training) position_ids = ( model_inputs_kwargs["position_ids"] if "position_ids" in model_inputs_kwargs else None ) with ( apply_generated_loras( self.base_model, generated_loras, self.hypernet.layer_indices, position_ids, self.training, ), apply_generated_layernorm( self.base_model, generated_layernorms, self.hypernet.layer_indices, position_ids, self.training, ), ): model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs) if return_generated_lora: return model_outputs, (generated_loras, generated_layernorms) else: return model_outputs @torch.inference_mode() def generate( self, ctx_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_length"]] = None, ctx_position_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, *model_inputs_args: Any, **model_inputs_kwargs: dict[str, Any], ): generated_loras = None generated_layernorms = None if ctx_ids is None and not self.use_base_input_as_ctx: 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: if self.use_base_input_as_ctx: ctx_ids = ( model_inputs_kwargs["input_ids"] if "input_ids" in model_inputs_kwargs else model_inputs_args[0] ) ctx_attn_mask = ( model_inputs_kwargs["attention_mask"] if "attention_mask" in model_inputs_kwargs else None ) ctx_position_ids = ( model_inputs_kwargs["position_ids"] if "position_ids" in model_inputs_kwargs else None ) # ctx_features = self.ctx_encoder( # input_ids=ctx_ids, attention_mask=ctx_attn_mask # ) # generated_loras, generated_layernorms = self.hypernet.generate_weights( # ctx_features, ctx_attn_mask # ) generated_loras, generated_layernorms = self.generate_weights( ctx_ids, ctx_attn_mask, ctx_position_ids ) # apply lora hook to the base model # self.apply_generated_loras(generated_loras) position_ids = ( model_inputs_kwargs["position_ids"] if "position_ids" in model_inputs_kwargs else None ) with ( apply_generated_loras( self.base_model, generated_loras, self.hypernet.layer_indices, position_ids, self.training, ), apply_generated_layernorm( self.base_model, generated_layernorms, self.hypernet.layer_indices, position_ids, self.training, ), ): model_outputs = self.base_model.generate( *model_inputs_args, **model_inputs_kwargs ) return model_outputs class ModulatedModelWithSharedInput(nn.Module): def __init__( self, modulated_model: ModulatedPretrainedModel, base_tokenizer: PreTrainedTokenizerBase, ctx_tokenizer: PreTrainedTokenizerBase, ctx_end_predicate: Optional[str] = None, remove_ctx_from_base_input: bool = False, ): super().__init__() self.modulated_model = modulated_model self.base_tokenizer = base_tokenizer self.ctx_tokenizer = ctx_tokenizer self.register_module("modulated_model", self.modulated_model) self.ctx_end_predicate = ctx_end_predicate self.remove_ctx_from_base_input = remove_ctx_from_base_input # delegate to self.modulated_model @property def config(self): return self.modulated_model.config @property def device(self): return self.modulated_model.device def tie_weights(self): self.modulated_model.base_model.tie_weights() def state_dict(self, *args, **kwargs): return self.modulated_model.state_dict(*args, **kwargs) def forward(self, *args, **kwargs): input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0] input_txts = self.base_tokenizer.batch_decode( input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True, ) ctx_txts = deepcopy(input_txts) if self.ctx_end_predicate: ctx_txts = [ txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts ] ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to( self.device ) ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask if self.remove_ctx_from_base_input: raise NotImplementedError("Not implemented") # input_txts = [ # txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts # ] # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True) # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask # kwargs["input_ids"] = input_ids # kwargs["attention_mask"] = attention_mask return self.modulated_model(ctx_ids, ctx_attn_mask, *args, **kwargs) def generate(self, *args, **kwargs): input_ids = kwargs["input_ids"] if "input_ids" in kwargs else args[0] input_txts = self.base_tokenizer.batch_decode( input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True, ) ctx_txts = deepcopy(input_txts) if self.ctx_end_predicate: ctx_txts = [ txt.split(self.ctx_end_predicate)[0].strip() for txt in input_txts ] ctx_inputs = self.ctx_tokenizer(ctx_txts, return_tensors="pt", padding=True).to( self.device ) ctx_ids, ctx_attn_mask = ctx_inputs.input_ids, ctx_inputs.attention_mask if self.remove_ctx_from_base_input: raise NotImplementedError("Not implemented") # input_txts = [ # txt.split(self.ctx_end_predicate)[1].strip() for txt in input_txts # ] # inputs = self.base_tokenizer(input_txts, return_tensors="pt", padding=True) # input_ids, attention_mask = inputs.input_ids, inputs.attention_mask # kwargs["input_ids"] = input_ids # kwargs["attention_mask"] = attention_mask return self.modulated_model.generate(ctx_ids, ctx_attn_mask, *args, **kwargs) @contextmanager def apply_generated_layernorm( base_model: nn.Module, generated_layernorms: Optional[dict[str, Float[Tensor, "bs n_layers h"]]] = None, layer_indices: Optional[Iterable[int]] = None, position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, training: bool = False, ): if generated_layernorms is None: yield base_model return try: hooks = [] for module_name in generated_layernorms: for layer_idx in layer_indices: hooks += add_generated_layernorm_hook( base_model, module_name, layer_idx, W=generated_layernorms[module_name][:, layer_idx], position_ids=position_ids, training=training, ) yield base_model finally: remove_hook_handles(hooks) @contextmanager def apply_generated_loras( base_model: nn.Module, generated_loras: Optional[dict] = None, layer_indices: Optional[Iterable[int]] = None, position_ids: Optional[Integer[Tensor, "bs seq_len"]] = 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: # TODO: handle sequence packing??? 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, position_ids=position_ids, training=training, ) yield base_model finally: remove_hook_handles(hooks) @contextmanager def apply_generated_loras_packed_sequence( base_model: nn.Module, generated_loras: Optional[dict] = None, layer_indices: Optional[Iterable[int]] = None, training: bool = False, ): pass # 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), ) print(base_model.modules_to_save) ctx_model_config = base_model.config 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_config, 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) state_dict = torch.load( "train_outputs/runs/Jan15_19-10-21_slurm0-a3nodeset-11_d1842c41/checkpoint-55000/pytorch_model.bin", ) print(state_dict.keys()) breakpoint() model.load_state_dict(state_dict) modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs) print(modelout) breakpoint()