diff --git a/intx_sft.py b/intx_sft.py index dc8fd2d..f12ab94 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -54,6 +54,7 @@ from ctx_to_lora.modeling.hypernet import ( ) from ctx_to_lora.trainer import train_model from ctx_to_lora.utils import ( + compile_linear, extract_cli_args, get_run_name, log_num_train_params, @@ -213,6 +214,7 @@ def main(): hypernet_config, ctx_encoder_args, ctx_args.use_kl_loss, + ctx_args.use_sequence_packing, ) training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef else: @@ -223,6 +225,7 @@ def main(): torch.load(ctx_args.from_pretrained_checkpoint, weights_only=False), train=True, use_flash_attn=model_args.use_flash_attn, + use_sequence_packing=ctx_args.use_sequence_packing, ) tokenizer = get_tokenizer(model.base_model.config.name_or_path, train=True) ctx_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path @@ -404,15 +407,23 @@ def main(): if isinstance(model, ModulatedPretrainedModel): logger.info("Applying liger-kernel to ModulatedPretrainedModel") if isinstance(model.base_model, PeftModel): - _apply_liger_kernel_to_instance(model=model.base_model.base_model.model) + base_model = model.base_model.base_model else: - _apply_liger_kernel_to_instance(model=model.base_model.model) + base_model = model.base_model + if ctx_name is not None: logger.info("Applying liger-kernel to ctx_encoder_model") - _apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model) + ctx_base_model = model.ctx_encoder.base_model + _apply_liger_kernel_to_instance(model=ctx_base_model) + compile_linear(ctx_base_model) + elif isinstance(model, PeftModel): logger.info("Applying liger-kernel to PeftModel") - _apply_liger_kernel_to_instance(model=model.base_model.model) + base_model = model.base_model + + _apply_liger_kernel_to_instance(model=base_model.model) + compile_linear(base_model) + base_model.loss_function = torch.compile(base_model.loss_function) if LOCAL_RANK == 0: wandb.init( diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index c4ec057..79bff30 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/src/ctx_to_lora/model_loading.py @@ -49,8 +49,8 @@ def get_tokenizer( padding_side = "left" if not train else "right" truncation_side = "left" - if peft_config: - model_name_or_path = peft_config.base_model_name_or_path + # if peft_config: + # model_name_or_path = peft_config.base_model_name_or_path if tokenizer_kwargs is None: tokenizer_kwargs = {} diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 1c02ce1..767ab90 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -3,6 +3,7 @@ from collections.abc import Iterable from contextlib import contextmanager from copy import deepcopy from dataclasses import dataclass +from functools import partial from typing import Any import torch @@ -42,10 +43,17 @@ from ctx_to_lora.model_loading import ( ) from ctx_to_lora.modeling.aggregator import AG, AggregatorConfig, get_aggregator_config from ctx_to_lora.modeling.ctx_encoder import EarlyExit +from ctx_to_lora.modeling.lora_layer import ( + apply_lora_to_layers, + lora_forward, + lora_forward_packed, +) from ctx_to_lora.utils import ( get_base_model, + get_layers, get_num_layers, get_peft_in_out_features, + get_peft_modules, ) logger = logging.getLogger() @@ -580,13 +588,16 @@ class ModulatedPretrainedModel(nn.Module): ctx_encoder_args: CtxEncoderArguments, # use_kl_loss is only used for training use_kl_loss: bool = False, + use_sequence_packing: bool = False, use_base_input_as_ctx: bool = False, ): 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_kl_loss = use_kl_loss + self.use_sequence_packing = use_sequence_packing self.use_base_input_as_ctx = use_base_input_as_ctx self.lora_id = 1 self.active_adapters = [] @@ -621,9 +632,31 @@ class ModulatedPretrainedModel(nn.Module): 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 if not self.use_sequence_packing else lora_forward_packed + ) + 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 @@ -780,7 +813,8 @@ class ModulatedPretrainedModel(nn.Module): **model_inputs_kwargs: dict[str, Any], ) -> tuple | ModelOutput: """Forward pass of the modulated model.""" - + # TODO: allow more than one LoRA per sample (multi-lora) + # TODO: allow multi queries per one LoRA (for efficient training) generated_loras = None generated_layernorms = None if ctx_ids is None and not self.use_base_input_as_ctx: @@ -820,23 +854,28 @@ class ModulatedPretrainedModel(nn.Module): 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) + # 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) + + apply_lora_to_layers( + self.base_model, self.hypernet.layer_indices, generated_loras, position_ids + ) + model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs) if return_generated_lora: return model_outputs, (generated_loras, generated_layernorms) @@ -886,7 +925,7 @@ class ModulatedPretrainedModel(nn.Module): @torch.inference_mode() def generate( self, - # TODO: allow more than one LoRA per sample + # TODO: allow more than one LoRA per sample (multi-lora) ctx_ids: Integer[Tensor, "bs ctx_length"] | None = None, ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None, ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None, @@ -1113,16 +1152,6 @@ def apply_generated_loras( remove_hook_handles(hooks) -@contextmanager -def apply_generated_loras_packed_sequence( - base_model: nn.Module, - generated_loras: dict | None = None, - layer_indices: Iterable[int] | None = None, - training: bool = False, -): - pass - - # needed for loading model from checkpoint # see https://github.com/huggingface/transformers/pull/34632 torch.serialization.add_safe_globals( diff --git a/src/ctx_to_lora/modeling/lora_layer.py b/src/ctx_to_lora/modeling/lora_layer.py new file mode 100644 index 0000000..109fae2 --- /dev/null +++ b/src/ctx_to_lora/modeling/lora_layer.py @@ -0,0 +1,171 @@ +from collections.abc import Iterable +from functools import partial +from operator import attrgetter + +import torch +import torch.nn.functional as F +from einops import einsum +from jaxtyping import Float, Integer +from torch import Tensor + +from ctx_to_lora.utils import get_layers + +# Patch corresponding layers of LLMs to include A and B lora matrices +""" +def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + self._check_forward_args(x, *args, **kwargs) + adapter_names = kwargs.pop("adapter_names", None) + + if self.disable_adapters: + if self.merged: + self.unmerge() + result = self.base_layer(x, *args, **kwargs) + elif adapter_names is not None: + result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **kwargs) + elif self.merged: + result = self.base_layer(x, *args, **kwargs) + else: + result = self.base_layer(x, *args, **kwargs) + torch_result_dtype = result.dtype + + lora_A_keys = self.lora_A.keys() + for active_adapter in self.active_adapters: + if active_adapter not in lora_A_keys: + continue + + lora_A = self.lora_A[active_adapter] + lora_B = self.lora_B[active_adapter] + dropout = self.lora_dropout[active_adapter] + scaling = self.scaling[active_adapter] + x = self._cast_input_dtype(x, lora_A.weight.dtype) + + if not self.use_dora[active_adapter]: + result = result + lora_B(lora_A(dropout(x))) * scaling + else: + if isinstance(dropout, nn.Identity) or not self.training: + base_result = result + else: + x = dropout(x) + base_result = None + + result = result + self.lora_magnitude_vector[active_adapter]( + x, + lora_A=lora_A, + lora_B=lora_B, + scaling=scaling, + base_layer=self.get_base_layer(), + base_result=base_result, + ) + + result = result.to(torch_result_dtype) + + return result +""" + + +def lora_forward( + x: Float[Tensor, "bs seq_len d_in"], + self, + # pre_modulated_forward_fn: callable[..., Float[Tensor, "bs seq_len d_out"]], + A: Float[Tensor, "bs r d_in"], + B: Float[Tensor, "bs d_out r"], + lora_dropout_p: float, + scaling: float, + *args, + **kwargs, +) -> Float[Tensor, "bs seq_len d_out"]: + """ + Forward pass of the LoRA layer. + # TODO: implement the packed version + # TODO: implement the one-lora-multi-query version + + Args: + x (torch.Tensor): Input tensor with shape (batch_size, input_dim). + A (torch.Tensor): LoRA matrix A with shape (bs, r, d_in). + B (torch.Tensor): LoRA matrix B with shape (bs, d_out, r). + + Returns: + torch.Tensor: Output tensor with shape (batch_size, output_dim). + """ + base_out = torch.nn.Linear.forward(self, x, *args, **kwargs) + delta_x = F.dropout(x, p=lora_dropout_p, training=self.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 + return base_out + delta_x + + +def lora_forward_packed( + x: Float[Tensor, "bs seq_len d_in"], + seq_lens: Integer[Tensor, "bs"], + self, + A: Float[Tensor, "bs r d_in"], + B: Float[Tensor, "bs d_out r"], + lora_dropout_p: float, + scaling: float, + *args, + **kwargs, +) -> Float[Tensor, "bs seq_len d_out"]: + """ + Forward pass of the LoRA layer for packed inputs. + + Args: + x (torch.Tensor): Input tensor with shape (batch_size, input_dim). + A (torch.Tensor): LoRA matrix A with shape (input_dim, lora_dim). + B (torch.Tensor): LoRA matrix B with shape (lora_dim, output_dim). + + Returns: + torch.Tensor: Output tensor with shape (batch_size, output_dim). + """ + # bs of x should be 1 in this case + base_out = torch.nn.Linear.forward(self, x, *args, **kwargs) + + delta_x = F.dropout(x, p=lora_dropout_p, training=self.training) + # [tot_seq_len, r, d_in] + repeated_A = A.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum()) + # [tot_seq_len, d_out, r] + repeated_B = B.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum()) + delta_x = einsum( + repeated_A, delta_x, "tot_len r d_in, bs tot_len d_in -> bs tot_len r" + ) + delta_x = einsum( + repeated_B, delta_x, "tot_len d_out r, bs tot_len r -> bs tot_len d_out" + ) + delta_x = delta_x * scaling + + return base_out + delta_x + + +def apply_lora_to_layers( + model: torch.nn.Module, + layer_indices: Iterable[int], + generated_loras: dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]], + position_ids: Integer[Tensor, "bs seq_len"] = None, +) -> None: + layers = get_layers(model) + if position_ids is not None: + position_ids = position_ids.squeeze() + seq_lens = position_ids[torch.where(position_ids == 0)[0][1:] - 1] + seq_lens = torch.cat( + [seq_lens, torch.tensor([position_ids[-1]], device=seq_lens.device)] + ) + seq_lens += 1 + for layer_idx in layer_indices: + layer = layers[layer_idx] + + for mname in generated_loras: + if mname in ["q_proj", "k_proj", "v_proj", "o_proj", "qkv_proj"]: + long_mname = f"self_attn.{mname}" + elif mname in ["down_proj", "up_proj", "gate_proj"]: + long_mname = f"mlp.{mname}" + module = attrgetter(long_mname)(layer) + module.forward = partial( + module.forward, + A=generated_loras[mname]["A"][:, layer_idx], + B=generated_loras[mname]["B"][:, layer_idx], + ) + if position_ids is not None: + module.forward = partial( + module.forward, + seq_lens=seq_lens, + ) diff --git a/src/ctx_to_lora/utils.py b/src/ctx_to_lora/utils.py index 16bc60e..5682847 100644 --- a/src/ctx_to_lora/utils.py +++ b/src/ctx_to_lora/utils.py @@ -161,6 +161,16 @@ def save_yaml(data, path): yaml.dump(data, file) +def get_peft_modules(model: PeftModel, peft_config: PeftConfig) -> list[dict[str, str]]: + return [ + {"name": name, "module": module} + for name, module in model.named_modules() + if name.split(".")[-1] in peft_config.target_modules + and isinstance(module, BaseTunerLayer) + and check_target_module_exists(peft_config, name) + ] + + def get_peft_in_out_features( model: PeftModel, peft_config: PeftConfig | None = None, @@ -169,11 +179,9 @@ def get_peft_in_out_features( return None, None 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 + for module_info in get_peft_modules(model, peft_config): + module_name = module_info["name"] + module = module_info["module"] # 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), ( @@ -239,6 +247,12 @@ def get_lora_module_names( return module_names +def compile_linear(model): + for name, module in model.named_modules(): + if isinstance(module, torch.nn.Linear): + module.compile() + + def clear_gpu(): gc.collect() torch.cuda.empty_cache()