import torch import torch.nn as nn from torch import Tensor import math import torch.nn.functional as F from transformers import AutoConfig, PretrainedConfig from jaxtyping import Float from dataclasses import asdict, dataclass from typing import List, Optional, Tuple, Dict import einops from .configIBA import MainConfig, HyperXSConfig, TrainingConfig def transpose(weight, fan_in_fan_out): return weight.T if fan_in_fan_out else weight class LoraLayer: def __init__( self, rank: int, train_cfg: TrainingConfig, # batch: int, lora_alpha: int, lora_dropout: float, ): self.rank = rank self.batch_train = train_cfg.per_device_train_batch_size self.batch_valid = train_cfg.per_device_eval_batch_size # self.batch = batch self.lora_alpha = lora_alpha # Optional dropout if lora_dropout > 0.0: self.lora_dropout = nn.Dropout(p=lora_dropout) else: self.lora_dropout = lambda x: x # Mark the weight as unmerged self.disable_adapters = False class LoraXSLinear(nn.Linear, LoraLayer): # Lora implemented in a dense layer def __init__( self, in_features: int, out_features: int, train_cfg: TrainingConfig, rank: int = 64, # batch: int = 32, lora_alpha: int = 1, lora_dropout: float = 0.0, fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) **kwargs, ): nn.Linear.__init__(self, in_features, out_features, **kwargs) LoraLayer.__init__(self, rank=rank, train_cfg=train_cfg, lora_alpha=lora_alpha, lora_dropout=lora_dropout) self.fan_in_fan_out = fan_in_fan_out # Actual trainable parameters if rank > 0: # self.register_buffer("lora_A", self.weight.new_zeros(in_features, rank), persistent=False) self.register_buffer("lora_A", torch.zeros([in_features, rank]), persistent=True) self.register_buffer("lora_B", torch.zeros([rank, out_features]), persistent=True) self.scaling = self.lora_alpha / self.rank # Freezing the pre-trained weight matrix self.weight.requires_grad = False self.lora_R = None # self.lora_A.weight.requires_grad = False # self.lora_B.weight.requires_grad = False if fan_in_fan_out: self.weight.data = self.weight.data.T self.reset_parameters() def reset_parameters(self): nn.Linear.reset_parameters(self) if hasattr(self, "lora_A"): # initialize A the same way as the default for nn.Linear and B to zero nn.init.kaiming_uniform_(self.lora_A, mode='fan_out', a=math.sqrt(5)) nn.init.kaiming_uniform_(self.lora_B, mode='fan_in', a=math.sqrt(5)) # def train(self, mode: bool = True): # nn.Linear.train(self, mode) def set_R(self, R: torch.Tensor): self.lora_R = R def decompose_weight_svd(self, rank): W = self.weight.data device, dtype = W.device, W.dtype #out_features, in_features = W.shape try: U, S, Vt = torch.linalg.svd(W,full_matrices=False) except torch.linalg.LinAlgError as e: print(f"SVD computation failed: {e}") return None, None # Set first r-rank columns U_r = U[:, :rank] # Shape: (d, r) S_r_values = S[:rank] sqrt_S_r_diag = torch.diag(torch.sqrt(S_r_values)) # Shape: (r, r) Vt_r = Vt[:rank, :] # Shape: (r, e) B = U_r @ sqrt_S_r_diag # Shape: (d, r) A = sqrt_S_r_diag @ Vt_r # Shape: (r, d) #return B.to(device, dtype), A.to(device, dtype) self.lora_A = A.T.to(device, dtype) self.lora_B = B.T.to(device, dtype) # Safer way to do with trainable params # with torch.no_grad(): # self.lora_A.T.weight.copy_(A.to(device, dtype)) # self.lora_B.T.weight.copy_(B.to(device, dtype)) def forward(self, x: torch.Tensor): previous_dtype = self.weight.dtype if self.disable_adapters: result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) elif self.rank > 0: result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if self.lora_R is not None: lora_R = self.lora_R result = result + (self.lora_dropout(x) @ self.lora_A) @ (lora_R @ self.lora_B) * self.scaling # else: # # unapplied layers else: result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias) if result.dtype != previous_dtype: result = result.to(previous_dtype) return result class HyperNetXSexp(nn.Module): def __init__( self, hyperxs_cfg: HyperXSConfig, hf_model_cfg: PretrainedConfig, ): super(HyperNetXSexp, self).__init__() self.n_modules = hyperxs_cfg.modules_per_layer # qkvo attn, up down gate mlp self.rank = hyperxs_cfg.lora_attn_dim # rank self.latent_feature_dim = hyperxs_cfg.latent_feature_dim # latent feature: embedding -> latent self.module_embed_dim = hyperxs_cfg.module_embed_dim self.layer_embed_dim = hyperxs_cfg.layer_embed_dim self.hyper_out = hyperxs_cfg.lora_attn_dim ** 2 # n_flat_indim = hf_model_cfg.hidden_size * hyperxs_cfg.n_cross_attn_tokens + self.module_embed_dim + self.layer_embed_dim # hyper_in_dim = n_flat_indim = self.latent_feature_dim * hyperxs_cfg.n_cross_attn_tokens + self.module_embed_dim + self.layer_embed_dim n_flat_outdim = hyperxs_cfg.out_proj_dim * hyperxs_cfg.n_cross_attn_tokens n_proj = 4 * n_flat_outdim self.latent_proj = nn.Linear(hf_model_cfg.hidden_size, self.latent_feature_dim) # rescale the embedđing first self.mixture = nn.Linear(n_flat_indim, n_flat_outdim) self.c_fc = nn.Linear(n_flat_outdim, n_proj) self.c_proj = nn.Linear(n_proj, self.hyper_out) self.act = nn.GELU() # Post-layer Normalization # self.ln_latent = nn.LayerNorm(self.latent_feature_dim, eps=hyperxs_cfg.layer_norm_epsilon) # self.ln_1 = nn.LayerNorm(n_flat_outdim, eps=hyperxs_cfg.layer_norm_epsilon) # self.ln_2 = nn.LayerNorm(n_proj, eps=hyperxs_cfg.layer_norm_epsilon) self.ln_latent = nn.LayerNorm(hf_model_cfg.hidden_size, eps=hyperxs_cfg.layer_norm_epsilon) self.ln_1 = nn.LayerNorm(n_flat_indim, eps=hyperxs_cfg.layer_norm_epsilon) self.ln_2 = nn.LayerNorm(n_flat_outdim, eps=hyperxs_cfg.layer_norm_epsilon) # A lookup table for each layer self.layer_embedding = nn.Embedding(hf_model_cfg.num_hidden_layers, self.layer_embed_dim) # Embedding for MLP self.module_embedding = nn.Embedding(self.n_modules, self.module_embed_dim) self.hyperxs_cfg = hyperxs_cfg self.hf_model_cfg = hf_model_cfg self.reset_parameters() def reset_parameters(self): # Initialize the MLP layers INIT_STD = 1e-3 nn.init.kaiming_normal_(self.latent_proj.weight, a=0, mode='fan_in', nonlinearity='leaky_relu') nn.init.constant_(self.latent_proj.bias, 0) nn.init.kaiming_normal_(self.mixture.weight, a=0, mode='fan_in', nonlinearity='leaky_relu') # nn.init.normal_(self.mixture.weight, mean=0.0, std=INIT_STD) nn.init.constant_(self.mixture.bias, 0) nn.init.kaiming_normal_(self.c_fc.weight, a=0, mode='fan_in', nonlinearity='leaky_relu') # nn.init.normal_(self.c_fc.weight, mean=0.0, std=INIT_STD) nn.init.constant_(self.c_fc.bias, 0) nn.init.normal_(self.layer_embedding.weight, mean=0.0, std=INIT_STD) # partly zeros for the last layer # nn.init.kaiming_normal_(self.c_proj.weight, a=0, mode='fan_in', nonlinearity='leaky_relu') nn.init.constant_(self.c_proj.weight, 0) nn.init.constant_(self.c_proj.bias, 0) # with torch.no_grad(): # # Get the dimensions for loraB and loraA per rank. [B_part, A_part] for each rank. # dim_b = self.outW[0] # dim_a = self.outW[1] # dim_per_rank = dim_b + dim_a # # It starts as all zeros, so the loraB part is already correct. # new_bias = torch.zeros_like(self.c_proj.bias) # # Reshape the flat bias vector into (rank, dim_per_rank) for easy manipulation. # new_bias_reshaped = new_bias.view(self.rank, dim_per_rank) # # Select the part of the bias that corresponds to loraA for all ranks. # # This is the slice from dim_b to the end for each rank. # bias_a_part = new_bias_reshaped[:, dim_b:] # # Initialize this loraA part with a small normal distribution. # # A small standard deviation is crucial to keep the initial LoRA adjustment small. # nn.init.kaiming_normal_(bias_a_part, a=0, mode='fan_in', nonlinearity='leaky_relu') #, mean=0.0, std=INIT_STD) # self.c_proj.bias.data.copy_(new_bias) def forward(self, x: Float[Tensor, 'b s f'], layer_idx) -> Float[Tensor, 'b r in out']: batch_size = x.shape[0] dtype_in = x.dtype x = x.to(self.latent_proj.weight.dtype) # preprocess x = self.ln_latent(x) x = self.latent_proj(x) # x = self.ln_latent(x) # flatten x = einops.rearrange(x, 'batch seq fea -> batch (seq fea)') # get weight from mlp_embedding module_embedding = self.module_embedding.weight # (n_mlp, embed_dim) # mlp_embedding = mlp_embedding[None, ...] module_embedding = module_embedding.expand(batch_size, -1, -1) x = x[:, None, ...] x = x.expand(-1, self.n_modules, -1) # Concatenate by the last dim & rearrange into 2D x = torch.cat((module_embedding, x), dim=-1) x = einops.rearrange(x, 'batch n_modules in_dim -> (batch n_modules) in_dim') # Add parameters to distinguish adapters if self.layer_embed_dim > 0: # Get the layer_embedding (1, embedding) -> (embedding) layer_embedding = self.layer_embedding(torch.tensor(layer_idx, device=x.device)) # Optimize the memory layer_embedding = layer_embedding.expand(batch_size, self.n_modules, -1) layer_embedding = einops.rearrange(layer_embedding, 'batch n_modules in_dim -> (batch n_modules) in_dim') x = torch.cat((layer_embedding, x), dim=-1) assert x.shape == (batch_size*self.n_modules, self.mixture.weight.data.shape[1]), 'Wrong at hypernetMLP.forward.x' # Post LayerNorm h = self.ln_1(x) h = self.mixture(x) # h = self.ln_1(h) h = self.act(h) # 2nd layer h = self.ln_2(h) h = self.c_fc(h) # h = self.ln_2(h) h = self.act(h) # 3rd layer h = self.c_proj(h) h = einops.rearrange(h, '(batch n_modules) (rank r) -> batch n_modules rank r', batch = batch_size, n_modules=self.n_modules, rank = self.rank, r = self.rank) h = h.to(dtype_in) return h def test_hypernet(): """ A simple test function for the HyperNetMLP class. Given empty B @ A """ mainCfg=MainConfig() print(mainCfg) hf_model_cfg = AutoConfig.from_pretrained( mainCfg.model.base_model_name ) print(hf_model_cfg) print("--- Starting HyperNetMLP Test ---") # 1. Define parameters for the test in_features = hf_model_cfg.hidden_size # 768 reduced_dim = 128 out_features = 256 batch_size = 27 rank = 30 outW = [768, 2*768] n_mlp=2 input_tensor = torch.randn(batch_size, mainCfg.hyperxs.n_cross_attn_tokens, in_features) model = HyperNetXSexp(mainCfg.hyperxs, hf_model_cfg) count_parameters(model) # print(model) output = model(input_tensor, layer_idx=torch.tensor(1, dtype=torch.long)) print('output shape', output.shape) B = output[:,1,:,:768] print('input shape', input_tensor.shape) print('output shape and sum of B', output.shape, output.sum(), B.sum()) if output.shape == (batch_size, n_mlp, rank, rank) and B.sum().item()==0: print("\n--- HyperNetMLP Test Passed Successfully! ✅ ---") def count_parameters(model:nn.Module): print(f'Counting params in {model.__class__.__name__}') total_params = 0 # Use a set to store the IDs of parameters that have already been counted counted_param_ids = set() print(f"{'Parameter Name':^60} | {'Shape':^20} | {'Num Params':^20}") print("-" * 110) for name, parameter in model.named_parameters(): if not parameter.requires_grad: continue # if not 'hypernet' in name or 'dummy' in name: # continue # Get the unique ID of the parameter tensor in memory param_id = id(parameter) if param_id in counted_param_ids: # Optional: print a message to verify that sharing is working print(f"Skipping shared parameter: {name}") continue counted_param_ids.add(param_id) shape = list(parameter.shape) # the number of parameters in this layer num_params = parameter.numel() # layer name and n_params # print(f"{name:<50} | {num_params:<10,}") # if 'hypernet' in name or 'dummy' in name: print(f"{name:<60} | {str(shape):<25} | {num_params:,}") total_params += num_params print(f"Model: {model.__class__.__name__} Total Trainable Params: {total_params:,}") return total_params if __name__ == "__main__": print("Hello world from iba_lora") mainCfg=MainConfig() # print(mainCfg) hf_model_cfg = AutoConfig.from_pretrained( mainCfg.model.base_model_name ) # print(hf_model_cfg) print('-'*50) test_hypernet()