feat: add _extract_lora_svd helper for LoRA extraction
Adds module-level _extract_lora_svd() that SVD-decomposes a weight delta into (lora_up, lora_down, alpha), supporting fixed-rank and auto-rank (energy threshold) modes with a singular value noise floor. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
02daae172e
commit
aae488bc00
@@ -3965,6 +3965,56 @@ class LoRAMergeFormula:
|
||||
return (output,)
|
||||
|
||||
|
||||
def _extract_lora_svd(delta: torch.Tensor, rank: int, rank_mode: str, energy_threshold: float):
|
||||
"""
|
||||
SVD-decompose a weight delta into (lora_up, lora_down, alpha).
|
||||
|
||||
Returns None if the delta is near-zero (layer unaffected by the LoRA).
|
||||
Returns (lora_up, lora_down, alpha) otherwise.
|
||||
|
||||
lora_up shape: (rows, r)
|
||||
lora_down shape: (r, cols)
|
||||
alpha = float(r) — effective scale = 1.0 when loaded at strength 1.0
|
||||
|
||||
Note: in auto mode, energy_threshold is measured against the energy of
|
||||
above-floor singular values only, not total delta energy.
|
||||
"""
|
||||
NEAR_ZERO_NORM = 1e-8
|
||||
SV_FLOOR_RATIO = 1e-4 # drop singular values below floor_ratio * S[0]
|
||||
|
||||
if delta.norm().item() < NEAR_ZERO_NORM:
|
||||
return None
|
||||
|
||||
# Ensure 2D: conv layers arrive as [C_out, C_in, kH, kW]
|
||||
if delta.ndim > 2:
|
||||
delta = delta.reshape(delta.shape[0], -1)
|
||||
|
||||
delta = delta.float()
|
||||
|
||||
U, S, Vh = torch.linalg.svd(delta, full_matrices=False)
|
||||
|
||||
# Apply singular value floor to exclude noise
|
||||
sv_floor = SV_FLOOR_RATIO * S[0].item()
|
||||
valid = S > sv_floor
|
||||
if not valid.any():
|
||||
return None
|
||||
U, S, Vh = U[:, valid], S[valid], Vh[valid, :]
|
||||
|
||||
# Determine rank r
|
||||
if rank_mode == "auto":
|
||||
energy_cumsum = S.pow(2).cumsum(0) / S.pow(2).sum()
|
||||
r = int((energy_cumsum < energy_threshold).sum().item()) + 1
|
||||
r = min(r, S.shape[0])
|
||||
else:
|
||||
r = min(rank, S.shape[0])
|
||||
|
||||
sqrt_s = S[:r].sqrt()
|
||||
lora_up = U[:, :r] * sqrt_s.unsqueeze(0) # (rows, r)
|
||||
lora_down = sqrt_s.unsqueeze(1) * Vh[:r, :] # (r, cols)
|
||||
|
||||
return lora_up, lora_down, float(r)
|
||||
|
||||
|
||||
class LoRAOptimizer(_LoRAMergeBase):
|
||||
"""
|
||||
Auto-optimizer that analyzes a LoRA stack (sign conflicts, magnitude
|
||||
|
||||
@@ -1148,5 +1148,16 @@ class TestExtractLoRAFromDelta(unittest.TestCase):
|
||||
self.assertLessEqual(alpha, 2.0)
|
||||
|
||||
|
||||
def test_conv_layer_reshaped_to_2d(self):
|
||||
"""4D conv delta is reshaped to 2D; lora_down has flat spatial dimension."""
|
||||
mod = _load_module()
|
||||
# Simulate a conv2d weight delta: (C_out=64, C_in=32, kH=3, kW=3)
|
||||
delta = torch.randn(64, 32, 3, 3)
|
||||
up, down, alpha = mod._extract_lora_svd(delta, rank=4, rank_mode="fixed", energy_threshold=0.99)
|
||||
self.assertEqual(up.shape[0], 64) # rows = C_out
|
||||
self.assertEqual(down.shape[1], 32 * 3 * 3) # cols = C_in * kH * kW
|
||||
self.assertEqual(up.shape[1], down.shape[0]) # inner dim consistent
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user