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

Merge branch 'main' of https://github.com/SakanaAI/ctx-to-lora

Browse source
This commit is contained in:
51616 2025-09-29 00:40:54 +09:00
commit 4b1e00fcce
2 changed files with 37 additions and 18 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -291,21 +291,24 @@ class HyperLoRA(nn.Module):
for m in self.target_modules
}
)
else:
self.bias_A = nn.ParameterDict(
# else:
# self.bias_A = nn.ParameterDict(
# {
# m: nn.Parameter(
# torch.zeros((self.n_layers, self.r, self.d_in[m]))
# )
# for m in self.target_modules
# }
# )
if self.config.use_bias:
self.bias_B = nn.ParameterDict(
{
m: nn.Parameter(
torch.zeros((self.n_layers, self.r, self.d_in[m]))
torch.zeros((self.n_layers, self.r, self.d_out[m]))
)
for m in self.target_modules
}
)
self.bias_B = nn.ParameterDict(
{
m: nn.Parameter(torch.zeros((self.n_layers, self.r, self.d_out[m])))
for m in self.target_modules
}
)
self.scaler_A = nn.ParameterDict(
{
@ -1028,7 +1031,9 @@ class ModulatedPretrainedModel(nn.Module):
generated_loras = combine_lora(
generated_loras,
n_ctx_chunks,
lora_bias=self.hypernet.get_head_bias(),
lora_bias=self.hypernet.get_head_bias()
if self.hypernet.config.use_bias
else None,
)
# input_ids in model_inputs_kwargs contains only
@ -1092,6 +1097,8 @@ class ModulatedPretrainedModel(nn.Module):
ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
n_ctx_chunks: Integer[Tensor, "n_ctx"] | None = None,
n_queries: Integer[Tensor, "n_ctx"] | None = None,
scalers: Float[Tensor, "n_ctx"] | None = None,
bias_scaler: float | None = None,
*model_inputs_args: Any,
**model_inputs_kwargs: dict[str, Any],
):
@ -1102,7 +1109,11 @@ class ModulatedPretrainedModel(nn.Module):
generated_loras = combine_lora(
generated_loras,
n_ctx_chunks,
lora_bias=self.hypernet.get_head_bias(),
lora_bias=self.hypernet.get_head_bias()
if self.hypernet.config.use_bias
else None,
scalers=scalers,
bias_scaler=bias_scaler,
)
# apply lora hook to the base model
@ -1196,7 +1207,9 @@ class ModulatedPretrainedModel(nn.Module):
generated_loras = self.combine_lora(
generated_loras,
n_ctx_chunks,
lora_bias=self.hypernet.get_head_bias(),
lora_bias=self.hypernet.get_head_bias()
if self.hypernet.config.use_bias
else None,
)
# # for generation the ctx_ids are batched not packed

18
src/ctx_to_lora/modeling/lora_merger.py
View file

@ -4,7 +4,7 @@ Utilities for merging / aggregating LoRA adapters coming from multiple chunks.
import torch
from einops import rearrange
from jaxtyping import Integer
from jaxtyping import Float, Integer
from torch import Tensor
@ -15,10 +15,13 @@ def compute_rank(n_lora, rank):
def combine_lora(
generated_loras: dict[str, dict[str, Tensor]],
n_chunks: Integer[Tensor, "n_ctx"],
lora_bias: dict[str, dict[str, Tensor]],
lora_bias: dict[str, dict[str, Tensor]] | None = None,
scalers: Float[Tensor, "n_ctx"] | None = None,
bias_scaler: float | None = None,
) -> dict[str, dict[str, Tensor]]:
total_chunks = int(n_chunks.sum())
if bias_scaler is None:
bias_scaler = 1
# Assume all modules share same base rank r
first_module = next(iter(generated_loras))
sampled_lora = generated_loras[first_module]["A"]
@ -33,11 +36,14 @@ def combine_lora(
rank_dim = 2
num_groups = len(n_chunks)
rank_per_group = (n_chunks * base_rank).tolist()
bias_tensor = None
for module_name, module_loras in generated_loras.items():
for matrix_key in ("A", "B"):
bias_tensor = lora_bias[module_name][matrix_key]
if lora_bias is not None:
bias_tensor = lora_bias[module_name][matrix_key]
loras = module_loras[matrix_key]
if (scalers is not None) and (matrix_key == "A"):
loras = loras * scalers[:, None, None, None]
flat_loras = rearrange(
loras, "tot_chunks n_layers r dim -> 1 n_layers (tot_chunks r) dim"
@ -64,7 +70,7 @@ def combine_lora(
# combined_rank, combined_rank + base_rank
# )
combined[g, :, combined_rank : combined_rank + base_rank, :] = (
bias_tensor
bias_tensor * bias_scaler
)
combined_loras[module_name][matrix_key] = combined