flash attn perceiver now workning

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)

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)

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)

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()

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(

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():