diff --git a/hyperlora/configs.py b/hyperlora/configs.py index 5ab95b4..9f98304 100644 --- a/hyperlora/configs.py +++ b/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 - ) -> List[dataclass]: + self, yaml_arg: str, other_args: Optional[list[str]] = None + ) -> 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. diff --git a/hyperlora/hooks.py b/hyperlora/hooks.py new file mode 100644 index 0000000..5b75df7 --- /dev/null +++ b/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) diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index b925085..000d4ea 100644 --- a/hyperlora/intx_sft.py +++ b/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__) diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index 28bdd8c..393c3d2 100644 --- a/hyperlora/modeling_utils.py +++ b/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 einops import rearrange, repeat +from typing import Any, Iterable, Optional, Tuple, Union 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 model_loading import get_lora_config, get_model_and_tokenizer -from utils import get_num_layers -from pooling import get_pooling_fn, POOL_FN +from utils import get_lora_module_names, get_num_layers, get_peft_in_out_features 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.ln = nn.LayerNorm(output_size) + 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 // 2), - nn.LayerNorm(output_size // 2), + nn.Embedding(num_layers, output_size // 4), + nn.LayerNorm(output_size // 4), ) self.module_embs = nn.Sequential( - nn.Embedding(num_modules, output_size // 2), - nn.LayerNorm(output_size // 2), + nn.Embedding(num_modules, output_size // 4), + nn.LayerNorm(output_size // 4), ) - self.mixer = Mixer(output_size * 2, output_size * 8, output_size * 2) - self.mlp = MLPResidualBlock(output_size * 2, output_size * 8, output_size * 2) + 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)) @@ -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.MAX: MaxPool, - # AGGREGATOR.LAST_TOKEN: LastTokenPool, - AGGREGATOR.POOLER: Pooler, - AGGREGATOR.PERCEIVER: Perceiver, + AGGREGATOR_TYPE.POOLER: Pooler, + AGGREGATOR_TYPE.PERCEIVER: Perceiver, } class HyperLoRA(nn.Module): def __init__( self, - lora_config: LoraConfig, - layer_indices: Integer[Tensor, "num_layers"], - aggregator: AGGREGATOR, - aggregator_kwargs: Optional[dict[str, Any]] = None, + # latent_size: int, + # lora_config: LoraConfig, + # layer_indices: Integer[Tensor, "num_layers"], + # 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.layer_indices = layer_indices + self.lora_config = config.lora_config - self.in_features = ... - self.out_features = ... + self.target_modules = self.lora_config.target_modules + 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] - # loras = self.layers(emb) - # return loras - return emb + # [bs, n_layers, n_modules, feature_dim] + emb = self.aggregator(features, attn_mask) + + # [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 = ... - logits = ... - # TODO: get ctx_features offline - features = self.get_ctx_features(ctx_ids, ctx_attn_mask) - generated_loras = self.hypernet(features, ctx_attn_mask) + generated_loras = self.hypernet.generate_loras(ctx_features, ctx_attn_mask) # apply lora hook to the base model - self.apply_lora_hook(generated_loras) - model_outputs = self.base_model(**model_inputs_kwargs) + # self.apply_generated_loras(generated_loras) + with apply_generated_loras( + self.base_model, + generated_loras, + self.hypernet.layer_indices, + self.training, + ): + model_outputs = self.base_model(**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() diff --git a/hyperlora/pooling.py b/hyperlora/pooling.py index 08d1bd5..ac6d171 100644 --- a/hyperlora/pooling.py +++ b/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 diff --git a/hyperlora/utils.py b/hyperlora/utils.py index 63db19e..41f7ba4 100644 --- a/hyperlora/utils.py +++ b/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