Perceiver blocks + layer-to-layer (#4)

scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh

@@ -31,6 +31,7 @@ uv run accelerate launch --main_process_port $port \
 --add_negative_prompt=False \
 --add_repeat_prompt=False \
 --use_sequence_packing=True \
+--max_qas_len=2048 \
 --max_packed_inp_len=8192 \
 --max_packed_ctx_len=16384 \
 --per_rank_gen=True \

scripts/short_ctx/gemma_qa_short_ctx_exp_self_attn.sh

@@ -31,6 +31,7 @@ uv run accelerate launch --main_process_port $port \
 --add_negative_prompt=False \
 --add_repeat_prompt=False \
 --use_sequence_packing=True \
+--max_qas_len=2048 \
 --max_packed_inp_len=8192 \
 --max_packed_ctx_len=16384 \
 --per_rank_gen=True \

scripts/short_ctx/gemma_qa_short_ctx_perceiver.sh

@@ -18,10 +18,9 @@ uv run accelerate launch --main_process_port $port \
 --gradient_accumulation_steps=2 \
 --per_device_eval_batch_size=64 \
 --target_modules=down_proj \
---num_self_attends_per_block=8 \
---n_cross_attn_layers=8 \
---concat_latents_context=True \
---num_latent_factor=1 \
+--num_blocks=4 \
+--num_self_attn_per_block=3 \
+--n_latent_queries=208 \
 --num_pre_head_layers=1 \
 --lora_r=8 \
 --eval_steps=1000 \
@@ -32,6 +31,7 @@ uv run accelerate launch --main_process_port $port \
 --add_negative_prompt=False \
 --add_repeat_prompt=False \
 --use_sequence_packing=True \
+--max_qas_len=2048 \
 --max_packed_inp_len=8192 \
 --max_packed_ctx_len=16384 \
 --per_rank_gen=True \

scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+#SBATCH --job-name=ctxlora
+#SBATCH --nodes=1
+#SBATCH --partition=a3
+#SBATCH --gpus=4
+#SBATCH --output=slurm_logs/%x-%j.out
+#SBATCH --error=slurm_logs/%x-%j.out
+
+port=$((10000 + ($SLURM_JOBID % 50000)))
+echo "Using port: $port"
+
+# --gradient_accumulation_steps=8 --gradient_clipping=1.0
+uv run accelerate launch --main_process_port $port \
+--num_processes=4 --gpu_ids all intx_sft.py $1 \
+--model_name_or_path=google/gemma-2-2b-it \
+--max_steps=12_000 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=4 \
+--per_device_eval_batch_size=32 \
+--target_modules=down_proj \
+--num_blocks=4 \
+--num_self_attn_per_block=3 \
+--ctx_encoder_type=per_layer_activations \
+--n_latent_queries=8 \
+--num_pre_head_layers=1 \
+--lora_r=8 \
+--eval_steps=1000 \
+--save_steps=1000 \
+--learning_rate=4e-5 \
+--lora_dropout=0.0 \
+--neftune_noise_alpha=5 \
+--add_negative_prompt=False \
+--add_repeat_prompt=False \
+--use_sequence_packing=True \
+--max_qas_len=2048 \
+--max_packed_inp_len=4096 \
+--max_packed_ctx_len=8192 \
+--per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=0.1 \
+--logging_steps=50

src/ctx_to_lora/configs.py

@@ -483,19 +483,12 @@ 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(
+    n_latent_queries: int = field(
         default=208,  # 26 * 8
         metadata={"help": "Number of latent queries of Perceiver."},
     )
@@ -504,27 +497,23 @@ class AggregatorArguments:
     #     metadata={"help": "Number of blocks for Perceiver."},
     # )
 
-    # misnomer
-    num_self_attends_per_block: int = field(
+    num_blocks: int = field(
+        default=8,
+        metadata={"help": "Number of blocks for Perceiver."},
+    )
+    num_self_attn_per_block: int = field(
         default=6,
-        metadata={"help": "Number of attn blocks for Perceiver."},
+        metadata={"help": "Number of self-attention layers per block for Perceiver."},
     )
-    n_cross_attn_layers: int = field(
-        default=6,
-        metadata={"help": "Number of cross-attention layers for Perceiver. "},
+    shared_weights: bool = field(
+        default=False,
+        metadata={"help": "Whether to share weights across blocks for Perceiver."},
     )
-    # self_attention_widening_factor: int = field(
-    #     default=4,
-    #     metadata={"help": "Self-attention widening factor for Perceiver."},
+
+    # decoder_depth: int = field(
+    #     default=1,
+    #     metadata={"help": "Decoder depth for Perceiver."},
     # )
-    # cross_attention_widening_factor: int = field(
-    #     default=4,
-    #     metadata={"help": "Cross-attention widening factor for Perceiver."},
-    # )
-    decoder_depth: int = field(
-        default=1,
-        metadata={"help": "Decoder depth for Perceiver."},
-    )
 
 
 # needed for loading model from checkpoint

src/ctx_to_lora/data/processing.py

@@ -665,9 +665,10 @@ def get_tokenized_dataset(
         num_proc=num_proc,
         **tokenize_kwargs,
     )
-    tokenized_ds = tokenized_ds.shuffle()
+
     if ("train" in split) and is_caching_enabled():
         tokenized_ds.save_to_disk(ds_path, num_proc=num_proc)
+        tokenized_ds = tokenized_ds.shuffle()
         # force reload from disk for fingerprint consistency
         tokenized_ds = datasets.load_from_disk(ds_path)
     return tokenized_ds
@@ -1006,6 +1007,8 @@ def split_too_long_qas(samples: dict[str, any], max_qas_len: int):
     # e.g., if max_qas_len = 512, and qas is 1024 tokens long,
     # we split it such that each sample has at most 512 tokens
     # and the ctx_ids and ctx_attn_mask are the same for all samples
+    if max_qas_len < 0:
+        return samples
     input_ids = samples["input_ids"]
     attention_mask = samples["attention_mask"]
     labels = samples["labels"]

src/ctx_to_lora/eval_utils.py

@@ -736,6 +736,7 @@ def evaluate(
 
     _get_tokenized_dataset = partial(
         get_tokenized_dataset,
+        max_qas_len=-1,
         base_model_max_len=model.base_model.config.max_position_embeddings,
         tokenizer=tokenizer,
         tokenizer_kwargs=tokenizer_kwargs,

src/ctx_to_lora/modeling/aggregator.py

@@ -2,10 +2,10 @@ import logging
 from dataclasses import dataclass
 from enum import Enum
 
-from einops import rearrange, unpack, repeat
+import torch
+from einops import rearrange, repeat, unpack
 from jaxtyping import Float, Integer
 from torch import Tensor, nn
-import torch
 from transformers import (
     PretrainedConfig,
     PreTrainedModel,
@@ -41,17 +41,18 @@ 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
+    # 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
+    lora_r: int
+    per_rank_gen: bool
 
-    layer_to_layer_ctx_encoder: bool = False
-    n_base_queries: int = 208
+    layer_to_layer_ctx_encoder: bool  #  = False
+    n_latent_queries: int  # = 208
+    num_blocks: int
+    num_self_attn_per_block: int
+    shared_weights: bool
 
 
 def get_aggregator_config(
@@ -92,8 +93,7 @@ class Perceiver(nn.Module):
         lora_r,
         num_latent_factor,  # unused
         layer_to_layer_ctx_encoder,
-        n_base_queries,
-        concat_latents_context,
+        n_latent_queries,
         *args,
         **kwargs,
     ):
@@ -109,22 +109,26 @@ class Perceiver(nn.Module):
         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,
-            n_cross_attn_layers=kwargs["n_cross_attn_layers"],
+            # resampler_depth=kwargs["num_self_attends_per_block"] + 1,
+            # n_cross_attn_layers=kwargs["n_cross_attn_layers"],
             # resampler_n_latents=n_output_queries * num_latent_factor,
-            resampler_n_latents=n_base_queries,
+            num_blocks=kwargs["num_blocks"],
+            num_self_attn_per_block=kwargs["num_self_attn_per_block"],
+            shared_weights=kwargs["shared_weights"],
+            n_latents=n_latent_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,
+            # resampler_depth=kwargs["decoder_depth"],
+            num_blocks=1,
+            num_self_attn_per_block=0,
+            shared_weights=False,
+            n_latents=n_output_queries,
             intermediate_size_factor=4,
             hidden_size=output_size,
             attn_implementation="flash_attention_2",
@@ -164,17 +168,17 @@ class Perceiver(nn.Module):
         if self.layer_to_layer:
             ctx_features = rearrange(
                 ctx_features,
-                "bs num_layers seq_len feature_dim -> bs (num_layers seq_len) feature_dim",
+                "1 num_layers seq_len feature_dim -> 1 (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",
+                    "bs seq_len -> (num_layers bs) seq_len",
                     num_layers=self.num_layers,
                 )
             ctx_position_ids = repeat(
                 ctx_position_ids,
-                "bs seq_len -> bs (num_layers seq_len)",
+                "1 seq_len -> 1 (num_layers seq_len)",
                 num_layers=self.num_layers,
             )
 
@@ -191,7 +195,7 @@ class Perceiver(nn.Module):
             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 bs) (per_layer_sz) d -> bs (num_layers per_layer_sz) d"),
                 num_layers=self.num_layers,
                 per_layer_sz=per_layer_size,
             )
@@ -338,7 +342,7 @@ if __name__ == "__main__":
         lora_r=lora_r,
         num_latent_factor=8,
         layer_to_layer_ctx_encoder=False,
-        n_base_queries=208,
+        n_latent_queries=208,
         num_self_attends_per_block=16,
         decoder_depth=1,
     ).to(device)
@@ -373,7 +377,7 @@ if __name__ == "__main__":
         lora_r=lora_r,
         num_latent_factor=8,
         layer_to_layer_ctx_encoder=True,
-        n_base_queries=208,
+        n_latent_queries=208,
         num_self_attends_per_block=16,
         decoder_depth=1,
     ).to(device)

src/ctx_to_lora/modeling/ctx_encoder.py

@@ -117,3 +117,50 @@ CTX_ENCODER_CLS = {
     CTX_ENCODER_TYPE.EMBED_ONLY: EmbeddingOnly,
     CTX_ENCODER_TYPE.PER_LAYER_ACTIVATIONS: PerLayerActivations,
 }
+
+if __name__ == "__main__":
+    from transformers import AutoModel, AutoTokenizer
+
+    from ctx_to_lora.configs import CtxEncoderArguments
+
+    # Test configuration
+    model_name = "google/gemma-2-2b-it"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.set_default_device(device)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    base_model = AutoModel.from_pretrained(model_name, device_map=device)
+
+    # Test input
+    test_text = ["Hello world", "this is a test.", "qwewq" * 50]
+    inputs = tokenizer(test_text)
+    position_ids = torch.cat(
+        [torch.arange(len(x)) for x in inputs["input_ids"]]
+    ).unsqueeze(0)
+    input_ids = torch.cat([torch.tensor(x) for x in inputs["input_ids"]]).unsqueeze(0)
+    inputs = {"input_ids": input_ids, "position_ids": position_ids}
+    print(inputs)
+
+    print(f"Testing all encoder classes with model: {model_name}")
+    print(f"Input shape: {inputs['input_ids'].shape}")
+    print("-" * 50)
+
+    # Test EarlyExit
+    # print("Testing EarlyExit...")
+    config = CtxEncoderArguments(layer_idx=2)
+    # early_exit_encoder = EarlyExit(base_model, config)
+    # early_exit_output = early_exit_encoder(**inputs)
+    # print(f"EarlyExit output shape: {early_exit_output.shape}")
+
+    # # Test EmbeddingOnly
+    # print("\nTesting EmbeddingOnly...")
+    # embed_only_encoder = EmbeddingOnly(base_model, config)
+    # embed_only_output = embed_only_encoder(**inputs)
+    # print(f"EmbeddingOnly output shape: {embed_only_output.shape}")
+
+    # Test PerLayerActivations
+    print("\nTesting PerLayerActivations...")
+    per_layer_encoder = PerLayerActivations(base_model, config)
+    per_layer_output = per_layer_encoder(**inputs)
+    print(f"PerLayerActivations output shape: {per_layer_output.shape}")
+
+    print("\nAll tests completed successfully!")

src/ctx_to_lora/modeling/idefics2.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 """PyTorch Idefics2 model."""
 
-from copy import deepcopy
 import math
 
 import torch
@@ -50,13 +49,13 @@ class Idefics2PerceiverConfig(PretrainedConfig):
             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):
+        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):
+        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):
+        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.
@@ -69,41 +68,39 @@ class Idefics2PerceiverConfig(PretrainedConfig):
     def __init__(
         self,
         input_size: int,
-        intermediate_size_factor: int = 1,
+        num_blocks: int,
+        num_self_attn_per_block: int,
+        shared_weights: bool,
+        intermediate_size_factor: int,
         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,
+        n_latents=64,
+        n_heads=16,
+        head_dim=128,
         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.num_blocks = num_blocks
+        self.num_self_attn_per_block = num_self_attn_per_block
+        self.shared_weights = shared_weights
+
         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.n_latents = n_latents
+        self.n_heads = n_heads
         self.num_key_value_heads = num_key_value_heads
-        self.resampler_head_dim = resampler_head_dim
+        self.head_dim = head_dim
         self.attention_dropout = attention_dropout
-        if self.num_key_value_heads > self.resampler_n_heads:
+        if self.num_key_value_heads > self.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}"
+                f" n_heads={self.n_heads}"
             )
         super().__init__(**kwargs)
 
