mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
self-attn and concat_latents_context working
This commit is contained in:
51616
parent
6ad1dadd3b
commit
9e12010fc5
5 changed files with 299 additions and 125 deletions
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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."},
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue