99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
# utils.py - REMOVE apply_density function (keep it only in methods.py)
|
|
import argparse
|
|
import torch
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
# ---------- Rank Dimension Detection ----------
|
|
def detect_rank_dimensions(shape, key):
|
|
"""
|
|
Smart detection of which dimensions are rank dimensions.
|
|
Uses shape heuristics first, key names as fallback.
|
|
"""
|
|
if len(shape) != 2:
|
|
return [] # Not a standard 2D LoRA weight
|
|
|
|
# Heuristic 1: Rank is usually the smaller dimension
|
|
# and model dimensions are typically >= 768
|
|
if shape[0] <= 256 and shape[1] >= 512:
|
|
return [0] # rank is dim 0
|
|
elif shape[1] <= 256 and shape[0] >= 512:
|
|
return [1] # rank is dim 1
|
|
|
|
# Heuristic 2: Very rectangular tensors (rank << model dim)
|
|
if shape[1] >= 4 * shape[0]:
|
|
return [0] # rank is dim 0
|
|
elif shape[0] >= 4 * shape[1]:
|
|
return [1] # rank is dim 1
|
|
|
|
# Fallback to key name hints
|
|
key_lower = key.lower()
|
|
if "down" in key_lower or "_a" in key_lower or ".a" in key_lower:
|
|
return [0] # Usually rank is dimension 0
|
|
elif "up" in key_lower or "_b" in key_lower or ".b" in key_lower:
|
|
return [1] # Usually rank is dimension 1
|
|
|
|
return [] # Couldn't determine
|
|
|
|
|
|
# ---------- Rank Projection ----------
|
|
def random_projection(out_rank, in_rank, device):
|
|
"""Create orthonormal projection matrix."""
|
|
mat = torch.randn(out_rank, in_rank, device=device)
|
|
q, _ = torch.linalg.qr(mat)
|
|
return q
|
|
|
|
|
|
def project_rank(t1, t2, key):
|
|
"""
|
|
Match tensor shapes via projection for rank dimensions,
|
|
padding for non-rank dimensions.
|
|
"""
|
|
s1, s2 = list(t1.shape), list(t2.shape)
|
|
|
|
if s1 == s2:
|
|
return t1, t2
|
|
|
|
max_dims = [max(a, b) for a, b in zip(s1, s2)]
|
|
|
|
def expand_tensor(t, target_shape, key):
|
|
result = t
|
|
|
|
# Get rank dimensions for this tensor
|
|
rank_dims = detect_rank_dimensions(t.shape, key)
|
|
|
|
for dim in range(len(target_shape)):
|
|
if result.shape[dim] < target_shape[dim]:
|
|
diff = target_shape[dim] - result.shape[dim]
|
|
|
|
# Check if this is a rank dimension
|
|
if dim in rank_dims:
|
|
# Project rank instead of padding (better quality)
|
|
proj = random_projection(
|
|
target_shape[dim],
|
|
result.shape[dim],
|
|
result.device,
|
|
)
|
|
|
|
if dim == 0:
|
|
# lora_down style: [rank, in_dim]
|
|
result = proj @ result
|
|
elif dim == 1:
|
|
# lora_up style: [out_dim, rank]
|
|
result = result @ proj.T
|
|
else:
|
|
# Higher dimensions (rare) - use padding
|
|
pad_shape = list(result.shape)
|
|
pad_shape[dim] = diff
|
|
padding = torch.zeros(pad_shape, device=result.device)
|
|
result = torch.cat([result, padding], dim=dim)
|
|
else:
|
|
# Non-rank dimension: use zero padding
|
|
pad_shape = list(result.shape)
|
|
pad_shape[dim] = diff
|
|
padding = torch.zeros(pad_shape, device=result.device)
|
|
result = torch.cat([result, padding], dim=dim)
|
|
|
|
return result
|
|
|
|
return expand_tensor(t1, max_dims, key), expand_tensor(t2, max_dims, key) |