From 9e12010fc584a16346d485a9aeaf02904ff0a2e1 Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 1 Jul 2025 15:07:06 +0000 Subject: [PATCH] self-attn and concat_latents_context working --- intx_sft.py | 6 +- src/ctx_to_lora/configs.py | 16 +- src/ctx_to_lora/modeling/aggregator.py | 18 +- src/ctx_to_lora/modeling/ctx_encoder.py | 2 +- src/ctx_to_lora/modeling/idefics2.py | 382 +++++++++++++++++------- 5 files changed, 299 insertions(+), 125 deletions(-) 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 11a867c..f9d54a2 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -476,10 +476,18 @@ class AggregatorArguments: # default=0.0, # metadata={"help": "Attention dropout probability for Perceiver."}, # ) - # num_latent_factor: int = field( - # default=8, - # metadata={"help": "Number of latent factors 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."}, diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py index 6984307..0cb744f 100644 --- a/src/ctx_to_lora/modeling/aggregator.py +++ b/src/ctx_to_lora/modeling/aggregator.py @@ -41,6 +41,7 @@ class AggregatorConfig: pooling_type: POOL_FN # perceiver + concat_latents_context: bool num_self_attends_per_block: int # = 16 decoder_depth: int # 1 = only cross-attention num_latent_factor: int = 8 @@ -94,6 +95,7 @@ class Perceiver(nn.Module): num_latent_factor, # unused layer_to_layer_ctx_encoder, n_base_queries, + concat_latents_context, *args, **kwargs, ): @@ -109,6 +111,7 @@ 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, @@ -119,6 +122,7 @@ class Perceiver(nn.Module): 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, @@ -175,13 +179,13 @@ class Perceiver(nn.Module): 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}" - ) + # 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: diff --git a/src/ctx_to_lora/modeling/ctx_encoder.py b/src/ctx_to_lora/modeling/ctx_encoder.py index 65f3965..3b037ac 100644 --- a/src/ctx_to_lora/modeling/ctx_encoder.py +++ b/src/ctx_to_lora/modeling/ctx_encoder.py @@ -53,7 +53,7 @@ def maybe_add_batch_dim(kwargs): class EarlyExit(nn.Module): def __init__(self, base_model: PreTrainedModel, config: CtxEncoderArguments): super().__init__() - base_model = get_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[: self.exit_layer] diff --git a/src/ctx_to_lora/modeling/idefics2.py b/src/ctx_to_lora/modeling/idefics2.py index b7635eb..0a53f8c 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,13 @@ class Idefics2PerceiverConfig(PretrainedConfig): resampler_head_dim=96, num_key_value_heads=4, attention_dropout=0.0, + is_cross_attn=True, + concat_latents_context=False, **kwargs, ): # for mlp + self.concat_latents_context = concat_latents_context + self.is_cross_attn = is_cross_attn self.input_size = input_size self.intermediate_size_factor = intermediate_size_factor # for perceiver @@ -157,9 +162,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 +358,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 +370,60 @@ 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"] + old_cur_len = 0 + cur_len = 0 + n_latents = latents.shape[1] + kv_inp = torch.empty( + 1, + cu_seq_lens_k[-1], + latents.shape[2], + dtype=context.dtype, + device=context.device, + ) + for i, (old_cu_len, cu_len) in enumerate( + zip(old_cu_seq_lens_k[1:], cu_seq_lens_k[1:]) + ): + ctx_len = old_cu_len - old_cur_len + old_end_idx = cur_len + ctx_len + kv_inp[0, cur_len:old_end_idx] = context[ + 0, old_cur_len:old_cu_len + ] + kv_inp[0, old_end_idx : old_end_idx + n_latents] = latents[i] + + cur_len = cu_len + old_cur_len = old_cu_len + 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 +521,7 @@ class Idefics2PerceiverFlashAttention2(Idefics2PerceiverAttention): IDEFICS2_PERCEIVER_ATTENTION_CLASSES = { - "eager": Idefics2PerceiverAttention, + # "eager": Idefics2PerceiverAttention, "flash_attention_2": Idefics2PerceiverFlashAttention2, } @@ -491,12 +533,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 +550,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 +564,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 +593,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 +650,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_x_attn_layers = 1 self.depth = config.resampler_depth + self.n_self_attn_layers = self.depth - self.n_x_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_x_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,38 +752,116 @@ 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 * (i + 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: - print(f"compressed_context.shape= {compressed_context.shape}") - print(f"context.shape= {context.shape}") - print( - f"position_ids.shape= {position_ids.shape if position_ids is not None else None}" - ) - print( - f"cu_seq_lens_q.shape= {cu_seq_lens_q.shape if cu_seq_lens_q is not None else None}" - ) - print( - f"cu_seq_lens_k.shape= {cu_seq_lens_k.shape if cu_seq_lens_k is not None else None}" - ) - print(f"max_length_q= {max_length_q}") - print(f"max_length_k= {max_length_k}") - 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_x_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] + + position_ids = torch.arange( + self.n_latents, device=position_ids.device, dtype=torch.int32 + ).repeat(1, bsz) + + for processor in self.layers[self.n_x_attn_layers :]: + cu_seq_lens_k = cu_seq_lens_q + max_length_k = max_length_q + # 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) @@ -719,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( @@ -775,78 +929,45 @@ __all__ = [ ] if __name__ == "__main__": - config = Idefics2PerceiverConfig( + encoder_config = Idefics2PerceiverConfig( + concat_latents_context=True, input_size=1234, resampler_n_latents=64, - resampler_depth=8, + resampler_depth=1 + 8, attn_implementation="flash_attention_2", ) decoder_config = Idefics2PerceiverConfig( - input_size=config.hidden_size, + concat_latents_context=True, + input_size=encoder_config.hidden_size, resampler_n_latents=16, resampler_depth=1, attn_implementation="flash_attention_2", ) - model = Idefics2Perceiver(config, decoder_config).to("cuda").to(torch.bfloat16) + model = ( + 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") - # # mask[..., -1] = 0 - # # print(mask) - # out = model( - # inp, - # # attention_mask=mask, - # position_ids=torch.cat( - # [ - # torch.arange(128, device="cuda"), - # torch.arange(64, device="cuda"), - # torch.arange(37, device="cuda"), - # torch.arange(128, device="cuda"), - # torch.arange(64, device="cuda"), - # torch.arange(37, device="cuda"), - # ], - # dim=-1, - # ).unsqueeze(0), - # # position_ids=torch.arange(128, device="cuda"), - # ) - # print(out) - # print(out.shape) - # print(model.encoder.latents.grad) - # print(model.decoder.latents.grad) - # out.mean().backward() - # print(model.encoder.latents.grad) - # print(model.decoder.latents.grad) - # breakpoint() - 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]) - ] + out = model( + inp, + # attention_mask=mask, + position_ids=torch.cat( + [ + torch.arange(128, device="cuda"), + torch.arange(64, device="cuda"), + torch.arange(37, device="cuda"), + torch.arange(128, device="cuda"), + torch.arange(64, device="cuda"), + torch.arange(37, device="cuda"), + ], + dim=-1, + ).unsqueeze(0), + # position_ids=torch.arange(128, device="cuda"), ) - print(out) print(out.shape) print(model.encoder.latents.grad) @@ -855,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()