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)