99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
#based off https://github.com/catid/dora/blob/main/dora.py
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from typing import TYPE_CHECKING, Union, List
|
|
|
|
from toolkit.network_mixins import ToolkitModuleMixin, ExtractableModuleMixin
|
|
|
|
if TYPE_CHECKING:
|
|
from toolkit.lora_special import LoRASpecialNetwork
|
|
|
|
# diffusers specific stuff
|
|
LINEAR_MODULES = [
|
|
'Linear',
|
|
'LoRACompatibleLinear'
|
|
# 'GroupNorm',
|
|
]
|
|
CONV_MODULES = [
|
|
'Conv2d',
|
|
'LoRACompatibleConv'
|
|
]
|
|
|
|
|
|
class DoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
|
|
# def __init__(self, d_in, d_out, rank=4, weight=None, bias=None):
|
|
def __init__(
|
|
self,
|
|
lora_name,
|
|
org_module: torch.nn.Module,
|
|
multiplier=1.0,
|
|
lora_dim=4,
|
|
alpha=1,
|
|
dropout=None,
|
|
rank_dropout=None,
|
|
module_dropout=None,
|
|
network: 'LoRASpecialNetwork' = None,
|
|
use_bias: bool = False,
|
|
**kwargs
|
|
):
|
|
self.can_merge_in = False
|
|
"""if alpha == 0 or None, alpha is rank (no scaling)."""
|
|
ToolkitModuleMixin.__init__(self, network=network)
|
|
torch.nn.Module.__init__(self)
|
|
self.lora_name = lora_name
|
|
self.scalar = torch.tensor(1.0)
|
|
|
|
self.lora_dim = lora_dim
|
|
|
|
if org_module.__class__.__name__ in CONV_MODULES:
|
|
raise NotImplementedError("Convolutional layers are not supported yet")
|
|
|
|
if type(alpha) == torch.Tensor:
|
|
alpha = alpha.detach().float().numpy() # without casting, bf16 causes error
|
|
alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
|
|
self.scale = alpha / self.lora_dim
|
|
# self.register_buffer("alpha", torch.tensor(alpha)) # 定数として扱える eng: treat as constant
|
|
|
|
self.multiplier: Union[float, List[float]] = multiplier
|
|
# wrap the original module so it doesn't get weights updated
|
|
self.org_module = [org_module]
|
|
self.dropout = dropout
|
|
self.rank_dropout = rank_dropout
|
|
self.module_dropout = module_dropout
|
|
self.is_checkpointing = False
|
|
|
|
# m = Magnitude column-wise across output dimension
|
|
self.magnitude = nn.Parameter(self.get_orig_weight().norm(p=2, dim=0, keepdim=True))
|
|
|
|
d_out = org_module.out_features
|
|
d_in = org_module.in_features
|
|
|
|
std_dev = 1 / torch.sqrt(torch.tensor(self.lora_dim).float())
|
|
self.lora_up = nn.Parameter(torch.randn(d_out, self.lora_dim) * std_dev)
|
|
self.lora_down = nn.Parameter(torch.zeros(self.lora_dim, d_in))
|
|
|
|
def apply_to(self):
|
|
self.org_forward = self.org_module[0].forward
|
|
self.org_module[0].forward = self.forward
|
|
# del self.org_module
|
|
|
|
def get_orig_weight(self):
|
|
return self.org_module[0].weight.data.detach()
|
|
|
|
def get_orig_bias(self):
|
|
if hasattr(self.org_module[0], 'bias') and self.org_module[0].bias is not None:
|
|
return self.org_module[0].bias.data.detach()
|
|
return None
|
|
|
|
def dora_forward(self, x, *args, **kwargs):
|
|
lora = torch.matmul(self.lora_up, self.lora_down)
|
|
adapted = self.get_orig_weight() + lora
|
|
column_norm = adapted.norm(p=2, dim=0, keepdim=True)
|
|
norm_adapted = adapted / column_norm
|
|
calc_weights = self.magnitude * norm_adapted
|
|
return F.linear(x, calc_weights, self.get_orig_bias())
|
|
|