From 9613c65c20647e3014b309b5eda98cff4ff066b0 Mon Sep 17 00:00:00 2001 From: Rujikorn Charakorn Date: Mon, 7 Jul 2025 17:31:07 +0900 Subject: [PATCH] Perceiver blocks + layer-to-layer (#4) --- .../gemma_qa_short_ctx_exp_default.sh | 1 + .../gemma_qa_short_ctx_exp_self_attn.sh | 1 + ...ttn.sh => gemma_qa_short_ctx_perceiver.sh} | 8 +- .../gemma_qa_short_ctx_perceiver_l2l.sh | 41 ++ src/ctx_to_lora/configs.py | 39 +- src/ctx_to_lora/data/processing.py | 5 +- src/ctx_to_lora/eval_utils.py | 1 + src/ctx_to_lora/modeling/aggregator.py | 58 +-- src/ctx_to_lora/modeling/ctx_encoder.py | 47 +++ src/ctx_to_lora/modeling/idefics2.py | 356 ++++++++---------- 10 files changed, 303 insertions(+), 254 deletions(-) rename scripts/short_ctx/{gemma_qa_short_ctx_exp_concat_latent_context_attn.sh => gemma_qa_short_ctx_perceiver.sh} (89%) create mode 100644 scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh diff --git a/scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh b/scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh index b78f1ee..8acb5ab 100644 --- a/scripts/short_ctx/gemma_qa_short_ctx_exp_default.sh +++ b/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 \ diff --git a/scripts/short_ctx/gemma_qa_short_ctx_exp_self_attn.sh b/scripts/short_ctx/gemma_qa_short_ctx_exp_self_attn.sh index 64b46b0..8c60d95 100644 --- a/scripts/short_ctx/gemma_qa_short_ctx_exp_self_attn.sh +++ b/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 \ diff --git a/scripts/short_ctx/gemma_qa_short_ctx_exp_concat_latent_context_attn.sh b/scripts/short_ctx/gemma_qa_short_ctx_perceiver.sh similarity index 89% rename from scripts/short_ctx/gemma_qa_short_ctx_exp_concat_latent_context_attn.sh rename to scripts/short_ctx/gemma_qa_short_ctx_perceiver.sh index 59d607c..c674f73 100644 --- a/scripts/short_ctx/gemma_qa_short_ctx_exp_concat_latent_context_attn.sh +++ b/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 \ diff --git a/scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh b/scripts/short_ctx/gemma_qa_short_ctx_perceiver_l2l.sh new file mode 100644 index 0000000..6df5dd4 --- /dev/null +++ b/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 diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index f3765b7..d656aae 100644 --- a/src/ctx_to_lora/configs.py +++ b/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 diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 9f22f35..b0bbd9e 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/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"] diff --git a/src/ctx_to_lora/eval_utils.py b/src/ctx_to_lora/eval_utils.py index 9214ce6..47203fa 100644 --- a/src/ctx_to_lora/eval_utils.py +++ b/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, diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py index c116f3a..cf05808 100644 --- a/src/ctx_to_lora/modeling/aggregator.py +++ b/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) diff --git a/src/ctx_to_lora/modeling/ctx_encoder.py b/src/ctx_to_lora/modeling/ctx_encoder.py index 3b037ac..fdb2c28 100644 --- a/src/ctx_to_lora/modeling/ctx_encoder.py +++ b/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!") diff --git a/src/ctx_to_lora/modeling/idefics2.py b/src/ctx_to_lora/modeling/idefics2.py index f1d013c..b07a508 100644 --- a/src/ctx_to_lora/modeling/idefics2.py +++ b/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)