From 16e493cfdfa31d8870c454e9c85e76f8fe22c4d3 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Thu, 13 Aug 2026 16:36:13 +0200 Subject: [PATCH] perf(lora): bound knee SVD probes --- nodes/dora_extract_wd.py | 85 +++++++++++++++++++-------- nodes/dora_learned_wd.py | 23 ++++++-- nodes/lora_extract_svd.py | 85 +++++++++++++++++++-------- nodes/text_encoder_extract.py | 24 +++++--- tests/test_knee_partial_svd.py | 101 +++++++++++++++++++++++++++++++++ 5 files changed, 259 insertions(+), 59 deletions(-) create mode 100644 tests/test_knee_partial_svd.py diff --git a/nodes/dora_extract_wd.py b/nodes/dora_extract_wd.py index b3039b2..db5fcc5 100644 --- a/nodes/dora_extract_wd.py +++ b/nodes/dora_extract_wd.py @@ -213,13 +213,34 @@ def _svd_extract_linear_lowrank( pass # Construct LoRA matrices - lora_up = U @ torch.diag(S) # [out_dim, rank] + lora_up = U * S.unsqueeze(0) # [out_dim, rank] lora_down = Vh # [rank, in_dim] - # Compute reconstruction diff - diff = weight - (lora_up @ lora_down) + return (lora_down, lora_up, None), "low rank" - return (lora_down, lora_up, diff), "low rank" + +def _svd_extract_knee_lowrank( + weight: torch.Tensor, + mode: str, + max_rank: int, + probe_offset: int, + niter: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Probe beyond the output rank and retry once when the knee hits the probe tail.""" + max_possible_rank = min(weight.shape) + output_cap = max(1, min(max_rank, max_possible_rank)) + offset = max(1, probe_offset) + probe_rank = min(max_possible_rank, output_cap + offset) + probe_limit = min(max_possible_rank, max(probe_rank, output_cap * 2)) + + while True: + U, S, V = torch.svd_lowrank(weight, q=probe_rank, niter=niter) + detected_rank = _compute_rank(S, mode, 0, None) + near_probe_tail = detected_rank >= max(1, probe_rank - offset) + if not near_probe_tail or probe_rank >= probe_limit: + rank = min(detected_rank, output_cap) + return U[:, :rank], S[:rank], V[:, :rank].T, rank + probe_rank = probe_limit def _svd_extract_linear( @@ -230,6 +251,7 @@ def _svd_extract_linear( max_rank: int = None, clamp_quantile: float = 0.99, niter: int = 2, + knee_probe_offset: int = 32, ) -> tuple: """ SVD decomposition for linear layers. @@ -250,14 +272,18 @@ def _svd_extract_linear( target_rank = min(target_rank, max_rank) return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter) - # For adaptive modes, use full SVD to analyze singular values - try: - U, S, Vh = linalg.svd(weight, full_matrices=False) - except Exception as e: - raise RuntimeError(f"SVD failed: {e}") + if mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None: + U, S, Vh, rank = _svd_extract_knee_lowrank( + weight, mode, max_rank, knee_probe_offset, niter + ) + else: + # Other adaptive modes require the full spectrum under their current contract. + try: + U, S, Vh = linalg.svd(weight, full_matrices=False) + except Exception as e: + raise RuntimeError(f"SVD failed: {e}") - # Compute rank based on mode - rank = _compute_rank(S, mode, mode_param, max_rank) + rank = _compute_rank(S, mode, mode_param, max_rank) # Check if decomposition is worthwhile if rank >= min(out_dim, in_dim): @@ -278,13 +304,10 @@ def _svd_extract_linear( # Construct LoRA matrices # lora_up = U @ diag(S), lora_down = Vh - lora_up = U @ torch.diag(S) # [out_dim, rank] + lora_up = U * S.unsqueeze(0) # [out_dim, rank] lora_down = Vh # [rank, in_dim] - # Compute reconstruction diff for optional sparse bias - diff = weight - (lora_up @ lora_down) - - return (lora_down, lora_up, diff), "low rank" + return (lora_down, lora_up, None), "low rank" def _svd_extract_conv( @@ -294,6 +317,7 @@ def _svd_extract_conv( device: str, max_rank: int = None, clamp_quantile: float = 0.99, + knee_probe_offset: int = 32, ) -> tuple: """ SVD decomposition for conv2d layers. @@ -310,7 +334,10 @@ def _svd_extract_conv( else: mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k] - result, mode_str = _svd_extract_linear(mat, mode, mode_param, device, max_rank, clamp_quantile) + result, mode_str = _svd_extract_linear( + mat, mode, mode_param, device, max_rank, clamp_quantile, + knee_probe_offset=knee_probe_offset, + ) if mode_str == "full": return weight, "full" @@ -363,6 +390,7 @@ def _extract_chunked_layer( mode_param: float, device: str, max_rank: int = None, + knee_probe_offset: int = 32, ) -> tuple: """Extract LoRA from fused layer by chunking.""" out_dim, in_dim = weight_diff.shape @@ -375,7 +403,10 @@ def _extract_chunked_layer( chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :] try: - result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank) + result, mode_str = _svd_extract_linear( + chunk, mode, mode_param, device, max_rank, + knee_probe_offset=knee_probe_offset, + ) if mode_str == "full": return None, None, 0 lora_down, lora_up, _ = result @@ -458,6 +489,7 @@ def extract_dora_from_files( lazy_load: bool = True, force_clear_cache: bool = True, glob_skip_patterns: bool = False, + knee_probe_offset: int = 32, ) -> None: """ Extract LoRA from difference between two models, writing incrementally to disk. @@ -602,11 +634,13 @@ def extract_dora_from_files( try: if is_conv: result, mode_str = _svd_extract_conv( - weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile + weight_diff, mode, conv_param, device, conv_max_rank, + clamp_quantile, knee_probe_offset ) else: result, mode_str = _svd_extract_linear( - weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter + weight_diff, mode, linear_param, device, linear_max_rank, + clamp_quantile, svd_niter, knee_probe_offset ) except Exception as e: # Try chunked extraction for large tensors @@ -615,7 +649,8 @@ def extract_dora_from_files( if num_chunks > 1: print(f"[DoRA Extract] Chunked: {key} ({num_chunks} chunks)") lora_up, lora_down, rank = _extract_chunked_layer( - weight_diff, num_chunks, mode, linear_param, device, linear_max_rank + weight_diff, num_chunks, mode, linear_param, device, + linear_max_rank, knee_probe_offset ) if lora_up is not None: layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() @@ -913,6 +948,8 @@ class DoRAExtractKnee(io.ComfyNode): *_get_model_inputs(), io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Knee detection method"), + io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, + tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."), io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum rank for linear layers"), io.Int.Input("conv_max_rank", default=128, min=1, max=16384, @@ -924,7 +961,8 @@ class DoRAExtractKnee(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, knee_method, knee_probe_offset, + linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: @@ -937,7 +975,8 @@ class DoRAExtractKnee(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns + glob_skip_patterns=glob_skip_patterns, + knee_probe_offset=knee_probe_offset, ) return io.NodeOutput(output_path) diff --git a/nodes/dora_learned_wd.py b/nodes/dora_learned_wd.py index 39981c2..e773c30 100644 --- a/nodes/dora_learned_wd.py +++ b/nodes/dora_learned_wd.py @@ -276,6 +276,7 @@ def _extract_chunked_learned_layer( early_stop_stall: int = 2000, early_stop_lr: float = 9.01e-9, layer_name: str = "Layer", + knee_probe_offset: int = 32, ) -> tuple: """Extract Learned DoRA from fused layer by chunking.""" out_dim, in_dim = tensor_base.shape @@ -296,7 +297,10 @@ def _extract_chunked_learned_layer( weight_diff = w_dir - chunk_base try: - result, mode_str = _svd_extract_linear(weight_diff, mode, mode_param, device, max_rank) + result, mode_str = _svd_extract_linear( + weight_diff, mode, mode_param, device, max_rank, + knee_probe_offset=knee_probe_offset, + ) if mode_str == "full": return None, None, 0 lora_down, lora_up, _ = result @@ -351,6 +355,7 @@ def extract_dora_learned_from_files( lazy_load: bool = True, force_clear_cache: bool = True, glob_skip_patterns: bool = False, + knee_probe_offset: int = 32, ) -> None: """ Extract Learned DoRA from difference between two models. @@ -465,11 +470,13 @@ def extract_dora_learned_from_files( # 1. Run Analytical SVD (Initialization) if is_conv: result, mode_str = _svd_extract_conv( - weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile + weight_diff, mode, conv_param, device, conv_max_rank, + clamp_quantile, knee_probe_offset ) else: result, mode_str = _svd_extract_linear( - weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter + weight_diff, mode, linear_param, device, linear_max_rank, + clamp_quantile, svd_niter, knee_probe_offset ) # 2. Run Optimization (Learned Rounding) @@ -492,7 +499,8 @@ def extract_dora_learned_from_files( tensor_base, tensor_ft, num_chunks, mode, linear_param, device, linear_max_rank, optimize_iters, learning_rate, optimizer_type, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, - layer_name=lora_name + layer_name=lora_name, + knee_probe_offset=knee_probe_offset, ) if lora_up is not None: layer_results[f"{lora_name}.lora_up.weight"] = lora_up.to(save_torch_dtype).cpu().contiguous() @@ -719,6 +727,7 @@ class DoRALearnedExtractKnee(io.ComfyNode): inputs=[ *_get_model_inputs(), io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."), + io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."), io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."), io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."), *_get_learned_inputs(), @@ -729,7 +738,8 @@ class DoRALearnedExtractKnee(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, knee_method, knee_probe_offset, + linear_max_rank, conv_max_rank, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: @@ -741,7 +751,8 @@ class DoRALearnedExtractKnee(io.ComfyNode): model_a_path, model_b_path, knee_method, 0, 0, device, save_dtype, output_path, optimize_iters, learning_rate, optimizer, lr_schedule, lr_patience, lr_factor, lr_cooldown, early_stop_loss, early_stop_stall, early_stop_lr, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + lazy_load=lazy_load, force_clear_cache=force_clear_cache, + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset ) return io.NodeOutput(output_path) diff --git a/nodes/lora_extract_svd.py b/nodes/lora_extract_svd.py index 594af99..0825eaa 100644 --- a/nodes/lora_extract_svd.py +++ b/nodes/lora_extract_svd.py @@ -214,13 +214,34 @@ def _svd_extract_linear_lowrank( pass # Construct LoRA matrices - lora_up = U @ torch.diag(S) # [out_dim, rank] + lora_up = U * S.unsqueeze(0) # [out_dim, rank] lora_down = Vh # [rank, in_dim] - # Compute reconstruction diff - diff = weight - (lora_up @ lora_down) + return (lora_down, lora_up, None), "low rank" - return (lora_down, lora_up, diff), "low rank" + +def _svd_extract_knee_lowrank( + weight: torch.Tensor, + mode: str, + max_rank: int, + probe_offset: int, + niter: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Probe beyond the output rank and retry once when the knee hits the probe tail.""" + max_possible_rank = min(weight.shape) + output_cap = max(1, min(max_rank, max_possible_rank)) + offset = max(1, probe_offset) + probe_rank = min(max_possible_rank, output_cap + offset) + probe_limit = min(max_possible_rank, max(probe_rank, output_cap * 2)) + + while True: + U, S, V = torch.svd_lowrank(weight, q=probe_rank, niter=niter) + detected_rank = _compute_rank(S, mode, 0, None) + near_probe_tail = detected_rank >= max(1, probe_rank - offset) + if not near_probe_tail or probe_rank >= probe_limit: + rank = min(detected_rank, output_cap) + return U[:, :rank], S[:rank], V[:, :rank].T, rank + probe_rank = probe_limit def _svd_extract_linear( @@ -231,6 +252,7 @@ def _svd_extract_linear( max_rank: int = None, clamp_quantile: float = 0.99, niter: int = 2, + knee_probe_offset: int = 32, ) -> tuple: """ SVD decomposition for linear layers. @@ -251,14 +273,18 @@ def _svd_extract_linear( target_rank = min(target_rank, max_rank) return _svd_extract_linear_lowrank(weight, target_rank, device, clamp_quantile, niter) - # For adaptive modes, use full SVD to analyze singular values - try: - U, S, Vh = linalg.svd(weight, full_matrices=False) - except Exception as e: - raise RuntimeError(f"SVD failed: {e}") + if mode in {"sv_knee", "sv_cumulative_knee"} and max_rank is not None: + U, S, Vh, rank = _svd_extract_knee_lowrank( + weight, mode, max_rank, knee_probe_offset, niter + ) + else: + # Other adaptive modes require the full spectrum under their current contract. + try: + U, S, Vh = linalg.svd(weight, full_matrices=False) + except Exception as e: + raise RuntimeError(f"SVD failed: {e}") - # Compute rank based on mode - rank = _compute_rank(S, mode, mode_param, max_rank) + rank = _compute_rank(S, mode, mode_param, max_rank) # Check if decomposition is worthwhile if rank >= min(out_dim, in_dim): @@ -279,13 +305,10 @@ def _svd_extract_linear( # Construct LoRA matrices # lora_up = U @ diag(S), lora_down = Vh - lora_up = U @ torch.diag(S) # [out_dim, rank] + lora_up = U * S.unsqueeze(0) # [out_dim, rank] lora_down = Vh # [rank, in_dim] - # Compute reconstruction diff for optional sparse bias - diff = weight - (lora_up @ lora_down) - - return (lora_down, lora_up, diff), "low rank" + return (lora_down, lora_up, None), "low rank" def _svd_extract_conv( @@ -295,6 +318,7 @@ def _svd_extract_conv( device: str, max_rank: int = None, clamp_quantile: float = 0.99, + knee_probe_offset: int = 32, ) -> tuple: """ SVD decomposition for conv2d layers. @@ -311,7 +335,10 @@ def _svd_extract_conv( else: mat = weight.reshape(out_ch, -1) # [out_ch, in_ch*k*k] - result, mode_str = _svd_extract_linear(mat, mode, mode_param, device, max_rank, clamp_quantile) + result, mode_str = _svd_extract_linear( + mat, mode, mode_param, device, max_rank, clamp_quantile, + knee_probe_offset=knee_probe_offset, + ) if mode_str == "full": return weight, "full" @@ -364,6 +391,7 @@ def _extract_chunked_layer( mode_param: float, device: str, max_rank: int = None, + knee_probe_offset: int = 32, ) -> tuple: """Extract LoRA from fused layer by chunking.""" out_dim, in_dim = weight_diff.shape @@ -376,7 +404,10 @@ def _extract_chunked_layer( chunk = weight_diff[i * chunk_size:(i + 1) * chunk_size, :] try: - result, mode_str = _svd_extract_linear(chunk, mode, mode_param, device, max_rank) + result, mode_str = _svd_extract_linear( + chunk, mode, mode_param, device, max_rank, + knee_probe_offset=knee_probe_offset, + ) if mode_str == "full": return None, None, 0 lora_down, lora_up, _ = result @@ -460,6 +491,7 @@ def extract_lora_from_files( force_clear_cache: bool = True, glob_skip_patterns: bool = False, include_1d_diffs: bool = False, + knee_probe_offset: int = 32, ) -> None: """ Extract LoRA from difference between two models, writing incrementally to disk. @@ -606,11 +638,13 @@ def extract_lora_from_files( try: if is_conv: result, mode_str = _svd_extract_conv( - weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile + weight_diff, mode, conv_param, device, conv_max_rank, + clamp_quantile, knee_probe_offset ) else: result, mode_str = _svd_extract_linear( - weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter + weight_diff, mode, linear_param, device, linear_max_rank, + clamp_quantile, svd_niter, knee_probe_offset ) except Exception as e: # Try chunked extraction for large tensors @@ -619,7 +653,8 @@ def extract_lora_from_files( if num_chunks > 1: print(f"[LoRA Extract] Chunked: {key} ({num_chunks} chunks)") lora_up, lora_down, rank = _extract_chunked_layer( - weight_diff, num_chunks, mode, linear_param, device, linear_max_rank + weight_diff, num_chunks, mode, linear_param, device, + linear_max_rank, knee_probe_offset ) if lora_up is not None: layer_results[f"{lora_name}.lora_B.weight"] = lora_up.to(layer_save_dtype).cpu().contiguous() @@ -922,6 +957,8 @@ class LoRAExtractKnee(io.ComfyNode): *_get_model_inputs(), io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Knee detection method"), + io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, + tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."), io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum rank for linear layers"), io.Int.Input("conv_max_rank", default=128, min=1, max=16384, @@ -933,7 +970,8 @@ class LoRAExtractKnee(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, knee_method, knee_probe_offset, + linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache, include_1d_diffs=False) -> io.NodeOutput: @@ -947,7 +985,8 @@ class LoRAExtractKnee(io.ComfyNode): device, save_dtype, output_path, linear_max_rank, conv_max_rank, clamp_quantile, min_diff, skip_patterns, mismatch_mode, chunk_large_layers, lazy_load=lazy_load, force_clear_cache=force_clear_cache, - glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs + glob_skip_patterns=glob_skip_patterns, include_1d_diffs=include_1d_diffs, + knee_probe_offset=knee_probe_offset, ) return io.NodeOutput(output_path) diff --git a/nodes/text_encoder_extract.py b/nodes/text_encoder_extract.py index 020433b..8c82375 100644 --- a/nodes/text_encoder_extract.py +++ b/nodes/text_encoder_extract.py @@ -82,6 +82,7 @@ def extract_te_from_files( lazy_load: bool = True, force_clear_cache: bool = True, glob_skip_patterns: bool = False, + knee_probe_offset: int = 32, ) -> None: """ Extract LoRA/DoRA from difference between two Text Encoder models, writing incrementally to disk. @@ -207,11 +208,13 @@ def extract_te_from_files( try: if is_conv: result, mode_str = _svd_extract_conv( - weight_diff, mode, conv_param, device, conv_max_rank, clamp_quantile + weight_diff, mode, conv_param, device, conv_max_rank, + clamp_quantile, knee_probe_offset ) else: result, mode_str = _svd_extract_linear( - weight_diff, mode, linear_param, device, linear_max_rank, clamp_quantile, svd_niter + weight_diff, mode, linear_param, device, linear_max_rank, + clamp_quantile, svd_niter, knee_probe_offset ) except Exception as e: # Try chunked extraction for large tensors @@ -220,7 +223,8 @@ def extract_te_from_files( if num_chunks > 1: print(f"[TE Extract] Chunked: {key} ({num_chunks} chunks)") lora_up, lora_down, rank = _extract_chunked_layer( - weight_diff, num_chunks, mode, linear_param, device, linear_max_rank + weight_diff, num_chunks, mode, linear_param, device, + linear_max_rank, knee_probe_offset ) if lora_up is not None: # PEFT suffixes @@ -451,6 +455,7 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode): inputs=[ *_get_te_model_inputs(), io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."), + io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."), io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."), io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."), *_get_common_inputs(), @@ -460,7 +465,8 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, knee_method, knee_probe_offset, + linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: @@ -474,7 +480,8 @@ class TextEncoderLoRAExtractKnee(io.ComfyNode): linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + lazy_load=lazy_load, force_clear_cache=force_clear_cache, + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset ) return io.NodeOutput(output_path) @@ -653,6 +660,7 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode): inputs=[ *_get_te_model_inputs(), io.Combo.Input("knee_method", options=["sv_knee", "sv_cumulative_knee"], default="sv_knee", tooltip="Detect the knee from raw singular values or their cumulative distribution."), + io.Int.Input("knee_probe_offset", default=32, min=1, max=4096, tooltip="Extra singular values probed beyond Max Rank to avoid detecting a false knee at the partial-spectrum boundary."), io.Int.Input("linear_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for linear layers."), io.Int.Input("conv_max_rank", default=128, min=1, max=16384, tooltip="Maximum extracted rank for convolution layers."), *_get_common_inputs(), @@ -662,7 +670,8 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode): ) @classmethod - def execute(cls, model_a, model_b, knee_method, linear_max_rank, conv_max_rank, + def execute(cls, model_a, model_b, knee_method, knee_probe_offset, + linear_max_rank, conv_max_rank, chunk_large_layers, clamp_quantile, min_diff, mismatch_mode, output_filename, save_dtype, device, skip_patterns, glob_skip_patterns, lazy_load, force_clear_cache) -> io.NodeOutput: @@ -676,7 +685,8 @@ class TextEncoderDoRAExtractKnee(io.ComfyNode): linear_max_rank=linear_max_rank, conv_max_rank=conv_max_rank, clamp_quantile=clamp_quantile, min_diff=min_diff, skip_patterns_str=skip_patterns, mismatch_mode=mismatch_mode, chunk_large_layers=chunk_large_layers, - lazy_load=lazy_load, force_clear_cache=force_clear_cache, glob_skip_patterns=glob_skip_patterns + lazy_load=lazy_load, force_clear_cache=force_clear_cache, + glob_skip_patterns=glob_skip_patterns, knee_probe_offset=knee_probe_offset ) return io.NodeOutput(output_path) diff --git a/tests/test_knee_partial_svd.py b/tests/test_knee_partial_svd.py new file mode 100644 index 0000000..87458a5 --- /dev/null +++ b/tests/test_knee_partial_svd.py @@ -0,0 +1,101 @@ +import importlib +import sys +import types +from pathlib import Path + +import pytest +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[1] +PACKAGE_NAME = "modelutils_knee_partial_svd_tests" + + +def _load_module(name): + if PACKAGE_NAME not in sys.modules: + package = types.ModuleType(PACKAGE_NAME) + package.__path__ = [str(REPO_ROOT)] + sys.modules[PACKAGE_NAME] = package + return importlib.import_module(f"{PACKAGE_NAME}.{name}") + + +@pytest.fixture(scope="module", params=["nodes.lora_extract_svd", "nodes.dora_extract_wd"]) +def extraction(request): + return _load_module(request.param) + + +def _fake_factors(weight, q, niter): + del niter + return ( + torch.eye(weight.shape[0], q), + torch.linspace(float(q), 1.0, q), + torch.eye(weight.shape[1], q), + ) + + +def test_knee_probe_uses_offset_and_expands_at_boundary(monkeypatch, extraction): + probes = [] + detected = iter([9, 5]) + + def fake_svd(weight, q, niter): + probes.append(q) + return _fake_factors(weight, q, niter) + + monkeypatch.setattr(extraction.torch, "svd_lowrank", fake_svd) + monkeypatch.setattr( + extraction, "_compute_rank", lambda singular_values, mode, param, cap: next(detected) + ) + + _, _, _, rank = extraction._svd_extract_knee_lowrank( + torch.eye(20), "sv_knee", max_rank=6, probe_offset=4, niter=2 + ) + + assert probes == [10, 12] + assert rank == 5 + + +def test_knee_probe_stops_when_knee_is_clear(monkeypatch, extraction): + probes = [] + + def fake_svd(weight, q, niter): + probes.append(q) + return _fake_factors(weight, q, niter) + + monkeypatch.setattr(extraction.torch, "svd_lowrank", fake_svd) + monkeypatch.setattr(extraction, "_compute_rank", lambda *args: 3) + + _, _, _, rank = extraction._svd_extract_knee_lowrank( + torch.eye(20), "sv_knee", max_rank=6, probe_offset=4, niter=2 + ) + + assert probes == [10] + assert rank == 3 + + +def test_lowrank_factor_construction_does_not_build_residual(monkeypatch, extraction): + monkeypatch.setattr(extraction.torch, "svd_lowrank", _fake_factors) + + result, mode = extraction._svd_extract_linear_lowrank( + torch.eye(6), rank=3, device="cpu", clamp_quantile=1.0, niter=2 + ) + down, up, residual = result + + assert mode == "low rank" + assert residual is None + assert torch.equal(up, torch.eye(6, 3) * torch.tensor([3.0, 2.0, 1.0])) + assert torch.equal(down, torch.eye(3, 6)) + + +def test_knee_control_is_adjacent_to_method(): + modules_and_nodes = [ + (_load_module("nodes.lora_extract_svd"), "LoRAExtractKnee"), + (_load_module("nodes.dora_extract_wd"), "DoRAExtractKnee"), + (_load_module("nodes.dora_learned_wd"), "DoRALearnedExtractKnee"), + (_load_module("nodes.text_encoder_extract"), "TextEncoderLoRAExtractKnee"), + (_load_module("nodes.text_encoder_extract"), "TextEncoderDoRAExtractKnee"), + ] + + for module, node_name in modules_and_nodes: + names = [item.id for item in getattr(module, node_name).define_schema().inputs] + method_index = names.index("knee_method") + assert names[method_index + 1] == "knee_probe_offset"