@@ -225,8 +222,8 @@ class Idefics2PerceiverAttention(nn.Module):
         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_heads = config.n_heads
+        self.head_dim = config.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
@@ -361,7 +358,6 @@ class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention):
         self,
         latents: 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,
@@ -376,37 +372,6 @@ class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention):
         query_states = self.q_proj(latents)
         if is_cross_attn:
             kv_inp = context
-            if concat_latents_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
 
@@ -526,14 +491,12 @@ IDEFICS2_PERCEIVER_ATTENTION_CLASSES = {
 
 
 class Idefics2PerceiverLayer(nn.Module):
-    def __init__(self, config, layer_idx: int):
+    def __init__(self, config, is_cross_attn: bool):
         super().__init__()
         self.hidden_size = config.hidden_size
-        self.n_latents = config.resampler_n_latents
-        self.depth = config.resampler_depth
+        self.n_latents = config.n_latents
         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.is_cross_attn = is_cross_attn
 
         self.input_latents_norm = Idefics2RMSNorm(
             self.hidden_size, eps=self.rms_norm_eps
@@ -545,7 +508,7 @@ class Idefics2PerceiverLayer(nn.Module):
         )
         self.self_attn = IDEFICS2_PERCEIVER_ATTENTION_CLASSES[
             config._attn_implementation
-        ](config, layer_idx=layer_idx)
+        ](config)
         self.post_attention_layernorm = Idefics2RMSNorm(
             self.hidden_size, eps=self.rms_norm_eps
         )
