fix(lora): support ai-toolkit DoRA magnitudes for Krea-2 LoRAs

[Pfannkuchensack]

Summary

fix: Krea-2 DoRA LoRAs trained with ai-toolkit fail to load with ValueError: Unsupported lora format: dict_keys(['to_gate.magnitude', 'to_k.magnitude', 'to_q.magnitude', 'to_v.magnitude']).

Why. ai-toolkit stores the DoRA magnitude as <layer>.magnitude, a suffix the Krea-2 converter's suffix table didn't know. The key fell through to the generic rsplit fallback, so every magnitude in an attention block got grouped into a bogus parent layer (...attn) whose sub-keys were to_q.magnitude, to_k.magnitude, … — exactly the dict shown in the error.

While fixing that, a second, pre-existing bug surfaced. Two mutually incompatible DoRA magnitude conventions exist, and they index different axes of the (out_features, in_features) weight:

Trainer	Key	Indexes	Shape
LyCORIS / kohya / OneTrainer	.dora_scale	in_features	(1, in)
PEFT / Diffusers / ai-toolkit	.lora_magnitude_vector.weight, .magnitude	out_features	(out,)

DoRALayer only implemented the LyCORIS math (norm over the output dim). The existing .lora_magnitude_vector.weight -> dora_scale mapping therefore applied a PEFT magnitude with LyCORIS math: silently wrong weights on square layers, and a broadcast error on non-square ones (e.g. Krea-2's ff.down, 6144 × 16384).

How.

• DoRALayer carries an explicit magnitude_is_out_dim flag instead of guessing the orientation from tensor shape (shape alone can't disambiguate a square layer). The out-dim branch normalizes over the input dim, matching PEFT/ai-toolkit; the LyCORIS path is untouched.
• any_lora_layer_from_state_dict also dispatches on the new dora_magnitude value key.
• The Krea-2 converter maps .magnitude and routes both PEFT-style magnitude keys to dora_magnitude.

The Krea-2 model-config probe already accepted these files, so no identification changes were needed.

Related Issues / Discussions

Closes #9515

Distinct from #9424 (Krea-2 LoKr/LyCORIS support), which is not addressed here.

QA Instructions

Test adapter: https://huggingface.co/khronex/krea2-dora-ehrmantraut-test (ai-toolkit 0.10.17, network.type: dora, native/ComfyUI key layout, 256 modules / 768 tensors)

End-to-end (this is the reproducer from the issue)

1. Install the LoRA above (it probes as lora / krea-2 / lycoris).
2. Generate with Krea-2-Raw + the LoRA applied at weight 1.0.
3. Before this PR: generation fails during text encoding with Unsupported lora format. After: the image generates and the trained subject is clearly present.

Verified on an RTX 4090, Krea-2-Raw (fp8_storage: true), 1024×1024, 28 steps, CFG 4.5, seed 1234, identical prompt with and without the LoRA. The subject identity changes as expected while the composition stays put. Server log for the LoRA run shows krea2_lora_loader executing, and zero Failed to find module for LoRA layer key warnings — all 256 layers patch onto real modules.

Structural check

All 256 converted patch keys were resolved against a meta-device Krea2Transformer2DModel (stock config): 0 missing modules, 0 shape mismatches, including magnitude length == out_features.

Numeric check

The merged weight was compared against ai-toolkit's own forward pass (ToolkitModuleMixin.forward + DoRAModule.apply_dora) for out < in, square, and out > in layers — max absolute deviation 1.4e-6.

Unit tests

pytest tests/backend/patches tests/backend/model_manager/configs/test_krea2_lora_config.py

New tests/backend/patches/layers/test_dora_layer.py covers both magnitude conventions (the out-dim test asserts equivalence with the ai-toolkit reference formula); test_krea2_lora_conversion_utils.py gains an ai-toolkit .magnitude regression test plus orientation assertions on the existing DoRA tests.

Merge Plan

Normal merge — backend only, no DB or schema changes.

Note the behavior change for adapters that were already loading: Krea-2 LoRAs with .lora_magnitude_vector.weight (PEFT/Diffusers DoRA) previously had their magnitude applied along the wrong axis. Square layers will now produce different — correct — output; non-square ones previously raised a broadcast error.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[lstein]

Approve

Reviewed adversarially at 37b3c3a323, merged onto origin/main in a scratch worktree (the branch is BEHIND). I went after this by execution rather than by reading, and I couldn't break it.

The orientation claim checks out independently

I didn't take the numbers in the PR description on faith — I pulled the safetensors header of the linked test adapter (768 tensors / 256 modules) and measured it directly:

• 256/256 magnitudes have shape [out_features]; zero have [in_features]. The PEFT/out-dim claim holds on real data.
• 152 of 256 modules are non-square, so the pre-existing .lora_magnitude_vector.weight -> dora_scale mapping was not a theoretical hazard.
• On origin/main that key set reproduces the reported failure verbatim: ValueError: Unsupported lora format: dict_keys(['to_gate.magnitude', 'to_k.magnitude', 'to_q.magnitude', 'to_v.magnitude']). On this branch it converts to 256 DoRALayers, all magnitude_is_out_dim=True.
• All 256 converted keys resolve against a meta-device stock Krea2Transformer2DModel (KREA2_TRANSFORMER_CONFIG): 0 missing modules, 0 shape mismatches, magnitude length == out_features everywhere, and no bogus .attn / .mlp parent groups.

Attacks on the math — all clean

• Out-dim branch vs an independently written PEFT/diffusers reference (m * (W + dV) / ||W + dV|| over dims >= 1): max abs deviation 1.2e-7 across (12,20), (16,16), (20,12), a 4-D conv, and with and without alpha.
• End-to-end through LayerPatcher.apply_smart_model_patches on a custom-wrapped nn.Linear (sidecar path — the base weight is never mutated, so a naive weight comparison is misleading here): 2.4e-7 vs the same reference for all three shapes, weight restored on context exit.
• The LyCORIS .dora_scale path is byte-identical to main across 7 cases including conv4d and explicit alpha — the "untouched" claim is verified, not just visually plausible.
• Magnitude stored as (out,), (out,1) or (1,out) all produce identical results; a wrong-length magnitude raises rather than silently broadcasting.
• to(dtype) preserves the flag; fp16 output finite; calc_size unaffected.
• The new kwarg is last and defaulted, so both existing DoRALayer(...) construction sites in test_all_custom_modules.py keep the LyCORIS behavior.
• There is no second DoRA implementation anywhere in the tree (grepped for direction_norm / DoRA) — the sidecar path goes through get_parameters, so the flag cannot be bypassed.
• The model-config probes (configs/lora.py, configs/main.py) ignore .magnitude entirely and still identify the file through its lora_A/lora_B pairs; is_state_dict_likely_krea2_lora returns True on the real key set. Agreed that no identification change was needed.
• Test sensitivity: reverting the three source files while keeping the tests fails all 5 new/changed assertions, so none of them are decorative.
• ruff check and ruff format --check clean; 799 tests pass in tests/backend/patches + tests/backend/model_manager/configs after merging main.

Merge note

main now carries #9449, whose per-module kohya gate uses _SUFFIX_TO_VALUE_KEY as its convertibility test — so adding .magnitude there changes that gate too. I merged and checked: clean merge, full suite green, and kohya-flattened .magnitude modules now convert into proper out-dim DoRA layers instead of the old warn-and-skip. Strictly better; nothing to do beyond the normal merge.

Non-blocking findings

1. KeyError instead of a readable error for an orphan magnitude (dora_layer.py:53). A module carrying .magnitude but no A/B — e.g. an ai-toolkit LoKr-DoRA file, which is #9424 territory — now groups dora_magnitude on its own; any_lora_layer_from_state_dict takes the DoRA branch and values["lora_up.weight"] raises an uncaught KeyError: 'lora_up.weight'. Pre-PR the same file produced the clear ValueError: Unsupported lora format. Which error surfaces first depends on group iteration order — I reproduced both. The shape of this is pre-existing for dora_scale, so it isn't a regression this PR introduces on its own, but gating the DoRA branch on lora_up.weight being present would fix both at once.
2. Both conventions on one layer silently drops one (dora_layer.py:52-53). dora_magnitude wins and dora_scale is discarded with no warning, because both are in handled_keys so warn_on_unhandled_keys stays quiet. Contrived, but a one-line warning is cheap.
3. A wrong-length magnitude surfaces as RuntimeError: shape '[12, 1]' is invalid for input of size 20 — no layer name, no hint about orientation. A named ValueError would save someone a bad afternoon.
4. Scope. Only the Krea-2 converter learns .magnitude. ai-toolkit DoRA adapters for Z-Image / Wan / Anima / FLUX still die with the same Unsupported lora format, since their _group_by_layer suffix tables don't know the suffix. Worth a follow-up issue rather than widening this PR.
5. anima_lora_conversion_utils._make_layer_patch strips only dora_scale from LoKr layers, not dora_magnitude. Harmless today (Anima never produces it), but a latent inconsistency the moment .magnitude is added there.
6. Stale comment: tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_all_custom_modules.py (~line 766) points at dora_layer.py:74-82 for the broadcast rationale; that branch has moved to roughly 106-115.

Nice piece of work — carrying the orientation explicitly rather than inferring it from shape is the right call, and the docstring explaining why is exactly what the next person will need.