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 into main

Browse source
This commit is contained in:
51616 2025-09-04 13:54:54 +00:00
commit 237b6dc6fa
4 changed files with 269 additions and 213 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
)

284
src/ctx_to_lora/trainer.py
View file

@ -13,14 +13,10 @@ from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
logger = logging.getLogger()
class DistillationTrainer(Trainer):
def __init__(self, *args, **kwargs):
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False)
super().__init__(*args, **kwargs)
class ModulatedModelTrainer(Trainer):
# modified from the base Trainer to support per-context average loss
def get_batch_samples_ctx(self, epoch_iterator, num_batches, device):
def get_batch_samples(self, epoch_iterator, num_batches, device):
# only used with `use_per_ctx_average_loss=True`
batch_samples = []
num_items_in_batch = None
@ -33,27 +29,17 @@ class DistillationTrainer(Trainer):
count_num_items_in_batch = (
len(batch_samples) > 0
and "labels" in batch_samples[0]
and "ctx_ids" in batch_samples[0]
and (
# num_items_in_batch is passed to model forward
# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3757
self.model_accepts_loss_kwargs
# num_items_in_batch is passed to compute_loss_func
# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3773
or self.compute_loss_func is not None
# num_items_in_batch is also verified if (self.model_accepts_loss_kwargs or self.compute_loss_func)
# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3790
)
and "n_ctx_chunks" in batch_samples[0]
)
if count_num_items_in_batch:
# For now we don't support object detection
try:
num_items_in_batch = sum(
[(batch["ctx_position_ids"] == 0).sum() for batch in batch_samples]
)
except (TypeError, AttributeError):
pass
num_items_in_batch = dict()
num_items_in_batch["ctx"] = torch.tensor(
sum([batch["n_ctx_chunks"].numel() for batch in batch_samples])
)
num_items_in_batch["labels"] = sum(
[(batch["labels"].ne(-100)).sum() for batch in batch_samples]
)
if num_items_in_batch is not None:
if self.args.average_tokens_across_devices:
@ -68,117 +54,138 @@ class DistillationTrainer(Trainer):
return batch_samples, num_items_in_batch
class DistillationTrainer(ModulatedModelTrainer):
def __init__(self, *args, **kwargs):
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False)
super().__init__(*args, **kwargs)
def compute_loss(
self, model, inputs, return_outputs=False, num_items_in_batch=None
):
if self.model_accepts_loss_kwargs:
loss_kwargs = {}
if num_items_in_batch is not None:
loss_kwargs["num_items_in_batch"] = num_items_in_batch
inputs = {**inputs, **loss_kwargs}
use_kl = "logprobs_vals" in inputs
if use_kl:
labels = inputs.pop("labels", None)
label_pos = torch.where(labels != -100)
# NOTE: the loss output from this fn will be ***added***
# meaning that we should always scale the loss wrt `num_items_in_batch`
# (average over the number of items in the accumulated batch)
labels = inputs.pop("labels", None)
label_pos = torch.where(labels != -100)
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
if not use_kl:
loss = outputs["loss"]
target_logp = inputs.pop("logprobs_vals").squeeze(0)
indices = inputs.pop("logprobs_indices").squeeze(0)
else:
target_logp = inputs.pop("logprobs_vals").squeeze(0)
indices = inputs.pop("logprobs_indices").squeeze(0)
assert label_pos[0].shape[0] == target_logp.shape[0], (
"Label positions and target log probabilities should have the same # tokens."
f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}"
)
assert label_pos[0].shape[0] == target_logp.shape[0], (
"Label positions and target log probabilities should have the same # tokens."
f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}"
##### KL loss
outputs_logits = outputs.logits[label_pos[0], label_pos[1] - 1] # shift back 1
teacher_logp = torch.full_like(outputs_logits, -torch.inf)
teacher_logp.scatter_(1, indices, target_logp)
# reduction = "batchmean" if num_items_in_batch is None else "sum"
p = teacher_logp.exp()
logq = nn.functional.log_softmax(outputs_logits, dim=-1)
loss = -torch.sum(p * logq, dim=-1)
if self.use_per_ctx_average_loss:
n_queries_per_ctx = inputs["n_queries"].tolist()
position_ids = inputs["position_ids"].squeeze(0)
# account only label positions
label_mask = labels.squeeze(0) != -100
label_pos_ids = label_mask * position_ids
label_pos_ids_diff = label_pos_ids.diff(
append=torch.tensor([0], device=position_ids.device)
)
start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0]
end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0]
label_seq_lens = end_label_pos - start_label_pos
cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0)
start_indices = torch.cat(
(
torch.tensor([0], device=cu_label_seq_lens.device),
cu_label_seq_lens[:-1],
)
)
##### KL loss
outputs_logits = outputs.logits[
label_pos[0], label_pos[1] - 1
] # shift back 1
teacher_logp = torch.full_like(outputs_logits, -torch.inf)
teacher_logp.scatter_(1, indices, target_logp)
reduction = "batchmean" if num_items_in_batch is None else "sum"
p = teacher_logp.exp()
logq = nn.functional.log_softmax(outputs_logits, dim=-1)
loss = -torch.sum(p * logq, dim=-1)
if self.use_per_ctx_average_loss:
n_queries_per_ctx = inputs["n_queries"].tolist()
position_ids = inputs["position_ids"].squeeze(0)
# account only label positions
label_mask = labels.squeeze(0) != -100
label_pos_ids = label_mask * position_ids
label_pos_ids_diff = label_pos_ids.diff(
append=torch.tensor([0], device=position_ids.device)
)
start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0]
end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0]
label_seq_lens = end_label_pos - start_label_pos
cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0)
start_indices = torch.cat(
(
torch.tensor([0], device=cu_label_seq_lens.device),
cu_label_seq_lens[:-1],
)
)
# these stack and split can be optimized but let's keep it simple
# mean across tokens of each q
qa_losses = torch.stack(
[
loss[start:end].mean()
for start, end in zip(start_indices, cu_label_seq_lens)
]
)
# mean across queries of each ctx
per_ctx_losses = [
ctx.mean() for ctx in torch.split(qa_losses, n_queries_per_ctx)
# these stack and split can be optimized but let's keep it simple
# mean across tokens of each q
qa_losses = torch.stack(
[
loss[start:end].mean()
for start, end in zip(start_indices, cu_label_seq_lens)
]
)
# mean across contexts
loss = torch.stack(per_ctx_losses)
# mean across queries of each ctx
per_ctx_losses = [
ctx.mean() for ctx in torch.split(qa_losses, n_queries_per_ctx)
]
if reduction == "batchmean":
loss = loss.mean()
elif reduction == "sum" and num_items_in_batch is not None:
loss = loss.sum() / num_items_in_batch
#####
# mean across contexts
loss = torch.stack(per_ctx_losses)
if num_items_in_batch is not None:
if self.use_per_ctx_average_loss:
loss = loss.sum() / num_items_in_batch["ctx"]
else:
loss = loss.sum() / num_items_in_batch["labels"]
else:
# eval
loss = loss.mean()
# if reduction == "batchmean":
# loss = loss.mean()
# elif reduction == "sum":
# # loss does not scale with grad acc
# # num_items_in_batch does
# # this works for both token-avg and ctx-avg
# # loss = loss.sum() / num_items_in_batch
# `num_items_in_batch` is # tokens if `args.use_ctx_average_loss=False``
# loss = loss.sum() / num_items_in_batch
#####
##### unpack gen lora dict and compute regularization loss
l1_norm = 0
n_modules = len(gen_loras)
for module, lora in gen_loras.items():
l1_norm += lora["A"].abs().mean() + lora["B"].abs().mean()
l1_norm += lora["A"].abs().sum(0).mean() + lora["B"].abs().sum(0).mean()
l1_norm /= n_modules
if num_items_in_batch is not None:
# during eval `num_items_in_batch` will be None
l1_norm /= num_items_in_batch["ctx"]
total_loss = loss + self.gen_lora_l1_reg_coef * l1_norm
#####
scaler = self.args.gradient_accumulation_steps
if self.args.average_tokens_across_devices and num_items_in_batch is not None:
total_loss *= self.accelerator.num_processes
scaler *= self.accelerator.num_processes
# rough estimate of the losses (we only log the values from one step)
if (self.state.global_step == 1 and self.args.logging_first_step) or (
self.args.logging_strategy == IntervalStrategy.STEPS
and self.state.global_step % self.state.logging_steps == 0
):
loss_field = "kl_loss" if use_kl else "ce_loss"
self.log({loss_field: loss.item(), "gen_lora_l1_norm": l1_norm.item()})
loss += self.gen_lora_l1_reg_coef * l1_norm
#####
# compensate `num_items_in_batch` division
self.log(
{
"kl_loss": loss.item() * scaler,
"gen_lora_l1_norm": l1_norm.item() * scaler,
}
)
if (
self.args.average_tokens_across_devices
and self.model_accepts_loss_kwargs
and num_items_in_batch is not None
):
loss *= self.accelerator.num_processes
return (loss, outputs) if return_outputs else loss
return (total_loss, outputs) if return_outputs else total_loss
class ModulatedModelTrainer(Trainer):
class CrossEntropyTrainer(ModulatedModelTrainer):
def __init__(self, *args, **kwargs):
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
super().__init__(*args, **kwargs)
@ -197,11 +204,11 @@ class ModulatedModelTrainer(Trainer):
labels = inputs.pop("labels")
else:
labels = None
if self.model_accepts_loss_kwargs:
loss_kwargs = {}
if num_items_in_batch is not None:
loss_kwargs["num_items_in_batch"] = num_items_in_batch
inputs = {**inputs, **loss_kwargs}
loss_kwargs = {}
if num_items_in_batch is not None:
loss_kwargs["num_items_in_batch"] = num_items_in_batch["labels"]
inputs = {**inputs, **loss_kwargs}
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
# Save past state if it exists
if self.args.past_index >= 0:
@ -216,7 +223,7 @@ class ModulatedModelTrainer(Trainer):
# User-defined compute_loss function
if self.compute_loss_func is not None:
loss = self.compute_loss_func(
outputs, labels, num_items_in_batch=num_items_in_batch
outputs, labels, num_items_in_batch=num_items_in_batch["labels"]
)
elif model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
loss = self.label_smoother(outputs, labels, shift_labels=True)
@ -236,25 +243,36 @@ class ModulatedModelTrainer(Trainer):
##### unpack gen lora dict and compute regularization loss
l1_norm = 0
n_modules = len(gen_loras)
for module, lora in gen_loras.items():
l1_norm += lora["A"].abs().mean() + lora["B"].abs().mean()
l1_norm += lora["A"].abs().sum(0).mean() + lora["B"].abs().sum(0).mean()
l1_norm /= n_modules
if num_items_in_batch is not None:
# during eval `num_items_in_batch` will be None
l1_norm /= num_items_in_batch["ctx"]
total_loss = loss + self.gen_lora_l1_reg_coef * l1_norm
#####
scaler = self.args.gradient_accumulation_steps
if self.args.average_tokens_across_devices and num_items_in_batch is not None:
total_loss *= self.accelerator.num_processes
scaler *= self.accelerator.num_processes
# rough estimate of the losses (we only log the values from one step)
if (self.state.global_step == 1 and self.args.logging_first_step) or (
self.args.logging_strategy == IntervalStrategy.STEPS
and self.state.global_step % self.state.logging_steps == 0
):
self.log({"ce_loss": loss.item(), "gen_lora_l1_norm": l1_norm.item()})
loss += self.gen_lora_l1_reg_coef * l1_norm
#####
# compensate `num_items_in_batch` division
self.log(
{
"ce_loss": loss.item() * scaler,
"gen_lora_l1_norm": l1_norm.item() * scaler,
}
)
if (
self.args.average_tokens_across_devices
and self.model_accepts_loss_kwargs
and num_items_in_batch is not None
):
loss *= self.accelerator.num_processes
return (loss, outputs) if return_outputs else loss
return (total_loss, outputs) if return_outputs else total_loss
def get_decay_parameter_names(model) -> list[str]:
@ -269,7 +287,7 @@ def get_decay_parameter_names(model) -> list[str]:
decay_parameters = get_parameter_names(
model,
[nn.Embedding, nn.LayerNorm],
["bias", "layernorm", "rmsnorm", "latents_q"],
["scaler", "bias", "layernorm", "rmsnorm", "latents_q"],
)
return decay_parameters
@ -299,8 +317,8 @@ def train_model(
is_modulated_model = isinstance(model, ModulatedPretrainedModel)
trainer_cls = Trainer
if is_modulated_model:
logger.info("Training with modulated model. Using ModulatedModelTrainer.")
trainer_cls = ModulatedModelTrainer
logger.info("Training with modulated model. Using CrossEntropyTrainer.")
trainer_cls = CrossEntropyTrainer
trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef
del training_args.gen_lora_l1_reg_coef
@ -319,8 +337,8 @@ def train_model(
training_args.per_device_train_batch_size = 128
trainer = trainer_cls(**trainer_kwargs)
if getattr(trainer, "use_per_ctx_average_loss", False):
trainer.get_batch_samples = trainer.get_batch_samples_ctx
# if getattr(trainer, "use_per_ctx_average_loss", False):
# trainer.get_batch_samples = trainer.get_batch_samples_ctx
# MONKEY PATCH: remove embedding layers from weight decay
trainer.get_decay_parameter_names = get_decay_parameter_names