492 lines
20 KiB
Python
492 lines
20 KiB
Python
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from safetensors.torch import load_file
|
|
import logging
|
|
|
|
# --- QLoRA Support (Optional) ---
|
|
try:
|
|
import bitsandbytes as bnb
|
|
except ImportError:
|
|
bnb = None
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ==========================================
|
|
# 1. Helper Functions
|
|
# ==========================================
|
|
|
|
def get_dequantized_weight(module):
|
|
if bnb and hasattr(module, "weight") and hasattr(module.weight, "quant_state"):
|
|
weight = module.weight
|
|
return bnb.functional.dequantize_4bit(weight, weight.quant_state).to(torch.float32)
|
|
elif bnb and module.__class__.__name__ == "Linear8bitLt":
|
|
if hasattr(module, "state") and hasattr(module.state, "SCB"):
|
|
try:
|
|
return module.weight.detach().float()
|
|
except:
|
|
return module.weight.float()
|
|
else:
|
|
return module.weight.float()
|
|
elif hasattr(module, "weight"):
|
|
return module.weight
|
|
return None
|
|
|
|
def compute_null_space_basis(weight, rank, device="cpu"):
|
|
is_conv = weight.dim() == 4
|
|
if is_conv:
|
|
out_ch, in_ch, k1, k2 = weight.shape
|
|
w_flat = weight.reshape(out_ch, -1)
|
|
else:
|
|
w_flat = weight
|
|
|
|
try:
|
|
U, S, Vh = torch.linalg.svd(w_flat.to(device, dtype=torch.float32), full_matrices=False)
|
|
except Exception as e:
|
|
print(f"SVD failed on device {device}, trying CPU. Error: {e}")
|
|
w_flat_cpu = w_flat.to("cpu", dtype=torch.float32)
|
|
U, S, Vh = torch.linalg.svd(w_flat_cpu, full_matrices=False)
|
|
U, S, Vh = U.to(device), S.to(device), Vh.to(device)
|
|
|
|
actual_rank = min(rank, U.shape[1], Vh.shape[0])
|
|
|
|
u_null = U[:, -actual_rank:].clone()
|
|
v_null = Vh[-actual_rank:, :].clone()
|
|
|
|
del U, S, Vh
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
return u_null, v_null
|
|
|
|
# ==========================================
|
|
# 2. NullUniversalModule (Inference Optimized)
|
|
# ==========================================
|
|
|
|
class NullUniversalModule(nn.Module):
|
|
def __init__(
|
|
self,
|
|
lora_name,
|
|
org_module,
|
|
mean_down, basis_down,
|
|
mean_up, basis_up,
|
|
multiplier=1.0,
|
|
total_rank=32,
|
|
loaded_null_down=None,
|
|
loaded_null_up=None,
|
|
loaded_scale=None,
|
|
storage_dtype=torch.float32,
|
|
use_dora=False,
|
|
shared_alpha_down=None,
|
|
num_experts=1,
|
|
):
|
|
super().__init__()
|
|
self.lora_name = lora_name
|
|
self.multiplier = multiplier
|
|
self.org_module = org_module
|
|
self.use_dora = use_dora
|
|
self.num_experts = num_experts
|
|
|
|
# Rank Splitting
|
|
self.r_half = total_rank // 2
|
|
if self.r_half < 1: self.r_half = 1
|
|
|
|
# Architecture Extraction
|
|
if isinstance(org_module, nn.Conv2d) or org_module.__class__.__name__ == "Conv2d":
|
|
self.is_conv = True
|
|
self.in_dim = org_module.in_channels
|
|
self.out_dim = org_module.out_channels
|
|
self.k_size = org_module.kernel_size
|
|
self.stride = org_module.stride
|
|
self.padding = org_module.padding
|
|
self.dilation = org_module.dilation
|
|
self.groups = org_module.groups
|
|
elif isinstance(org_module, nn.Linear) or (hasattr(org_module, "in_features") and hasattr(org_module, "out_features")):
|
|
self.is_conv = False
|
|
self.in_dim = org_module.in_features
|
|
self.out_dim = org_module.out_features
|
|
self.k_size = (1, 1)
|
|
self.stride = (1, 1); self.padding = (0, 0); self.dilation = (1, 1); self.groups = 1
|
|
else:
|
|
# Fallback
|
|
self.is_conv = False
|
|
w = get_dequantized_weight(org_module)
|
|
if w is not None:
|
|
if w.dim() == 4:
|
|
self.is_conv = True
|
|
self.out_dim, in_dim_per_group, k1, k2 = w.shape
|
|
self.groups = getattr(org_module, "groups", 1)
|
|
self.in_dim = in_dim_per_group * self.groups
|
|
self.k_size = (k1, k2)
|
|
self.stride = getattr(org_module, "stride", (1, 1))
|
|
self.padding = getattr(org_module, "padding", (0, 0))
|
|
self.dilation = getattr(org_module, "dilation", (1, 1))
|
|
else:
|
|
self.out_dim, self.in_dim = w.shape[:2]
|
|
self.groups = 1
|
|
self.k_size = (1, 1)
|
|
self.stride = (1, 1); self.padding = (0, 0); self.dilation = (1, 1)
|
|
else:
|
|
raise ValueError(f"Could not infer dimensions for module: {lora_name}")
|
|
|
|
# Cap r_half by maximum possible rank
|
|
if self.is_conv:
|
|
max_rank = min(self.out_dim, (self.in_dim // self.groups) * self.k_size[0] * self.k_size[1])
|
|
else:
|
|
max_rank = min(self.out_dim, self.in_dim)
|
|
|
|
if self.r_half > max_rank:
|
|
self.r_half = max_rank
|
|
|
|
# --- Universal Basis (For Trainable Path) ---
|
|
if mean_down.dim() == 1:
|
|
mean_down_reshaped = mean_down.unsqueeze(0)
|
|
else:
|
|
mean_down_reshaped = mean_down.view(1, -1)
|
|
|
|
basis_down_t = basis_down[:, :self.r_half].t()
|
|
fused_down = torch.cat([basis_down_t, mean_down_reshaped], dim=0)
|
|
|
|
self.register_buffer("fused_down", fused_down.to(storage_dtype).contiguous())
|
|
self.register_buffer("mean_up", mean_up.to(storage_dtype).contiguous())
|
|
self.register_buffer("basis_up", basis_up[:, :self.r_half].to(storage_dtype).contiguous())
|
|
|
|
# --- Trainable Parameters ---
|
|
if shared_alpha_down is not None:
|
|
self.alpha_down = shared_alpha_down
|
|
self.is_alpha_down_shared = True
|
|
else:
|
|
self.alpha_down = nn.Parameter(torch.empty(self.r_half, self.r_half))
|
|
self.is_alpha_down_shared = False
|
|
|
|
self.alpha_up = nn.Parameter(torch.zeros(self.num_experts, self.r_half, self.r_half))
|
|
|
|
if self.num_experts > 1:
|
|
self.gate = nn.Linear(self.in_dim, self.num_experts)
|
|
else:
|
|
self.gate = None
|
|
self.register_buffer("expert_mask", torch.ones(1))
|
|
|
|
# --- Scale ---
|
|
if self.use_dora:
|
|
if loaded_scale is not None and loaded_scale.shape[0] == self.out_dim:
|
|
self.m = nn.Parameter(loaded_scale)
|
|
else:
|
|
self.m = nn.Parameter(torch.ones(self.out_dim))
|
|
self.s = None
|
|
else:
|
|
if loaded_scale is not None:
|
|
if loaded_scale.shape[0] > self.r_half * 2:
|
|
self.s = nn.Parameter(loaded_scale[:self.r_half * 2])
|
|
else:
|
|
self.s = nn.Parameter(loaded_scale)
|
|
else:
|
|
self.s = nn.Parameter(torch.ones(self.r_half * 2))
|
|
|
|
# --- Null Space Basis (Frozen) ---
|
|
loaded_from_cache = False
|
|
if loaded_null_down is not None and loaded_null_up is not None:
|
|
expected_in = self.in_dim if not self.is_conv else ((self.in_dim // self.groups) * self.k_size[0] * self.k_size[1])
|
|
expected_out = self.out_dim
|
|
|
|
if loaded_null_down.shape[1] != expected_in or loaded_null_up.shape[0] != expected_out:
|
|
if loaded_null_down.shape[0] == expected_out and loaded_null_up.shape[1] == expected_in:
|
|
loaded_null_down, loaded_null_up = loaded_null_up, loaded_null_down
|
|
|
|
if loaded_null_down.shape[1] == expected_in and loaded_null_up.shape[0] == expected_out:
|
|
self.register_buffer("null_down", loaded_null_down[:self.r_half, :].to(storage_dtype).contiguous())
|
|
self.register_buffer("null_up", loaded_null_up[:, :self.r_half].to(storage_dtype).contiguous())
|
|
loaded_from_cache = True
|
|
else:
|
|
print(f"Warning: SVD cache mismatch for {lora_name}. Recomputing.")
|
|
|
|
if not loaded_from_cache:
|
|
weight_for_svd = get_dequantized_weight(org_module)
|
|
if weight_for_svd is None:
|
|
raise ValueError(f"Failed to extract weights from {lora_name} for SVD.")
|
|
|
|
# Move to CPU for SVD to save VRAM if needed, or stick to GPU
|
|
device = weight_for_svd.device
|
|
with torch.no_grad():
|
|
u_null, v_null = compute_null_space_basis(weight_for_svd, self.r_half, device=device)
|
|
|
|
self.register_buffer("null_down", v_null.to(storage_dtype).contiguous())
|
|
self.register_buffer("null_up", u_null.to(storage_dtype).contiguous())
|
|
del weight_for_svd
|
|
|
|
def apply_to(self):
|
|
self.org_forward = self.org_module.forward
|
|
self.org_module.forward = self.forward
|
|
|
|
def restore(self):
|
|
if hasattr(self, 'org_forward'):
|
|
self.org_module.forward = self.org_forward
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
return self._forward(x)
|
|
|
|
def _forward(self, x):
|
|
dtype = x.dtype
|
|
device = x.device
|
|
|
|
# --- Cross-Path Architecture (Sandwich) ---
|
|
frozen_down = self.null_down.to(device=device, dtype=dtype)
|
|
frozen_up = self.null_up.to(device=device, dtype=dtype)
|
|
|
|
w_fused = self.fused_down.to(device=device, dtype=dtype)
|
|
a_down = self.alpha_down[:self.r_half, :self.r_half].to(device=device, dtype=dtype)
|
|
|
|
# Prepare Gate
|
|
gate_scores = None
|
|
if self.num_experts > 1:
|
|
if self.is_conv:
|
|
gate_input = x.mean(dim=(2, 3))
|
|
else:
|
|
gate_input = x.mean(dim=1) if x.dim() == 3 else x
|
|
|
|
self.gate.to(device)
|
|
gate_logits = self.gate(gate_input.to(self.gate.weight.dtype))
|
|
gate_scores = F.softmax(gate_logits, dim=-1)
|
|
|
|
# Prepare Scaling
|
|
s1, s2 = 1.0, 1.0
|
|
if not self.use_dora:
|
|
s = self.s.to(device=device, dtype=dtype)
|
|
s1 = s[:self.r_half].mean()
|
|
s2 = s[self.r_half:].mean()
|
|
|
|
# === Path 1: Frozen Down -> Train Up (MoE) ===
|
|
if self.is_conv:
|
|
fd = frozen_down.view(self.r_half, self.in_dim // self.groups, *self.k_size)
|
|
h1 = F.conv2d(x, fd, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)
|
|
else:
|
|
h1 = F.linear(x, frozen_down)
|
|
|
|
if not self.use_dora: h1 = h1 * s1
|
|
|
|
# MoE Mixing
|
|
a_up_experts = self.alpha_up.to(device=device, dtype=dtype)
|
|
if self.num_experts > 1:
|
|
if self.is_conv:
|
|
b, r, h, w = h1.shape
|
|
# gate_scores: (B, E), Experts: (E, R, R) -> (B, R, R)
|
|
mixed_weight = torch.einsum('be, eoi -> boi', gate_scores.to(dtype=dtype), a_up_experts)
|
|
h1_flat = h1.view(b, r, -1).permute(0, 2, 1) # (B, HW, R)
|
|
h1_mixed_flat = torch.bmm(h1_flat, mixed_weight.transpose(1, 2))
|
|
h1_mixed = h1_mixed_flat.permute(0, 2, 1).view(b, r, h, w)
|
|
else:
|
|
mixed_weight = torch.einsum('be, eoi -> boi', gate_scores.to(dtype=dtype), a_up_experts)
|
|
if h1.dim() == 2:
|
|
h1_mixed = torch.bmm(h1.unsqueeze(1), mixed_weight.transpose(1, 2)).squeeze(1)
|
|
else:
|
|
h1_mixed = torch.bmm(h1, mixed_weight.transpose(1, 2))
|
|
else:
|
|
a_up = a_up_experts[0]
|
|
if self.is_conv:
|
|
h1_mixed = F.conv2d(h1, a_up.view(self.r_half, self.r_half, 1, 1))
|
|
else:
|
|
h1_mixed = F.linear(h1, a_up)
|
|
|
|
# Basis Projection
|
|
b_up = self.basis_up.to(device=device, dtype=dtype)
|
|
m_up = self.mean_up.to(device=device, dtype=dtype)
|
|
|
|
if self.is_conv:
|
|
path1 = F.conv2d(h1_mixed, b_up.view(self.out_dim, self.r_half, 1, 1))
|
|
h1_sum = h1.sum(dim=1, keepdim=True)
|
|
path1_bias = F.conv2d(h1_sum, m_up.view(self.out_dim, 1, 1, 1))
|
|
path1 = path1 + path1_bias
|
|
else:
|
|
path1 = F.linear(h1_mixed, b_up)
|
|
h1_sum = h1.sum(dim=-1, keepdim=True)
|
|
path1_bias = F.linear(h1_sum, m_up.view(-1, 1))
|
|
path1 = path1 + path1_bias
|
|
|
|
|
|
# === Path 2: Train Down -> Frozen Up ===
|
|
if self.is_conv:
|
|
h_fused = F.conv2d(x, w_fused.view(self.r_half + 1, self.in_dim // self.groups, *self.k_size), stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)
|
|
h_in = h_fused.narrow(1, 0, self.r_half)
|
|
bias_down = h_fused.narrow(1, self.r_half, 1)
|
|
h2 = F.conv2d(h_in, a_down.view(self.r_half, self.r_half, 1, 1))
|
|
h2 = h2.add_(bias_down)
|
|
else:
|
|
h_fused = F.linear(x, w_fused)
|
|
if x.dim() == 2:
|
|
h_in = h_fused.narrow(1, 0, self.r_half)
|
|
bias_down = h_fused.narrow(1, self.r_half, 1)
|
|
else:
|
|
h_in = h_fused.narrow(2, 0, self.r_half)
|
|
bias_down = h_fused.narrow(2, self.r_half, 1)
|
|
h2 = F.linear(h_in, a_down)
|
|
h2 = h2.add_(bias_down)
|
|
|
|
if not self.use_dora: h2 = h2 * s2
|
|
|
|
if self.is_conv:
|
|
fu = frozen_up.view(self.out_dim, self.r_half, 1, 1)
|
|
path2 = F.conv2d(h2, fu)
|
|
else:
|
|
path2 = F.linear(h2, frozen_up)
|
|
|
|
lx = path1 + path2
|
|
|
|
if self.use_dora:
|
|
m = self.m.to(device=device, dtype=dtype)
|
|
if self.is_conv:
|
|
lx = lx * m.view(1, -1, 1, 1)
|
|
else:
|
|
lx = lx * m
|
|
|
|
org_out = self.org_forward(x)
|
|
return org_out + lx * self.multiplier
|
|
|
|
|
|
class NullUniversalNetwork(nn.Module):
|
|
def __init__(self, unet, multiplier=1.0, use_dora=False, enable_effilora=False, num_experts=4):
|
|
super().__init__()
|
|
self.multiplier = multiplier
|
|
self.use_dora = use_dora
|
|
self.enable_effilora = enable_effilora
|
|
self.num_experts = num_experts if enable_effilora else 1
|
|
self.modules_dict = nn.ModuleDict()
|
|
self.unet = unet
|
|
self.shared_alpha_down = None
|
|
|
|
def load_basis_and_init(self, basis_path, weights_path=None, total_rank=32, storage_dtype=torch.float32):
|
|
if not os.path.exists(basis_path):
|
|
raise FileNotFoundError(f"Basis file not found: {basis_path}")
|
|
|
|
print(f"Loading Basis from: {basis_path}")
|
|
basis_state_dict = load_file(basis_path)
|
|
|
|
if self.enable_effilora:
|
|
r_half = total_rank // 2
|
|
self.shared_alpha_down = nn.Parameter(torch.empty(r_half, r_half))
|
|
print(f"EffiLoRA Enabled: Shared Alpha Down (Rank={r_half}, Experts={self.num_experts})")
|
|
|
|
loaded_weights = {}
|
|
if weights_path and os.path.exists(weights_path):
|
|
if weights_path.endswith(".safetensors"):
|
|
loaded_weights = load_file(weights_path)
|
|
else:
|
|
loaded_weights = torch.load(weights_path, map_location="cpu")
|
|
|
|
keys = set()
|
|
for k in basis_state_dict.keys():
|
|
if k.endswith(".basis"):
|
|
keys.add(k.rsplit(".", 1)[0])
|
|
|
|
layer_groups = {}
|
|
for k in keys:
|
|
if "lora_down" in k:
|
|
base_name = k.split(".lora_down")[0]
|
|
if base_name not in layer_groups: layer_groups[base_name] = {}
|
|
layer_groups[base_name]["down"] = k
|
|
elif "lora_up" in k:
|
|
base_name = k.split(".lora_up")[0]
|
|
if base_name not in layer_groups: layer_groups[base_name] = {}
|
|
layer_groups[base_name]["up"] = k
|
|
|
|
count = 0
|
|
for base_name, parts in layer_groups.items():
|
|
if "down" not in parts or "up" not in parts: continue
|
|
|
|
# Skip Text Encoder
|
|
if "te1" in base_name or "te2" in base_name or "text_encoder" in base_name:
|
|
continue
|
|
|
|
target_module = self.find_target_module(base_name)
|
|
if target_module is not None:
|
|
mean_d = basis_state_dict[parts["down"] + ".mean"]
|
|
basis_d = basis_state_dict[parts["down"] + ".basis"]
|
|
mean_u = basis_state_dict[parts["up"] + ".mean"]
|
|
basis_u = basis_state_dict[parts["up"] + ".basis"]
|
|
|
|
safe_name = base_name.replace(".", "_")
|
|
|
|
def find_key(dct, suffix):
|
|
k1 = f"{safe_name}.{suffix}"
|
|
k2 = f"modules_dict.{safe_name}.{suffix}"
|
|
return dct.get(k1, dct.get(k2, None))
|
|
|
|
l_null_d, l_null_u, l_scale = None, None, None
|
|
if loaded_weights:
|
|
l_null_d = find_key(loaded_weights, "null_down")
|
|
l_null_u = find_key(loaded_weights, "null_up")
|
|
if self.use_dora: l_scale = find_key(loaded_weights, "m")
|
|
else: l_scale = find_key(loaded_weights, "s")
|
|
|
|
u_module = NullUniversalModule(
|
|
safe_name, target_module,
|
|
mean_d, basis_d, mean_u, basis_u,
|
|
multiplier=self.multiplier,
|
|
total_rank=total_rank,
|
|
loaded_null_down=l_null_d,
|
|
loaded_null_up=l_null_u,
|
|
loaded_scale=l_scale,
|
|
storage_dtype=storage_dtype,
|
|
use_dora=self.use_dora,
|
|
shared_alpha_down=self.shared_alpha_down if self.enable_effilora else None,
|
|
num_experts=self.num_experts,
|
|
)
|
|
|
|
if loaded_weights:
|
|
ad = find_key(loaded_weights, "alpha_down")
|
|
au = find_key(loaded_weights, "alpha_up")
|
|
gate_w = find_key(loaded_weights, "gate.weight")
|
|
gate_b = find_key(loaded_weights, "gate.bias")
|
|
|
|
if self.enable_effilora:
|
|
shared_ad = loaded_weights.get("shared_alpha_down", None)
|
|
if shared_ad is not None: self.shared_alpha_down.data.copy_(shared_ad)
|
|
elif ad is not None and count == 0: self.shared_alpha_down.data.copy_(ad)
|
|
else:
|
|
if ad is not None: u_module.alpha_down.data.copy_(ad)
|
|
|
|
if au is not None:
|
|
if au.shape == u_module.alpha_up.shape: u_module.alpha_up.data.copy_(au)
|
|
elif au.dim() == 2 and u_module.alpha_up.dim() == 3: u_module.alpha_up.data[0].copy_(au)
|
|
|
|
if u_module.gate is not None and gate_w is not None:
|
|
u_module.gate.weight.data.copy_(gate_w)
|
|
if gate_b is not None: u_module.gate.bias.data.copy_(gate_b)
|
|
|
|
self.modules_dict[safe_name] = u_module
|
|
count += 1
|
|
|
|
del basis_state_dict
|
|
if loaded_weights: del loaded_weights
|
|
if torch.cuda.is_available(): torch.cuda.empty_cache()
|
|
print(f"Applied Null-Universal LoRA to {count} modules.")
|
|
|
|
def find_target_module(self, lora_key):
|
|
search_key = lora_key
|
|
root = None
|
|
if search_key.startswith("lora_unet_"):
|
|
root = self.unet
|
|
search_key = search_key.replace("lora_unet_", "")
|
|
elif search_key.startswith("unet."):
|
|
root = self.unet
|
|
search_key = search_key.replace("unet.", "")
|
|
|
|
if root is None: return None
|
|
target_name_converted = search_key.replace(".", "_")
|
|
for name, module in root.named_modules():
|
|
if name.replace(".", "_") == target_name_converted: return module
|
|
return None
|
|
|
|
def apply_to(self):
|
|
for module in self.modules_dict.values():
|
|
module.apply_to()
|
|
|
|
def set_multiplier(self, multiplier):
|
|
for module in self.modules_dict.values():
|
|
module.multiplier = multiplier
|
|
|
|
def restore(self):
|
|
for module in self.modules_dict.values():
|
|
module.restore()
|
|
|