mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
initial one-lora-multi-q model interface
This commit is contained in:
51616
parent
3daf068d4a
commit
a37998a015
2 changed files with 62 additions and 106 deletions
|
|
@ -805,16 +805,16 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
|
||||
def forward(
|
||||
self,
|
||||
ctx_ids: Integer[Tensor, "bs ctx_len"] | None = None,
|
||||
ctx_attn_mask: Integer[Tensor, "bs ctx_len"] | None = None,
|
||||
ctx_position_ids: Integer[Tensor, "bs ctx_len"] | None = None,
|
||||
ctx_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None,
|
||||
ctx_attn_mask: Integer[Tensor, "n_ctx ctx_len"] | None = None,
|
||||
ctx_position_ids: Integer[Tensor, "n_ctx ctx_len"] | None = None,
|
||||
n_queries: Integer[Tensor, "n_ctx"] | None = None,
|
||||
return_generated_lora: bool | None = False,
|
||||
*model_inputs_args: Any,
|
||||
**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:
|
||||
|
|
@ -872,8 +872,24 @@ class ModulatedPretrainedModel(nn.Module):
|
|||
# ):
|
||||
# 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
|
||||
)
|
||||
|
||||
apply_lora_to_layers(
|
||||
self.base_model, self.hypernet.layer_indices, generated_loras, position_ids
|
||||
self.base_model,
|
||||
self.hypernet.layer_indices,
|
||||
generated_loras,
|
||||
n_queries,
|
||||
position_ids,
|
||||
)
|
||||
model_outputs = self.base_model(*model_inputs_args, **model_inputs_kwargs)
|
||||
|
||||
|
|
|
|||
|
|
@ -10,121 +10,61 @@ 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"],
|
||||
x: Float[Tensor, "tot_q seq_len d_in"],
|
||||
n_qs: Integer[Tensor, "n_ctx"],
|
||||
tot_q: int,
|
||||
A: Float[Tensor, "n_ctx r d_in"],
|
||||
B: Float[Tensor, "n_ctx d_out r"],
|
||||
lora_dropout_p: float,
|
||||
scaling: float,
|
||||
self,
|
||||
*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
|
||||
) -> Float[Tensor, "tot_q seq_len d_out"]:
|
||||
# A: [n_ctx, r, d_in] -> [tot_q, r, d_in]
|
||||
A = A.repeat_interleave(n_qs, dim=0, output_size=tot_q)
|
||||
# B: [n_ctx, d_out, r] -> [tot_q, d_out, r]
|
||||
B = B.repeat_interleave(n_qs, dim=0, output_size=tot_q)
|
||||
|
||||
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 = einsum(A, delta_x, "tot_q r d_in, tot_q s_len d_in -> tot_q s_len r")
|
||||
delta_x = einsum(B, delta_x, "tot_q d_out r, tot_q s_len r -> tot_q s_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"],
|
||||
x: Float[Tensor, "1 tot_len d_in"],
|
||||
n_qs: Integer[Tensor, "n_ctx"],
|
||||
tot_q: int,
|
||||
seq_lens: Integer[Tensor, "tot_q"],
|
||||
tot_len: int,
|
||||
A: Float[Tensor, "n_ctx r d_in"],
|
||||
B: Float[Tensor, "n_ctx d_out r"],
|
||||
lora_dropout_p: float,
|
||||
scaling: float,
|
||||
self,
|
||||
*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).
|
||||
"""
|
||||
) -> Float[Tensor, "1 tot_len d_out"]:
|
||||
# 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())
|
||||
# # [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())
|
||||
# TODO: is ther a better way to implement this?
|
||||
repeated_A = A.repeat_interleave(n_qs, dim=0, output_size=tot_q)
|
||||
repeated_A = repeated_A.repeat_interleave(seq_lens, dim=0, output_size=tot_len)
|
||||
|
||||
repeated_B = B.repeat_interleave(n_qs, dim=0, output_size=tot_q)
|
||||
repeated_B = repeated_B.repeat_interleave(seq_lens, dim=0, output_size=tot_len)
|
||||
|
||||
delta_x = einsum(
|
||||
repeated_A, delta_x, "tot_len r d_in, bs tot_len d_in -> bs tot_len r"
|
||||
)
|
||||
|
|
@ -139,7 +79,8 @@ def lora_forward_packed(
|
|||
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 _"]]],
|
||||
generated_loras: dict[str, dict[str, Float[Tensor, "n_ctx n_layers r _"]]],
|
||||
n_qs: Integer[Tensor, "n_ctx"],
|
||||
position_ids: Integer[Tensor, "bs seq_len"] = None,
|
||||
) -> None:
|
||||
layers = get_layers(model)
|
||||
|
|
@ -150,6 +91,8 @@ def apply_lora_to_layers(
|
|||
[seq_lens, torch.tensor([position_ids[-1]], device=seq_lens.device)]
|
||||
)
|
||||
seq_lens += 1
|
||||
tot_len = seq_lens.sum().item()
|
||||
tot_q = n_qs.sum().item()
|
||||
for layer_idx in layer_indices:
|
||||
layer = layers[layer_idx]
|
||||
|
||||
|
|
@ -159,13 +102,10 @@ def apply_lora_to_layers(
|
|||
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],
|
||||
)
|
||||
A = generated_loras[mname]["A"][:, layer_idx]
|
||||
B = generated_loras[mname]["B"][:, layer_idx]
|
||||
module.forward = partial(module.forward, n_qs=n_qs, tot_q=tot_q, A=A, B=B)
|
||||
if position_ids is not None:
|
||||
module.forward = partial(
|
||||
module.forward,
|
||||
seq_lens=seq_lens,
|
||||
module.forward, seq_lens=seq_lens, tot_len=tot_len
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue