Buckets:
| # 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 | |
| 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) | |
Xet Storage Details
- Size:
- 17.6 kB
- Xet hash:
- 61ff3ad82a2b07aef90208f9769b4fb6ddb8c840e9fa0bee81c307d4e00ddc59
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.