From c84ff102cbd60eff54743a9d2887e7c97a13f55c Mon Sep 17 00:00:00 2001 From: silveroxides Date: Thu, 13 Aug 2026 07:27:27 +0200 Subject: [PATCH] fix(lora): preserve per-role dtypes --- nodes/lora_merger.py | 60 +++++++++++++++++++++++--------- tests/test_lora_merger_async.py | 61 +++++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 17 deletions(-) diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index a6be4a1..751a76a 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -108,17 +108,20 @@ def _build_layer_map(lora_infos: List[Dict], base_model_path: Optional[str]) -> return layer_map -def _logical_layer_dtypes(indices, lora_infos, handlers): - source_dtypes = [] +def _logical_output_dtypes(indices, lora_infos, handlers): + pair_dtypes = [] + direct_dtypes = {role: [] for role in _DIRECT_MERGE_ROLES} for idx, orig_block in indices: block_keys = lora_infos[idx]["pairs"][orig_block] - for name in ( - "down", "up", "alpha", "mid", "reshape", "dora_scale", - "diff", "diff_b", "w_norm", "b_norm", "set_weight", - ): + for name in ("down", "up"): if name in block_keys: - source_dtypes.append(handlers[idx].get_dtype(block_keys[name])) - return source_dtypes + pair_dtypes.append(handlers[idx].get_dtype(block_keys[name])) + for role in _DIRECT_MERGE_ROLES: + if role in block_keys: + direct_dtypes[role].append( + handlers[idx].get_dtype(block_keys[role]) + ) + return pair_dtypes, direct_dtypes def _merge_direct_weighted(tensors): @@ -236,14 +239,16 @@ def _build_merge_work_units( units.append(unit) continue + pair_dtypes, direct_dtypes = _logical_output_dtypes( + indices, lora_infos, handlers + ) unit = { "kind": "merge", "core": core, "indices": list(indices), "entries": [], - "source_dtypes": _logical_layer_dtypes( - indices, lora_infos, handlers - ), + "pair_source_dtypes": pair_dtypes, + "direct_source_dtypes": direct_dtypes, } for idx, block_name in indices: block_keys = lora_infos[idx]["pairs"][block_name] @@ -610,7 +615,9 @@ def _process_merge_unit_on_device( ) output_rank = maximum_rank - layer_dtype = select_output_dtype(unit["source_dtypes"], save_dtype) + layer_dtype = select_output_dtype( + unit["pair_source_dtypes"], save_dtype + ) outputs.extend([ ( canonical_lora_key(unit["core"], "down"), @@ -655,7 +662,7 @@ def _process_merge_unit_on_device( generator, ) direct_dtype = select_output_dtype( - unit["source_dtypes"], + unit["direct_source_dtypes"][role], save_dtype, is_1d_diff=merged_direct.ndim == 1, ) @@ -754,7 +761,8 @@ def _run_multi_lora_merge( if verbose: print( f"[{operation}] LoRA {idx + 1} ({info['name']}): " - f"{len(pairs)} layers, ranks={_format_rank_distribution(ranks)}, " + f"{len(pairs)} target groups, " + f"factor-pair ranks={_format_rank_distribution(ranks)}, " f"format={info['format']['format']}" ) @@ -764,6 +772,17 @@ def _run_multi_lora_merge( units, stream_keys = _build_merge_work_units( layer_map, lora_infos, handlers, include_1d_diffs ) + if verbose: + coverage = Counter(len(indices) for indices in layer_map.values()) + coverage_text = ", ".join( + f"{contributors}-input={count}" + for contributors, count in sorted(coverage.items()) + ) + print( + f"[{operation}] Merge plan: {len(layer_map)} unique target groups; " + f"contributor coverage: {coverage_text or 'none'}; " + f"execution work units: {len(units)}" + ) cursors = [ _AsyncTensorCursor( handler, @@ -792,7 +811,11 @@ def _run_multi_lora_merge( ) as writer: with torch.no_grad(): for work_index, unit in enumerate( - tqdm(units, desc=f"Merging layers ({strategy})", unit="units") + tqdm( + units, + desc=f"Merging target groups ({strategy})", + unit="units", + ) ): if unit["kind"] == "raw_copy": for entry in unit["entries"]: @@ -856,9 +879,12 @@ def _run_multi_lora_merge( for cursor in cursors: cursor.finish() if verbose: - print(f"[{operation}] Output ranks: {_format_rank_distribution(output_ranks)}") + print( + f"[{operation}] Output factor-pair ranks: " + f"{_format_rank_distribution(output_ranks)}" + ) print(f"[{operation}] CUDA OOM CPU fallbacks: {fallback_count}") - print(f"[{operation}] Output state dict has {tensor_count} tensors") + print(f"[{operation}] Output file contains {tensor_count} tensors") print(f"[{operation}] Saved merged LoRA to {output_path}") return output_path finally: diff --git a/tests/test_lora_merger_async.py b/tests/test_lora_merger_async.py index 4528789..76f9510 100644 --- a/tests/test_lora_merger_async.py +++ b/tests/test_lora_merger_async.py @@ -233,6 +233,67 @@ def test_nonconcatenate_variants_apply_source_alpha_once( ) +@pytest.mark.parametrize( + "strategy", ["concatenate", "weighted_sum", "dare", "enhanced_dare"] +) +def test_fp32_alpha_does_not_promote_bf16_factors( + monkeypatch, tmp_path, merger, strategy +): + _patch_runtime(monkeypatch, tmp_path, merger) + first = tmp_path / f"bf16_alpha_first_{strategy}.safetensors" + second = tmp_path / f"bf16_alpha_second_{strategy}.safetensors" + tensors = { + "diffusion_model.layer.lora_A.weight": torch.ones( + (1, 2), dtype=torch.bfloat16 + ), + "diffusion_model.layer.lora_B.weight": torch.ones( + (2, 1), dtype=torch.bfloat16 + ), + "diffusion_model.layer.alpha": torch.tensor(1.0, dtype=torch.float32), + } + _write_uel(first, tensors) + _write_uel(second, tensors) + common = [str(first), str(second)], [1.0, 1.0] + + if strategy in {"concatenate", "weighted_sum"}: + output = merger.merge_multi_loras( + *common, + strategy, + "cpu", + torch.bfloat16, + f"bf16_alpha_{strategy}", + verbose=False, + ) + elif strategy == "dare": + output = merger.merge_multi_loras_dare( + *common, + 0.0, + 0.0, + 17, + "cpu", + torch.bfloat16, + "bf16_alpha_dare", + verbose=False, + ) + else: + output = merger.merge_multi_loras_dare_enhanced( + *common, + 1.0, + 1.0, + 0.0, + 0.0, + 17, + "cpu", + torch.bfloat16, + "bf16_alpha_enhanced", + verbose=False, + ) + + merged = _read_uel(output) + assert merged["diffusion_model.layer.lora_A.weight"].dtype == torch.bfloat16 + assert merged["diffusion_model.layer.lora_B.weight"].dtype == torch.bfloat16 + + def test_all_normal_inputs_use_async_uel_and_release_once( monkeypatch, tmp_path, merger ):