215 lines
6.8 KiB
Python
215 lines
6.8 KiB
Python
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)
|