@@ -593,7 +556,6 @@ 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,
@@ -646,37 +608,58 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
 
     def __init__(self, config) -> None:
         super().__init__(config)
+        self.num_blocks = config.num_blocks
+        self.num_self_attn_per_block = config.num_self_attn_per_block
+        self.shared_weights = config.shared_weights
         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.n_latents = config.n_latents
         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
-        if not self.concat_latents_context:
-            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)
-            ]
+        # First block
+        assert config.num_blocks > 0
+        first_x_attn = [Idefics2PerceiverLayer(config, is_cross_attn=True)]
+        first_self_attn_block = [
+            Idefics2PerceiverLayer(config, is_cross_attn=False)
+            for _ in range(config.num_self_attn_per_block)
+        ]
+
+        self.layers = nn.ModuleList(first_x_attn + first_self_attn_block)
+        for _ in range(1, config.num_blocks):
+            # cross-attention at the beginning of each block
+            if self.shared_weights:
+                x_attn = first_x_attn[0]
+            else:
+                x_attn = Idefics2PerceiverLayer(config, is_cross_attn=True)
+            self.layers.append(x_attn)
+
+            # self-attention
+            for i in range(config.num_self_attn_per_block):
+                if self.shared_weights:
+                    self_attn = first_self_attn_block[i]
+                else:
+                    self_attn = Idefics2PerceiverLayer(config, is_cross_attn=False)
+                self.layers.append(self_attn)
+
+        # self.layers = nn.ModuleList(self.layers)
+
+        # 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)
 
