| from functools import cache |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from .base import LycorisBaseModule |
| from ..logging import logger |
|
|
|
|
| @cache |
| def log_bypass_override(): |
| return logger.warning( |
| "Automatic Bypass-Mode detected in algo=full, " |
| "override with bypass_mode=False since algo=full not support bypass mode. " |
| "If you are using quantized model which require bypass mode, please don't use algo=full. " |
| ) |
|
|
|
|
| class FullModule(LycorisBaseModule): |
| name = "full" |
| support_module = { |
| "linear", |
| "conv1d", |
| "conv2d", |
| "conv3d", |
| } |
| weight_list = ["diff", "diff_b"] |
| weight_list_det = ["diff"] |
|
|
| def __init__( |
| self, |
| lora_name, |
| org_module: nn.Module, |
| multiplier=1.0, |
| lora_dim=4, |
| alpha=1, |
| dropout=0.0, |
| rank_dropout=0.0, |
| module_dropout=0.0, |
| use_tucker=False, |
| use_scalar=False, |
| rank_dropout_scale=False, |
| bypass_mode=None, |
| **kwargs, |
| ): |
| org_bypass = bypass_mode |
| super().__init__( |
| lora_name, |
| org_module, |
| multiplier, |
| dropout, |
| rank_dropout, |
| module_dropout, |
| rank_dropout_scale, |
| bypass_mode, |
| ) |
| if bypass_mode and org_bypass is None: |
| self.bypass_mode = False |
| log_bypass_override() |
|
|
| if self.module_type not in self.support_module: |
| raise ValueError(f"{self.module_type} is not supported in Full algo.") |
|
|
| if self.is_quant: |
| raise ValueError( |
| "Quant Linear is not supported and meaningless in Full algo." |
| ) |
|
|
| if self.bypass_mode: |
| raise ValueError("bypass mode is not supported in Full algo.") |
|
|
| self.weight = nn.Parameter(torch.zeros_like(org_module.weight)) |
| if org_module.bias is not None: |
| self.bias = nn.Parameter(torch.zeros_like(org_module.bias)) |
| else: |
| self.bias = None |
| self.is_diff = True |
| self._org_weight = [self.org_module[0].weight.data.cpu().clone()] |
| if self.org_module[0].bias is not None: |
| self.org_bias = [self.org_module[0].bias.data.cpu().clone()] |
| else: |
| self.org_bias = None |
|
|
| @classmethod |
| def make_module_from_state_dict(cls, lora_name, orig_module, diff, diff_b): |
| module = cls( |
| lora_name, |
| orig_module, |
| 1, |
| ) |
| module.weight.copy_(diff) |
| if diff_b is not None: |
| if orig_module.bias is not None: |
| module.bias.copy_(diff_b) |
| else: |
| module.bias = nn.Parameter(diff_b) |
| module.is_diff = True |
| return module |
|
|
| @property |
| def org_weight(self): |
| return self._org_weight[0] |
|
|
| @org_weight.setter |
| def org_weight(self, value): |
| self.org_module[0].weight.data.copy_(value) |
|
|
| def apply_to(self, **kwargs): |
| self.org_forward = self.org_module[0].forward |
| self.org_module[0].forward = self.forward |
| self.weight.data.add_(self.org_module[0].weight.data) |
| self._org_weight = [self.org_module[0].weight.data.cpu().clone()] |
| delattr(self.org_module[0], "weight") |
| if self.org_module[0].bias is not None: |
| self.bias.data.add_(self.org_module[0].bias.data) |
| self.org_bias = [self.org_module[0].bias.data.cpu().clone()] |
| delattr(self.org_module[0], "bias") |
| else: |
| self.org_bias = None |
| self.is_diff = False |
|
|
| def restore(self): |
| self.org_module[0].forward = self.org_forward |
| self.org_module[0].weight = nn.Parameter(self._org_weight[0]) |
| if self.org_bias is not None: |
| self.org_module[0].bias = nn.Parameter(self.org_bias[0]) |
|
|
| def custom_state_dict(self): |
| sd = {"diff": self.weight.data.cpu() - self._org_weight[0]} |
| if self.bias is not None: |
| sd["diff_b"] = self.bias.data.cpu() - self.org_bias[0] |
| return sd |
|
|
| def load_weight_prehook( |
| self, |
| state_dict, |
| prefix, |
| local_metadata, |
| strict, |
| missing_keys, |
| unexpected_keys, |
| error_msgs, |
| ): |
| diff_weight = state_dict.pop(f"{prefix}diff") |
| state_dict[f"{prefix}weight"] = diff_weight + self.weight.data.to(diff_weight) |
| if f"{prefix}diff_b" in state_dict: |
| diff_bias = state_dict.pop(f"{prefix}diff_b") |
| state_dict[f"{prefix}bias"] = diff_bias + self.bias.data.to(diff_bias) |
|
|
| def make_weight(self, scale=1, device=None): |
| drop = ( |
| torch.rand(self.dim, device=device) > self.rank_dropout |
| if self.rank_dropout and self.training |
| else 1 |
| ) |
| if drop != 1 or scale != 1 or self.is_diff: |
| diff_w, diff_b = self.get_diff_weight(scale, device=device) |
| weight = self.org_weight + diff_w * drop |
| if self.org_bias is not None: |
| bias = self.org_bias + diff_b * drop |
| else: |
| bias = None |
| else: |
| weight = self.weight |
| bias = self.bias |
| return weight, bias |
|
|
| def get_diff_weight(self, multiplier=1, shape=None, device=None): |
| if self.is_diff: |
| diff_b = None |
| if self.bias is not None: |
| diff_b = self.bias * multiplier |
| return self.weight * multiplier, diff_b |
| org_weight = self.org_module[0].weight.to(device, dtype=self.weight.dtype) |
| diff = self.weight.to(device) - org_weight |
| diff_b = None |
| if shape: |
| diff = diff.view(shape) |
| if self.bias is not None: |
| org_bias = self.org_module[0].bias.to(device, dtype=self.bias.dtype) |
| diff_b = self.bias.to(device) - org_bias |
| if device is not None: |
| diff = diff.to(device) |
| if self.bias is not None: |
| diff_b = diff_b.to(device) |
| if multiplier != 1: |
| diff = diff * multiplier |
| if diff_b is not None: |
| diff_b = diff_b * multiplier |
| return diff * multiplier, diff_b |
|
|
| def get_merged_weight(self, multiplier=1, shape=None, device=None): |
| weight, bias = self.make_weight(multiplier, device) |
| if shape is not None: |
| weight = weight.view(shape) |
| if bias is not None: |
| bias = bias.view(shape[0]) |
| return weight, bias |
|
|
| def forward(self, x: torch.Tensor, *args, **kwargs): |
| if ( |
| self.module_dropout |
| and self.training |
| and torch.rand(1) < self.module_dropout |
| ): |
| original = True |
| else: |
| original = False |
| if original: |
| return self.org_forward(x) |
| scale = self.multiplier |
| weight, bias = self.make_weight(scale, x.device) |
| kw_dict = self.kw_dict | {"weight": weight, "bias": bias} |
| return self.op(x, **kw_dict) |
|
|
