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

AB scaler + use bias for generation

Browse source
This commit is contained in:
51616 2025-09-04 21:15:42 +09:00
parent 3364a25591
commit 41a1cfb90a
3 changed files with 118 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

11
.github/instructions/.instructions.md vendored Normal file
View file

@ -0,0 +1,11 @@
---
applyTo: '**'
---
When asked for a change or implementation, only work on the specific task I have given you with the most concise and elegant solution that changes as little code as possible.
When reviewing, first check the code diff for potential bugs and inconsistencies.
After that, check if the change would cause any bugs somewhere else in the codebase.
You should give detailed, professional, and nicely formatted explanation.
Always think slowly and carefully before doing anything. Be objective.

4
README.md
View file

@ -102,6 +102,10 @@ gsutil -m rsync -r data/raw_datasets/self_gen gs://ctx-to-lora/data/raw_datasets
# downloading from gcp bucket to login node
mkdir -p data/raw_datasets/self_gen
gsutil -m rsync -r gs://ctx-to-lora/data/raw_datasets/self_gen data/raw_datasets/self_gen
# self-gen eval
mkdir -p data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/
gsutil -m cp -r gs://ctx-to-lora/data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/*level_0_val*.parquet data/raw_datasets/self_gen/google/gemma-2-2b-it_temp_0.0_closed_qa_prob_0.0/fw_qa_v2/min_0_to_2000/train/
```
### Upload/download checkpoints

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

@ -306,6 +306,19 @@ class HyperLoRA(nn.Module):
}
)
self.scaler_A = nn.ParameterDict(
{
m: nn.Parameter(torch.ones((1, self.n_layers, self.r, 1)))
for m in self.target_modules
}
)
self.scaler_B = nn.ParameterDict(
{
m: nn.Parameter(torch.zeros((1, self.n_layers, self.r, 1)))
for m in self.target_modules
}
)
if self.config.use_light_weight_lora:
# light-weight lora projection (per layer, per module)
# self.pre_lora_projection = nn.ParameterDict(
@ -596,6 +609,10 @@ class HyperLoRA(nn.Module):
)
# transpose B
# B = rearrange(B, "bs n_layers r d_out -> bs n_layers d_out r")
# apparently doing A * self.scaler_A is slow due to broadcasting
A = torch.einsum("ijkl,ijkl->ijkl", A, self.scaler_A[module])
B = torch.einsum("ijkl,ijkl->ijkl", B, self.scaler_B[module])
if self.config.use_light_weight_lora:
# A = einsum(
# self.pre_lora_projection[module],
@ -690,6 +707,7 @@ class ModulatedPretrainedModel(nn.Module):
self.ctx_encoder_args = ctx_encoder_args
self.use_base_input_as_ctx = use_base_input_as_ctx
self.use_sequence_packing = use_sequence_packing
self.model_accepts_loss_kwargs = True
self.active_adapters = []
self.register_module("base_model", base_model)
@ -975,56 +993,59 @@ class ModulatedPretrainedModel(nn.Module):
ctx_ids, ctx_attn_mask, ctx_position_ids
)
generated_loras = combine_lora(
generated_loras,
n_ctx_chunks,
lora_bias=self.hypernet.get_head_bias(),
)
if generated_loras is not None:
generated_loras = combine_lora(
generated_loras,
n_ctx_chunks,
lora_bias=self.hypernet.get_head_bias(),
)
# input_ids in model_inputs_kwargs contains only
# prompt + response (for hypernet training)
position_ids = (
model_inputs_kwargs["position_ids"]
if "position_ids" in model_inputs_kwargs
else None
)
# with (
# apply_generated_loras(
# self.base_model,
# generated_loras,
# self.hypernet.layer_indices,
# position_ids,
# self.training,
# ),
# apply_generated_layernorm(
# self.base_model,
# generated_layernorms,
# self.hypernet.layer_indices,
# position_ids,
# self.training,
# ),
# ):
# model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs)
# input_ids in model_inputs_kwargs contains only
# prompt + response (for hypernet training)
position_ids = (
model_inputs_kwargs["position_ids"]
if "position_ids" in model_inputs_kwargs
else None
)
# with (
# apply_generated_loras(
# self.base_model,
# generated_loras,
# self.hypernet.layer_indices,
# position_ids,
# self.training,
# ),
# apply_generated_layernorm(
# self.base_model,
# generated_layernorms,
# self.hypernet.layer_indices,
# position_ids,
# self.training,
# ),
# ):
# model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs)
if n_queries is None:
if ctx_position_ids is None:
n_queries = torch.ones(
ctx_ids.shape[0], dtype=torch.int32, device=self.device
)
else:
# quite redundant (we do cu_seqlens many places)
# TODO: compute cu_seqlens here and propagate that
n_queries = torch.ones(
(ctx_position_ids == 0).sum(), dtype=torch.int32, device=self.device
)
if n_queries is None:
if ctx_position_ids is None:
n_queries = torch.ones(
ctx_ids.shape[0], dtype=torch.int32, device=self.device
)
else:
# quite redundant (we do cu_seqlens many places)
# TODO: compute cu_seqlens here and propagate that
n_queries = torch.ones(
(ctx_position_ids == 0).sum(),
dtype=torch.int32,
device=self.device,
)
apply_lora_to_layers(
self.base_model,
self.hypernet.layer_indices,
generated_loras,
n_queries,
position_ids,
)
apply_lora_to_layers(
self.base_model,
self.hypernet.layer_indices,
generated_loras,
n_queries,
position_ids,
)
model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs)
if return_generated_lora:
@ -1133,42 +1154,44 @@ class ModulatedPretrainedModel(nn.Module):
ctx_ids, ctx_attn_mask, ctx_position_ids
)
# generated_loras = combine_lora(
# generated_loras,
# n_chunks=n_ctx_chunks,
# aggregation="sum",
# lora_bias=self.hypernet.get_head_bias(),
# )
if generated_loras is not None:
generated_loras = combine_lora(
generated_loras,
n_ctx_chunks,
lora_bias=self.hypernet.get_head_bias(),
)
# # for generation the ctx_ids are batched not packed
# ctx_ids_list = torch.split(ctx_ids, n_ctx_chunks, dim=0)
# # for generation the ctx_ids are batched not packed
# ctx_ids_list = torch.split(ctx_ids, n_ctx_chunks, dim=0)
# apply lora hook to the base model
# TODO: we dont this position_ids for generation?
position_ids = (
model_inputs_kwargs["position_ids"]
if "position_ids" in model_inputs_kwargs
else None
)
if n_queries is None:
if ctx_position_ids is None:
n_queries = torch.ones(
ctx_ids.shape[0], dtype=torch.int32, device=self.device
)
else:
# quite redundant (we do cu_seqlens many places)
# TODO: compute cu_seqlens here and propagate that
n_queries = torch.ones(
(ctx_position_ids == 0).sum(), dtype=torch.int32, device=self.device
)
# apply lora hook to the base model
# TODO: we dont this position_ids for generation?
position_ids = (
model_inputs_kwargs["position_ids"]
if "position_ids" in model_inputs_kwargs
else None
)
if n_queries is None:
if ctx_position_ids is None:
n_queries = torch.ones(
ctx_ids.shape[0], dtype=torch.int32, device=self.device
)
else:
# quite redundant (we do cu_seqlens many places)
# TODO: compute cu_seqlens here and propagate that
n_queries = torch.ones(
(ctx_position_ids == 0).sum(),
dtype=torch.int32,
device=self.device,
)
apply_lora_to_layers(
self.base_model,
self.hypernet.layer_indices,
generated_loras,
n_queries,
position_ids,
)
apply_lora_to_layers(
self.base_model,
self.hypernet.layer_indices,
generated_loras,
n_queries,
position_ids,
)
model_outputs = self.base_model.generate(
*model_inputs_args, **model_inputs_kwargs
)