From 4cac283561b72cbca44b22c06878c99e1701d052 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Sun, 12 Jul 2026 21:35:11 +0100 Subject: [PATCH] feat(lora): slice per-output dora_scale on fused-qkv targets LyCORIS wd=True saves a dora_scale companion for LoRA/LoHA/LoKR; on fused BFL targets the chunk paths passed it through whole, so apply failed with a shape mismatch and the module was dropped. Per-output magnitudes (wd_on_out=True, the default) partition exactly with the fused rows; per-input magnitudes couple the chunks through shared column norms and have no exact split. - slice per-output dora_scale rows with the chunk in the LoRA, LoKR and LoHA loaders - skip per-input DoRA on fused targets with a specific warning - cover sliced and skipped orientations in the offline suite --- modules/lora/native_adapter.py | 83 +++++++++++++++++++++++++----- test/test-flux2-native-adapters.py | 62 ++++++++++++++++++++++ 2 files changed, 132 insertions(+), 13 deletions(-) diff --git a/modules/lora/native_adapter.py b/modules/lora/native_adapter.py index 72a330aa4..dd9285c5f 100644 --- a/modules/lora/native_adapter.py +++ b/modules/lora/native_adapter.py @@ -407,27 +407,55 @@ def resolve_group_targets(resolve_targets, prefix_used, base): # slices and Tucker-decomposed LoHAs on fused targets are skipped with a # warning (no slice variant exists, and Tucker keys cannot arise on Linear # layers per LyCORIS upstream — see network_hada.NetworkModuleHadaChunk). +# +# DoRA on fused targets (LoRA/LoKR/LoHA): per-output dora_scale rows are +# sliced with the chunk via slice_dora_scale (exact, row norms are +# independent); per-input dora_scale couples the chunks through shared +# column norms and the group is skipped with a warning. # - OFT/BOFT: no chunk variant exists; fused targets are skipped with a # warning. Discrimination is by ``oft_blocks.ndim`` (3-D OFT, 4-D BOFT), # mirroring upstream LyCORIS ``algo_check``. -def _slice_lora_chunk(w, chunk: ChunkSpec): - """Return a shallow copy of ``w`` with ``lora_up.weight`` sliced per ``chunk``. +def slice_chunk_rows(t, chunk: ChunkSpec): + """Slice dim 0 of ``t`` per ``chunk``. Equal-chunks form uses ``torch.chunk`` (faster for the symmetric case); row-range form uses tensor slicing for arbitrary partitions. """ - up = w["lora_up.weight"] if chunk.is_equal_chunks: - sliced = torch.chunk(up, chunk.total, dim=0)[chunk.idx].contiguous() - else: - sliced = up[chunk.start:chunk.end].contiguous() + return torch.chunk(t, chunk.total, dim=0)[chunk.idx].contiguous() + return t[chunk.start:chunk.end].contiguous() + + +def _slice_lora_chunk(w, chunk: ChunkSpec): + """Return a shallow copy of ``w`` with ``lora_up.weight`` sliced per ``chunk``.""" out = dict(w) - out["lora_up.weight"] = sliced + out["lora_up.weight"] = slice_chunk_rows(w["lora_up.weight"], chunk) return out +def slice_dora_scale(w, chunk: ChunkSpec, fused_out): + """Slice a per-output ``dora_scale`` to the chunk rows; ``None`` if unsliceable. + + LyCORIS ``wd_on_out=True`` (the LoKr/LoHA default) stores per-output + magnitudes of shape ``(out, 1)`` (1-D ``(out,)`` also seen); the rows + partition exactly with the fused weight, so the chunk slice preserves the + DoRA math (row norms are independent across rows). ``wd_on_out=False`` + stores per-input magnitudes whose column norms span every fused row, + coupling the chunks; no exact per-chunk equivalent exists and the caller + must skip the group. + """ + ds = w.get("dora_scale") + if ds is None: + return w + if ds.ndim >= 1 and ds.shape[0] == fused_out: + out = dict(w) + out["dora_scale"] = slice_chunk_rows(ds, chunk) + return out + return None + + def try_load_lora(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=(), @@ -454,6 +482,7 @@ def try_load_lora(name, network_on_disk, lora_scale, *, unmapped = 0 mismatch = 0 + skipped = 0 for (prefix, base), w in groups.items(): if "lora_down.weight" not in w or "lora_up.weight" not in w: continue @@ -465,7 +494,14 @@ def try_load_lora(name, network_on_disk, lora_scale, *, unmapped += 1 continue - target_w = _slice_lora_chunk(w, chunk) if chunk is not None else w + target_w = w + if chunk is not None: + target_w = _slice_lora_chunk(w, chunk) + target_w = slice_dora_scale(target_w, chunk, fused_out=w["lora_up.weight"].shape[0]) + if target_w is None: + log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} per-input DoRA on fused target skipped (unsupported)') + skipped += 1 + continue if not shapes_match(sd_module, target_w["lora_down.weight"], target_w["lora_up.weight"]): log.warning( @@ -480,14 +516,14 @@ def try_load_lora(name, network_on_disk, lora_scale, *, nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=target_w, sd_module=sd_module) net.modules[network_key] = network_lora.NetworkModuleLora(net, nw) - return finalize_network(net, name, "LoRA", lora_scale, t0, unmapped=unmapped, mismatch=mismatch) + return finalize_network(net, name, "LoRA", lora_scale, t0, unmapped=unmapped, mismatch=mismatch, skipped=skipped) def try_load_lokr(name, network_on_disk, lora_scale, *, resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=(), network_prefix=NETWORK_PREFIX_DEFAULT, - arch_name="generic"): # pylint: disable=unused-argument + arch_name="generic"): """Generic LoKR loader. Stores only the compact LoKR factors and dispatches to @@ -512,6 +548,7 @@ def try_load_lokr(name, network_on_disk, lora_scale, *, ) unmapped = 0 + skipped = 0 for (prefix, base), w in groups.items(): has_1 = "lokr_w1" in w or ("lokr_w1_a" in w and "lokr_w1_b" in w) has_2 = "lokr_w2" in w or ("lokr_w2_a" in w and "lokr_w2_b" in w) @@ -524,7 +561,20 @@ def try_load_lokr(name, network_on_disk, lora_scale, *, if sd_module is None: unmapped += 1 continue - nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) + target_w = w + if chunk is not None: + # Kron rows = w1 rows * w2 rows; the tucker w2_a orientation + # differs but tucker is conv-only and chunks are Linear-only. + w1 = w.get("lokr_w1") + w2 = w.get("lokr_w2") + w1_rows = w1.shape[0] if w1 is not None else w["lokr_w1_a"].shape[0] + w2_rows = w2.shape[0] if w2 is not None else w["lokr_w2_a"].shape[0] + target_w = slice_dora_scale(w, chunk, fused_out=w1_rows * w2_rows) + if target_w is None: + log.warning(f'Network load: type=LoKR name="{name}" arch={arch_name} key={network_key} per-input DoRA on fused target skipped (unsupported)') + skipped += 1 + continue + nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=target_w, sd_module=sd_module) if chunk is None: net.modules[network_key] = network_lokr.NetworkModuleLokr(net, nw) elif chunk.is_equal_chunks: @@ -532,7 +582,7 @@ def try_load_lokr(name, network_on_disk, lora_scale, *, else: net.modules[network_key] = network_lokr.NetworkModuleLokrSliceChunk(net, nw, chunk.start, chunk.end) - return finalize_network(net, name, "LoKR", lora_scale, t0, unmapped=unmapped) + return finalize_network(net, name, "LoKR", lora_scale, t0, unmapped=unmapped, skipped=skipped) def try_load_loha(name, network_on_disk, lora_scale, *, @@ -581,7 +631,14 @@ def try_load_loha(name, network_on_disk, lora_scale, *, if sd_module is None: unmapped += 1 continue - nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=w, sd_module=sd_module) + target_w = w + if chunk is not None: + target_w = slice_dora_scale(w, chunk, fused_out=w["hada_w1_a"].shape[0]) + if target_w is None: + log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={network_key} per-input DoRA on fused target skipped (unsupported)') + skipped += 1 + continue + nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=target_w, sd_module=sd_module) if chunk is None: net.modules[network_key] = network_hada.NetworkModuleHada(net, nw) elif chunk.is_equal_chunks: diff --git a/test/test-flux2-native-adapters.py b/test/test-flux2-native-adapters.py index 162774d42..41e63f10b 100644 --- a/test/test-flux2-native-adapters.py +++ b/test/test-flux2-native-adapters.py @@ -329,6 +329,31 @@ def sd_lora_with_dora_scale(): } +def sd_lora_dora_fused_qkv(): + """Kohya LoRA on fused img_attn.qkv with per-output dora_scale (LyCORIS wd_on_out=True).""" + return { + 'lora_unet_double_blocks_0_img_attn_qkv.lora_down.weight': torch.randn(RANK_LORA, HIDDEN), + 'lora_unet_double_blocks_0_img_attn_qkv.lora_up.weight': torch.randn(3 * QKV_OUT, RANK_LORA), + 'lora_unet_double_blocks_0_img_attn_qkv.alpha': torch.tensor(float(RANK_LORA)), + 'lora_unet_double_blocks_0_img_attn_qkv.dora_scale': torch.rand(3 * QKV_OUT, 1) + 0.5, + } + + +def sd_lokr_dora_fused_qkv(per_input=False): + """Kohya LoKR on fused img_attn.qkv with a dora_scale companion. + + ``per_input=True`` stores the wd_on_out=False orientation ``(1, in)``, + which has no exact per-chunk split and must be skipped by the loader. + """ + sd = dict(sd_lokr_kohya_qkv()) + if per_input: + ds = torch.rand(1, HIDDEN) + 0.5 + else: + ds = torch.rand(3 * QKV_OUT, 1) + 0.5 + sd['lora_unet_double_blocks_0_img_attn_qkv.dora_scale'] = ds + return sd + + def sd_lokr_bfl_proj(): """BFL-format LoKR on a non-fused proj target.""" return { @@ -912,6 +937,40 @@ def test_lora_dora_threading(): return True +def test_lora_dora_fused_qkv_sliced(): + """Per-output dora_scale is sliced with the up-weight chunk and applies cleanly.""" + net = _load_via(F.try_load_lora, sd_lora_dora_fused_qkv()) + assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}' + for nk, mod in net.modules.items(): + assert mod.dora_scale is not None and mod.dora_scale.shape[0] == QKV_OUT, \ + f'{nk}: dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}' + updown, _ex_bias = mod.calc_updown(mod.sd_module.weight) + assert_shape(updown, mod.sd_module.weight.shape, label=nk) + assert_finite(updown, label=nk) + return True + + +def test_lokr_dora_fused_qkv_sliced(): + """Per-output dora_scale rides the kron chunk; each module sees its own rows.""" + net = _load_via(F.try_load_lokr, sd_lokr_dora_fused_qkv()) + assert net is not None and len(net.modules) == 3, f'got {net.modules if net else None}' + for nk, mod in net.modules.items(): + assert isinstance(mod, network_lokr.NetworkModuleLokrChunk), f'{nk}: type={type(mod).__name__}' + assert mod.dora_scale is not None and mod.dora_scale.shape[0] == QKV_OUT, \ + f'{nk}: dora_scale shape {tuple(mod.dora_scale.shape) if mod.dora_scale is not None else None}' + updown, _ex_bias = mod.calc_updown(mod.sd_module.weight) + assert_shape(updown, mod.sd_module.weight.shape, label=nk) + assert_finite(updown, label=nk) + return True + + +def test_dora_per_input_fused_skipped(): + """wd_on_out=False dora_scale on a fused target cannot be split; the group is skipped.""" + net = _load_via(F.try_load_lokr, sd_lokr_dora_fused_qkv(per_input=True)) + assert net is None, f'expected no modules, got {net.modules if net else None}' + return True + + def test_lokr_bfl_non_fused(): net = _load_via(F.try_load_lokr, sd_lokr_bfl_proj()) assert net is not None and len(net.modules) == 1 @@ -1332,6 +1391,9 @@ def run_tests(): test_lora_peft_saved_dreambooth_style, test_lora_peft_saved_diffusers_style, test_lora_dora_threading, + test_lora_dora_fused_qkv_sliced, + test_lokr_dora_fused_qkv_sliced, + test_dora_per_input_fused_skipped, test_lokr_bfl_non_fused, test_lokr_kohya_fused_qkv_chunked, test_lokr_simpletuner_lycoris_format,