From 7f8e5035e91b60df6297b8d116ffbbbb6678ecb0 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Wed, 24 Jun 2026 01:41:25 +0200 Subject: [PATCH] perf(merge): skip dense+compression SVD for single-LoRA groups MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Layers only one LoRA touches have no conflict to sparsify and no second diff to align, so TIES/DARE/refinement are no-ops on them. Previously, turning on sparsification or refinement disabled the exact low-rank fast path for EVERY group — so single-LoRA layers needlessly materialized a full dense diff and ran a compression SVD (the 99%-of-Pass-2 cost on the Krea 2 softwatercolor merge, where each layer only one LoRA touches paid a large GPU SVD). Take the exact low-rank fast path for single-LoRA groups regardless of those toggles, emitting the LoRA's native factors directly. Size-safe: a single LoRA's compress_rank is max(rank, 64) >= its native rank, so compression never shrank it anyway (only padded rank<64 up to 64) — bypassing is always equal-or-smaller, and peak VRAM drops too (no dense tensor is built). Multi-LoRA groups are unchanged. No version bump. Co-Authored-By: Claude Opus 4.8 --- lora_optimizer.py | 16 ++++++-- tests/test_lora_optimizer.py | 77 ++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 3 deletions(-) diff --git a/lora_optimizer.py b/lora_optimizer.py index efd7896..1fd4e65 100644 --- a/lora_optimizer.py +++ b/lora_optimizer.py @@ -7701,11 +7701,21 @@ class LoRAOptimizer(_LoRAMergeBase): _group_patch_cache["misses"] += 1 linear_stats = None + # Single-LoRA groups (layers only one LoRA touches) have no conflict + # to sparsify and no second diff to align, so TIES/DARE/refinement + # are no-ops on them. Take the exact low-rank fast path for them even + # when those global toggles are on — this skips the wasteful + # dense-materialize + compression SVD and emits the LoRA's native + # factors directly. Size-safe: a single LoRA's compress_rank is + # max(rank, 64) >= its native rank, so compression never shrinks it + # (only pads rank<64 ones to 64) — bypassing is always equal-or-smaller. + _single_lora_group = pf_n_loras == 1 # exactly one contributor + _linear_quality_ok = (sparsification == "disabled" + and merge_refinement == "none") if (pf_mode in ("weighted_sum", "weighted_average", "normalize") - and sparsification == "disabled" - and merge_refinement == "none" and not has_virtual_loras - and not stack_has_preserve): + and not stack_has_preserve + and (_single_lora_group or _linear_quality_ok)): _t_lin = _prof_t() if _merge_prof is not None else 0.0 linear_patch_info = self._build_exact_linear_patch( target_group, active_loras, raw_n, pf_mode, diff --git a/tests/test_lora_optimizer.py b/tests/test_lora_optimizer.py index 7db2e34..8c0f6c5 100644 --- a/tests/test_lora_optimizer.py +++ b/tests/test_lora_optimizer.py @@ -4213,6 +4213,83 @@ class TestGlobalModeMergesWithDeclaredMode(unittest.TestCase): self.assertEqual(decisions["alias_b"], "weighted_sum", decisions) +@unittest.skipIf(torch is None, "torch is not installed in this environment") +class TestSingleLoraSkipsCompression(unittest.TestCase): + """Single-LoRA groups (layers only one LoRA touches) take the exact + low-rank fast path even when sparsification/refinement is on — those are + conflict/multi-diff ops that are no-ops on a lone LoRA. This avoids the + wasteful dense-materialize + compression SVD and emits native factors + (a LoRAAdapter), while genuine multi-LoRA groups still go dense+compress.""" + + def _run(self, **kwargs): + class FakePatcher: + def __init__(self, model): + self.model = model + + def clone(self): + return FakePatcher(self.model) + + def add_patches(self, patches, strength=1.0, strength_clip=None): + return list(patches.keys()) + + inner = types.SimpleNamespace( + layer=types.SimpleNamespace(weight=torch.zeros(32, 32)), + layer2=types.SimpleNamespace(weight=torch.zeros(32, 32))) + model = FakePatcher(inner) + + def make_lora(seed, scale, prefixes): + g = torch.Generator().manual_seed(seed) + d = {} + for prefix in prefixes: + d[f"{prefix}.lora_up.weight"] = torch.randn(32, 4, generator=g) * scale + d[f"{prefix}.lora_down.weight"] = torch.randn(4, 32, generator=g) + d[f"{prefix}.alpha"] = torch.tensor(4.0) + return d + + stack = [ + # alias_a is shared (2 LoRAs -> conflict); alias_b is single-LoRA + {"name": "A", "lora": make_lora(1, 0.1, ("alias_a", "alias_b")), "strength": 1.0}, + {"name": "B", "lora": make_lora(2, 0.1, ("alias_a",)), "strength": 0.8}, + ] + opt = lora_optimizer.LoRAOptimizer() + opt._get_model_keys = lambda m: {"alias_a": "layer.weight", "alias_b": "layer2.weight"} + _, _, _, _, lora_data = opt.optimize_merge( + model, stack, 1.0, cache_patches="disabled", **kwargs) + patches = {} + for k, v in lora_data["model_patches"].items(): + patches[k[0] if isinstance(k, tuple) else k] = v + return patches + + def _is_native_rank4(self, patch): + # The fast path emits the LoRA's native rank-4 factors untouched. + # The dense+compress path (DARE breaks the low-rank structure, then SVD + # re-fits) lands at a higher rank or stays a dense ("diff",) tuple. + return (isinstance(patch, lora_optimizer.LoRAAdapter) + and patch.weights[1].shape[0] == 4) + + def test_single_lora_skips_compression_under_sparsification(self): + # sparsification on + compression on: the single-LoRA group must STILL + # take the low-rank fast path (no dense diff, no sparsification, no + # compression SVD) — proven by its native rank 4 surviving intact; the + # multi-LoRA group must NOT (it genuinely needs deconfliction). + patches = self._run(sparsification="dare", sparsification_density=0.9, + patch_compression="smart") + self.assertTrue(self._is_native_rank4(patches["layer2.weight"]), + "single-LoRA group should emit native rank-4 factors") + self.assertFalse(self._is_native_rank4(patches["layer.weight"]), + "multi-LoRA group should still go dense+compress") + + def test_single_lora_lowrank_path_is_size_safe(self): + # The bypassed single-LoRA patch is stored at the LoRA's native rank (4), + # never padded up to the rank-64 compression floor. + patches = self._run(sparsification="dare", sparsification_density=0.9, + patch_compression="smart") + adapter = patches["layer2.weight"] + mat_up, mat_down = adapter.weights[0], adapter.weights[1] + self.assertEqual(mat_down.shape[0], 4) # native rank, not 64 + self.assertEqual(mat_up.shape[1], 4) + + @unittest.skipIf(torch is None, "torch is not installed in this environment") class TestBatchedKarcher(unittest.TestCase): """The batched Karcher implementation must match the per-unit reference