add hyperlora w/ mean pooler + biggest head

hyperlora/configs.py

@@ -1,14 +1,13 @@
 import dataclasses
 import os
 import sys
-import yaml
 from dataclasses import dataclass, field
 from enum import Enum, auto
 from typing import Any, Dict, List, Literal, NewType, Optional, Tuple
 
+import yaml
 from transformers import MODEL_FOR_CAUSAL_LM_MAPPING, HfArgumentParser
 
-
 MODEL_CONFIG_CLASSES = list(MODEL_FOR_CAUSAL_LM_MAPPING.keys())
 MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
 
@@ -18,8 +17,8 @@ DataClassType = NewType("DataClassType", Any)
 
 class ArgumentParser(HfArgumentParser):
     def parse_yaml_and_args(
-        self, yaml_arg: str, other_args: Optional[List[str]] = None
+        self, yaml_arg: str, other_args: Optional[list[str]] = None
-    ) -> List[dataclass]:
+    ) -> list[dataclass]:
         """
         Parse a YAML file and overwrite the default/loaded values with the values provided to the command line.
 
@@ -84,7 +83,7 @@ class ArgumentParser(HfArgumentParser):
 
         return outputs
 
-    def parse(self) -> DataClassType | Tuple[DataClassType]:
+    def parse(self) -> DataClassType | tuple[DataClassType]:
         if len(sys.argv) == 2 and sys.argv[1].endswith(".yaml"):
             # If we pass only one argument to the script and it's the path to a YAML file,
             # let's parse it to get our arguments.

hyperlora/hooks.py

@@ -0,0 +1,185 @@
+import logging
+from operator import attrgetter
+from typing import Callable, Iterable, Optional
+
+import torch
+import torch.nn.functional as F
+from einops import einsum
+from jaxtyping import Float
+from torch import Tensor
+from torch.utils.hooks import RemovableHandle
+from utils import get_layers
+
+logger = logging.getLogger(__name__)
+
+
+def remove_hook_handles(handles: list[RemovableHandle]) -> None:
+    """
+    Removes all hook handles to detach them from the model.
+
+    Args:
+        handles (list[RemovableHandle]): A list of hook handles to be removed.
+    """
+    for handle in handles:
+        if handle is not None:
+            handle.remove()
+
+
+def apply_hook_to_model(
+    model: torch.nn.Module,
+    module_name: str,
+    pre_hook: Optional[Callable] = None,
+    post_hook: Optional[Callable] = None,
+) -> torch.nn.Module:
+    if not ((pre_hook is not None) or (post_hook is not None)):
+        raise ValueError("No hooks provided. Nothing to apply.")
+
+    module = attrgetter(module_name)(model)
+
+    if pre_hook is not None:
+        pre_hook_handle = module.register_forward_pre_hook(pre_hook)
+    if post_hook is not None:
+        post_hook_handle = module.register_forward_hook(post_hook)
+    return pre_hook_handle, post_hook_handle
+
+
+def apply_hook_to_layer(
+    layer: torch.nn.Module,
+    mname: str,
+    pre_hook: Optional[Callable] = None,
+    post_hook: Optional[Callable] = None,
+) -> tuple[Optional[RemovableHandle], Optional[RemovableHandle]]:
+    """
+    Applies pre and/or post hooks to a specific layer in the model.
+
+    Args:
+        layer (torch.nn.Module): The layer to which hooks will be applied.
+        mname (str): The name of the module within the layer.
+        pre_hook (Optional[Callable], optional): Function to be called before the forward pass.
+        post_hook (Optional[Callable], optional): Function to be called after the forward pass.
+
+    Returns:
+        Tuple[Optional[RemovableHandle], Optional[RemovableHandle]]:
+            Handles for the pre and post hooks.
+    """
+    if not ((pre_hook is not None) or (post_hook is not None)):
+        raise ValueError("No hooks provided. Nothing to apply.")
+
+    pre_hook_handle = None
+    post_hook_handle = None
+
+    if mname in ["q_proj", "k_proj", "v_proj", "o_proj", "qkv_proj"]:
+        mname = f"self_attn.{mname}"
+    module = attrgetter(mname)(layer)
+
+    if pre_hook is not None:
+        pre_hook_handle = module.register_forward_pre_hook(pre_hook)
+    if post_hook is not None:
+        post_hook_handle = module.register_forward_hook(post_hook)
+    return pre_hook_handle, post_hook_handle
+
+
+def apply_hook_to_layers(
+    model: torch.nn.Module,
+    module_names: list[str],
+    layer_indices: Iterable[int],
+    pre_hook: Optional[Callable] = None,
+    post_hook: Optional[Callable] = None,
+) -> list[RemovableHandle]:
+    """
+    Applies custom hooks to specified layers and modules in the model.
+
+    Args:
+        model (torch.nn.Module): The model to which hooks will be applied.
+        module_names (list[str]): Names of the modules to hook.
+        layer_indices (Iterable[int]): Indices of the layers to hook.
+        pre_hook (Optional[Callable], optional): Function to be called before the forward pass.
+        post_hook (Optional[Callable], optional): Function to be called after the forward pass.
+
+    Returns:
+        list[RemovableHandle]: A list of hook handles.
+    """
+
+    layers = get_layers(model)
+    out_handles = []
+    c = 0
+    for layer_idx in layer_indices:
+        layer = layers[layer_idx]
+        for mname in module_names:
+            handles = apply_hook_to_layer(
+                layer,
+                mname,
+                pre_hook=pre_hook,
+                post_hook=post_hook,
+            )
+            out_handles += handles
+            if handles[0] is not None or handles[1] is not None:
+                c += 1
+    if c == 0:
+        raise ValueError(
+            "No forward hooks applied. Something might be wrong. "
+            "Check if the module names are correct."
+        )
+    return out_handles
+
+
+def add_generated_lora_hook(
+    model: torch.nn.Module,
+    module_name: str,
+    layer_index: int,
+    A: Float[Tensor, "bs r d_in"],
+    B: Float[Tensor, "bs d_out r"],
+    scaling: float,
+    input_dropout: float,
+    training: bool,
+) -> list[RemovableHandle]:
+    """
+    Adds LoRA hooks to the specified modules and layers of the model.
+
+    Args:
+        model (torch.nn.Module): The model to which hooks will be applied.
+        module_name (str): Name of the module to hook.
+        A (Tensor): Weight tensor A.
+        B (Tensor): Weight tensor B.
+        scaling (float): Scaling factor.
+        input_dropout (float): Dropout rate for input.
+        training (bool): Flag indicating if the model is in training mode.
+
+    Returns:
+        list[RemovableHandle]: A list of hook handles.
+    """
+
+    def lora_hook(
+        module: torch.nn.Module,
+        args: tuple | Float[Tensor, "bs seq_len d_in"],
+        output: Float[Tensor, "bs seq_len d_out"],
+    ) -> Float[Tensor, "bs seq_len d_out"]:
+
+        if isinstance(output, tuple):
+            model_out = output[0]
+        else:
+            model_out = output
+
+        x = args[0].to(A.dtype)  # [bs, seq_len, d_in]
+
+        delta_x = F.dropout(x, input_dropout, training)
+        delta_x = einsum(A, delta_x, "bs r d_in, bs seq_len d_in -> bs seq_len r")
+        delta_x = einsum(B, delta_x, "bs d_out r, bs seq_len r -> bs seq_len d_out")
+        delta_x = delta_x * scaling
+
+        # # A and B repeat for each input token
+        # lora_A = A.repeat_interleave(seq_len, dim=0)
+        # lora_B = B.repeat_interleave(seq_len, dim=0)
+        # x = x.reshape(bs * seq_len, 1, -1)
+        # delta_x = (
+        #     torch.bmm(torch.bmm(F.dropout(x, input_dropout, training), lora_A), lora_B)
+        #     * scaling
+        # )
+
+        newoutput = model_out + delta_x.to(model_out.dtype)
+        if isinstance(output, tuple):
+            return (newoutput, *output[1:])
+        else:
+            return newoutput
+
+    return apply_hook_to_layers(model, [module_name], [layer_index], post_hook=lora_hook)

hyperlora/intx_sft.py

@@ -5,7 +5,6 @@ import time
 
 import numpy as np
 import torch
-from configs import CtxTrainingArguments, ExperimentSetup, LoRAArguments, ModelArguments
 from data_utils import (
     convert_ctx_prompt_response_to_messages,
     get_preprocessing_fn,
@@ -13,8 +12,6 @@ from data_utils import (
     tokenize_chat_messages,
 )
 from datasets import load_dataset
-
-from configs import ArgumentParser
 from model_loading import get_lora_config, get_model_and_tokenizer
 from modeling_utils import ModulatedPretrainedModel
 from training_utils import TRAINING_TASK, train_model
@@ -28,6 +25,14 @@ from transformers import (
 )
 from utils import log_num_train_params
 
+from configs import (
+    ArgumentParser,
+    CtxTrainingArguments,
+    ExperimentSetup,
+    LoRAArguments,
+    ModelArguments,
+)
+
 logger = logging.getLogger(__name__)
 
 

hyperlora/modeling_utils.py

@@ -1,26 +1,82 @@
 import logging
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum
 from functools import partial
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Iterable, Optional, Tuple, Union
-from einops import rearrange, repeat
 
 import torch
+from einops import rearrange, repeat, unpack
+from einops.layers.torch import EinMix as Mix
+from hooks import add_generated_lora_hook, remove_hook_handles
 from jaxtyping import Float, Integer
+from model_loading import get_lora_config, get_model_and_tokenizer
 from peft import LoraConfig
+from pooling import POOL_FN, get_pooling_fn
 from torch import Tensor, nn
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import ModelOutput
-
+from utils import get_lora_module_names, get_num_layers, get_peft_in_out_features
-from model_loading import get_lora_config, get_model_and_tokenizer
-from utils import get_num_layers
-from pooling import get_pooling_fn, POOL_FN
 
 logger = logging.getLogger(__name__)
 
 
+AGGREGATOR_TYPE = Enum("AGGREGATOR_TYPE", ["POOLER", "PERCEIVER"])
+
+
+@dataclass
+class AggregatorConfig:
+    feature_size: int
+    num_layers: int
+    num_modules: int
+    pooling_type: POOL_FN
+
+
+def get_aggregator_config(
+    model: PreTrainedModel,
+    pooling_type: POOL_FN = POOL_FN.MEAN,
+):
+    lora_config = model.peft_config["default"]
+    return AggregatorConfig(
+        feature_size=model.config.hidden_size,
+        num_layers=get_num_layers(model),
+        num_modules=len(lora_config.target_modules),
+        pooling_type=pooling_type,
+    )
+
+
+@dataclass
+class HypernetConfig:
+    latent_size: int
+    lora_config: LoraConfig
+    module_names: dict[str, list[str]]
+    layer_indices: Iterable[int]
+    feature_sizes: tuple[dict[str, int], dict[str, int]]
+    aggregator_type: AGGREGATOR_TYPE
+    aggregator_config: AggregatorConfig
+
+
+def get_hypernet_config(model: PreTrainedModel, latent_size: int = 256):
+    lora_config = model.peft_config["default"]
+    indices = torch.arange(get_num_layers(model), device=model.device)
+    return HypernetConfig(
+        latent_size=latent_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_type=AGGREGATOR_TYPE.POOLER,
+        aggregator_config=get_aggregator_config(model, POOL_FN.MEAN),
+    )
+
+
 # TODO: implement Perceiver
-class Perceiver(nn.Module): ...
+class Perceiver(nn.Module):
+    """perceiver w/ bottleneck size = n_modules * n_layers"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        pass
 
 
 class Mixer(nn.Module):
