434 lines
17 KiB
Python
434 lines
17 KiB
Python
# based heavily on https://github.com/KohakuBlueleaf/LyCORIS/blob/eb460098187f752a5d66406d3affade6f0a07ece/lycoris/modules/lokr.py
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from toolkit.network_mixins import ToolkitModuleMixin
|
|
|
|
from typing import TYPE_CHECKING, Union, List
|
|
|
|
from optimum.quanto import QBytesTensor, QTensor
|
|
from torchao.dtypes import AffineQuantizedTensor
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
from toolkit.lora_special import LoRASpecialNetwork
|
|
|
|
|
|
def factorization(dimension: int, factor: int = -1) -> tuple[int, int]:
|
|
'''
|
|
return a tuple of two value of input dimension decomposed by the number closest to factor
|
|
second value is higher or equal than first value.
|
|
|
|
In LoRA with Kroneckor Product, first value is a value for weight scale.
|
|
secon value is a value for weight.
|
|
|
|
Becuase of non-commutative property, A(kron)B != B(kron)A. Meaning of two matrices is slightly different.
|
|
|
|
examples)
|
|
factor
|
|
-1 2 4 8 16 ...
|
|
127 -> 127, 1 127 -> 127, 1 127 -> 127, 1 127 -> 127, 1 127 -> 127, 1
|
|
128 -> 16, 8 128 -> 64, 2 128 -> 32, 4 128 -> 16, 8 128 -> 16, 8
|
|
250 -> 125, 2 250 -> 125, 2 250 -> 125, 2 250 -> 125, 2 250 -> 125, 2
|
|
360 -> 45, 8 360 -> 180, 2 360 -> 90, 4 360 -> 45, 8 360 -> 45, 8
|
|
512 -> 32, 16 512 -> 256, 2 512 -> 128, 4 512 -> 64, 8 512 -> 32, 16
|
|
1024 -> 32, 32 1024 -> 512, 2 1024 -> 256, 4 1024 -> 128, 8 1024 -> 64, 16
|
|
'''
|
|
|
|
if factor > 0 and (dimension % factor) == 0:
|
|
m = factor
|
|
n = dimension // factor
|
|
return m, n
|
|
if factor == -1:
|
|
factor = dimension
|
|
m, n = 1, dimension
|
|
length = m + n
|
|
while m < n:
|
|
new_m = m + 1
|
|
while dimension % new_m != 0:
|
|
new_m += 1
|
|
new_n = dimension // new_m
|
|
if new_m + new_n > length or new_m > factor:
|
|
break
|
|
else:
|
|
m, n = new_m, new_n
|
|
if m > n:
|
|
n, m = m, n
|
|
return m, n
|
|
|
|
|
|
def make_weight_cp(t, wa, wb):
|
|
rebuild2 = torch.einsum('i j k l, i p, j r -> p r k l',
|
|
t, wa, wb) # [c, d, k1, k2]
|
|
return rebuild2
|
|
|
|
|
|
def make_kron(w1, w2, scale):
|
|
if len(w2.shape) == 4:
|
|
w1 = w1.unsqueeze(2).unsqueeze(2)
|
|
w2 = w2.contiguous()
|
|
rebuild = torch.kron(w1, w2)
|
|
|
|
return rebuild*scale
|
|
|
|
|
|
class LokrModule(ToolkitModuleMixin, nn.Module):
|
|
def __init__(
|
|
self,
|
|
lora_name,
|
|
org_module: nn.Module,
|
|
multiplier=1.0,
|
|
lora_dim=4,
|
|
alpha=1,
|
|
dropout=0.,
|
|
rank_dropout=0.,
|
|
module_dropout=0.,
|
|
use_cp=False,
|
|
decompose_both=False,
|
|
network: 'LoRASpecialNetwork' = None,
|
|
factor: int = -1, # factorization factor
|
|
**kwargs,
|
|
):
|
|
""" if alpha == 0 or None, alpha is rank (no scaling). """
|
|
ToolkitModuleMixin.__init__(self, network=network)
|
|
torch.nn.Module.__init__(self)
|
|
factor = int(factor)
|
|
self.lora_name = lora_name
|
|
self.lora_dim = lora_dim
|
|
self.cp = False
|
|
self.use_w1 = False
|
|
self.use_w2 = False
|
|
self.can_merge_in = True
|
|
# only the plain Linear path (no cp/Conv2d) gets the factorized forward;
|
|
# Conv2d and the cp branch keep using get_weight()/make_kron() below.
|
|
self._fast_linear = False
|
|
|
|
# avoid the weight property on quantized OstrisLinear: it dequantizes the
|
|
# whole weight just to answer .shape
|
|
if getattr(org_module, "is_ostris_quantized", False):
|
|
self.shape = torch.Size((org_module.out_features, org_module.in_features))
|
|
else:
|
|
self.shape = org_module.weight.shape
|
|
if org_module.__class__.__name__ == 'Conv2d':
|
|
in_dim = org_module.in_channels
|
|
k_size = org_module.kernel_size
|
|
out_dim = org_module.out_channels
|
|
|
|
in_m, in_n = factorization(in_dim, factor)
|
|
out_l, out_k = factorization(out_dim, factor)
|
|
# ((a, b), (c, d), *k_size)
|
|
shape = ((out_l, out_k), (in_m, in_n), *k_size)
|
|
|
|
self.cp = use_cp and k_size != (1, 1)
|
|
if decompose_both and lora_dim < max(shape[0][0], shape[1][0])/2:
|
|
self.lokr_w1_a = nn.Parameter(
|
|
torch.empty(shape[0][0], lora_dim))
|
|
self.lokr_w1_b = nn.Parameter(
|
|
torch.empty(lora_dim, shape[1][0]))
|
|
else:
|
|
self.use_w1 = True
|
|
self.lokr_w1 = nn.Parameter(torch.empty(
|
|
shape[0][0], shape[1][0])) # a*c, 1-mode
|
|
|
|
if lora_dim >= max(shape[0][1], shape[1][1])/2:
|
|
self.use_w2 = True
|
|
self.lokr_w2 = nn.Parameter(torch.empty(
|
|
shape[0][1], shape[1][1], *k_size))
|
|
elif self.cp:
|
|
self.lokr_t2 = nn.Parameter(torch.empty(
|
|
lora_dim, lora_dim, shape[2], shape[3]))
|
|
self.lokr_w2_a = nn.Parameter(
|
|
torch.empty(lora_dim, shape[0][1])) # b, 1-mode
|
|
self.lokr_w2_b = nn.Parameter(
|
|
torch.empty(lora_dim, shape[1][1])) # d, 2-mode
|
|
else: # Conv2d not cp
|
|
# bigger part. weight and LoRA. [b, dim] x [dim, d*k1*k2]
|
|
self.lokr_w2_a = nn.Parameter(
|
|
torch.empty(shape[0][1], lora_dim))
|
|
self.lokr_w2_b = nn.Parameter(torch.empty(
|
|
lora_dim, shape[1][1]*shape[2]*shape[3]))
|
|
# w1 (kron) (w2_a x w2_b) = (a, b)(kron)((c, dim)x(dim, d*k1*k2)) = (a, b)(kron)(c, d*k1*k2) = (ac, bd*k1*k2)
|
|
|
|
self.op = F.conv2d
|
|
self.extra_args = {
|
|
"stride": org_module.stride,
|
|
"padding": org_module.padding,
|
|
"dilation": org_module.dilation,
|
|
"groups": org_module.groups
|
|
}
|
|
|
|
else: # Linear
|
|
in_dim = org_module.in_features
|
|
out_dim = org_module.out_features
|
|
|
|
in_m, in_n = factorization(in_dim, factor)
|
|
out_l, out_k = factorization(out_dim, factor)
|
|
# ((a, b), (c, d)), out_dim = a*c, in_dim = b*d
|
|
shape = ((out_l, out_k), (in_m, in_n))
|
|
# stash factor pair for the factorized (no full-kron) forward path
|
|
self._in_m, self._in_n = in_m, in_n
|
|
self._out_l, self._out_k = out_l, out_k
|
|
self._fast_linear = True
|
|
|
|
# smaller part. weight scale
|
|
if decompose_both and lora_dim < max(shape[0][0], shape[1][0])/2:
|
|
self.lokr_w1_a = nn.Parameter(
|
|
torch.empty(shape[0][0], lora_dim))
|
|
self.lokr_w1_b = nn.Parameter(
|
|
torch.empty(lora_dim, shape[1][0]))
|
|
else:
|
|
self.use_w1 = True
|
|
self.lokr_w1 = nn.Parameter(torch.empty(
|
|
shape[0][0], shape[1][0])) # a*c, 1-mode
|
|
|
|
if lora_dim < max(shape[0][1], shape[1][1])/2:
|
|
# bigger part. weight and LoRA. [b, dim] x [dim, d]
|
|
self.lokr_w2_a = nn.Parameter(
|
|
torch.empty(shape[0][1], lora_dim))
|
|
self.lokr_w2_b = nn.Parameter(
|
|
torch.empty(lora_dim, shape[1][1]))
|
|
# w1 (kron) (w2_a x w2_b) = (a, b)(kron)((c, dim)x(dim, d)) = (a, b)(kron)(c, d) = (ac, bd)
|
|
else:
|
|
self.use_w2 = True
|
|
self.lokr_w2 = nn.Parameter(
|
|
torch.empty(shape[0][1], shape[1][1]))
|
|
|
|
self.op = F.linear
|
|
self.extra_args = {}
|
|
|
|
self.dropout = dropout
|
|
if dropout:
|
|
print("[WARN]LoKr haven't implemented normal dropout yet.")
|
|
self.rank_dropout = rank_dropout
|
|
self.module_dropout = module_dropout
|
|
|
|
if isinstance(alpha, torch.Tensor):
|
|
alpha = float(alpha.detach().float().item())
|
|
alpha = lora_dim if alpha is None or alpha == 0 else alpha
|
|
if self.use_w2 and self.use_w1:
|
|
# use scale = 1
|
|
alpha = lora_dim
|
|
self._set_runtime_scale(float(alpha) / self.lora_dim)
|
|
self.register_buffer('alpha', torch.tensor(alpha)) # treat as constant
|
|
|
|
if self.use_w2:
|
|
torch.nn.init.constant_(self.lokr_w2, 0)
|
|
else:
|
|
if self.cp:
|
|
torch.nn.init.kaiming_uniform_(self.lokr_t2, a=math.sqrt(5))
|
|
torch.nn.init.kaiming_uniform_(self.lokr_w2_a, a=math.sqrt(5))
|
|
torch.nn.init.constant_(self.lokr_w2_b, 0)
|
|
|
|
if self.use_w1:
|
|
torch.nn.init.kaiming_uniform_(self.lokr_w1, a=math.sqrt(5))
|
|
else:
|
|
torch.nn.init.kaiming_uniform_(self.lokr_w1_a, a=math.sqrt(5))
|
|
torch.nn.init.kaiming_uniform_(self.lokr_w1_b, a=math.sqrt(5))
|
|
|
|
self.multiplier = multiplier
|
|
self.org_module = [org_module]
|
|
weight = make_kron(
|
|
self.lokr_w1 if self.use_w1 else self.lokr_w1_a@self.lokr_w1_b,
|
|
(self.lokr_w2 if self.use_w2
|
|
else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.cp
|
|
else self.lokr_w2_a@self.lokr_w2_b),
|
|
self.multiplier * self.scale
|
|
)
|
|
assert torch.sum(torch.isnan(weight)) == 0, "weight is nan"
|
|
|
|
# Same as locon.py
|
|
def apply_to(self):
|
|
self.org_forward = self.org_module[0].forward
|
|
self.org_module[0].forward = self.forward
|
|
|
|
def get_weight(self, orig_weight=None):
|
|
weight = make_kron(
|
|
self.lokr_w1 if self.use_w1 else self.lokr_w1_a@self.lokr_w1_b,
|
|
(self.lokr_w2 if self.use_w2
|
|
else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.cp
|
|
else self.lokr_w2_a@self.lokr_w2_b),
|
|
self._runtime_scale
|
|
)
|
|
if orig_weight is not None:
|
|
weight = weight.reshape(orig_weight.shape)
|
|
if self.training and self.rank_dropout:
|
|
drop = torch.rand(weight.size(0)) < self.rank_dropout
|
|
weight *= drop.view(-1, [1] *
|
|
len(weight.shape[1:])).to(weight.device)
|
|
return weight
|
|
|
|
@torch.no_grad()
|
|
def merge_in(self, merge_weight=1.0):
|
|
if not self.can_merge_in:
|
|
return
|
|
|
|
# extract weight from org_module
|
|
org_sd = self.org_module[0].state_dict()
|
|
# todo find a way to merge in weights when doing quanto quantized model
|
|
if 'weight._data' in org_sd:
|
|
# quanto quantized weight
|
|
return
|
|
|
|
weight_key = "weight"
|
|
from toolkit.util.quantize import is_quantized_tensor
|
|
org_weight = self.org_module[0].weight
|
|
is_ao_quantized = is_quantized_tensor(org_weight)
|
|
orig_dtype = org_weight.dtype
|
|
# dequantize torchao weights so the delta can be merged in full precision
|
|
weight = (org_weight.dequantize() if is_ao_quantized else org_weight).float()
|
|
|
|
scale = self.scale
|
|
# handle trainable scaler method locon does
|
|
if hasattr(self, 'scalar'):
|
|
scale = scale * self.scalar
|
|
|
|
lokr_weight = self.get_weight(weight)
|
|
|
|
merged_weight = (
|
|
weight
|
|
+ (lokr_weight * merge_weight).to(weight.device, dtype=weight.dtype)
|
|
)
|
|
|
|
# write the merged weight back, re-quantizing if the original was torchao quantized so the
|
|
# model stays quantized across continuous merge/reset cycles
|
|
if is_ao_quantized:
|
|
from toolkit.util.quantize import get_torchao_config, requantize_module_weight
|
|
requantize_module_weight(
|
|
self.org_module[0], merged_weight, orig_dtype, get_torchao_config(self._get_base_qtype())
|
|
)
|
|
else:
|
|
org_sd[weight_key] = merged_weight.to(orig_dtype)
|
|
self.org_module[0].load_state_dict(org_sd)
|
|
|
|
def get_orig_weight(self, device):
|
|
weight = self.org_module[0].weight
|
|
if weight.device != device:
|
|
weight = weight.to(device)
|
|
if isinstance(weight, QTensor) or isinstance(weight, QBytesTensor):
|
|
return weight.dequantize().data.detach()
|
|
elif isinstance(weight, AffineQuantizedTensor):
|
|
return weight.dequantize().data.detach()
|
|
else:
|
|
return weight.data.detach()
|
|
|
|
def get_orig_bias(self, device):
|
|
if hasattr(self.org_module[0], 'bias') and self.org_module[0].bias is not None:
|
|
bias = self.org_module[0].bias
|
|
if bias.device != device:
|
|
bias = bias.to(device)
|
|
if isinstance(bias, QTensor) or isinstance(bias, QBytesTensor):
|
|
return bias.dequantize().data.detach()
|
|
elif isinstance(bias, AffineQuantizedTensor):
|
|
return bias.dequantize().data.detach()
|
|
else:
|
|
return self.org_module[0].bias.data.detach()
|
|
return None
|
|
|
|
def _get_delta_factors(self):
|
|
"""(A, w2_or_w2a, w2_b_or_None) without ever combining the two kron
|
|
factors into a full-size matrix."""
|
|
A = self.lokr_w1 if self.use_w1 else self.lokr_w1_a @ self.lokr_w1_b # (out_l, in_m), always small
|
|
if self.use_w2:
|
|
return A, self.lokr_w2, None
|
|
else:
|
|
return A, self.lokr_w2_a, self.lokr_w2_b
|
|
|
|
def _call_forward_fast_linear(self, x):
|
|
"""Factorized LoKr delta for the plain Linear case (quantized or not).
|
|
Never materializes a (out_dim, in_dim) tensor: uses the kron
|
|
mixed-product identity plus the existing low-rank factoring of w2, so
|
|
both compute and the backward-pass gradient are O(rank) instead of
|
|
O(out_dim * in_dim)."""
|
|
if isinstance(x, QTensor) or isinstance(x, QBytesTensor):
|
|
x = x.dequantize()
|
|
orig_dtype = x.dtype
|
|
|
|
if getattr(self.org_module[0], "is_ostris_quantized", False):
|
|
# org_forward here does raw CUDA stream/event orchestration and
|
|
# in-place buffer swapping (manager_modules.py _mm_forward) to
|
|
# move the quantized weight onto the device -- that can't be
|
|
# traced by inductor, so break the graph at exactly this call.
|
|
base_out = torch._dynamo.disable(self.org_forward)(x)
|
|
else:
|
|
base_out = self.org_forward(x)
|
|
|
|
A, w2a, w2b = self._get_delta_factors()
|
|
# match the base path's compute dtype (usually bf16) rather than
|
|
# upcasting the (much larger) activation tensor to the fp32 master
|
|
# params -- keeps every intermediate here at the same footprint the
|
|
# old full-kron GEMM had, instead of doubling it.
|
|
compute_dtype = base_out.dtype
|
|
x_ = x.to(compute_dtype) if x.dtype != compute_dtype else x
|
|
A = A.to(compute_dtype)
|
|
w2a = w2a.to(compute_dtype)
|
|
if w2b is not None:
|
|
w2b = w2b.to(compute_dtype)
|
|
|
|
X = x_.unflatten(-1, (self._in_m, self._in_n))
|
|
|
|
if w2b is None:
|
|
# use_w2: w2a is the full (out_k, in_n) factor -- still far smaller
|
|
# than the full kron product, so a single einsum here is fine.
|
|
tmp = torch.einsum('...qs,os->...qo', X, w2a) # (..., in_m, out_k)
|
|
else:
|
|
# low-rank w2 = w2a @ w2b, rank << out_k, in_n: fold the rank
|
|
# factor through first so we never touch an (out_k, in_n) tensor.
|
|
tmp = torch.einsum('...qs,rs->...qr', X, w2b) # (..., in_m, rank)
|
|
tmp = torch.einsum('...qr,or->...qo', tmp, w2a) # (..., in_m, out_k)
|
|
|
|
# scale folded into A (not applied to the reduction output) to avoid
|
|
# an inductor lowering bug under torch.compile
|
|
delta = torch.einsum('...qo,pq->...po', tmp, A * self._runtime_scale) # (..., out_l, out_k)
|
|
delta = delta.flatten(-2, -1)
|
|
|
|
if self.training and self.rank_dropout:
|
|
# equivalent to the old row-mask on the full weight: zeroing a
|
|
# weight ROW is the same as zeroing that output channel post-matmul.
|
|
drop = torch.rand(delta.size(-1), device=delta.device) < self.rank_dropout
|
|
delta = delta * drop.to(delta.dtype)
|
|
|
|
multiplier = torch.mean(self.network_ref().torch_multiplier).to(compute_dtype)
|
|
delta = delta * multiplier
|
|
|
|
return (base_out + delta).to(orig_dtype)
|
|
|
|
def _call_forward(self, x):
|
|
if self._fast_linear:
|
|
return self._call_forward_fast_linear(x)
|
|
|
|
# legacy path: Conv2d / cp branch. Still materializes the full kron
|
|
# product -- TODO: extend the factorized path to cover these.
|
|
if isinstance(x, QTensor) or isinstance(x, QBytesTensor):
|
|
x = x.dequantize()
|
|
|
|
orig_dtype = x.dtype
|
|
|
|
orig_weight = self.get_orig_weight(x.device)
|
|
lokr_weight = self.get_weight(orig_weight).to(dtype=orig_weight.dtype)
|
|
multiplier = self.network_ref().torch_multiplier
|
|
|
|
if x.dtype != orig_weight.dtype:
|
|
x = x.to(dtype=orig_weight.dtype)
|
|
|
|
# we do not currently support split batch multipliers for lokr. Just do a mean
|
|
multiplier = torch.mean(multiplier)
|
|
|
|
weight = (
|
|
orig_weight
|
|
+ lokr_weight * multiplier
|
|
)
|
|
bias = self.get_orig_bias(x.device)
|
|
if bias is not None:
|
|
bias = bias.to(weight.device, dtype=weight.dtype)
|
|
output = self.op(
|
|
x,
|
|
weight.view(self.shape),
|
|
bias,
|
|
**self.extra_args
|
|
)
|
|
return output.to(orig_dtype)
|