diff --git a/intx_sft.py b/intx_sft.py index 0a6869c..9fed971 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -538,7 +538,10 @@ def main(): if isinstance(model, ModulatedPretrainedModel): logger.info("Applying liger-kernel to ModulatedPretrainedModel") - _apply_liger_kernel_to_instance(model=model.base_model.base_model.model) + if isinstance(model.base_model, PeftModel): + _apply_liger_kernel_to_instance(model=model.base_model.base_model.model) + else: + _apply_liger_kernel_to_instance(model=model.base_model.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) diff --git a/src/ctx_to_lora/hooks.py b/src/ctx_to_lora/hooks.py index a86a361..b9be4a9 100644 --- a/src/ctx_to_lora/hooks.py +++ b/src/ctx_to_lora/hooks.py @@ -5,7 +5,7 @@ from typing import Callable, Iterable, Optional import torch import torch.nn.functional as F from einops import einsum -from jaxtyping import Float +from jaxtyping import Float, Integer from torch import Tensor from torch.utils.hooks import RemovableHandle from ctx_to_lora.utils import get_layers @@ -133,7 +133,8 @@ def add_generated_lora_hook( B: Float[Tensor, "bs d_out r"], scaling: float, input_dropout: float, - training: bool, + position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, + training: bool = False, ) -> list[RemovableHandle]: """ Adds LoRA hooks to the specified modules and layers of the model. @@ -184,7 +185,51 @@ def add_generated_lora_hook( else: return newoutput - return apply_hook_to_layers(model, [module_name], [layer_index], post_hook=lora_hook) + 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 + + def lora_hook_packed_sequence( + 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 + + # bs of x should be 1 in this case + x = args[0].to(A.dtype) # [1, tot_seq_len, d_in] + # print(f"seq_lens: {seq_lens}") + # print(f"sum of seq_lens: {seq_lens.sum()}") + # print(f"tot_len: {x.shape[1]}") + delta_x = F.dropout(x, input_dropout, training) + n_seq = A.shape[0] + # [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 + newoutput = model_out + delta_x.to(model_out.dtype) + if isinstance(output, tuple): + return (newoutput, *output[1:]) + else: + return newoutput + + hook_fn = lora_hook if position_ids is None else lora_hook_packed_sequence + + return apply_hook_to_layers(model, [module_name], [layer_index], post_hook=hook_fn) def add_generated_layernorm_hook( @@ -193,6 +238,7 @@ def add_generated_layernorm_hook( layer_index: int, W: Float[Tensor, "bs hidden_size"], training: bool, + position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, ) -> list[RemovableHandle]: """ Adds layer normalization hooks to specified modules and layers. @@ -203,6 +249,7 @@ def add_generated_layernorm_hook( layer_index (int): Index of layer to modify W (Tensor): Learned weight tensor of shape [batch_size, hidden_size] training (bool): Whether model is in training mode + position_ids (Optional[Tensor]): Position IDs for packed sequence processing Returns: list[RemovableHandle]: Hook handles for removal @@ -213,7 +260,6 @@ def add_generated_layernorm_hook( args: tuple, output: Float[Tensor, "bs seq_len hidden_size"], ) -> Float[Tensor, "bs seq_len hidden_size"]: - # For models that return tuples from layernorm (e.g., some attention implementations) if isinstance(output, tuple): main_output = output[0] rest = output[1:] @@ -229,6 +275,37 @@ def add_generated_layernorm_hook( return (new_output, *rest) if rest else new_output - return apply_hook_to_layers( - model, [module_name], [layer_index], post_hook=layernorm_hook - ) + 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 + + def layernorm_hook_packed_sequence( + module: torch.nn.Module, + args: tuple, + output: Float[Tensor, "bs seq_len hidden_size"], + ) -> Float[Tensor, "bs seq_len hidden_size"]: + if isinstance(output, tuple): + main_output = output[0] + rest = output[1:] + else: + main_output = output + rest = None + + # [1, tot_seq_len, d_in] + x = args[0].to(W.dtype) + # Repeat weights for each sequence + + repeated_W = W.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum()) + # Apply learned weights to layernorm output + scaled_output = x * repeated_W.unsqueeze(0) + new_output = main_output + scaled_output.to(output.dtype) + + return (new_output, *rest) if rest else new_output + + hook_fn = layernorm_hook if position_ids is None else layernorm_hook_packed_sequence + + return apply_hook_to_layers(model, [module_name], [layer_index], post_hook=hook_fn) diff --git a/src/ctx_to_lora/model_loading.py b/src/ctx_to_lora/model_loading.py index 9e5e20d..a85d950 100644 --- a/src/ctx_to_lora/model_loading.py +++ b/src/ctx_to_lora/model_loading.py @@ -179,7 +179,7 @@ def get_model( def get_lora_config(model_dir, **kwargs): - if "target_modules" not in kwargs: + if "target_modules" not in kwargs or kwargs["target_modules"] is None: logger.info("No target modules specified for LoRA.") return None r = kwargs.pop("lora_r", 8) diff --git a/src/ctx_to_lora/modeling_idefics2.py b/src/ctx_to_lora/modeling_idefics2.py new file mode 100644 index 0000000..72016a5 --- /dev/null +++ b/src/ctx_to_lora/modeling_idefics2.py @@ -0,0 +1,816 @@ +# coding=utf-8 +# Copyright 2024 the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Idefics2 model.""" + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import CrossEntropyLoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.generation import GenerationMixin +from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask +from transformers.modeling_outputs import BaseModelOutput, ModelOutput +from transformers.modeling_utils import PreTrainedModel, ALL_ATTENTION_FUNCTIONS +from transformers.utils import ( + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_2_available, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from transformers.models.auto import AutoModel +from transformers.configuration_utils import PretrainedConfig +from transformers.models.idefics2.configuration_idefics2 import Idefics2Config + + +if is_flash_attn_2_available(): + from transformers.modeling_flash_attention_utils import _flash_attention_forward + +logger = logging.get_logger(__name__) + + +class Idefics2PerceiverConfig(PretrainedConfig): + r""" + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the perceiver block. + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + resampler_n_latents (`int`, *optional*, defaults to 64): + Number of latent embeddings to resample ("compress") the input sequence to (usually < 128). + resampler_depth (`int`, *optional*, defaults to 3): + Depth of the Perceiver Resampler (Transformer w/ cross attention). Should be shallow (<= 3). + resampler_n_heads (`int`, *optional*, defaults to 16): + Number of heads in each Transformer block (for multi-headed self-attention). + resampler_head_dim (`int`, *optional*, defaults to 96): + Dimensionality of each head projection in the Transformer block. + num_key_value_heads (`int`, *optional*, defaults to 4): + Number of key-value heads in the perceiver attention block. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + """ + + model_type = "idefics2_perceiver" + + def __init__( + self, + input_size: int, + intermediate_size_factor: int = 1, + hidden_act="silu", + hidden_size=4096, + rms_norm_eps=1e-06, + resampler_n_latents=64, + resampler_depth=3, + resampler_n_heads=16, + resampler_head_dim=96, + num_key_value_heads=4, + attention_dropout=0.0, + **kwargs, + ): + # for mlp + self.input_size = input_size + self.intermediate_size_factor = intermediate_size_factor + # for perceiver + self.hidden_act = hidden_act + self.hidden_size = hidden_size + self.rms_norm_eps = rms_norm_eps + self.resampler_n_latents = resampler_n_latents + self.resampler_depth = resampler_depth + self.resampler_n_heads = resampler_n_heads + self.num_key_value_heads = num_key_value_heads + self.resampler_head_dim = resampler_head_dim + self.attention_dropout = attention_dropout + if self.num_key_value_heads > self.resampler_n_heads: + raise ValueError( + f"num_key_value_heads={self.num_key_value_heads} must be less than or equal to" + f" resampler_n_heads={self.resampler_n_heads}" + ) + super().__init__(**kwargs) + + +class Idefics2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + output_size: int, + hidden_act: str, + ): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, output_size, bias=False) + self.act_fn = ACT2FN[hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +IDEFICS2_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`Idefics2Config`] or [`Idefics2VisionConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare Idefics2 Model outputting raw hidden-states without any specific head on top.", + IDEFICS2_START_DOCSTRING, +) +class Idefics2PreTrainedModel(PreTrainedModel): + config_class = Idefics2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = [ + "Idefics2VisionAttention", + "Idefics2MLP", + "Idefics2PerceiverLayer", + "Idefics2DecoderLayer", + ] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + + def _init_weights(self, module): + std = ( + self.config.architecturesinitializer_range + if hasattr(self.config, "initializer_range") + else self.config.initializer_range + ) + + if hasattr(module, "class_embedding"): + module.class_embedding.data.normal_(mean=0.0, std=std) + + if isinstance(module, (nn.Linear, nn.Conv2d)): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +# Copied from transformers.models.llama.modeling_llama.repeat_kv +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Idefics2 +class Idefics2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Idefics2RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Idefics2PerceiverAttention(nn.Module): + def __init__(self, config, layer_idx: Optional[int] = None) -> None: + """Perceiver Cross-Attention Module --> let long-form inputs be `context`, resampled embeddings be `latents`""" + super().__init__() + self.config = config + self.layer_idx = None + self.hidden_size = config.hidden_size + self.num_heads = config.resampler_n_heads + self.head_dim = config.resampler_head_dim + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.attention_dropout = config.attention_dropout + + self.q_proj = nn.Linear( + self.hidden_size, self.num_heads * self.head_dim, bias=False + ) + self.k_proj = nn.Linear( + self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False + ) + self.v_proj = nn.Linear( + self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False + ) + self.o_proj = nn.Linear( + self.num_heads * self.head_dim, self.hidden_size, bias=False + ) + + self.is_causal = False + + def forward( + self, + latents: torch.Tensor, + context: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: bool = False, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + """ + Runs Perceiver Self-Attention, with special (context, latents) appended along the `seq` dimension! + + Args: + latents (`torch.Tensor`): Tensor of shape [bsz, n_latents, embed_dim] representing fixed length latents to compress to. + context (`torch.Tensor`): Tensor of shape [bsz, seq, embed_dim] representing long-form context to resample. + attention_mask (`torch.Tensor`, *optional*): Tensor of shape [bsz, 1, seq, n_latents] representing attention mask. + position_ids (`torch.LongTensor`, *optional*): Tensor of shape [bsz, seq] representing position indices of each input token. + past_key_value (`Tuple[torch.Tensor]`, *optional*): Tuple of tensors containing cached key and value states. + output_attentions (`bool`, *optional*, defaults to `False`): Whether to return attention weights. + use_cache (`bool`, *optional*, defaults to `False`): Whether to use past_key_value for caching. + """ + bsz, q_len, _ = latents.size() + kv_seq_len = q_len + context.size()[1] + + hidden_states = torch.concat([context, latents], dim=-2) + + query_states = self.q_proj(latents) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view( + bsz, q_len, self.num_heads, self.head_dim + ).transpose(1, 2) + key_states = key_states.view( + bsz, kv_seq_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = value_states.view( + bsz, kv_seq_len, self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + past_key_value = getattr(self, "past_key_value", past_key_value) + + if past_key_value is not None: + key_states, value_states = past_key_value.update( + key_states, value_states, self.layer_idx + ) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul( + query_states, key_states.transpose(2, 3) + ) / math.sqrt(self.head_dim) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax( + attn_weights, dim=-1, dtype=torch.float32 + ).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +# NO LONGER EXIST Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2 with MistralAttention->Idefics2PerceiverAttention,MistralFlashAttention->Idefics2PerceiverFlashAttention,Mistral->Idefics2 +# TODO cyril: modular +class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention): + """ + Idefics2 flash attention module. This module inherits from `Idefics2PerceiverAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + # Ignore copy + def forward( + self, + latents: torch.Tensor, + context: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = latents.size() + # kv_seq_len = q_len + context.size()[1] + kv_seq_len = context.size()[1] + + # Query, Key, Value Projections --> Note that in Flamingo, latents are *concatenated* with context prior to attn! + # Note: This results in queries w/ `seq = n_latents`, and keys, values with `seq = len(context) + n_latents` + query_states = self.q_proj(latents) + # key_states = self.k_proj(torch.cat([context, latents], dim=-2)) + # value_states = self.v_proj(torch.cat([context, latents], dim=-2)) + key_states = self.k_proj(context) + value_states = self.v_proj(context) + + # query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) + query_states = query_states.view( + *latents.shape[:2], self.num_heads, self.head_dim + ) + key_states = key_states.view( + *context.shape[:2], self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + value_states = value_states.view( + *context.shape[:2], self.num_key_value_heads, self.head_dim + ).transpose(1, 2) + + # kv_seq_len = key_states.shape[-2] + # if past_key_value is not None: + # kv_seq_len += past_key_value[0].shape[-2] + + # if past_key_value is not None: + # # Activate slicing cache only if the config has a value `sliding_windows` attribute + # if ( + # hasattr(self.config, "sliding_window") + # and kv_seq_len > self.config.sliding_window + # ): + # slicing_tokens = kv_seq_len - self.config.sliding_window + + # past_key = past_key_value[0] + # past_value = past_key_value[1] + + # past_key = past_key[:, :, slicing_tokens:, :].contiguous() + # past_value = past_value[:, :, slicing_tokens:, :].contiguous() + + # if past_key.shape[-2] != self.config.sliding_window - 1: + # raise ValueError( + # "past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1," + # f" head_dim`), got {past_key.shape}" + # ) + + # past_key_value = (past_key, past_value) + + # if attention_mask is not None: + # attention_mask = attention_mask[:, slicing_tokens:] + # attention_mask = torch.cat( + # [attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1 + # ) + + # key_states = torch.cat([past_key_value[0], key_states], dim=2) + # value_states = torch.cat([past_key_value[1], value_states], dim=2) + + past_key_value = (key_states, value_states) if use_cache else None + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + dropout_rate = 0.0 if not self.training else self.attention_dropout + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + # Reashape to the expected shape for Flash Attention + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + dropout=dropout_rate, + position_ids=position_ids, + sliding_window=None, + is_causal=self.is_causal, + use_top_left_mask=self._flash_attn_uses_top_left_mask, + **kwargs, + ) + + attn_output = attn_output.reshape( + bsz, q_len, self.num_heads * self.head_dim + ).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +IDEFICS2_PERCEIVER_ATTENTION_CLASSES = { + "eager": Idefics2PerceiverAttention, + "flash_attention_2": Idefics2PerceiverFlashAttention2, +} + + +class Idefics2PerceiverLayer(nn.Module): + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.n_latents = config.resampler_n_latents + self.depth = config.resampler_depth + self.rms_norm_eps = config.rms_norm_eps + + self.input_latents_norm = Idefics2RMSNorm( + self.hidden_size, eps=self.rms_norm_eps + ) + self.input_context_norm = Idefics2RMSNorm( + self.hidden_size, eps=self.rms_norm_eps + ) + self.self_attn = IDEFICS2_PERCEIVER_ATTENTION_CLASSES[ + config._attn_implementation + ](config, layer_idx=layer_idx) + self.post_attention_layernorm = Idefics2RMSNorm( + self.hidden_size, eps=self.rms_norm_eps + ) + self.mlp = Idefics2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.hidden_size * 4, + output_size=config.hidden_size, + hidden_act=config.hidden_act, + ) + + def forward( + self, + latents: torch.Tensor, + context: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + latents (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + context (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): attention mask of size + `(batch, sequence_length)` where padding elements are indicated by 0. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + """ + residual = latents + + latents = self.input_latents_norm(latents) + context = self.input_context_norm(context) + + latents, self_attn_weights, present_key_value = self.self_attn( + latents=latents, + context=context, + attention_mask=attention_mask, + position_ids=position_ids, + **kwargs, + ) + latents = residual + latents + residual = latents + + latents = self.post_attention_layernorm(latents) + latents = self.mlp(latents) + latents = residual + latents + + outputs = (latents,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +IDEFICS2_INPUTS_DOCSTRING = r""" + Args: + context (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_dim)`): + The hidden states of the image after vision encoder and modality projection. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) +""" + + +@add_start_docstrings( + "Idefics2 perceiver resampler model that performs `depth` blocks of cross-attention with a fixed ", + "`n_latents` inputs to decrease embedding sequence length. The Resampler acts as a form of learned pooling and ", + "is derived from [Perceiver: General Perception with Iterative Attention](https://arxiv.org/abs/2103.03206)", + IDEFICS2_START_DOCSTRING, +) +class Idefics2PerceiverResampler(Idefics2PreTrainedModel): + _supports_sdpa = False + config_class = Idefics2PerceiverConfig + + def __init__(self, config) -> None: + super().__init__(config) + self.hidden_size = config.hidden_size + self.hidden_act = config.hidden_act + self.n_latents = config.resampler_n_latents + self.depth = config.resampler_depth + self.rms_norm_eps = config.rms_norm_eps + + # Create Latents for Perceiver + self.latents = nn.Parameter(torch.randn(self.n_latents, self.hidden_size)) + + # Create Transformer Blocks + self.layers = nn.ModuleList( + [Idefics2PerceiverLayer(config, idx) for idx in range(self.depth)] + ) + self.norm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps) + + self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" + assert self._use_flash_attention_2 + + def forward( + self, + context: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + # seq embed -> bsz seq embed + if position_ids is None: + bsz = context.shape[0] + else: + # flattened packed sequence + bsz = torch.where(position_ids == 0, 1, 0).sum() + + latents = self.latents.unsqueeze(0).expand((bsz, *self.latents.size())) + + # latent_attention_mask = torch.ones( + # (attention_mask.size(0), latents.size(1)), + # dtype=attention_mask.dtype, + # device=attention_mask.device, + # ) + # attention_mask = torch.cat([attention_mask, latent_attention_mask], dim=-1) + attention_mask = ( + _prepare_4d_attention_mask( + attention_mask, latents.dtype, tgt_len=self.n_latents + ) + if not self._use_flash_attention_2 + else attention_mask + ) + + compressed_context = latents + + cu_seq_lens_q = torch.tensor( + [self.n_latents] * (bsz + 1), device=context.device, dtype=torch.int32 + ) * torch.arange(bsz + 1, device=context.device, dtype=torch.int32) + max_length_q = self.n_latents + # cu_seq_lens_k = None + # max_length_k = None + if attention_mask is not None: + logger.warning_once("Using attention mask for resampler") + seq_lens_k = attention_mask.sum(dim=-1, dtype=torch.int32) + cu_seq_lens_k = torch.cumsum(seq_lens_k, dim=0, dtype=torch.int32) + max_length_k = seq_lens_k.max().item() + + elif position_ids is not None: + logger.warning_once("Using position ids for resampler") + # [bsz + 1] + # cu_seq_lens_q = torch.tensor( + # [self.n_latents] * (bsz + 1), device=context.device, dtype=torch.int32 + # ) * torch.arange(bsz + 1, device=context.device, dtype=torch.int32) + # max_length_q = self.n_latents + + position_ids = position_ids.flatten() + indices = torch.arange( + position_ids.size(0), device=position_ids.device, dtype=torch.int32 + ) + # [bsz + 1] + cu_seq_lens_k = torch.cat( + ( + indices[position_ids == 0], + torch.tensor( + position_ids.size(), + device=position_ids.device, + dtype=torch.int32, + ), + ) + ) + + max_length_k = position_ids.max() + 1 + else: + raise ValueError("either position_ids or attention_mask is required") + + for perceiver_layer in self.layers: + layer_outputs = perceiver_layer( + compressed_context, + context, + attention_mask=None, + position_ids=position_ids, + past_key_value=None, + output_attentions=False, + use_cache=False, + cu_seq_lens_q=cu_seq_lens_q, + cu_seq_lens_k=cu_seq_lens_k, + max_length_q=max_length_q, + max_length_k=max_length_k, + ) + + compressed_context = layer_outputs[0] + + compressed_context = self.norm(compressed_context) + + return compressed_context + + +class Idefics2Perceiver(Idefics2PreTrainedModel): + def __init__( + self, config: Idefics2PerceiverConfig, decoder_config: Idefics2PerceiverConfig + ): + super().__init__(config) + self.modality_projection = Idefics2MLP( + hidden_size=config.input_size, + intermediate_size=config.intermediate_size_factor * config.input_size, + output_size=config.hidden_size, + hidden_act=config.hidden_act, + ) + self.encoder = Idefics2PerceiverResampler._from_config(config) + self.decoder = Idefics2PerceiverResampler._from_config(decoder_config) + + def forward( + self, + context: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + ): + if position_ids is None: + bsz = context.shape[0] + else: + bsz = torch.where(position_ids == 0, 1, 0).sum() + projected_inputs = self.modality_projection(context) + + # [bsz, n_latents, dim] + latents = self.encoder( + context=projected_inputs, + attention_mask=attention_mask, + position_ids=position_ids, + ) + + # if position_ids is None: + # # padded + # # attn_mask = torch.ones( + # # (bsz, self.encoder.n_latents), + # # device=context.device, + # # dtype=torch.int32, + # # ) + # outputs = self.decoder(latents) + # else: + # packed sequence + latent_position_ids = torch.arange( + self.encoder.n_latents, device=context.device + ).unsqueeze(0) + latent_position_ids = torch.tile(latent_position_ids, (1, bsz)) + outputs = self.decoder(latents, position_ids=latent_position_ids) + + return outputs + + +__all__ = [ + "Idefics2Perceiver", +] + +if __name__ == "__main__": + config = Idefics2PerceiverConfig( + input_size=1234, + resampler_n_latents=64, + resampler_depth=8, + attn_implementation="flash_attention_2", + ) + decoder_config = Idefics2PerceiverConfig( + input_size=config.hidden_size, + resampler_n_latents=16, + resampler_depth=1, + attn_implementation="flash_attention_2", + ) + model = Idefics2Perceiver(config, decoder_config).to("cuda").to(torch.bfloat16) + print(model) + inp = torch.rand(1, 192 + 37, 1234, dtype=torch.bfloat16, device="cuda") + inp = torch.cat([inp, inp], dim=1) + # mask = torch.ones(1, 128, dtype=torch.bfloat16, device="cuda") + # mask[..., -1] = 0 + # print(mask) + out = model( + inp, + # attention_mask=mask, + position_ids=torch.cat( + [ + torch.arange(128, device="cuda"), + torch.arange(64, device="cuda"), + torch.arange(37, device="cuda"), + torch.arange(128, device="cuda"), + torch.arange(64, device="cuda"), + torch.arange(37, device="cuda"), + ], + dim=-1, + ).unsqueeze(0), + # position_ids=torch.arange(128, device="cuda"), + ) + print(out) + print(out.shape) + print(model.encoder.latents.grad) + print(model.decoder.latents.grad) + out.mean().backward() + print(model.encoder.latents.grad) + print(model.decoder.latents.grad) + breakpoint() diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 9eebd67..a951612 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -57,6 +57,7 @@ from ctx_to_lora.utils import ( get_peft_in_out_features, get_base_model, ) +from ctx_to_lora.modeling_idefics2 import Idefics2PerceiverConfig, Idefics2Perceiver logger = logging.getLogger() @@ -164,36 +165,54 @@ class Perceiver(nn.Module): super().__init__() self.num_layers = num_layers self.num_modules = num_modules - self.config = PerceiverConfig( - d_model=feature_size, # + num_bands - num_latents=num_layers * num_modules * num_latent_factor, - d_latents=output_size, - # attention_probs_dropout_prob=0.0, - # num_blocks=8, - # num_self_attends_per_block=6, - # self_attention_widening_factor=4, - **kwargs, - ) - decoder = PerceiverBasicDecoder( - self.config, - output_num_channels=output_size, - output_index_dims=num_layers * num_modules, - num_channels=output_size, - final_project=False, - trainable_position_encoding_kwargs=dict( - num_channels=output_size, - index_dims=num_layers * num_modules, - ), - ) + # self.config = PerceiverConfig( + # d_model=feature_size, # + num_bands + # num_latents=num_layers * num_modules * num_latent_factor, + # d_latents=output_size, + # # attention_probs_dropout_prob=0.0, + # # num_blocks=8, + # # num_self_attends_per_block=6, + # # self_attention_widening_factor=4, + # **kwargs, + # ) + # decoder = PerceiverBasicDecoder( + # self.config, + # output_num_channels=output_size, + # output_index_dims=num_layers * num_modules, + # num_channels=output_size, + # final_project=False, + # trainable_position_encoding_kwargs=dict( + # num_channels=output_size, + # index_dims=num_layers * num_modules, + # ), + # ) - self.perceiver = PerceiverModel(self.config, decoder=decoder) + # self.perceiver = PerceiverModel(self.config, decoder=decoder) + self.config = Idefics2PerceiverConfig( + input_size=feature_size, + # the first layer is xattn + resampler_depth=kwargs["num_self_attends_per_block"] + 1, + resampler_n_latents=num_layers * num_modules * num_latent_factor, + intermediate_size_factor=4, + hidden_size=output_size, + attn_implementation="flash_attention_2", + ) + self.decoder_config = Idefics2PerceiverConfig( + input_size=output_size, + resampler_depth=1, + resampler_n_latents=num_layers * num_modules, + hidden_size=output_size, + attn_implementation="flash_attention_2", + ) + self.perceiver = Idefics2Perceiver(self.config, self.decoder_config) def forward( self, ctx_features: Float[Tensor, "bs seq_len feature_dim"], ctx_attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, + ctx_position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, ): - x = self.perceiver(ctx_features, ctx_attn_mask).logits + x = self.perceiver(ctx_features, ctx_attn_mask, ctx_position_ids) x = rearrange( x, "bs (n_layers n_modules) d -> bs n_layers n_modules d", @@ -342,10 +361,18 @@ def maybe_add_batch_dim(kwargs): try: batched_input = False batched_attn_mask = False - if "input_ids" in kwargs and len(kwargs["input_ids"].shape) == 1: + if ( + "input_ids" in kwargs + and kwargs["input_ids"] is not None + and len(kwargs["input_ids"].shape) == 1 + ): kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0) batched_input = True - if "attention_mask" in kwargs and len(kwargs["attention_mask"].shape) == 1: + if ( + "attention_mask" in kwargs + and kwargs["attention_mask"] is not None + and len(kwargs["attention_mask"].shape) == 1 + ): kwargs["attention_mask"] = kwargs["attention_mask"].unsqueeze(0) batched_attn_mask = True yield batched_input, batched_attn_mask @@ -370,7 +397,7 @@ class EarlyExit(nn.Module): def config(self): return self.base_model.config - @torch.inference_mode() + @torch.no_grad() def forward(self, **kwargs): # if len(kwargs["input_ids"].shape) == 1: # kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0) @@ -461,8 +488,8 @@ class HyperLoRA(nn.Module): output_size=self.config.latent_size, dropout_rate=getattr(self.config, "dropout_rate", 0), ) - self.pre_layers_norm = nn.LayerNorm(self.config.latent_size) - self.norm = nn.LayerNorm(self.config.latent_size) + # self.pre_layers_norm = nn.LayerNorm(self.config.latent_size) + # self.norm = nn.LayerNorm(self.config.latent_size) if self.extra_modules: self.extra_layers = MLPResidualBlock( input_size=self.config.latent_size, @@ -470,8 +497,8 @@ class HyperLoRA(nn.Module): output_size=self.config.latent_size, dropout_rate=getattr(self.config, "dropout_rate", 0), ) - self.pre_extra_layers_norm = nn.LayerNorm(self.config.latent_size) - self.extra_norm = nn.LayerNorm(self.config.latent_size) + # self.pre_extra_layers_norm = nn.LayerNorm(self.config.latent_size) + # self.extra_norm = nn.LayerNorm(self.config.latent_size) if self.config.use_light_weight_lora: # light-weight lora projection (per module) @@ -610,10 +637,11 @@ class HyperLoRA(nn.Module): self, features: Float[Tensor, "bs seq_len feature_dim"], attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, + position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, ): # [bs, n_layers, n_total_modules, feature_dim] - emb = self.aggregator(features.to(torch.float32), attn_mask) + emb = self.aggregator(features, attn_mask, position_ids) lora_emb, extra_emb = unpack( emb, [[self.num_modules], [self.num_extra_modules]], @@ -621,17 +649,13 @@ class HyperLoRA(nn.Module): ) # [bs, n_layers, n_modules, r, max_in_d_outim] - lora_emb = self.norm(self.layers(self.pre_layers_norm(lora_emb))) + lora_emb = self.layers(lora_emb) flat_loras = self.head(lora_emb) flat_layernorms = None if self.num_extra_modules: # [bs, n_layers, n_extra_modules, base_hidden_size] - emb = self.extra_norm( - self.extra_layers( - self.pre_extra_layers_norm(extra_emb[:, :, self.num_modules :]) - ) - ) + extra_emb = self.extra_layers(extra_emb) flat_layernorms = self.extra_head(extra_emb) return flat_loras, flat_layernorms @@ -640,8 +664,12 @@ class HyperLoRA(nn.Module): self, features: Float[Tensor, "bs seq_len feature_dim"], attn_mask: Optional[Integer[Tensor, "bs seq_len"]] = None, + position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, ): - flat_loras, flat_layernorms = self.forward(features, attn_mask) + flat_loras, flat_layernorms = self.forward(features, attn_mask, position_ids) + # print(f"flat_loras: {flat_loras.shape}") + # if flat_layernorms is not None: + # print(f"flat_layernorms: {flat_layernorms.shape}") return self._to_lora_dict(flat_loras), self._to_layernorm_dict(flat_layernorms) @@ -735,7 +763,7 @@ class ModulatedPretrainedModel(nn.Module): self.base_model = self.base_model.to(*args, **kwargs) self.ctx_encoder = self.ctx_encoder.to(*args, **kwargs) # self.hypernet = self.hypernet.to(*args, **kwargs) - self.hypernet.to(torch.float32) + # self.hypernet.to(torch.float32) return self # Delegate to base_model @@ -865,21 +893,62 @@ class ModulatedPretrainedModel(nn.Module): def generate_weights( self, ctx_ids: Integer[Tensor, "bs ctx_len"], - ctx_attn_mask: Integer[Tensor, "bs ctx_len"], + ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_len"]] = None, + ctx_position_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, *args: Any, **kwargs: Any, ): - with torch.inference_mode(): + with torch.no_grad(): + # TODO: for modernbert ctx_encoder pass + # `cu_seq_len` and `max_seq_len` to the forward call ctx_features = self.ctx_encoder( - input_ids=ctx_ids, attention_mask=ctx_attn_mask, *args, **kwargs + input_ids=ctx_ids, + attention_mask=ctx_attn_mask, + position_ids=ctx_position_ids, + *args, + **kwargs, ) - return self.hypernet.generate_weights(ctx_features, ctx_attn_mask) + # print(f"padded ctx_features: {ctx_features.shape}") + # if ctx_attn_mask is not None: + # print(f"padded ctx_attn_mask: {ctx_attn_mask.shape}") + # if ctx_position_ids is not None: + # print(f"padded ctx_position_ids: {ctx_position_ids.shape}") + + return self.hypernet.generate_weights( + ctx_features, ctx_attn_mask, ctx_position_ids + ) + + # @torch.inference_mode() + # def _unpack(self, ctx_position_ids, ctx_features): + # # [1, len, d] -> [len, d] + # ctx_features = ctx_features.squeeze(0) + # lengths = torch.where(ctx_position_ids == 0)[1] + # total_len = torch.tensor([len(ctx_features)], device=self.device) + # lengths = torch.cat([lengths, total_len]) + # # compute the difference between the lengths + # lens = torch.diff(lengths).unsqueeze(1) + + # ctx_features = unpack(ctx_features, lens, "* d") + # ctx_attn_mask = [torch.ones(len(x), device=self.device) for x in ctx_features] + # ctx_features = torch.nn.utils.rnn.pad_sequence( + # ctx_features, + # batch_first=True, + # padding_value=0, + # ) + # ctx_attn_mask = torch.nn.utils.rnn.pad_sequence( + # ctx_attn_mask, + # batch_first=True, + # padding_value=0, + # ) + + # return ctx_features, ctx_attn_mask def forward( self, # ctx_features: Optional[Float[Tensor, "bs ctx_length feature_dim"]] = None, ctx_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_len"]] = None, + ctx_position_ids: Optional[Integer[Tensor, "bs ctx_len"]] = None, # chat_ids: Optional[Integer[Tensor, "bs chat_len"]] = None, # chat_attn_mask: Optional[Integer[Tensor, "bs chat_len"]] = None, # chat_labels: Optional[Integer[Tensor, "bs chat_len"]] = None, @@ -889,6 +958,7 @@ class ModulatedPretrainedModel(nn.Module): """Forward pass of the modulated model.""" generated_loras = None + generated_layernorms = None if ctx_ids is None and not self.use_base_input_as_ctx: logger.warning( ( @@ -913,78 +983,48 @@ class ModulatedPretrainedModel(nn.Module): if "attention_mask" in model_inputs_kwargs else None ) - with torch.inference_mode(): - ctx_features = self.ctx_encoder( - input_ids=ctx_ids, attention_mask=ctx_attn_mask + ctx_position_ids = ( + model_inputs_kwargs["position_ids"] + if "position_ids" in model_inputs_kwargs + else None ) - generated_loras, generated_layernorms = self.hypernet.generate_weights( - ctx_features, ctx_attn_mask + # with torch.inference_mode(): + # ctx_features = self.ctx_encoder( + # input_ids=ctx_ids, + # attention_mask=ctx_attn_mask, + # position_ids=ctx_position_ids, + # ) + # generated_loras, generated_layernorms = self.hypernet.generate_weights( + # ctx_features, ctx_attn_mask + # ) + generated_loras, generated_layernorms = self.generate_weights( + ctx_ids, ctx_attn_mask, ctx_position_ids ) - # compute kl loss - # - compute logits from the base model from tokenized chat [bs, chat_len, vocab_size] - # - compute logits from model w/ generated loras [bs, answer_len, vocab_size] - # - select only the position of the label tokens (but we're actually not using the labels themselves) - # - compute kl loss between the two logits - if self.use_kl_loss and "labels" in model_inputs_kwargs: - raise NotImplementedError("Not implemented") - # bs = chat_ids.shape[0] - # labels = model_inputs_kwargs.pop("labels") # [bs, prompt_res_len] - # first_label_pos = torch.argmax((labels != -100).float(), dim=-1) - # # just a place holder, we dont use labels as groundtruth anyway - # labels[torch.arange(bs), first_label_pos - 1] = 1 - # # same for chat_labels - # first_label_pos = torch.argmax((chat_labels != -100).float(), dim=-1) - # chat_labels[torch.arange(bs), first_label_pos - 1] = 1 - - # with torch.no_grad(): - # chat_outputs = self.base_model( - # input_ids=chat_ids, attention_mask=chat_attn_mask - # ) - - # with apply_generated_loras( - # self.base_model, - # generated_loras, - # self.hypernet.layer_indices, - # self.training, - # ): - # model_outputs = self.base_model( - # *model_inputs_args, **model_inputs_kwargs - # ) - # pred_logits = model_outputs.logits # [bs, prompt_res_len, vocab_size] - - # # [bs, ctx_prompt_res_len, vocab_size] - # base_chat_logits = chat_outputs.logits - # base_chat_logits = base_chat_logits[torch.where(chat_labels != -100)] - # pred_logits = pred_logits[torch.where(labels != -100)] - # kl_loss = F.kl_div( - # F.log_softmax(pred_logits, dim=-1), - # F.softmax(base_chat_logits, dim=-1), - # reduction="none", - # ) - # # only compute kl loss on tokens where labels != -100 - # kl_loss = kl_loss.sum(dim=-1).mean() - # model_outputs = ModelOutput(loss=kl_loss, **model_outputs) - else: - # input_ids in model_inputs_kwargs contains only - # prompt + response (for hypernet training) - with ( - apply_generated_loras( - self.base_model, - generated_loras, - self.hypernet.layer_indices, - self.training, - ), - apply_generated_layernorm( - self.base_model, - generated_layernorms, - self.hypernet.layer_indices, - self.training, - ), - ): - model_outputs = self.base_model( - *model_inputs_args, **model_inputs_kwargs - ) + # input_ids in model_inputs_kwargs contains only + # prompt + response (for hypernet training) + position_ids = ( + model_inputs_kwargs["position_ids"] + 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) return model_outputs @@ -993,10 +1033,12 @@ class ModulatedPretrainedModel(nn.Module): self, ctx_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, ctx_attn_mask: Optional[Integer[Tensor, "bs ctx_length"]] = None, + ctx_position_ids: Optional[Integer[Tensor, "bs ctx_length"]] = None, *model_inputs_args: Any, **model_inputs_kwargs: dict[str, Any], ): generated_loras = None + generated_layernorms = None if ctx_ids is None and not self.use_base_input_as_ctx: logger.warning( ( @@ -1020,26 +1062,41 @@ class ModulatedPretrainedModel(nn.Module): if "attention_mask" in model_inputs_kwargs else None ) - ctx_features = self.ctx_encoder( - input_ids=ctx_ids, attention_mask=ctx_attn_mask - ) - generated_loras, generated_layernorms = self.hypernet.generate_weights( - ctx_features, ctx_attn_mask + ctx_position_ids = ( + model_inputs_kwargs["position_ids"] + if "position_ids" in model_inputs_kwargs + else None + ) + # ctx_features = self.ctx_encoder( + # input_ids=ctx_ids, attention_mask=ctx_attn_mask + # ) + # generated_loras, generated_layernorms = self.hypernet.generate_weights( + # ctx_features, ctx_attn_mask + # ) + generated_loras, generated_layernorms = self.generate_weights( + ctx_ids, ctx_attn_mask, ctx_position_ids ) # apply lora hook to the base model # self.apply_generated_loras(generated_loras) + position_ids = ( + model_inputs_kwargs["position_ids"] + 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, ), ): @@ -1141,6 +1198,7 @@ def apply_generated_layernorm( base_model: nn.Module, generated_layernorms: Optional[dict[str, Float[Tensor, "bs n_layers h"]]] = None, layer_indices: Optional[Iterable[int]] = None, + position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, training: bool = False, ): if generated_layernorms is None: @@ -1156,6 +1214,7 @@ def apply_generated_layernorm( module_name, layer_idx, W=generated_layernorms[module_name][:, layer_idx], + position_ids=position_ids, training=training, ) yield base_model @@ -1168,6 +1227,7 @@ def apply_generated_loras( base_model: nn.Module, generated_loras: Optional[dict] = None, layer_indices: Optional[Iterable[int]] = None, + position_ids: Optional[Integer[Tensor, "bs seq_len"]] = None, training: bool = False, ): if generated_loras is None: @@ -1187,6 +1247,7 @@ def apply_generated_loras( 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, + position_ids=position_ids, training=training, ) @@ -1195,6 +1256,16 @@ def apply_generated_loras( remove_hook_handles(hooks) +@contextmanager +def apply_generated_loras_packed_sequence( + base_model: nn.Module, + generated_loras: Optional[dict] = None, + layer_indices: Optional[Iterable[int]] = 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/utils.py b/src/ctx_to_lora/utils.py index e30b2a9..9eacbb5 100644 --- a/src/ctx_to_lora/utils.py +++ b/src/ctx_to_lora/utils.py @@ -162,7 +162,7 @@ def get_peft_in_out_features( ) -> tuple[dict[str, int], dict[str, int]]: if peft_config is None: - peft_config = model.peft_config["default"] + return None, None in_features = dict() out_features = dict() for module_name, module in model.named_modules():