import logging from collections.abc import Iterable from contextlib import contextmanager from dataclasses import dataclass from functools import partial from typing import Any, Literal import torch from einops import rearrange, unpack from einops.layers.torch import EinMix as Mix from jaxtyping import Float, Integer from peft import ( LoraConfig, LoraRuntimeConfig, PeftConfig, PeftModel, ) from peft.tuners._buffer_dict import BufferDict from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists from peft.utils import PeftType, TaskType from torch import Tensor, nn from transformers import ( PretrainedConfig, PreTrainedModel, ) from transformers.modeling_outputs import ModelOutput from transformers.models.modernbert.modeling_modernbert import ModernBertModel from ctx_to_lora.configs import ( AggregatorArguments, CtxEncoderArguments, HypernetArguments, ) 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_model, ) from ctx_to_lora.modeling.aggregator import ( AGGREGATOR_CLS, AggregatorConfig, get_aggregator_config, ) from ctx_to_lora.modeling.ctx_encoder import CTX_ENCODER_CLS, CTX_ENCODER_TYPE from ctx_to_lora.modeling.lora_layer import ( apply_lora_to_layers, lora_forward, lora_forward_packed, ) from ctx_to_lora.modeling.lora_merger import combine_lora from ctx_to_lora.utils import ( get_layers, get_num_layers, get_peft_in_out_features, get_peft_modules, ) logger = logging.getLogger() @dataclass class HypernetConfig: latent_size: int use_light_weight_lora: bool light_weight_latent_size: int per_rank_gen: bool use_per_rank_bias: bool per_layer_processing: bool use_token_mixing: bool num_pre_head_layers: int dropout_rate: float lora_config: LoraConfig extra_modules: list[str] | None 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, ctx_encoder_args: CtxEncoderArguments, ): 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, 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, ctx_encoder_args.ctx_encoder_type == CTX_ENCODER_TYPE.PER_LAYER_ACTIVATIONS, hypernet_args.latent_size, num_modules, num_extra_modules, lora_config.r, hypernet_args.per_rank_gen, aggregator_args, ), ) 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 ResMLPBlock(nn.Module): def __init__( self, input_size: int, hidden_size: int, output_size: int, dropout_rate: float = 0, ): super().__init__() layers = [] layers = [ nn.LayerNorm(input_size), nn.Dropout(dropout_rate), nn.Linear(input_size, hidden_size), nn.SiLU(), nn.Dropout(dropout_rate), nn.Linear(hidden_size, output_size), nn.LayerNorm(output_size), ] self.mlp = nn.Sequential(*layers) def forward(self, x): return x + self.mlp(x) class ResMLPBlockPerLayer(nn.Module): def __init__( self, n_layers: int, input_size: int, hidden_size: int, output_size: int, ): super().__init__() layers = [ nn.LayerNorm(input_size), Mix( "bs n_layers n_modules r d_in -> bs n_layers n_modules r d_hid", weight_shape="n_layers d_in d_hid", bias_shape="n_layers d_hid", n_layers=n_layers, d_in=input_size, d_hid=hidden_size, ), nn.SiLU(), Mix( "bs n_layers n_modules r d_hid -> bs n_layers n_modules r d_out", weight_shape="n_layers d_hid d_out", bias_shape="n_layers d_out", n_layers=n_layers, d_hid=hidden_size, d_out=output_size, ), nn.LayerNorm(output_size), ] self.layers = nn.Sequential(*layers) def forward(self, x): return x + self.layers(x) 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): # TODO: add a way to handle dependencies between chunks self.agg_config = self.config.aggregator_config self.aggregator = AGGREGATOR_CLS[self.agg_config.aggregator_type]( **vars(self.agg_config) ) self.lora_config = self.config.lora_config self.r = self.lora_config.r self.target_modules = ( tuple(sorted(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: if self.config.per_layer_processing: layers = [ ResMLPBlockPerLayer( self.n_layers, self.d_latent, self.d_latent * 4, self.d_latent, ) for _ in range(self.config.num_pre_head_layers) ] else: layers = [ ResMLPBlock( 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(self.config.num_pre_head_layers) ] self.layers = nn.Sequential(*layers) d_lora = max(self.d_in[m] + self.d_out[m] for m in self.target_modules) # self.bias_a = nn.ParameterDict( # { # m: nn.Parameter( # torch.normal( # 0, 1 / self.d_in[m], (self.n_layers, self.r, self.d_in[m]) # ) # ) # for m in self.target_modules # } # ) # self.bias_b = nn.ParameterDict( # { # m: nn.Parameter(torch.zeros((self.n_layers, self.r, self.d_out[m]))) # for m in self.target_modules # } # ) if self.config.use_light_weight_lora: # light-weight lora projection (per layer, 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 # } # ) self.pre_lora_projection = nn.ParameterDict( { k: Mix( "bs n_layers r d_latent -> bs n_layers r d_in", weight_shape="n_layers d_latent d_in", bias_shape="n_layers d_in", n_layers=len(self.layer_indices), d_latent=self.config.light_weight_latent_size, d_in=self.d_in[k], ) for k in self.target_modules } ) with torch.no_grad(): for param in self.pre_lora_projection.values(): for i in range(param.weight.shape[0]): nn.init.orthogonal_( param.weight[i], # 1 / self.config.light_weight_latent_size**0.5, ) param.bias.data[:] = 0 # 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 # } # ) self.post_lora_projection = nn.ParameterDict( { k: Mix( "bs n_layers d_latent r -> bs n_layers d_out r", weight_shape="n_layers d_latent d_out", bias_shape="n_layers d_out", n_layers=len(self.layer_indices), d_latent=self.config.light_weight_latent_size, d_out=self.d_out[k], ) for k in self.target_modules } ) with torch.no_grad(): for param in self.post_lora_projection.values(): for i in range(param.weight.shape[0]): nn.init.orthogonal_( param.weight[i], # 1 / self.config.light_weight_latent_size**0.5, ) param.bias.data[:] = 0 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: if self.config.use_per_rank_bias: self.head = Mix( "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", weight_shape="n_layers 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="n_layers d_latent d_lora", # bias_shape=None, # no bias bias_shape="n_layers 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.use_per_rank_bias: # self.head = Mix( # "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", # weight_shape="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: # self.head = Mix( # "bs n_layers n_modules r d_latent -> bs n_layers n_modules r d_lora", # weight_shape="d_latent d_lora", # # bias_shape=None, # no bias # bias_shape="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: if self.config.use_per_rank_bias: 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 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_layers n_modules d_latent d_lora", # bias_shape=None, # no bias bias_shape="n_layers n_modules 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 d_latent d_lora", # # bias_shape=None, # no bias # bias_shape="n_modules 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, # ) # else: # 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( # # *[ # # ResMLPBlock( # # 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(self.config.num_pre_head_layers) # # ] # # ) # self.extra_layers = ResMLPBlock( # 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 get_head_bias(self): # bias_dict = dict() # for module in self.target_modules: # bias_A = self.bias_a[module][None, ...] # bias_B = self.bias_b[module][None, ...] # # transpose B # bias_B = rearrange(bias_B, "bs n_layers r d_out -> bs n_layers d_out r") # bias_dict[module] = dict(A=bias_A, B=bias_B) # return bias_dict 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", # ) A = self.pre_lora_projection[module](A) B = self.post_lora_projection[module](B) 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: Integer[Tensor, "bs seq_len"] | None = None, position_ids: Integer[Tensor, "bs seq_len"] | None = None, ): # [bs, n_layers x n_total_modules x r, feature_dim] with torch.autocast(device_type="cuda", dtype=torch.bfloat16): # OMG!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! 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) norm_lora_emb = lora_emb / torch.norm(lora_emb, dim=-1, keepdim=True) flat_loras = self.head(norm_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: Integer[Tensor, "bs seq_len"] | None = None, position_ids: Integer[Tensor, "bs seq_len"] | None = None, ): flat_loras, flat_layernorms = self.forward(features, attn_mask, position_ids) return self._to_lora_dict(flat_loras), self._to_layernorm_dict(flat_layernorms) class ModulatedPretrainedModel(nn.Module): def __init__( self, base_model: PeftModel, hypernet_config: HypernetConfig, ctx_encoder_args: CtxEncoderArguments, use_base_input_as_ctx: bool = False, # need non-packed inputs for generation use_sequence_packing: bool = True, lora_aggregation: Literal["mean", "sum"] = "sum", ): assert not use_base_input_as_ctx super().__init__() self.device = base_model.device self.peft_config = base_model.peft_config["default"] self.hypernet_config = hypernet_config self.ctx_encoder_args = ctx_encoder_args self.use_base_input_as_ctx = use_base_input_as_ctx self.use_sequence_packing = use_sequence_packing self.lora_aggregation = lora_aggregation self.lora_id = 1 self.active_adapters = [] self.register_module("base_model", base_model) self._init_model() self._bias_hyper_init() @classmethod def from_state_dict( cls, state_dict: dict, train: bool = True, base_model_kwargs: dict = None, use_flash_attn: bool = True, **kwargs: Any, ): lora_config = state_dict["hypernet_config"].lora_config print(f"lora_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, model_kwargs=base_model_kwargs, use_flash_attn=use_flash_attn, ) hypernet_config = state_dict["hypernet_config"] if getattr(hypernet_config, "num_pre_head_layers", None) is None: hypernet_config.num_pre_head_layers = 4 if getattr(hypernet_config, "use_per_rank_bias", None) is None: hypernet_config.use_per_rank_bias = False ctx_encoder_args = state_dict["ctx_encoder_args"] model = cls(base_model, hypernet_config, ctx_encoder_args, **kwargs) model.load_state_dict(state_dict) return model def _init_model(self): # disable adapter of the base model # this only works with LoRA(?) # we disable to avoid peft lora computation self.base_model.disable_adapter_layers() self.hypernet = ( HyperLoRA(self.hypernet_config).to(self.device).to(torch.float32) ) layers = get_layers(self.base_model) lora_forward_fn = ( lora_forward_packed if self.use_sequence_packing else lora_forward ) for layer_idx in self.hypernet.layer_indices: for module_info in get_peft_modules(layers[layer_idx], self.peft_config): name = module_info["name"] module = module_info["module"] logger.debug(f"Applying LoRA forward to {name}") module.forward = partial( lora_forward_fn, self=module, lora_dropout_p=self.peft_config.lora_dropout, scaling=self.peft_config.lora_alpha, ) 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 = CTX_ENCODER_CLS[self.ctx_encoder_args.ctx_encoder_type]( encoder_model, self.ctx_encoder_args ) # delegate to base_model @property def config(self): return self.base_model.config @property def generation_config(self): return self.base_model.generation_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 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: if not self.hypernet.config.use_per_rank_bias: 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) # scale the biases by the output size # e.g., rank=16 gives bigger (2x?) gradient magnitudes at initialization # compared to rank=8 r = self.hypernet_config.lora_config.r bias_cat = bias_cat / r self.hypernet.head.bias.data[..., i, :, : bias_cat.shape[1]] = bias_cat 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) 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): 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() def remove_compile_prefix(sd: dict[str, Tensor]) -> dict[str, Tensor]: COMPILED_PREFIX = "_orig_mod." for k in list(sd.keys()): if k.startswith(COMPILED_PREFIX): sd[k[len(COMPILED_PREFIX) :]] = sd.pop(k) return sd load_result = self.hypernet.load_state_dict( remove_compile_prefix(state_dict), strict=True, # , *args, **kwargs ) logger.info(f"load result: {load_result}") return load_result def generate_weights( self, ctx_ids: Integer[Tensor, "bs ctx_len"], ctx_attn_mask: Integer[Tensor, "bs ctx_len"] | None = None, ctx_position_ids: Integer[Tensor, "bs ctx_len"] | None = 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, ) 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) return self.hypernet.generate_weights( ctx_features, ctx_attn_mask, ctx_position_ids ) def forward( self, ctx_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None, ctx_attn_mask: Integer[Tensor, "n_ctx ctx_len"] | None = None, ctx_position_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None, n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None, n_queries: Integer[Tensor, "n_ctx"] | None = None, return_generated_lora: bool | None = False, *model_inputs_args: Any, **model_inputs_kwargs: dict[str, Any], ) -> tuple | ModelOutput: """Forward pass of the modulated model.""" # TODO: allow more than one LoRA per sample (multi-lora) 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, ) ) 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 ) generated_loras, generated_layernorms = self.generate_weights( ctx_ids, ctx_attn_mask, ctx_position_ids ) generated_loras = combine_lora( generated_loras, n_ctx_chunks, aggregation=self.lora_aggregation, # lora_bias=self.hypernet.get_head_bias(), ) # 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 n_queries is None: if ctx_position_ids is None: n_queries = torch.ones( ctx_ids.shape[0], dtype=torch.int32, device=self.device ) else: # quite redundant (we do cu_seqlens many places) # TODO: compute cu_seqlens here and propagate that n_queries = torch.ones( (ctx_position_ids == 0).sum(), dtype=torch.int32, device=self.device ) apply_lora_to_layers( self.base_model, self.hypernet.layer_indices, generated_loras, n_queries, position_ids, ) 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_with_multi_loras( self, ctx_ids: Integer[Tensor, "bs ctx_length"], ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None, ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None, n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None, n_queries: Integer[Tensor, "n_ctx"] | None = None, lora_aggregation: Literal["mean", "sum"] = "sum", *model_inputs_args: Any, **model_inputs_kwargs: dict[str, Any], ): generated_loras, _ = self.generate_weights( ctx_ids, ctx_attn_mask, ctx_position_ids ) generated_loras = combine_lora( generated_loras, n_ctx_chunks, aggregation=lora_aggregation, # lora_bias=self.hypernet.get_head_bias(), ) # apply lora hook to the base model position_ids = ( model_inputs_kwargs["position_ids"] if "position_ids" in model_inputs_kwargs else None ) if n_queries is None: # assumes each group has one q n_queries = torch.ones( len(n_ctx_chunks), dtype=torch.int32, device=self.device ) apply_lora_to_layers( self.base_model, self.hypernet.layer_indices, generated_loras, n_queries, position_ids, # TODO: remove (cant be used with generation?) ) model_outputs = self.base_model.generate( *model_inputs_args, **model_inputs_kwargs ) return model_outputs @torch.inference_mode() def generate( self, # TODO: allow more than one LoRA per sample (multi-lora) ctx_ids: Integer[Tensor, "n_chunks ctx_length"] | None = None, ctx_attn_mask: Integer[Tensor, "n_chunks ctx_length"] | None = None, ctx_position_ids: Integer[Tensor, "n_chunks ctx_length"] | None = None, n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None, n_queries: Integer[Tensor, "n_ctx"] | None = 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.active_adapters 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, ) ) if ctx_ids is None and self.active_adapters: logger.info( ( "*" * 100, "\n\nUsing active LoRAs for generation\n\n", "*" * 100, ) ) 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 ) generated_loras, generated_layernorms = self.generate_weights( ctx_ids, ctx_attn_mask, ctx_position_ids ) # generated_loras = combine_lora( # generated_loras, # n_chunks=n_ctx_chunks, # aggregation="sum", # lora_bias=self.hypernet.get_head_bias(), # ) # # for generation the ctx_ids are batched not packed # ctx_ids_list = torch.split(ctx_ids, n_ctx_chunks, dim=0) # apply lora hook to the base model # TODO: we dont this position_ids for generation? position_ids = ( model_inputs_kwargs["position_ids"] if "position_ids" in model_inputs_kwargs else None ) if n_queries is None: if ctx_position_ids is None: n_queries = torch.ones( ctx_ids.shape[0], dtype=torch.int32, device=self.device ) else: # quite redundant (we do cu_seqlens many places) # TODO: compute cu_seqlens here and propagate that n_queries = torch.ones( (ctx_position_ids == 0).sum(), dtype=torch.int32, device=self.device ) apply_lora_to_layers( self.base_model, self.hypernet.layer_indices, generated_loras, n_queries, position_ids, ) model_outputs = self.base_model.generate( *model_inputs_args, **model_inputs_kwargs ) return model_outputs @contextmanager def apply_generated_layernorm( base_model: nn.Module, generated_layernorms: dict[str, Float[Tensor, "bs n_layers h"]] | None = None, layer_indices: Iterable[int] | None = None, position_ids: Integer[Tensor, "bs seq_len"] | None = 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: dict | None = None, layer_indices: Iterable[int] | None = None, position_ids: Integer[Tensor, "bs seq_len"] | None = 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, position_ids=position_ids, 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), # ) # 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", # weights_only=False, # ) # print(state_dict.keys()) # breakpoint() # model.load_state_dict(state_dict) # modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs) # print(modelout) # breakpoint()