From 7365b84225960292907c177e2baafdee623b2fb9 Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 1 Jul 2025 14:49:59 +0900 Subject: [PATCH 1/2] l2l prototype (not complete) --- src/ctx_to_lora/configs.py | 21 ++- src/ctx_to_lora/modeling/aggregator.py | 163 +++++++++++++++++++++++- src/ctx_to_lora/modeling/ctx_encoder.py | 62 ++++++++- src/ctx_to_lora/modeling/hypernet.py | 21 ++- 4 files changed, 247 insertions(+), 20 deletions(-) diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index afe8a06..11a867c 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -435,7 +435,16 @@ class CtxEncoderArguments: default=None, metadata={"help": "Context encoder model name or path."}, ) - # TODO: allow using activations from multiple layers? + ctx_encoder_type: Literal[ + "embedding_only", "per_layer_activations", "early_exit" + ] = field( + default="early_exit", + metadata={ + "help": "Context encoder type. " + "Options: 'embedding_only', 'per_layer_activations', 'early_exit'." + }, + ) + # used only with `early_exit` type layer_idx: int | None = field( default=None, metadata={ @@ -467,9 +476,13 @@ 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."}, + # 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, diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py index 118b40c..63aa66e 100644 --- a/src/ctx_to_lora/modeling/aggregator.py +++ b/src/ctx_to_lora/modeling/aggregator.py @@ -2,7 +2,8 @@ import logging from dataclasses import dataclass from enum import Enum -from einops import rearrange, unpack +import torch +from einops import rearrange, repeat, unpack from jaxtyping import Float, Integer from torch import Tensor, nn from transformers import ( @@ -43,13 +44,20 @@ class AggregatorConfig: num_self_attends_per_block: int # = 16 decoder_depth: int # 1 = only cross-attention num_latent_factor: int = 8 + n_base_queries: int = 208 lora_r: int = 8 per_rank_gen: bool = False + # layer to layer generation + # maps activations of layer L to LoRA of that layer + layer_to_layer_ctx_encoder: bool = False + n_ctx_model_layers: int = -1 + def get_aggregator_config( model: PreTrainedModel, ctx_encoder_model_config: PretrainedConfig, + layer_to_layer_ctx_encoder: bool, output_size: int, num_modules: int, num_extra_modules: int, @@ -65,6 +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), ) @@ -81,7 +91,9 @@ 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, *args, **kwargs, ): @@ -92,11 +104,15 @@ 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.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( 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, + # 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", @@ -110,14 +126,73 @@ class Perceiver(nn.Module): attn_implementation="flash_attention_2", ) self.perceiver = Idefics2Perceiver(self.config, self.decoder_config) + if self.layer_to_layer: + # doesn't work + self.perceiver.forward = torch.vmap( + self.perceiver.forward, + in_dims=(1, None, None), + out_dims=1, + ) 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)", + ) + ctx_position_ids = repeat( + ctx_position_ids, + "bs seq_len -> bs (num_layers seq_len)", + ) + x = self.perceiver(ctx_features, ctx_attn_mask, ctx_position_ids) + + if self.layer_to_layer: + x = rearrange( + x, + ( + "(bs num_layers) (n_modules r) seq_len d -> " + "bs (num_layers n_modules r) d" + ), + bs=1, + num_layers=self.num_layers, + n_modules=self.num_modules, + r=self.r, + ) lora_x, extra_x = unpack( x, [ @@ -228,7 +303,85 @@ 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 + + # 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, + ) + + # Test with 3D input (standard case) + print("Testing with 3D input...") + ctx_features = torch.randn(batch_size, seq_len, feature_size) + ctx_attn_mask = torch.ones(batch_size, seq_len, dtype=torch.long) + ctx_position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) + + 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, + ) + + # Test with 4D input for layer-to-layer + ctx_features_4d = torch.randn(batch_size, num_layers, seq_len, feature_size) + + 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..65f3965 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 + 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 From fe7863c41b18bc9d7e175c773abf14922300404f Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 1 Jul 2025 05:57:16 +0000 Subject: [PATCH 2/2] idefics2 attn bug found --- src/ctx_to_lora/modeling/aggregator.py | 50 ++++++++-------- src/ctx_to_lora/modeling/idefics2.py | 82 +++++++++++++++++++++----- 2 files changed, 93 insertions(+), 39 deletions(-) diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py index 63aa66e..6984307 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 -import torch -from einops import rearrange, repeat, unpack +from einops import rearrange, unpack, repeat from jaxtyping import Float, Integer from torch import Tensor, nn +import torch from transformers import ( PretrainedConfig, PreTrainedModel, @@ -104,6 +104,7 @@ 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 @@ -126,13 +127,6 @@ class Perceiver(nn.Module): attn_implementation="flash_attention_2", ) self.perceiver = Idefics2Perceiver(self.config, self.decoder_config) - if self.layer_to_layer: - # doesn't work - self.perceiver.forward = torch.vmap( - self.perceiver.forward, - in_dims=(1, None, None), - out_dims=1, - ) def forward( self, @@ -172,26 +166,31 @@ class Perceiver(nn.Module): 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 -> (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) (n_modules r) seq_len d -> " - "bs (num_layers n_modules r) d" - ), - bs=1, + ("(bs num_layers) (per_layer_sz) d -> bs (num_layers per_layer_sz) d"), num_layers=self.num_layers, - n_modules=self.num_modules, - r=self.r, + per_layer_sz=per_layer_size, ) lora_x, extra_x = unpack( x, @@ -321,6 +320,9 @@ if __name__ == "__main__": lora_r = 8 batch_size = 2 seq_len = 100 + device = "cuda" + + torch.set_default_dtype(torch.bfloat16) # Create Perceiver instance perceiver = Perceiver( @@ -336,13 +338,13 @@ if __name__ == "__main__": 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(batch_size, seq_len, feature_size) - ctx_attn_mask = torch.ones(batch_size, seq_len, dtype=torch.long) - ctx_position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1) + 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) @@ -371,10 +373,12 @@ if __name__ == "__main__": 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(batch_size, num_layers, seq_len, feature_size) + 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( diff --git a/src/ctx_to_lora/modeling/idefics2.py b/src/ctx_to_lora/modeling/idefics2.py index 0f90430..b7635eb 100644 --- a/src/ctx_to_lora/modeling/idefics2.py +++ b/src/ctx_to_lora/modeling/idefics2.py @@ -683,6 +683,19 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel): 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, @@ -775,28 +788,65 @@ if __name__ == "__main__": attn_implementation="flash_attention_2", ) model = Idefics2Perceiver(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") + # # 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 = 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"), + 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)