diff --git a/intx_sft.py b/intx_sft.py index 1a3ee0b..a4575f1 100755 --- a/intx_sft.py +++ b/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 = ( diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index ad62730..f3765b7 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -442,7 +442,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={ @@ -474,17 +483,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, diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py index d33982f..de51670 100644 --- a/src/ctx_to_lora/modeling/aggregator.py +++ b/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,9 +41,12 @@ 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 @@ -53,6 +57,7 @@ class AggregatorConfig: 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, @@ -68,6 +73,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), ) @@ -84,7 +91,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, ): @@ -95,16 +105,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, @@ -116,11 +134,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, [ @@ -231,7 +306,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!") diff --git a/src/ctx_to_lora/modeling/ctx_encoder.py b/src/ctx_to_lora/modeling/ctx_encoder.py index a946664..3b037ac 100644 --- a/src/ctx_to_lora/modeling/ctx_encoder.py +++ b/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, +} diff --git a/src/ctx_to_lora/modeling/hypernet.py b/src/ctx_to_lora/modeling/hypernet.py index ce06953..ab1aadc 100644 --- a/src/ctx_to_lora/modeling/hypernet.py +++ b/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.ctx_encoder import EarlyExit +from ctx_to_lora.modeling.aggregator import ( + 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( - get_base_model(encoder_model), self.ctx_encoder_args.layer_idx + self.ctx_encoder = CTX_ENCODER_CLS[self.ctx_encoder_args.ctx_encoder_type]( + encoder_model, self.ctx_encoder_args ) # delegate to base_model diff --git a/src/ctx_to_lora/modeling/idefics2.py b/src/ctx_to_lora/modeling/idefics2.py index 0f90430..f1d013c 100644 --- a/src/ctx_to_lora/modeling/idefics2.py +++ b/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] - - # 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` + # kv_seq_len = context.size()[1] query_states = self.q_proj(latents) - # key_states = self.k_proj(torch.cat([context, latents], dim=-2)) - # value_states = self.v_proj(torch.cat([context, latents], dim=-2)) - key_states = self.k_proj(context) - value_states = self.v_proj(context) + 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 + + 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.hidden_size, eps=self.rms_norm_eps + self.input_context_norm = ( + 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( - [Idefics2PerceiverLayer(config, idx) for idx in range(self.depth)] - ) + 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) + ] + + 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: - layer_outputs = perceiver_layer( - compressed_context, - context, - attention_mask=None, - position_ids=position_ids, - past_key_value=None, - output_attentions=False, - use_cache=False, - cu_seq_lens_q=cu_seq_lens_q, - cu_seq_lens_k=cu_seq_lens_k, - max_length_q=max_length_q, - max_length_k=max_length_k, - ) + 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, + ) - 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, - intermediate_size=config.intermediate_size_factor * config.input_size, - output_size=config.hidden_size, - hidden_act=config.hidden_act, + hidden_size=encoder_config.input_size, + intermediate_size=encoder_config.intermediate_size_factor + * encoder_config.input_size, + 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) - print(model) + 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()