diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index 767ab90..4ffe384 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/src/ctx_to_lora/modeling/hypernet.py @@ -805,16 +805,16 @@ class ModulatedPretrainedModel(nn.Module): def forward( self, - ctx_ids: Integer[Tensor, "bs ctx_len"] | None = None, - ctx_attn_mask: Integer[Tensor, "bs ctx_len"] | None = None, - ctx_position_ids: Integer[Tensor, "bs ctx_len"] | None = None, + 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_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) - # 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: @@ -872,8 +872,24 @@ class ModulatedPretrainedModel(nn.Module): # ): # 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, position_ids + self.base_model, + self.hypernet.layer_indices, + generated_loras, + n_queries, + position_ids, ) model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs) diff --git a/src/ctx_to_lora/modeling/lora_layer.py b/src/ctx_to_lora/modeling/lora_layer.py index 109fae2..2a44852 100644 --- a/src/ctx_to_lora/modeling/lora_layer.py +++ b/src/ctx_to_lora/modeling/lora_layer.py @@ -10,121 +10,61 @@ 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"], + x: Float[Tensor, "tot_q seq_len d_in"], + n_qs: Integer[Tensor, "n_ctx"], + tot_q: int, + A: Float[Tensor, "n_ctx r d_in"], + B: Float[Tensor, "n_ctx d_out r"], lora_dropout_p: float, scaling: float, + self, *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 +) -> Float[Tensor, "tot_q seq_len d_out"]: + # A: [n_ctx, r, d_in] -> [tot_q, r, d_in] + A = A.repeat_interleave(n_qs, dim=0, output_size=tot_q) + # B: [n_ctx, d_out, r] -> [tot_q, d_out, r] + B = B.repeat_interleave(n_qs, dim=0, output_size=tot_q) - 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 = einsum(A, delta_x, "tot_q r d_in, tot_q s_len d_in -> tot_q s_len r") + delta_x = einsum(B, delta_x, "tot_q d_out r, tot_q s_len r -> tot_q s_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"], + x: Float[Tensor, "1 tot_len d_in"], + n_qs: Integer[Tensor, "n_ctx"], + tot_q: int, + seq_lens: Integer[Tensor, "tot_q"], + tot_len: int, + A: Float[Tensor, "n_ctx r d_in"], + B: Float[Tensor, "n_ctx d_out r"], lora_dropout_p: float, scaling: float, + self, *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). - """ +) -> Float[Tensor, "1 tot_len d_out"]: # 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()) + # # [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()) + # TODO: is ther a better way to implement this? + repeated_A = A.repeat_interleave(n_qs, dim=0, output_size=tot_q) + repeated_A = repeated_A.repeat_interleave(seq_lens, dim=0, output_size=tot_len) + + repeated_B = B.repeat_interleave(n_qs, dim=0, output_size=tot_q) + repeated_B = repeated_B.repeat_interleave(seq_lens, dim=0, output_size=tot_len) + delta_x = einsum( repeated_A, delta_x, "tot_len r d_in, bs tot_len d_in -> bs tot_len r" ) @@ -139,7 +79,8 @@ def lora_forward_packed( 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 _"]]], + generated_loras: dict[str, dict[str, Float[Tensor, "n_ctx n_layers r _"]]], + n_qs: Integer[Tensor, "n_ctx"], position_ids: Integer[Tensor, "bs seq_len"] = None, ) -> None: layers = get_layers(model) @@ -150,6 +91,8 @@ def apply_lora_to_layers( [seq_lens, torch.tensor([position_ids[-1]], device=seq_lens.device)] ) seq_lens += 1 + tot_len = seq_lens.sum().item() + tot_q = n_qs.sum().item() for layer_idx in layer_indices: layer = layers[layer_idx] @@ -159,13 +102,10 @@ def apply_lora_to_layers( 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], - ) + A = generated_loras[mname]["A"][:, layer_idx] + B = generated_loras[mname]["B"][:, layer_idx] + module.forward = partial(module.forward, n_qs=n_qs, tot_q=tot_q, A=A, B=B) if position_ids is not None: module.forward = partial( - module.forward, - seq_lens=seq_lens, + module.forward, seq_lens=seq_lens, tot_len=tot_len )