-            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"
@@ -751,117 +734,95 @@ 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")
 
-        if self.concat_latents_context:
-            old_cu_seq_lens_k = cu_seq_lens_k
-            cu_seq_lens_k = cu_seq_lens_k + cu_seq_lens_q
-            max_length_k = max_length_q + max_length_k
-            # print("Using concat latents + context 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}")
-            for layer in self.layers:
-                layer_outputs = 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,
-                    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,
-                )
+        x_attn_kwargs = dict(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            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,
+        )
+        self_attn_mask = (
+            torch.ones(
+                (bsz, self.n_latents), device=context.device, dtype=torch.float32
+            )
+            if attention_mask is not None
+            else None
+        )
+        self_attn_position_ids = (
+            torch.arange(
+                self.n_latents, device=position_ids.device, dtype=torch.int32
+            ).repeat(1, bsz)
+            if position_ids is not None
+            else None
+        )
+        self_attn_kwargs = dict(
+            attention_mask=self_attn_mask,
+            position_ids=self_attn_position_ids,
+            cu_seq_lens_q=cu_seq_lens_q,
+            cu_seq_lens_k=cu_seq_lens_q,
+            max_length_q=max_length_q,
+            max_length_k=max_length_q,
+        )
+        for i, layer in enumerate(self.layers):
+            inp_kwargs = dict(
+                latents=compressed_context,
+                context=context,
+                past_key_value=None,
+                output_attentions=False,
+                use_cache=False,
+            )
+            if layer.is_cross_attn:
+                attn_kwargs = {**inp_kwargs, **x_attn_kwargs}
+            else:
+                attn_kwargs = {**inp_kwargs, **self_attn_kwargs}
 
