Merge pull request #3 from SakanaAI:layer-to-layer

Enhance context encoder with new configurations and fix attention bugs

intx_sft.py

@@ -199,7 +199,11 @@ def main():
         logger.info("Using HyperLoRA")
         if not ctx_args.from_pretrained_checkpoint:
             hypernet_config = get_hypernet_config(
-                model, ctx_encoder_model_config, hypernet_args, aggregator_args
+                model,
+                ctx_encoder_model_config,
+                hypernet_args,
+                aggregator_args,
+                ctx_encoder_args,
             )
             if ctx_encoder_args.layer_idx is None:
                 ctx_encoder_args.layer_idx = (

src/ctx_to_lora/configs.py

@@ -435,7 +435,16 @@ class CtxEncoderArguments:
         default=None,
         metadata={"help": "Context encoder model name or path."},
     )
-    # TODO: allow using activations from multiple layers?
+    ctx_encoder_type: Literal["embed_only", "per_layer_activations", "early_exit"] = (
+        field(
+            default="early_exit",
+            metadata={
+                "help": "Context encoder type. "
+                "Options: 'embed_only', 'per_layer_activations', 'early_exit'."
+            },
+        )
+    )
+    # used only with `early_exit` type
     layer_idx: int | None = field(
         default=None,
         metadata={
@@ -467,17 +476,35 @@ class AggregatorArguments:
     #     default=0.0,
     #     metadata={"help": "Attention dropout probability for Perceiver."},
     # )
+    concat_latents_context: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to concatenate latents and context for Perceiver. "
+            "If False, only context is used. Used only when `use_self_attn` is False."
+        },
+    )
+    # not used (keep herer for backward compatibility)
     num_latent_factor: int = field(
         default=8,
         metadata={"help": "Number of latent factors for Perceiver."},
     )
+    n_base_queries: int = field(
+        default=208,  # 26 * 8
+        metadata={"help": "Number of latent queries of Perceiver."},
+    )
     # num_blocks: int = field(
     #     default=8,
     #     metadata={"help": "Number of blocks for Perceiver."},
     # )
+
+    # misnomer
     num_self_attends_per_block: int = field(
         default=6,
-        metadata={"help": "Number of self-attends per block for Perceiver."},
+        metadata={"help": "Number of attn blocks for Perceiver."},
+    )
+    n_cross_attn_layers: int = field(
+        default=6,
+        metadata={"help": "Number of cross-attention layers for Perceiver. "},
     )
     # self_attention_widening_factor: int = field(
     #     default=4,

src/ctx_to_lora/modeling/aggregator.py

@@ -2,9 +2,10 @@ import logging
 from dataclasses import dataclass
 from enum import Enum
 
-from einops import rearrange, unpack
+from einops import rearrange, unpack, repeat
 from jaxtyping import Float, Integer
 from torch import Tensor, nn
