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

per module head

Browse source
This commit is contained in:
51616 2024-12-27 14:17:12 +00:00
parent 786bb44499
commit ecab2d1d40
1 changed files with 25 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

37
hyperlora/modeling_utils.py
View file

@ -90,7 +90,7 @@ class HypernetConfig:
def get_hypernet_config( def get_hypernet_config(
model: PreTrainedModel, model: PreTrainedModel,
latent_size: int = 256, latent_size: int = 512,
aggregator_type: AGGREGATOR_TYPE = AGGREGATOR_TYPE.POOLER, aggregator_type: AGGREGATOR_TYPE = AGGREGATOR_TYPE.POOLER,
): ):
lora_config = model.peft_config["default"] lora_config = model.peft_config["default"]
@ -107,6 +107,9 @@ def get_hypernet_config(
class Perceiver(nn.Module): class Perceiver(nn.Module):
# TODO: we could expand the latent query size to be much bigger
# and then do cross-attn with an output query with size [n_layers x n_modules]
# https://huggingface.co/docs/transformers/en/model_doc/perceiver#transformers.models.perceiver.modeling_perceiver.PerceiverBasicDecoder
"""perceiver w/ bottleneck size = n_modules * n_layers""" """perceiver w/ bottleneck size = n_modules * n_layers"""
def __init__( def __init__(
@ -356,25 +359,35 @@ class HyperLoRA(nn.Module):
# TODO: check initialization of the head # TODO: check initialization of the head
# default values are prob. way too big # default values are prob. way too big
#
# we can separate the modules while vectorizing by
# making modules have the same output size then slicing
# the output to the correct size
# TODO: could be even more efficient if we use lightweight LoRA
# ie. project the input to a smaller subspace (w/ same size) for all modules
self.head = Mix( self.head = Mix(
"bs n_layers n_modules d -> bs n_layers r out_d", "bs n_layers n_modules d -> bs n_layers n_modules r out_d",
weight_shape="n_modules d r out_d", weight_shape="d r out_d",
bias_shape=None, # no bias bias_shape=None, # no bias
n_modules=len(self.target_modules),
d=config.latent_size, d=config.latent_size,
r=config.lora_config.r, r=config.lora_config.r,
out_d=sum(self.in_d[m] + self.out_d[m] for m in self.target_modules), out_d=max(self.in_d[m] + self.out_d[m] for m in self.target_modules),
) )
def _to_lora_dict( def _to_lora_dict(
self, flat_loras: Float[Tensor, "bs n_layers r _"] self, flat_loras: Float[Tensor, "bs n_layers n_modules r max_io_dim"]
) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]: ) -> dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]]:
# list of [bs, n_layers, r, in_out_dim] # # list of [bs, n_layers, r, in_out_dim]
# and in_out_dim might vary across modules # # and in_out_dim might vary across modules
# loras = unpack(
# flat_loras,
# [[self.in_d[m] + self.out_d[m]] for m in self.target_modules],
# "bs n_layers r *",
# )
loras = unpack( loras = unpack(
flat_loras, flat_loras,
[[self.in_d[m] + self.out_d[m]] for m in self.target_modules], [[] for _ in range(len(self.target_modules))],
"bs n_layers r *", "bs n_layers * r max_io_dim",
) )
# dict of {module: # dict of {module:
@ -383,7 +396,7 @@ class HyperLoRA(nn.Module):
lora_dict = dict() lora_dict = dict()
for module, lora in zip(self.target_modules, loras): for module, lora in zip(self.target_modules, loras):
A, B = unpack( A, B = unpack(
lora, lora[..., : self.in_d[module] + self.out_d[module]],
[[self.in_d[module]], [self.out_d[module]]], [[self.in_d[module]], [self.out_d[module]]],
"bs n_layers r *", "bs n_layers r *",
) )
@ -402,7 +415,7 @@ class HyperLoRA(nn.Module):
# [bs, n_layers, n_modules, feature_dim] # [bs, n_layers, n_modules, feature_dim]
emb = self.aggregator(features, attn_mask) emb = self.aggregator(features, attn_mask)
# [bs, n_layers, r, in_out_dim_mod1 + in_out_dim_mod2 + ...] # [bs, n_layers, n_modules, r, max_in_out_dim]
flat_loras = self.head(self.layers(emb)) flat_loras = self.head(self.layers(emb))
return flat_loras return flat_loras