From ae7f742e08f8ef32927592fedcddf49f4644fe20 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Fri, 4 Sep 2026 23:20:46 +0200 Subject: [PATCH] feat(lora): add factorized Delta CWB output --- docs/cwb_delta_lora_mergers.md | 5 ++ nodes/cwb_delta_lora_merger.py | 77 +++++++++++++++++++++++++---- tests/test_cwb_delta_lora_merger.py | 58 +++++++++++++++++++--- 3 files changed, 125 insertions(+), 15 deletions(-) diff --git a/docs/cwb_delta_lora_mergers.md b/docs/cwb_delta_lora_mergers.md index bf27f5c..d90bd33 100644 --- a/docs/cwb_delta_lora_mergers.md +++ b/docs/cwb_delta_lora_mergers.md @@ -13,6 +13,11 @@ Missing-layer handling can preserve LoRA 1, include explicit zero contributors, or abort. Excluded layers preserve LoRA 1 and discarded layers are omitted. Outputs contain only canonical `.diff` tensors with alpha already normalized. +`factorized_lora` output factorizes each completed merged layer immediately and +writes canonical `lora_A` and `lora_B` tensors at the selected maximum rank. It +never writes or retains a complete full-difference model. One-dimensional +differences remain `.diff` tensors. + Processing streams one logical layer from each input at a time and writes each completed layer incrementally. CUDA out-of-memory retries only the failed layer on CPU. The destination is replaced only after the complete output is written. diff --git a/nodes/cwb_delta_lora_merger.py b/nodes/cwb_delta_lora_merger.py index 0ec844c..f1cf247 100644 --- a/nodes/cwb_delta_lora_merger.py +++ b/nodes/cwb_delta_lora_merger.py @@ -147,6 +147,29 @@ def _factor_delta(down, up, layer): ) +def _factorize_delta(delta, rank): + if delta.ndim not in (2, 4): + return None + shape = delta.shape + matrix = delta if delta.ndim == 2 else delta.reshape(shape[0], -1) + target_rank = max(1, min(int(rank), min(matrix.shape))) + if target_rank == min(matrix.shape): + if matrix.shape[0] <= matrix.shape[1]: + up = torch.eye(matrix.shape[0], device=matrix.device, dtype=matrix.dtype) + down = matrix + else: + up = matrix + down = torch.eye(matrix.shape[1], device=matrix.device, dtype=matrix.dtype) + else: + up, singular_values, right = torch.svd_lowrank(matrix, q=target_rank, niter=2) + up = up * singular_values.unsqueeze(0) + down = right.t() + if delta.ndim == 4: + down = down.reshape(down.shape[0], *shape[1:]) + up = up.unsqueeze(-1).unsqueeze(-1) + return down, up + + def _inspect_inputs(handlers): infos = [] for index, handler in enumerate(handlers): @@ -281,7 +304,8 @@ def _to_processing_device(tensor, device): def _process_layer_on_device( - unit, loaded, settings: CWBSettings, save_dtype, device, diagnostics + unit, loaded, settings: CWBSettings, save_dtype, device, diagnostics, + output_representation, output_rank, ): deltas = [] try: @@ -321,7 +345,15 @@ def _process_layer_on_device( save_dtype, is_1d_diff=merged.ndim == 1, ) - return merged.to(dtype=dtype, device="cpu").contiguous(), dtype + if output_representation == "factorized_lora": + factors = _factorize_delta(merged, output_rank) + if factors is not None: + down, up = factors + return [ + (canonical_lora_key(unit.core, "down"), down.to(dtype=dtype, device="cpu").contiguous()), + (canonical_lora_key(unit.core, "up"), up.to(dtype=dtype, device="cpu").contiguous()), + ], dtype + return [(unit.output_key, merged.to(dtype=dtype, device="cpu").contiguous())], dtype finally: deltas.clear() @@ -338,12 +370,16 @@ def _release_cuda_oom(): torch.cuda.empty_cache() -def _process_layer(unit, loaded, settings, save_dtype, requested_device, diagnostics): +def _process_layer( + unit, loaded, settings, save_dtype, requested_device, diagnostics, + output_representation="full_difference", output_rank=384, +): device = torch.device(requested_device) attempt_diagnostics = CWBDiagnostics() try: output, dtype = _process_layer_on_device( - unit, loaded, settings, save_dtype, device, attempt_diagnostics + unit, loaded, settings, save_dtype, device, attempt_diagnostics, + output_representation, output_rank, ) diagnostics.absorb_runtime(attempt_diagnostics) return output, dtype, False @@ -353,7 +389,8 @@ def _process_layer(unit, loaded, settings, save_dtype, requested_device, diagnos _release_cuda_oom() retry_diagnostics = CWBDiagnostics() output, dtype = _process_layer_on_device( - unit, loaded, settings, save_dtype, torch.device("cpu"), retry_diagnostics + unit, loaded, settings, save_dtype, torch.device("cpu"), retry_diagnostics, + output_representation, output_rank, ) diagnostics.absorb_runtime(retry_diagnostics) return output, dtype, True @@ -370,6 +407,8 @@ class DeltaCWBLoRAMergerLogic: mismatch_mode, output_filename, save_dtype, + output_representation, + output_rank, process_device, exclude_patterns, discard_patterns, @@ -416,7 +455,10 @@ class DeltaCWBLoRAMergerLogic: "merge_method": "delta_cwb", "source_models": ",".join(source_names), "alpha_normalized": "true", - "output_representation": "full_direct_difference", + "output_representation": ( + "factorized_lora" if output_representation == "factorized_lora" + else "full_direct_difference" + ), "cwb_consensus_type": settings.consensus_type, "cwb_alignment_method": settings.alignment_method, "cwb_alignment_threshold": str(settings.alignment_threshold), @@ -442,21 +484,23 @@ class DeltaCWBLoRAMergerLogic: loaded[(source.source_index, key)] = cursors[ source.source_index ].take(key) - output, dtype, used_cpu = _process_layer( + output_items, dtype, used_cpu = _process_layer( unit, loaded, settings, save_dtype, process_device, diagnostics, + output_representation, + output_rank, ) - writer.write_batch([(unit.output_key, output)]) + writer.write_batch(output_items) report.merged_layers += int(not unit.preserve_anchor) report.anchor_preserved_layers += int(unit.preserve_anchor) report.zero_filled_contributors += unit.zero_contributors report.cuda_cpu_fallbacks += int(used_cpu) report.record_dtype(dtype) - del output + del output_items finally: for source_index, key in list(loaded): del loaded[(source_index, key)] @@ -496,6 +540,8 @@ class DeltaCWBLoRAMergerLogic: kwargs["mismatch_mode"], kwargs["output_filename"], requested_dtype, + kwargs["output_representation"], + kwargs["output_rank"], kwargs["process_device"], kwargs["exclude_patterns"], kwargs["discard_patterns"], @@ -551,6 +597,19 @@ def _fixed_inputs(count, default_filename): default="bf16", tooltip="Requested saved precision. Participating FP32 inputs preserve FP32.", ), + io.Combo.Input( + "output_representation", + options=["full_difference", "factorized_lora"], + default="full_difference", + tooltip="Save exact full differences or factorize each merged layer directly into compact LoRA factors.", + ), + io.Int.Input( + "output_rank", + default=384, + min=1, + max=16384, + tooltip="Maximum rank used when Output Representation is factorized_lora.", + ), io.Combo.Input( "process_device", options=["cuda", "cpu"], diff --git a/tests/test_cwb_delta_lora_merger.py b/tests/test_cwb_delta_lora_merger.py index 6afcf3e..7c09d98 100644 --- a/tests/test_cwb_delta_lora_merger.py +++ b/tests/test_cwb_delta_lora_merger.py @@ -160,6 +160,8 @@ def test_streamed_merge_normalizes_alpha_once_and_supports_mixed_rank( "skip", "dense/mixed_rank", torch.float32, + "full_difference", + 384, "cpu", "", "", @@ -196,7 +198,8 @@ def test_existing_diff_and_1d_output_remain_fp32(monkeypatch, tmp_path, delta_cw filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute( [str(first), str(second)], [first.name, second.name], _settings(delta_cwb), - "balanced_mean", False, "skip", "one_dim", torch.float16, "cpu", + "balanced_mean", False, "skip", "one_dim", torch.float16, + "full_difference", 384, "cpu", "", "", False, False, ) output = _read(tmp_path / "loras" / filename)[key] @@ -204,6 +207,42 @@ def test_existing_diff_and_1d_output_remain_fp32(monkeypatch, tmp_path, delta_cw torch.testing.assert_close(output, torch.tensor([2.0, 4.0])) +def test_factorized_output_writes_compact_factors_without_dense_tensor( + monkeypatch, tmp_path, delta_cwb +): + first = tmp_path / "factor_first.safetensors" + second = tmp_path / "factor_second.safetensors" + tensors = { + "diffusion_model.layer.lora_A.weight": torch.tensor([[1.0, 2.0]]), + "diffusion_model.layer.lora_B.weight": torch.tensor([[3.0], [6.0]]), + } + _write(first, tensors) + _write(second, tensors) + monkeypatch.setattr(delta_cwb.folder_paths, "models_dir", str(tmp_path)) + monkeypatch.setattr(delta_cwb, "prepare_for_large_operation", lambda *args: None) + monkeypatch.setattr(delta_cwb, "cleanup_after_operation", lambda: None) + + filename, _ = delta_cwb.DeltaCWBLoRAMergerLogic.execute( + [str(first), str(second)], [first.name, second.name], _settings(delta_cwb), + "balanced_mean", False, "skip", "factorized", torch.float32, + "factorized_lora", 1, "cpu", "", "", False, False, + ) + output_path = tmp_path / "loras" / filename + output = _read(output_path) + assert set(output) == { + "diffusion_model.layer.lora_A.weight", + "diffusion_model.layer.lora_B.weight", + } + assert "diffusion_model.layer.diff" not in output + reconstructed = ( + output["diffusion_model.layer.lora_B.weight"] + @ output["diffusion_model.layer.lora_A.weight"] + ) + torch.testing.assert_close(reconstructed, torch.tensor([[3.0, 6.0], [6.0, 12.0]])) + with MemoryEfficientSafeOpen(str(output_path), low_memory=True) as handler: + assert handler.metadata()["output_representation"] == "factorized_lora" + + def test_atomic_writer_preserves_destination_on_cwb_failure( monkeypatch, tmp_path, delta_cwb ): @@ -250,7 +289,8 @@ def test_atomic_writer_preserves_destination_on_cwb_failure( with pytest.raises(RuntimeError, match="planned failure"): delta_cwb.DeltaCWBLoRAMergerLogic.execute( [str(first), str(second)], [first.name, second.name], _settings(delta_cwb), - "balanced_mean", False, "skip", "atomic", torch.float32, "cpu", + "balanced_mean", False, "skip", "atomic", torch.float32, + "full_difference", 384, "cpu", "", "", False, False, ) assert destination.read_bytes() == b"existing destination" @@ -268,7 +308,10 @@ def test_complete_layer_cuda_oom_retries_on_cpu(monkeypatch, delta_cwb): ) calls = [] - def process(_unit, _loaded, _settings, _dtype, device, _diagnostics): + def process( + _unit, _loaded, _settings, _dtype, device, _diagnostics, + _output_representation, _output_rank, + ): calls.append(device.type) if device.type == "cuda": raise torch.cuda.OutOfMemoryError("CUDA out of memory") @@ -305,7 +348,8 @@ def test_exact_node_schema_contract(delta_cwb): fixed_ids = [value.id for value in schemas[0].inputs] assert fixed_ids == [ "lora_1", "lora_2", "cwb_preset", "cwb_config", "mismatch_mode", - "output_filename", "save_dtype", "process_device", "exclude_patterns", + "output_filename", "save_dtype", "output_representation", "output_rank", + "process_device", "exclude_patterns", "discard_patterns", "glob_patterns", "force_clear_cache", ] multi_ids = [value.id for value in schemas[2].inputs] @@ -313,11 +357,13 @@ def test_exact_node_schema_contract(delta_cwb): assert schemas[0].inputs[2].default == "balanced_mean" assert schemas[0].inputs[4].default == "skip" assert schemas[0].inputs[6].default == "bf16" - assert schemas[0].inputs[7].default == "cuda" + assert schemas[0].inputs[7].default == "full_difference" + assert schemas[0].inputs[8].default == 384 + assert schemas[0].inputs[9].default == "cuda" assert schemas[0].inputs[-2].default is False assert schemas[0].inputs[-1].default is True assert "agreement" in schemas[0].inputs[2].tooltip - assert "Only a failed CUDA layer retries on CPU" in schemas[0].inputs[7].tooltip + assert "Only a failed CUDA layer retries on CPU" in schemas[0].inputs[9].tooltip def test_invalid_regex_and_unselected_multi_input_fail(delta_cwb):