+import torch
 from transformers import (
     PretrainedConfig,
     PreTrainedModel,
@@ -40,16 +41,25 @@ class AggregatorConfig:
     pooling_type: POOL_FN
 
     # perceiver
+    concat_latents_context: bool
+    n_cross_attn_layers: int  # number of cross-attention layers
     num_self_attends_per_block: int  # = 16
     decoder_depth: int  # 1 = only cross-attention
     num_latent_factor: int = 8
+    n_base_queries: int = 208
     lora_r: int = 8
     per_rank_gen: bool = False
 
+    # layer to layer generation
+    # maps activations of layer L to LoRA of that layer
+    layer_to_layer_ctx_encoder: bool = False
+    n_ctx_model_layers: int = -1
+
 
 def get_aggregator_config(
     model: PreTrainedModel,
     ctx_encoder_model_config: PretrainedConfig,
+    layer_to_layer_ctx_encoder: bool,
     output_size: int,
     num_modules: int,
     num_extra_modules: int,
@@ -65,6 +75,8 @@ def get_aggregator_config(
         num_extra_modules=num_extra_modules,
         lora_r=lora_r,
         per_rank_gen=per_rank_gen,
+        layer_to_layer_ctx_encoder=layer_to_layer_ctx_encoder,
+        n_ctx_model_layers=ctx_encoder_model_config.num_hidden_layers,
         **vars(aggregator_args),
     )
 
@@ -81,7 +93,10 @@ class Perceiver(nn.Module):
         num_extra_modules,
         per_rank_gen,
         lora_r,
-        num_latent_factor,
+        num_latent_factor,  # unused
+        layer_to_layer_ctx_encoder,
+        n_base_queries,
+        concat_latents_context,
         *args,
         **kwargs,
     ):
@@ -92,16 +107,24 @@ class Perceiver(nn.Module):
         self.per_rank_gen = per_rank_gen
         self.r = lora_r if self.per_rank_gen else 1
         n_output_queries = num_layers * (num_modules * self.r + num_extra_modules)
+        self.tot_output_size = n_output_queries
+        self.layer_to_layer = layer_to_layer_ctx_encoder
+        if self.layer_to_layer:
+            n_output_queries = num_modules * self.r + num_extra_modules
         self.config = Idefics2PerceiverConfig(
+            concat_latents_context=concat_latents_context,
             input_size=feature_size,
             # the first layer is xattn
             resampler_depth=kwargs["num_self_attends_per_block"] + 1,
-            resampler_n_latents=n_output_queries * num_latent_factor,
+            n_cross_attn_layers=kwargs["n_cross_attn_layers"],
+            # resampler_n_latents=n_output_queries * num_latent_factor,
+            resampler_n_latents=n_base_queries,
             intermediate_size_factor=4,
             hidden_size=output_size,
             attn_implementation="flash_attention_2",
         )
         self.decoder_config = Idefics2PerceiverConfig(
+            concat_latents_context=concat_latents_context,
             input_size=output_size,
             resampler_depth=kwargs["decoder_depth"],
             resampler_n_latents=n_output_queries,
@@ -113,11 +136,68 @@ class Perceiver(nn.Module):
 
     def forward(
         self,
-        ctx_features: Float[Tensor, "bs seq_len feature_dim"],
+        ctx_features: Float[Tensor, "bs seq_len feature_dim"]
+        | Float[Tensor, "bs x seq_len feature_dim"],
         ctx_attn_mask: Integer[Tensor, "bs seq_len"] | None = None,
         ctx_position_ids: Integer[Tensor, "bs seq_len"] | None = None,
     ):
+        # if len(ctx_features.shape) == 3:
+        #     # ["bs seq_len feature_dim"] to
+        #     # ["bs n_output_queries feature_dim"]
+        #     x = self.perceiver(ctx_features, ctx_attn_mask, ctx_position_ids)
+        # elif len(ctx_features.shape) == 4:
+        #     # ["bs n_ctx_model_layers seq_len feature_dim"] to
+        #     # to ["bs n_ctx_model_layers n_output_queries feature_dim"]
+        #     x = self.perceiver(ctx_features, ctx_attn_mask, ctx_position_ids)
+        #     x = rearrange(
+        #         x,
+        #         (
+        #             "bs n_ctx_model_layers (n_modules r) seq_len d -> "
+        #             "bs (n_ctx_model_layers n_modules r) d"
+        #         ),
+        #         n_ctx_model_layers=self.n_ctx_model_layers,
+        #         n_modules=self.num_modules,
+        #         r=self.r,
+        #     )
+        # else:
+        #     raise ValueError(
+        #         f"ctx_features should be 3D or 4D tensor, got {ctx_features.shape}"
+        #     )
+
+        if self.layer_to_layer:
+            ctx_features = rearrange(
+                ctx_features,
+                "bs num_layers seq_len feature_dim -> bs (num_layers seq_len) feature_dim",
+            )
+            if ctx_attn_mask is not None:
+                ctx_attn_mask = repeat(
+                    ctx_attn_mask,
+                    "bs seq_len -> (bs num_layers) seq_len",
+                    num_layers=self.num_layers,
+                )
+            ctx_position_ids = repeat(
+                ctx_position_ids,
+                "bs seq_len -> bs (num_layers seq_len)",
+                num_layers=self.num_layers,
+            )
+
+        # print(f"ctx_features shape: {ctx_features.shape}")
+        # print(
+        #     f"ctx_attn_mask shape: {ctx_attn_mask.shape if ctx_attn_mask is not None else None}"
+        # )
+        # print(
+        #     f"ctx_position_ids shape: {ctx_position_ids.shape if ctx_position_ids is not None else None}"
+        # )
         x = self.perceiver(ctx_features, ctx_attn_mask, ctx_position_ids)
+
+        if self.layer_to_layer:
+            per_layer_size = self.num_modules * self.r + self.num_extra_modules
+            x = rearrange(
+                x,
+                ("(bs num_layers) (per_layer_sz) d -> bs (num_layers per_layer_sz) d"),
+                num_layers=self.num_layers,
+                per_layer_sz=per_layer_size,
+            )
         lora_x, extra_x = unpack(
             x,
             [
@@ -228,7 +308,90 @@ class Perceiver(nn.Module):
 #         return self.mlp(self.mixer(emb))
 
 
-AG = {
+AGGREGATOR_CLS = {
     # AGGREGATOR_TYPE.POOLER: Pooler,
     AGGREGATOR_TYPE.PERCEIVER: Perceiver,
 }
+
+if __name__ == "__main__":
+    # Test the Perceiver class
+    print("Testing Perceiver...")
+
+    # Test parameters
+    feature_size = 768
+    output_size = 512
+    num_layers = 12
+    num_modules = 4
+    num_extra_modules = 2
+    lora_r = 8
+    batch_size = 2
+    seq_len = 100
+    device = "cuda"
+
+    torch.set_default_dtype(torch.bfloat16)
+
+    # Create Perceiver instance
+    perceiver = Perceiver(
+        feature_size=feature_size,
+        output_size=output_size,
+        num_layers=num_layers,
+        num_modules=num_modules,
+        num_extra_modules=num_extra_modules,
+        per_rank_gen=True,
+        lora_r=lora_r,
+        num_latent_factor=8,
+        layer_to_layer_ctx_encoder=False,
+        n_base_queries=208,
+        num_self_attends_per_block=16,
+        decoder_depth=1,
+    ).to(device)
+
+    # Test with 3D input (standard case)
+    print("Testing with 3D input...")
+    ctx_features = torch.randn(1, batch_size * seq_len, feature_size).to(device)
+    ctx_attn_mask = None  # torch.ones(batch_size, seq_len, dtype=torch.long).to(device)
+    ctx_position_ids = torch.arange(seq_len).repeat(batch_size).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        lora_x, extra_x = perceiver(ctx_features, ctx_attn_mask, ctx_position_ids)
+
+    print(f"LoRA output shape: {lora_x.shape}")
+    print(f"Extra output shape: {extra_x.shape}")
+    print(
+        f"Expected LoRA shape: ({batch_size}, {num_layers}, {num_modules}, {lora_r}, {output_size})"
+    )
+    print(
+        f"Expected extra shape: ({batch_size}, {num_layers}, {num_extra_modules}, {output_size})"
+    )
+
+    # Test with layer-to-layer configuration
+    print("\nTesting with layer-to-layer configuration...")
+    perceiver_l2l = Perceiver(
+        feature_size=feature_size,
+        output_size=output_size,
+        num_layers=num_layers,
+        num_modules=num_modules,
+        num_extra_modules=num_extra_modules,
+        per_rank_gen=True,
+        lora_r=lora_r,
+        num_latent_factor=8,
+        layer_to_layer_ctx_encoder=True,
+        n_base_queries=208,
+        num_self_attends_per_block=16,
+        decoder_depth=1,
+    ).to(device)
+
+    # Test with 4D input for layer-to-layer
+    ctx_features_4d = torch.randn(1, num_layers, batch_size * seq_len, feature_size).to(
+        device
+    )
+
+    with torch.no_grad():
+        lora_x_l2l, extra_x_l2l = perceiver_l2l(
+            ctx_features_4d, ctx_attn_mask, ctx_position_ids
+        )
+
+    print(f"Layer-to-layer LoRA output shape: {lora_x_l2l.shape}")
+    print(f"Layer-to-layer extra output shape: {extra_x_l2l.shape}")
+
+    print("All tests passed!")

src/ctx_to_lora/modeling/ctx_encoder.py

@@ -1,10 +1,14 @@
 import logging
 from contextlib import contextmanager
+from enum import Enum
 
 import torch
 from torch import nn
 from transformers import PreTrainedModel
 
+from ctx_to_lora.configs import CtxEncoderArguments
+from ctx_to_lora.utils import get_base_model
+
 logger = logging.getLogger()
 
 
@@ -47,13 +51,14 @@ def maybe_add_batch_dim(kwargs):
 
 
 class EarlyExit(nn.Module):
-    def __init__(self, base_model: PreTrainedModel, exit_layer: int):
+    def __init__(self, base_model: PreTrainedModel, config: CtxEncoderArguments):
         super().__init__()
-        self.base_model = base_model
+        self.base_model = base_model = get_base_model(base_model)
+        self.exit_layer = config.layer_idx
         if "gte" in base_model.config.name_or_path:
-            self.base_model.encoder.layer = base_model.encoder.layer[:exit_layer]
+            self.base_model.encoder.layer = base_model.encoder.layer[: self.exit_layer]
         else:
-            self.base_model.layers = base_model.layers[:exit_layer]
+            self.base_model.layers = base_model.layers[: self.exit_layer]
 
     @property
     def config(self):
@@ -63,3 +68,52 @@ class EarlyExit(nn.Module):
     def forward(self, **kwargs):
         model_outputs = self.base_model(**kwargs)
         return model_outputs.last_hidden_state
+
+
+class EmbeddingOnly(nn.Module):
+    def __init__(self, base_model: PreTrainedModel, config: CtxEncoderArguments):
+        super().__init__()
+        self.base_model = base_model
+
+    @property
+    def config(self):
+        return self.base_model.config
+
+    @torch.no_grad()
+    def forward(self, **kwargs):
+        kwargs["output_hidden_states"] = True  # Force output of hidden states
+        outputs = self.base_model(**kwargs)
+        # Return the embeddings only
+        return outputs.hidden_states[0]  # The first hidden state is the embeddings
+
+
+class PerLayerActivations(nn.Module):
+    def __init__(self, base_model: PreTrainedModel, config: CtxEncoderArguments):
+        super().__init__()
+        self.base_model = base_model
+
+    @property
+    def config(self):
+        return self.base_model.config
+
+    @torch.no_grad()
+    def forward(self, **kwargs):
+        kwargs["output_hidden_states"] = True  # Force output of hidden states
+        outputs = self.base_model(**kwargs)
+        # Return all layers' activations except the last one
+        # from embeddings to the input of the last attn block
+        # Shape: (batch_size, num_layers, seq_len, hidden_size)
+        return torch.stack(outputs.hidden_states[:-1], dim=1)
+
+
+class CTX_ENCODER_TYPE(str, Enum):
+    EARLY_EXIT = "early_exit"
+    EMBED_ONLY = "embed_only"
+    PER_LAYER_ACTIVATIONS = "per_layer_activations"
+
+
+CTX_ENCODER_CLS = {
+    CTX_ENCODER_TYPE.EARLY_EXIT: EarlyExit,
+    CTX_ENCODER_TYPE.EMBED_ONLY: EmbeddingOnly,
+    CTX_ENCODER_TYPE.PER_LAYER_ACTIVATIONS: PerLayerActivations,
+}

src/ctx_to_lora/modeling/hypernet.py

@@ -41,15 +41,18 @@ from ctx_to_lora.hooks import (
 from ctx_to_lora.model_loading import (
     get_model,
 )
-from ctx_to_lora.modeling.aggregator import AG, AggregatorConfig, get_aggregator_config
+from ctx_to_lora.modeling.aggregator import (
-from ctx_to_lora.modeling.ctx_encoder import EarlyExit
+    AGGREGATOR_CLS,
+    AggregatorConfig,
+    get_aggregator_config,
+)
+from ctx_to_lora.modeling.ctx_encoder import CTX_ENCODER_CLS, CTX_ENCODER_TYPE
 from ctx_to_lora.modeling.lora_layer import (
     apply_lora_to_layers,
     lora_forward,
     lora_forward_packed,
 )
 from ctx_to_lora.utils import (
-    get_base_model,
     get_layers,
     get_num_layers,
     get_peft_in_out_features,
@@ -84,6 +87,7 @@ def get_hypernet_config(
     ctx_encoder_model_config: PretrainedConfig,
     hypernet_args: HypernetArguments,
     aggregator_args: AggregatorArguments,
+    ctx_encoder_args: CtxEncoderArguments,
 ):
     num_modules = 0
     lora_config = getattr(model, "peft_config", None)
@@ -101,6 +105,7 @@ def get_hypernet_config(
         aggregator_config=get_aggregator_config(
             model,
             ctx_encoder_model_config,
+            ctx_encoder_args.ctx_encoder_type == CTX_ENCODER_TYPE.PER_LAYER_ACTIVATIONS,
             hypernet_args.latent_size,
             num_modules,
             num_extra_modules,
@@ -222,7 +227,9 @@ class HyperLoRA(nn.Module):
 
     def _init_model(self):
         self.agg_config = self.config.aggregator_config
-        self.aggregator = AG[self.agg_config.aggregator_type](**vars(self.agg_config))
+        self.aggregator = AGGREGATOR_CLS[self.agg_config.aggregator_type](
+            **vars(self.agg_config)
+        )
 
         self.lora_config = self.config.lora_config
 
@@ -546,7 +553,7 @@ class HyperLoRA(nn.Module):
         attn_mask: Integer[Tensor, "bs seq_len"] | None = None,
         position_ids: Integer[Tensor, "bs seq_len"] | None = None,
     ):
-        # [bs, n_layers, n_total_modules, feature_dim]
+        # [bs, n_layers x n_total_modules x r, feature_dim]
         lora_emb, extra_emb = self.aggregator(features, attn_mask, position_ids)
         # lora_emb, extra_emb = unpack(
         #     emb,
@@ -673,8 +680,8 @@ class ModulatedPretrainedModel(nn.Module):
             requires_grad=False,
             use_flash_attn=base_model_attn_impl == "flash_attention_2",
         )
-        self.ctx_encoder = EarlyExit(
+        self.ctx_encoder = CTX_ENCODER_CLS[self.ctx_encoder_args.ctx_encoder_type](
-            get_base_model(encoder_model), self.ctx_encoder_args.layer_idx
+            encoder_model, self.ctx_encoder_args
         )
 
     # delegate to base_model

src/ctx_to_lora/modeling/idefics2.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """PyTorch Idefics2 model."""
 
+from copy import deepcopy
 import math
 
 import torch
@@ -78,9 +79,15 @@ class Idefics2PerceiverConfig(PretrainedConfig):
         resampler_head_dim=96,
         num_key_value_heads=4,
         attention_dropout=0.0,
+        is_cross_attn=True,
+        n_cross_attn_layers=1,
+        concat_latents_context=False,
         **kwargs,
     ):
         # for mlp
+        self.concat_latents_context = concat_latents_context
+        self.is_cross_attn = is_cross_attn
+        self.n_cross_attn_layers = n_cross_attn_layers
         self.input_size = input_size
         self.intermediate_size_factor = intermediate_size_factor
         # for perceiver
@@ -157,9 +164,9 @@ class Idefics2PreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         std = (
-            self.config.architecturesinitializer_range
+            self.config.initializer_range
             if hasattr(self.config, "initializer_range")
-            else self.config.initializer_range
+            else 0.02
         )
 
         if hasattr(module, "class_embedding"):
@@ -353,7 +360,9 @@ class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention):
     def forward(
         self,
         latents: torch.Tensor,
-        context: torch.Tensor,
+        is_cross_attn: bool,
+        concat_latents_context: bool,
+        context: torch.Tensor | None = None,
         attention_mask: torch.LongTensor | None = None,
         position_ids: torch.LongTensor | None = None,
         past_key_value: Cache | None = None,
@@ -363,25 +372,57 @@ class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention):
     ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
         bsz, q_len, _ = latents.size()
         # kv_seq_len = q_len + context.size()[1]
-        kv_seq_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))
+        if is_cross_attn:
-        # value_states = self.v_proj(torch.cat([context, latents], dim=-2))
+            kv_inp = context
-        key_states = self.k_proj(context)
+            if concat_latents_context:
-        value_states = self.v_proj(context)
+                # 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`
+                if context.shape[0] != latents.shape[0]:
+                    # context: [1, tot_len, dim]
+                    # latents: [bs, n_latents, dim]
+                    # slice context sample by sample and concat with latents
+                    old_cu_seq_lens_k = kwargs.pop("old_cu_seq_lens_k")
+                    cu_seq_lens_k = kwargs["cu_seq_lens_k"]
+                    kv_inp = torch.empty(
+                        1,
+                        cu_seq_lens_k[-1],
+                        latents.shape[2],
+                        dtype=context.dtype,
+                        device=context.device,
+                    )
+                    # Compute context lengths and indices
+                    ctx_lens = old_cu_seq_lens_k[1:] - old_cu_seq_lens_k[:-1]
+                    ctx_start = old_cu_seq_lens_k[:-1]
+                    ctx_end = old_cu_seq_lens_k[1:]
+                    kv_start = cu_seq_lens_k[:-1]
+                    kv_end = cu_seq_lens_k[1:]
+                    # Fill context slices
+                    for i in range(bsz):
+                        kv_inp[0, kv_start[i] : kv_start[i] + ctx_lens[i]] = context[
+                            0, ctx_start[i] : ctx_end[i]
+                        ]
+                        kv_inp[0, kv_start[i] + ctx_lens[i] : kv_end[i]] = latents[i]
+                else:
+                    kv_inp = torch.cat([context, latents], dim=-2)
+
+        else:
+            kv_inp = latents
+
+        key_states = self.k_proj(kv_inp)
+        value_states = self.v_proj(kv_inp)
 
         # 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
+            *kv_inp.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
+            *kv_inp.shape[:2], self.num_key_value_heads, self.head_dim
         ).transpose(1, 2)
 
         # kv_seq_len = key_states.shape[-2]
@@ -479,7 +520,7 @@ class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention):
 
 
 IDEFICS2_PERCEIVER_ATTENTION_CLASSES = {
-    "eager": Idefics2PerceiverAttention,
+    # "eager": Idefics2PerceiverAttention,
     "flash_attention_2": Idefics2PerceiverFlashAttention2,
 }
 
@@ -491,12 +532,16 @@ class Idefics2PerceiverLayer(nn.Module):
         self.n_latents = config.resampler_n_latents
         self.depth = config.resampler_depth
         self.rms_norm_eps = config.rms_norm_eps
+        self.is_cross_attn = config.is_cross_attn
+        self.concat_latents_context = config.concat_latents_context
 
         self.input_latents_norm = Idefics2RMSNorm(
             self.hidden_size, eps=self.rms_norm_eps
         )
-        self.input_context_norm = Idefics2RMSNorm(
+        self.input_context_norm = (
-            self.hidden_size, eps=self.rms_norm_eps
+            Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+            if self.is_cross_attn
+            else torch.nn.Identity()
         )
         self.self_attn = IDEFICS2_PERCEIVER_ATTENTION_CLASSES[
             config._attn_implementation
@@ -504,6 +549,10 @@ class Idefics2PerceiverLayer(nn.Module):
         self.post_attention_layernorm = Idefics2RMSNorm(
             self.hidden_size, eps=self.rms_norm_eps
         )
+        self.pre_ff_layernorm = Idefics2RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.post_ff_layernorm = Idefics2RMSNorm(
+            self.hidden_size, eps=self.rms_norm_eps
+        )
         self.mlp = Idefics2MLP(
             hidden_size=config.hidden_size,
             intermediate_size=config.hidden_size * 4,
@@ -514,7 +563,7 @@ class Idefics2PerceiverLayer(nn.Module):
     def forward(
         self,
         latents: torch.Tensor,
-        context: torch.Tensor,
+        context: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
         position_ids: torch.LongTensor | None = None,
         past_key_value: tuple[torch.Tensor] | None = None,
@@ -543,16 +592,21 @@ class Idefics2PerceiverLayer(nn.Module):
 
         latents, self_attn_weights, present_key_value = self.self_attn(
             latents=latents,
+            is_cross_attn=self.is_cross_attn,
+            concat_latents_context=self.concat_latents_context,
             context=context,
             attention_mask=attention_mask,
             position_ids=position_ids,
             **kwargs,
         )
+        latents = self.post_attention_layernorm(latents)
         latents = residual + latents
         residual = latents
 
-        latents = self.post_attention_layernorm(latents)
+        # latents = self.post_attention_layernorm(latents)
+        latents = self.pre_ff_layernorm(latents)
         latents = self.mlp(latents)
+        latents = self.post_ff_layernorm(latents)
         latents = residual + latents
 
         outputs = (latents,)
@@ -595,16 +649,34 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
         self.hidden_size = config.hidden_size
         self.hidden_act = config.hidden_act
         self.n_latents = config.resampler_n_latents
+        self.n_cross_attn_layers = config.n_cross_attn_layers
         self.depth = config.resampler_depth
+        self.n_self_attn_layers = self.depth - self.n_cross_attn_layers
         self.rms_norm_eps = config.rms_norm_eps
+        self.concat_latents_context = config.concat_latents_context
 
         # Create Latents for Perceiver
         self.latents = nn.Parameter(torch.randn(self.n_latents, self.hidden_size))
 
         # Create Transformer Blocks
-        self.layers = nn.ModuleList(
+        if not self.concat_latents_context:
-            [Idefics2PerceiverLayer(config, idx) for idx in range(self.depth)]
+            x_attn_config = deepcopy(config)
-        )
+            x_attn_config.is_cross_attn = True
+            self.layers = [
+                Idefics2PerceiverLayer(x_attn_config, idx)
+                for idx in range(self.n_cross_attn_layers)
+            ]
+
+            config.is_cross_attn = False
+            self.layers += [
+                Idefics2PerceiverLayer(config, idx)
+                for idx in range(self.n_self_attn_layers)
+            ]
+            self.layers = nn.ModuleList(self.layers)
+        else:
+            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"
@@ -679,25 +751,117 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
             )
 
             max_length_k = position_ids.max() + 1
+            if self.concat_latents_context:
+                last_pos_ids = torch.arange(
+                    position_ids[-1] + 1 + self.n_latents,
+                    device=position_ids.device,
+                    dtype=position_ids.dtype,
+                )
+                new_position_ids = torch.empty(
+                    cu_seq_lens_k[-1] + self.n_latents * bsz,
+                    dtype=position_ids.dtype,
+                    device=position_ids.device,
+                )
+
+                cur_idx = 0
+                for i, idx in enumerate(indices[position_ids == 0][1:]):
+                    l = idx - cur_idx
+
+                    new_cur_idx = cur_idx + self.n_latents * i
+                    n_new_ids = l + self.n_latents
+                    new_ids = torch.arange(
+                        n_new_ids,
+                        device=position_ids.device,
+                        dtype=position_ids.dtype,
+                    )
+                    new_position_ids[new_cur_idx : new_cur_idx + n_new_ids] = new_ids
+                    cur_idx = idx
+                new_cur_idx = cur_idx + self.n_latents * (bsz - 1)
+                new_position_ids[new_cur_idx:] = last_pos_ids
+                position_ids = new_position_ids
+
         else:
             raise ValueError("either position_ids or attention_mask is required")
 
-        for perceiver_layer in self.layers:
+        if self.concat_latents_context:
-            layer_outputs = perceiver_layer(
+            old_cu_seq_lens_k = cu_seq_lens_k
-                compressed_context,
+            cu_seq_lens_k = cu_seq_lens_k + cu_seq_lens_q
-                context,
+            max_length_k = max_length_q + max_length_k
-                attention_mask=None,
+            # print("Using concat latents + context for resampler")
-                position_ids=position_ids,
+            # print(f"updated position_ids: {position_ids}")
-                past_key_value=None,
+            # print(f"updated cu_seq_lens_k: {cu_seq_lens_k}")
-                output_attentions=False,
+            # print(f"updated max_length_k: {max_length_k}")
-                use_cache=False,
+            # print(f"cu_seq_lens_q: {cu_seq_lens_q}")
-                cu_seq_lens_q=cu_seq_lens_q,
+            # print(f"max_length_q: {max_length_q}")
-                cu_seq_lens_k=cu_seq_lens_k,
+            for layer in self.layers:
-                max_length_q=max_length_q,
+                layer_outputs = layer(
-                max_length_k=max_length_k,
+                    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,
+                    old_cu_seq_lens_k=old_cu_seq_lens_k,
+                )
+                compressed_context = layer_outputs[0]
+        else:
+            for encoder in self.layers[: self.n_cross_attn_layers]:
+                # print("Using cross-attention for resampler")
+                # print(f"position_ids: {position_ids}")
+                # print(f"cu_seq_lens_q: {cu_seq_lens_q}")
+                # print(f"max_length_q: {max_length_q}")
+                # print(f"cu_seq_lens_k: {cu_seq_lens_k}")
+                # print(f"max_length_k: {max_length_k}")
+                layer_outputs = encoder(
+                    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 = layer_outputs[0]
+
+            if self.n_self_attn_layers:
+                position_ids = torch.arange(
+                    self.n_latents, device=position_ids.device, dtype=torch.int32
+                ).repeat(1, bsz)
+                cu_seq_lens_k = cu_seq_lens_q
+                max_length_k = max_length_q
+
+            for processor in self.layers[self.n_cross_attn_layers :]:
+                # print(f"Using self-attention for resampler")
+                # print(f"updated position_ids: {position_ids}")
+                # print(f"updated cu_seq_lens_k: {cu_seq_lens_k}")
+                # print(f"updated max_length_k: {max_length_k}")
+                # print(f"cu_seq_lens_q: {cu_seq_lens_q}")
+                # print(f"max_length_q: {max_length_q}")
+                layer_outputs = processor(
+                    compressed_context,
+                    context=None,
+                    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)
 
@@ -706,16 +870,19 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
 
 class Idefics2Perceiver(Idefics2PreTrainedModel):
     def __init__(
-        self, config: Idefics2PerceiverConfig, decoder_config: Idefics2PerceiverConfig
+        self,
+        encoder_config: Idefics2PerceiverConfig,
+        decoder_config: Idefics2PerceiverConfig,
     ):
-        super().__init__(config)
+        super().__init__(encoder_config)
         self.modality_projection = Idefics2MLP(
-            hidden_size=config.input_size,
+            hidden_size=encoder_config.input_size,
-            intermediate_size=config.intermediate_size_factor * config.input_size,
+            intermediate_size=encoder_config.intermediate_size_factor
-            output_size=config.hidden_size,
+            * encoder_config.input_size,
-            hidden_act=config.hidden_act,
+            output_size=encoder_config.hidden_size,
+            hidden_act=encoder_config.hidden_act,
         )
-        self.encoder = Idefics2PerceiverResampler._from_config(config)
+        self.encoder = Idefics2PerceiverResampler._from_config(encoder_config)
         self.decoder = Idefics2PerceiverResampler._from_config(decoder_config)
 
     def forward(
@@ -762,20 +929,24 @@ __all__ = [
 ]
 
 if __name__ == "__main__":
-    config = Idefics2PerceiverConfig(
+    encoder_config = Idefics2PerceiverConfig(
+        concat_latents_context=True,
         input_size=1234,
         resampler_n_latents=64,
-        resampler_depth=8,
+        resampler_depth=1 + 8,
         attn_implementation="flash_attention_2",
     )
     decoder_config = Idefics2PerceiverConfig(
-        input_size=config.hidden_size,
+        concat_latents_context=True,
+        input_size=encoder_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)
+    model = (
-    print(model)
+        Idefics2Perceiver(encoder_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")
@@ -805,3 +976,40 @@ if __name__ == "__main__":
     print(model.encoder.latents.grad)
     print(model.decoder.latents.grad)
     breakpoint()
+
+    # inp = torch.rand(1, 192 + 37, 1234, dtype=torch.bfloat16, device="cuda")
+    # inp = torch.cat([inp, inp], dim=1)
+    # inp = inp.unsqueeze(1)
+    # # mask = torch.ones(1, 128, dtype=torch.bfloat16, device="cuda")
+    # # mask[..., -1] = 0
+    # # print(mask)
+    # out = torch.stack(
+    #     [
+    #         model(
+    #             inp[:, i],
+    #             # 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"),
+    #         )
+    #         for i in range(inp.shape[1])
+    #     ]
+    # )
+
+    # 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()