@@ -82,19 +138,22 @@ class Pooler(nn.Module):
         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)
+        self.feature_proj = nn.Linear(feature_size, output_size // 2)
-        self.ln = nn.LayerNorm(output_size)
+        self.ln = nn.LayerNorm(output_size // 2)
         self.layer_embs = nn.Sequential(
-            nn.Embedding(num_layers, output_size // 2),
+            nn.Embedding(num_layers, output_size // 4),
-            nn.LayerNorm(output_size // 2),
+            nn.LayerNorm(output_size // 4),
         )
         self.module_embs = nn.Sequential(
-            nn.Embedding(num_modules, output_size // 2),
+            nn.Embedding(num_modules, output_size // 4),
-            nn.LayerNorm(output_size // 2),
+            nn.LayerNorm(output_size // 4),
         )
-        self.mixer = Mixer(output_size * 2, output_size * 8, output_size * 2)
+        self.mixer = Mixer(output_size, output_size * 4, output_size)
-        self.mlp = MLPResidualBlock(output_size * 2, output_size * 8, output_size * 2)
+        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))
@@ -110,7 +169,7 @@ class Pooler(nn.Module):
         x = self.ln(self.feature_proj(self.pool_fn(features, attn_mask).float()))
         x = repeat(
             x,
-            "bs d -> bs n_modules n_layers d",
+            "bs d -> bs n_layers n_modules d",
             n_modules=self.num_modules,
             n_layers=self.num_layers,
         )
@@ -118,7 +177,7 @@ class Pooler(nn.Module):
         layer_embs = self.layer_embs(self.layer_indices)  # [num_layers, d]
         layer_embs = repeat(
             layer_embs,
-            "n_layers d -> bs n_modules n_layers d",
+            "n_layers d -> bs n_layers n_modules d",
             bs=bs,
             n_modules=self.num_modules,
         )
@@ -126,7 +185,7 @@ class Pooler(nn.Module):
         module_embs = self.module_embs(self.module_indices)  # [num_modules, d]
         module_embs = repeat(
             module_embs,
-            "n_modules d -> bs n_modules n_layers d",
+            "n_modules d -> bs n_layers n_modules d",
             bs=bs,
             n_layers=self.num_layers,
         )
@@ -135,53 +194,112 @@ class Pooler(nn.Module):
         return self.mlp(self.mixer(emb))
 
 
-AGGREGATOR = Enum("AGGREGATOR", ["POOLER", "PERCEIVER"])
-
 AGGREGATOR_CLS = {
-    # AGGREGATOR.MEAN: MeanPool,
+    AGGREGATOR_TYPE.POOLER: Pooler,
-    # AGGREGATOR.MAX: MaxPool,
+    AGGREGATOR_TYPE.PERCEIVER: Perceiver,
-    # AGGREGATOR.LAST_TOKEN: LastTokenPool,
-    AGGREGATOR.POOLER: Pooler,
-    AGGREGATOR.PERCEIVER: Perceiver,
 }
 
 
 class HyperLoRA(nn.Module):
     def __init__(
         self,
-        lora_config: LoraConfig,
+        # latent_size: int,
-        layer_indices: Integer[Tensor, "num_layers"],
+        # lora_config: LoraConfig,
-        aggregator: AGGREGATOR,
+        # layer_indices: Integer[Tensor, "num_layers"],
-        aggregator_kwargs: Optional[dict[str, Any]] = None,
+        # in_out_features: dict,
+        # aggregator_type: AGGREGATOR_TYPE,
+        # aggregator_kwargs: dict,
+        config: HypernetConfig,
     ):
         super().__init__()
-        # TODO: aggregator should output
-        # [bs, n_modules, n_layers, 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
 
         # NOTE: this class then only handles the output space of the hypernet
         # e.g., shared_AB_head, per_rank_gen, etc.
-        self.aggregator = AGGREGATOR_CLS[aggregator](**(aggregator_kwargs or {}))
+        # TODO: add different output spaces
 
-        self.lora_config = lora_config
+        # 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.aggregator = AGGREGATOR_CLS[config.aggregator_type](
+            **vars(config.aggregator_config),
+            output_size=config.latent_size,
+        )
 
-        self.target_modules = lora_config.target_modules
+        self.lora_config = config.lora_config
-        self.layer_indices = layer_indices
 
-        self.in_features = ...
+        self.target_modules = self.lora_config.target_modules
-        self.out_features = ...
+        self.layer_indices = config.layer_indices
+
+        # TODO: add lightweight LoRA i.e., a projection layer of the input of LoRA
+        # have to also change in_d and out_d accordingly
+        self.in_d, self.out_d = config.feature_sizes
+
+        # TODO: add different output spaces
+
+        self.layers = MLPResidualBlock(
+            input_size=config.latent_size,
+            hidden_size=config.latent_size * 4,
+            output_size=config.latent_size,
+        )
+
+        self.head = Mix(
+            "bs n_layers n_modules d -> bs n_layers r out_d",
+            weight_shape="n_modules d r out_d",
+            bias_shape=None,  # no bias
+            n_modules=len(self.target_modules),
+            d=config.latent_size,
+            r=config.lora_config.r,
+            out_d=sum(self.in_d[m] + self.out_d[m] for m in self.target_modules),
+        )
+
+    def _to_lora_dict(
+        self, flat_loras: Float[Tensor, "bs n_layers r _"]
+    ) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]:
+        # list of [bs, n_layers, r, in_out_dim]
+        # and in_out_dim might vary across modules
+        loras = unpack(
+            flat_loras,
+            [[self.in_d[m] + self.out_d[m]] for m in self.target_modules],
+            "bs n_layers r *",
+        )
+
+        # dict of {module:
+        #   {A: [bs, n_layers, r, in_dim],
+        #    B: [bs, n_layers, r, out_dim]}}
+        lora_dict = dict()
+        for module, lora in zip(self.target_modules, loras):
+            A, B = unpack(
+                lora,
+                [[self.in_d[module]], [self.out_d[module]]],
+                "bs n_layers r *",
+            )
+            # transpose B
+            B = rearrange(B, "bs n_layers r d_out -> bs n_layers d_out r")
+            lora_dict[module] = dict(A=A, B=B)
+
+        return lora_dict
 
     def forward(
         self,
         features: Float[Tensor, "bs seq_len feature_dim"],
         attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None,
-    ) -> Float[Tensor, "bs num_modules num_layers lora_r in_out_dim"]:
+    ):
 
-        emb = self.aggregator(features, attn_mask)  # [bs, n_features, feature_dim]
+        # [bs, n_layers, n_modules, feature_dim]
-        # loras = self.layers(emb)
+        emb = self.aggregator(features, attn_mask)
-        # return loras
+
-        return emb
+        # [bs, n_layers, r, in_out_dim_mod1 + in_out_dim_mod2 + ...]
+        flat_loras = self.head(self.layers(emb))
+
+        return flat_loras
+
+    def generate_loras(
+        self,
+        features: Float[Tensor, "bs seq_len feature_dim"],
+        attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None,
+    ):
+        flat_loras = self.forward(features, attn_mask)
+        return self._to_lora_dict(flat_loras)
 
 
 class ModulatedPretrainedModel(nn.Module):
@@ -214,13 +332,13 @@ class ModulatedPretrainedModel(nn.Module):
 
     def forward(
         self,
-        ctx_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None,
+        ctx_features: Optional[Float[Tensor, "bs ctx_length feature_dim"]] = None,
         ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_length"]] = None,
         **model_inputs_kwargs: dict[str, Any],
     ) -> Union[tuple, ModelOutput]:
         """Forward pass of the modulated model."""
 
-        if ctx_ids is None:
+        if ctx_features is None:
             logger.warning(
                 "No context ids provided, using the base model for the forward pass"
             )
@@ -228,21 +346,47 @@ class ModulatedPretrainedModel(nn.Module):
             # model_outputs.generated_loras = None
             return model_outputs
 
-        loss = ...
+        generated_loras = self.hypernet.generate_loras(ctx_features, ctx_attn_mask)
-        logits = ...
-        # TODO: get ctx_features offline
-        features = self.get_ctx_features(ctx_ids, ctx_attn_mask)
-        generated_loras = self.hypernet(features, ctx_attn_mask)
 
         # apply lora hook to the base model
-        self.apply_lora_hook(generated_loras)
+        # self.apply_generated_loras(generated_loras)
-        model_outputs = self.base_model(**model_inputs_kwargs)
+        with apply_generated_loras(
+            self.base_model,
+            generated_loras,
+            self.hypernet.layer_indices,
+            self.training,
+        ):
+            model_outputs = self.base_model(**model_inputs_kwargs)
         # model_outputs.generated_loras = generated_loras
 
         return model_outputs
 
-    def apply_lora_hook(self, generated_loras):
+
-        pass
+@contextmanager
+def apply_generated_loras(
+    base_model: nn.Module,
+    generated_loras: dict,
+    layer_indices: Iterable[int],
+    training: bool = False,
+):
+    try:
+        hooks = []
+        for module_name in generated_loras:
+            for layer_idx in layer_indices:
+                hooks += add_generated_lora_hook(
+                    base_model,
+                    module_name,
+                    layer_idx,
+                    A=generated_loras[module_name]["A"][:, layer_idx],
+                    B=generated_loras[module_name]["B"][:, layer_idx],
+                    scaling=base_model.peft_config["default"].lora_alpha,
+                    input_dropout=base_model.peft_config["default"].lora_dropout,
+                    training=training,
+                )
+
+        yield base_model
+    finally:
+        remove_hook_handles(hooks)
 
 
 if __name__ == "__main__":
@@ -256,20 +400,19 @@ if __name__ == "__main__":
         peft_config=get_lora_config(model_name),
     )
     print(base_model)
-    peft_config = base_model.peft_config["default"]
+    # lora_config = base_model.peft_config["default"]
+    # in_d, out_d = get_peft_in_out_features(base_model, peft_config=lora_config)
+
+    # hypernet = HyperLoRA(
+    #     lora_config,
+    #     torch.arange(get_num_layers(base_model), device=base_model.device),
+    #     in_out_features={"in": in_d, "out": out_d},
+    #     aggregator_type=AGGREGATOR_TYPE.POOLER,
+    #     aggregator_config=get_aggregator_config(base_model, POOL_FN.MEAN),
+    # ).to(base_model.device)
+
+    hypernet = HyperLoRA(get_hypernet_config(base_model)).to(base_model.device)
 
-    hypernet = HyperLoRA(
-        peft_config,
-        torch.arange(get_num_layers(base_model), device=base_model.device),
-        aggregator=AGGREGATOR.POOLER,
-        aggregator_kwargs={
-            "feature_size": base_model.config.hidden_size,
-            "output_size": 128,
-            "num_layers": get_num_layers(base_model),
-            "num_modules": len(peft_config.target_modules),
-            "pooling_type": POOL_FN.MEAN,
-        },
-    ).to(base_model.device)
     model = ModulatedPretrainedModel(base_model, hypernet).to(base_model.device)
     print(model)
 
@@ -279,7 +422,7 @@ if __name__ == "__main__":
     ctx_attn_mask = ctx_inputs["attention_mask"]
     print(ctx_features.shape)
 
-    hypernet.eval()
+    model.eval()
     agg_features = hypernet.aggregator(ctx_features, ctx_attn_mask)
     print(agg_features)
     print(agg_features.shape)
@@ -287,4 +430,14 @@ if __name__ == "__main__":
     hnetout = hypernet(ctx_features, ctx_attn_mask)
     print(hnetout)
     print(hnetout.shape)
+
+    prompt_msg = "hello"
+    prompt_inputs = tokenizer(prompt_msg, return_tensors="pt").to(model.device)
+
+    basemodelout = model.base_model(**prompt_inputs)
+    print(basemodelout)
+
+    modelout = model(ctx_features, ctx_attn_mask, **prompt_inputs)
+    print(modelout)
+
     breakpoint()

hyperlora/pooling.py

@@ -1,8 +1,8 @@
-from typing import Optional
-from jaxtyping import Float, Integer
 from enum import Enum
+from typing import Optional
 
 import torch
+from jaxtyping import Float, Integer
 from torch import Tensor
 from torch.nn import functional as F
 

hyperlora/utils.py

@@ -1,8 +1,31 @@
 import logging
+from contextlib import contextmanager
+from typing import Iterable, Optional
+
+import torch
+from peft import PeftConfig, PeftModel
+from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists
+from peft.utils import get_peft_model_state_dict
 
 logger = logging.getLogger(__name__)
 
 
+# taken from https://discuss.pytorch.org/t/opinion-eval-should-be-a-context-manager/18998/3
+@contextmanager
+def evaluating(*models):
+    """Temporarily switch to evaluation mode."""
+    is_training = [model.training if model is not None else False for model in models]
+    try:
+        for model in models:
+            if model is not None:
+                model.eval()
+        yield models
+    finally:
+        for model, training in zip(models, is_training):
+            if model is not None:
+                model.train(training)
+
+
 def get_layers(model):
     if hasattr(model, "model"):
         return get_layers(model.model)
@@ -36,3 +59,82 @@ def log_num_train_params(model):
         f"|| all params: {num_total_params:,d} "
         f"|| trainable%: {100 * num_trainable_params / num_total_params:.4f}"
     )
+
+
+def get_peft_in_out_features(
+    model: PeftModel,
+    peft_config: Optional[PeftConfig] = None,
+) -> tuple[dict[str, int], dict[str, int]]:
+
+    if peft_config is None:
+        peft_config = model.peft_config["default"]
+    in_features = dict()
+    out_features = dict()
+    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, torch.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
+
+        if name not in in_features:
+            in_features[name] = module.in_features
+            out_features[name] = module.out_features
+        else:
+            # assumes each module has the same input and output features
+            assert in_features[name] == module.in_features
+            assert out_features[name] == module.out_features
+
+    return in_features, out_features
+
+
+def generated_lora_to_state_dict(
+    lora_dict: dict,
+    module_names: dict,
+    target_modules: list[str],
+    layer_indices: Iterable[int],
+) -> dict:
+    lora_state_dict = dict()
+    for target_module in target_modules:
+        for layer_idx in layer_indices:
+            for module_name in module_names[target_module][layer_idx]:
+                if "lora_A" in module_name:
+                    lora_state_dict[module_name] = (
+                        lora_dict[target_module]["A"][layer_idx].cpu().contiguous()
+                    )
+                elif "lora_B" in module_name:
+                    lora_state_dict[module_name] = (
+                        lora_dict[target_module]["B"][layer_idx].cpu().contiguous()
+                    )
+                else:
+                    raise ValueError(f"Unexpected module name: {module_name}")
+    return lora_state_dict
+
+
+def get_lora_module_names(
+    model: PeftModel,
+    target_modules: list[str],
+    layer_indices: Iterable[int],
+) -> dict[str, list[str]]:
+    module_names = {
+        target_module: [[] for _ in range(len(layer_indices))]
+        for target_module in target_modules
+    }
+    for k in get_peft_model_state_dict(model):
+        if "lora" not in k:
+            continue
+        layer_idx = int(k.split("layers.")[-1].split(".")[0])
+        if layer_idx in layer_indices:
+            for target_module in target_modules:
+                if target_module in k:
+                    module_names[target_module][layer_idx].append(k)
+                    break
+    return module_names