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

l2l prototype (not complete)

Browse source
This commit is contained in:
51616 2025-07-01 14:49:59 +09:00
parent 92ba1ca76c
commit bf88de287d
4 changed files with 247 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

21
src/ctx_to_lora/configs.py
View file

@ -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,

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

@ -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!")

62
src/ctx_to_lora/modeling/ctx_encoder.py
View file

@ -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,
}

21
src/ctx_to_lora/modeling/hypernet.py
View file

@ -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