apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-26 17:11:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

n_cross_attn_layers args

Browse source
This commit is contained in:
51616 2025-07-01 15:19:10 +00:00
parent 9e12010fc5
commit a5142b8fcf
3 changed files with 16 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

8
src/ctx_to_lora/configs.py
View file

@ -496,9 +496,15 @@ class AggregatorArguments:
# default=8,
# metadata={"help": "Number of blocks for Perceiver."},
# )
# misnomer
num_self_attends_per_block: int = field(
default=6,
metadata={"help": "Number of self-attends per block for Perceiver."},
metadata={"help": "Number of attn blocks for Perceiver."},
)
n_cross_attn_layers: int = field(
default=6,
metadata={"help": "Number of cross-attention layers for Perceiver. "},
)
# self_attention_widening_factor: int = field(
# default=4,

2
src/ctx_to_lora/modeling/aggregator.py
View file

@ -42,6 +42,7 @@ class AggregatorConfig:
# perceiver
concat_latents_context: bool
n_cross_attn_layers: int # number of cross-attention layers
num_self_attends_per_block: int # = 16
decoder_depth: int # 1 = only cross-attention
num_latent_factor: int = 8
@ -115,6 +116,7 @@ class Perceiver(nn.Module):
input_size=feature_size,
# the first layer is xattn
resampler_depth=kwargs["num_self_attends_per_block"] + 1,
n_cross_attn_layers=kwargs["n_cross_attn_layers"],
# resampler_n_latents=n_output_queries * num_latent_factor,
resampler_n_latents=n_base_queries,
intermediate_size_factor=4,

12
src/ctx_to_lora/modeling/idefics2.py
View file

@ -80,12 +80,14 @@ class Idefics2PerceiverConfig(PretrainedConfig):
num_key_value_heads=4,
attention_dropout=0.0,
is_cross_attn=True,
n_cross_attn_layers=1,
concat_latents_context=False,
**kwargs,
):
# for mlp
self.concat_latents_context = concat_latents_context
self.is_cross_attn = is_cross_attn
self.n_cross_attn_layers = n_cross_attn_layers
self.input_size = input_size
self.intermediate_size_factor = intermediate_size_factor
# for perceiver
@ -650,9 +652,9 @@ 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.n_cross_attn_layers = config.n_cross_attn_layers
self.depth = config.resampler_depth
self.n_self_attn_layers = self.depth - self.n_x_attn_layers
self.n_self_attn_layers = self.depth - self.n_cross_attn_layers
self.rms_norm_eps = config.rms_norm_eps
self.concat_latents_context = config.concat_latents_context
@ -665,7 +667,7 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
x_attn_config.is_cross_attn = True
self.layers = [
Idefics2PerceiverLayer(x_attn_config, idx)
for idx in range(self.n_x_attn_layers)
for idx in range(self.n_cross_attn_layers)
]
config.is_cross_attn = False
@ -811,7 +813,7 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
)
compressed_context = layer_outputs[0]
else:
for encoder in self.layers[: self.n_x_attn_layers]:
for encoder in self.layers[: self.n_cross_attn_layers]:
# print("Using cross-attention for resampler")
# print(f"position_ids: {position_ids}")
# print(f"cu_seq_lens_q: {cu_seq_lens_q}")
@ -838,7 +840,7 @@ class Idefics2PerceiverResampler(Idefics2PreTrainedModel):
self.n_latents, device=position_ids.device, dtype=torch.int32
).repeat(1, bsz)
for processor in self.layers[self.n_x_attn_layers :]:
for processor in self.layers[self.n_cross_attn_layers :]:
cu_seq_lens_k = cu_seq_lens_q
max_length_k = max_length_q
# print(f"Using self-attention for resampler")