-                compressed_context = layer_outputs[0]
+            layer_outputs = layer(**attn_kwargs)
+            compressed_context = layer_outputs[0]
+        # for encoder in self.layers[: self.n_cross_attn_layers]:
+        #     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,
+        #     )
 
-            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
+        #     compressed_context = layer_outputs[0]
 
-            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,
-                )
+        # 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
 
-                compressed_context = layer_outputs[0]
+        # for processor in self.layers[self.n_cross_attn_layers :]:
+        #     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)
 
@@ -930,23 +891,25 @@ __all__ = [
 
 if __name__ == "__main__":
     encoder_config = Idefics2PerceiverConfig(
-        concat_latents_context=True,
         input_size=1234,
-        resampler_n_latents=64,
-        resampler_depth=1 + 8,
+        n_latents=64,
+        num_blocks=8,
+        num_self_attn_per_block=6,
+        shared_weights=False,
         attn_implementation="flash_attention_2",
     )
     decoder_config = Idefics2PerceiverConfig(
-        concat_latents_context=True,
         input_size=encoder_config.hidden_size,
-        resampler_n_latents=16,
-        resampler_depth=1,
+        n_latents=16,
+        num_blocks=1,
+        num_self_attn_per_block=0,
+        shared_weights=False,
         attn_implementation="flash_attention_2",
     )
     model = (
         Idefics2Perceiver(encoder_config, decoder_config).to("cuda").to(torch.bfloat16)
     )
-    # print(model)
+    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")
@@ -975,7 +938,6 @@ if __name__ == "__main__":
     out.mean().backward()
     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)