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