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

compile linear + loss + swtich from hook to lora_forward

Browse source
This commit is contained in:
51616 2025-06-20 23:41:32 +09:00
parent 36d78c827c
commit a36bd7943c
5 changed files with 265 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

19
intx_sft.py
View file

@ -54,6 +54,7 @@ from ctx_to_lora.modeling.hypernet import (
)
from ctx_to_lora.trainer import train_model
from ctx_to_lora.utils import (
compile_linear,
extract_cli_args,
get_run_name,
log_num_train_params,
@ -213,6 +214,7 @@ def main():
hypernet_config,
ctx_encoder_args,
ctx_args.use_kl_loss,
ctx_args.use_sequence_packing,
)
training_args.gen_lora_l1_reg_coef = ctx_args.gen_lora_l1_reg_coef
else:
@ -223,6 +225,7 @@ def main():
torch.load(ctx_args.from_pretrained_checkpoint, weights_only=False),
train=True,
use_flash_attn=model_args.use_flash_attn,
use_sequence_packing=ctx_args.use_sequence_packing,
)
tokenizer = get_tokenizer(model.base_model.config.name_or_path, train=True)
ctx_name = model.ctx_encoder_args.ctx_encoder_model_name_or_path
@ -404,15 +407,23 @@ def main():
if isinstance(model, ModulatedPretrainedModel):
logger.info("Applying liger-kernel to ModulatedPretrainedModel")
if isinstance(model.base_model, PeftModel):
_apply_liger_kernel_to_instance(model=model.base_model.base_model.model)
base_model = model.base_model.base_model
else:
_apply_liger_kernel_to_instance(model=model.base_model.model)
base_model = model.base_model
if ctx_name is not None:
logger.info("Applying liger-kernel to ctx_encoder_model")
_apply_liger_kernel_to_instance(model=model.ctx_encoder.base_model)
ctx_base_model = model.ctx_encoder.base_model
_apply_liger_kernel_to_instance(model=ctx_base_model)
compile_linear(ctx_base_model)
elif isinstance(model, PeftModel):
logger.info("Applying liger-kernel to PeftModel")
_apply_liger_kernel_to_instance(model=model.base_model.model)
base_model = model.base_model
_apply_liger_kernel_to_instance(model=base_model.model)
compile_linear(base_model)
base_model.loss_function = torch.compile(base_model.loss_function)
if LOCAL_RANK == 0:
wandb.init(

4
src/ctx_to_lora/model_loading.py
View file

@ -49,8 +49,8 @@ def get_tokenizer(
padding_side = "left" if not train else "right"
truncation_side = "left"
if peft_config:
model_name_or_path = peft_config.base_model_name_or_path
# if peft_config:
# model_name_or_path = peft_config.base_model_name_or_path
if tokenizer_kwargs is None:
tokenizer_kwargs = {}

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

@ -3,6 +3,7 @@ from collections.abc import Iterable
from contextlib import contextmanager
from copy import deepcopy
from dataclasses import dataclass
from functools import partial
from typing import Any
import torch
@ -42,10 +43,17 @@ from ctx_to_lora.model_loading import (
)
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.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,
get_peft_modules,
)
logger = logging.getLogger()
@ -580,13 +588,16 @@ class ModulatedPretrainedModel(nn.Module):
ctx_encoder_args: CtxEncoderArguments,
# use_kl_loss is only used for training
use_kl_loss: bool = False,
use_sequence_packing: bool = False,
use_base_input_as_ctx: bool = False,
):
super().__init__()
self.device = base_model.device
self.peft_config = base_model.peft_config["default"]
self.hypernet_config = hypernet_config
self.ctx_encoder_args = ctx_encoder_args
self.use_kl_loss = use_kl_loss
self.use_sequence_packing = use_sequence_packing
self.use_base_input_as_ctx = use_base_input_as_ctx
self.lora_id = 1
self.active_adapters = []
@ -621,9 +632,31 @@ class ModulatedPretrainedModel(nn.Module):
return model
def _init_model(self):
# disable adapter of the base model
# this only works with LoRA(?)
# we disable to avoid peft lora computation
self.base_model.disable_adapter_layers()
self.hypernet = (
HyperLoRA(self.hypernet_config).to(self.device).to(torch.float32)
)
layers = get_layers(self.base_model)
lora_forward_fn = (
lora_forward if not self.use_sequence_packing else lora_forward_packed
)
for layer_idx in self.hypernet.layer_indices:
for module_info in get_peft_modules(layers[layer_idx], self.peft_config):
name = module_info["name"]
module = module_info["module"]
logger.debug(f"Applying LoRA forward to {name}")
module.forward = partial(
lora_forward_fn,
self=module,
lora_dropout_p=self.peft_config.lora_dropout,
scaling=self.peft_config.lora_alpha,
)
ctx_model_name = self.ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_model_name is None:
ctx_model_name = self.base_model.config.name_or_path
@ -780,7 +813,8 @@ class ModulatedPretrainedModel(nn.Module):
**model_inputs_kwargs: dict[str, Any],
) -> tuple | ModelOutput:
"""Forward pass of the modulated model."""
# TODO: allow more than one LoRA per sample (multi-lora)
# TODO: allow multi queries per one LoRA (for efficient training)
generated_loras = None
generated_layernorms = None
if ctx_ids is None and not self.use_base_input_as_ctx:
@ -820,23 +854,28 @@ class ModulatedPretrainedModel(nn.Module):
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)
# 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)
apply_lora_to_layers(
self.base_model, self.hypernet.layer_indices, generated_loras, position_ids
)
model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs)
if return_generated_lora:
return model_outputs, (generated_loras, generated_layernorms)
@ -886,7 +925,7 @@ class ModulatedPretrainedModel(nn.Module):
@torch.inference_mode()
def generate(
self,
# TODO: allow more than one LoRA per sample
# TODO: allow more than one LoRA per sample (multi-lora)
ctx_ids: Integer[Tensor, "bs ctx_length"] | None = None,
ctx_attn_mask: Integer[Tensor, "bs ctx_length"] | None = None,
ctx_position_ids: Integer[Tensor, "bs ctx_length"] | None = None,
@ -1113,16 +1152,6 @@ def apply_generated_loras(
remove_hook_handles(hooks)
@contextmanager
def apply_generated_loras_packed_sequence(
base_model: nn.Module,
generated_loras: dict | None = None,
layer_indices: Iterable[int] | None = None,
training: bool = False,
):
pass
# needed for loading model from checkpoint
# see https://github.com/huggingface/transformers/pull/34632
torch.serialization.add_safe_globals(

171
src/ctx_to_lora/modeling/lora_layer.py Normal file
View file

@ -0,0 +1,171 @@
from collections.abc import Iterable
from functools import partial
from operator import attrgetter
import torch
import torch.nn.functional as F
from einops import einsum
from jaxtyping import Float, Integer
from torch import Tensor
from ctx_to_lora.utils import get_layers
# Patch corresponding layers of LLMs to include A and B lora matrices
"""
def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
self._check_forward_args(x, *args, **kwargs)
adapter_names = kwargs.pop("adapter_names", None)
if self.disable_adapters:
if self.merged:
self.unmerge()
result = self.base_layer(x, *args, **kwargs)
elif adapter_names is not None:
result = self._mixed_batch_forward(x, *args, adapter_names=adapter_names, **kwargs)
elif self.merged:
result = self.base_layer(x, *args, **kwargs)
else:
result = self.base_layer(x, *args, **kwargs)
torch_result_dtype = result.dtype
lora_A_keys = self.lora_A.keys()
for active_adapter in self.active_adapters:
if active_adapter not in lora_A_keys:
continue
lora_A = self.lora_A[active_adapter]
lora_B = self.lora_B[active_adapter]
dropout = self.lora_dropout[active_adapter]
scaling = self.scaling[active_adapter]
x = self._cast_input_dtype(x, lora_A.weight.dtype)
if not self.use_dora[active_adapter]:
result = result + lora_B(lora_A(dropout(x))) * scaling
else:
if isinstance(dropout, nn.Identity) or not self.training:
base_result = result
else:
x = dropout(x)
base_result = None
result = result + self.lora_magnitude_vector[active_adapter](
x,
lora_A=lora_A,
lora_B=lora_B,
scaling=scaling,
base_layer=self.get_base_layer(),
base_result=base_result,
)
result = result.to(torch_result_dtype)
return result
"""
def lora_forward(
x: Float[Tensor, "bs seq_len d_in"],
self,
# pre_modulated_forward_fn: callable[..., Float[Tensor, "bs seq_len d_out"]],
A: Float[Tensor, "bs r d_in"],
B: Float[Tensor, "bs d_out r"],
lora_dropout_p: float,
scaling: float,
*args,
**kwargs,
) -> Float[Tensor, "bs seq_len d_out"]:
"""
Forward pass of the LoRA layer.
# TODO: implement the packed version
# TODO: implement the one-lora-multi-query version
Args:
x (torch.Tensor): Input tensor with shape (batch_size, input_dim).
A (torch.Tensor): LoRA matrix A with shape (bs, r, d_in).
B (torch.Tensor): LoRA matrix B with shape (bs, d_out, r).
Returns:
torch.Tensor: Output tensor with shape (batch_size, output_dim).
"""
base_out = torch.nn.Linear.forward(self, x, *args, **kwargs)
delta_x = F.dropout(x, p=lora_dropout_p, training=self.training)
delta_x = einsum(A, delta_x, "bs r d_in, bs seq_len d_in -> bs seq_len r")
delta_x = einsum(B, delta_x, "bs d_out r, bs seq_len r -> bs seq_len d_out")
delta_x = delta_x * scaling
return base_out + delta_x
def lora_forward_packed(
x: Float[Tensor, "bs seq_len d_in"],
seq_lens: Integer[Tensor, "bs"],
self,
A: Float[Tensor, "bs r d_in"],
B: Float[Tensor, "bs d_out r"],
lora_dropout_p: float,
scaling: float,
*args,
**kwargs,
) -> Float[Tensor, "bs seq_len d_out"]:
"""
Forward pass of the LoRA layer for packed inputs.
Args:
x (torch.Tensor): Input tensor with shape (batch_size, input_dim).
A (torch.Tensor): LoRA matrix A with shape (input_dim, lora_dim).
B (torch.Tensor): LoRA matrix B with shape (lora_dim, output_dim).
Returns:
torch.Tensor: Output tensor with shape (batch_size, output_dim).
"""
# bs of x should be 1 in this case
base_out = torch.nn.Linear.forward(self, x, *args, **kwargs)
delta_x = F.dropout(x, p=lora_dropout_p, training=self.training)
# [tot_seq_len, r, d_in]
repeated_A = A.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum())
# [tot_seq_len, d_out, r]
repeated_B = B.repeat_interleave(seq_lens, dim=0, output_size=seq_lens.sum())
delta_x = einsum(
repeated_A, delta_x, "tot_len r d_in, bs tot_len d_in -> bs tot_len r"
)
delta_x = einsum(
repeated_B, delta_x, "tot_len d_out r, bs tot_len r -> bs tot_len d_out"
)
delta_x = delta_x * scaling
return base_out + delta_x
def apply_lora_to_layers(
model: torch.nn.Module,
layer_indices: Iterable[int],
generated_loras: dict[str, dict[str, Float[Tensor, "bs n_layers r _"]]],
position_ids: Integer[Tensor, "bs seq_len"] = None,
) -> None:
layers = get_layers(model)
if position_ids is not None:
position_ids = position_ids.squeeze()
seq_lens = position_ids[torch.where(position_ids == 0)[0][1:] - 1]
seq_lens = torch.cat(
[seq_lens, torch.tensor([position_ids[-1]], device=seq_lens.device)]
)
seq_lens += 1
for layer_idx in layer_indices:
layer = layers[layer_idx]
for mname in generated_loras:
if mname in ["q_proj", "k_proj", "v_proj", "o_proj", "qkv_proj"]:
long_mname = f"self_attn.{mname}"
elif mname in ["down_proj", "up_proj", "gate_proj"]:
long_mname = f"mlp.{mname}"
module = attrgetter(long_mname)(layer)
module.forward = partial(
module.forward,
A=generated_loras[mname]["A"][:, layer_idx],
B=generated_loras[mname]["B"][:, layer_idx],
)
if position_ids is not None:
module.forward = partial(
module.forward,
seq_lens=seq_lens,
)

24
src/ctx_to_lora/utils.py
View file

@ -161,6 +161,16 @@ def save_yaml(data, path):
yaml.dump(data, file)
def get_peft_modules(model: PeftModel, peft_config: PeftConfig) -> list[dict[str, str]]:
return [
{"name": name, "module": module}
for name, module in model.named_modules()
if name.split(".")[-1] in peft_config.target_modules
and isinstance(module, BaseTunerLayer)
and check_target_module_exists(peft_config, name)
]
def get_peft_in_out_features(
model: PeftModel,
peft_config: PeftConfig | None = None,
@ -169,11 +179,9 @@ def get_peft_in_out_features(
return None, None
in_features = dict()
out_features = dict()
for module_name, module in model.named_modules():
if not check_target_module_exists(peft_config, module_name):
continue
if not isinstance(module, BaseTunerLayer):
continue
for module_info in get_peft_modules(model, peft_config):
module_name = module_info["name"]
module = module_info["module"]
# support just Linear layer for now
# all modules should be a leave module that is Linear layer
assert isinstance(module.base_layer, torch.nn.Linear), (
@ -239,6 +247,12 @@ def get_lora_module_names(
return module_names
def compile_linear(model):
for name, module in model.named_modules():
if isinstance(module, torch.nn.Linear):
module.compile()
def clear_gpu():
gc.collect()
torch.cuda.empty_cache()