From aae488bc00f65d35fe7e68a37e1574417c185c95 Mon Sep 17 00:00:00 2001
From: Ethanfel <ethan.fel@ts-pc.fr>
Date: Tue, 31 Mar 2026 23:23:45 +0200
Subject: [PATCH] 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>
---
 lora_optimizer.py            | 50 ++++++++++++++++++++++++++++++++++++
 tests/test_lora_optimizer.py | 11 ++++++++
 2 files changed, 61 insertions(+)

diff --git a/lora_optimizer.py b/lora_optimizer.py
index a3c1b7b..d505f6a 100644
--- a/lora_optimizer.py
+++ b/lora_optimizer.py
@@ -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
diff --git a/tests/test_lora_optimizer.py b/tests/test_lora_optimizer.py
index b9ca727..063b2b5 100644
--- a/tests/test_lora_optimizer.py
+++ b/tests/test_lora_optimizer.py
@@ -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()