Skip to content

ai-toolkit DoRA (.magnitude) adapters fail to load for every base except Krea-2 #9558

Description

@lstein

Summary

#9517 taught the Krea-2 LoRA converter that ai-toolkit writes the DoRA magnitude as a bare <layer>.magnitude suffix (PEFT/out-dim orientation), and gave DoRALayer an explicit magnitude_is_out_dim flag. Every other architecture's converter still has a fixed suffix table that doesn't know .magnitude, so ai-toolkit DoRA adapters for those bases fail to load with the same error #9517 fixed for Krea-2:

ValueError: Unsupported lora format: dict_keys(['to_gate.magnitude', 'to_k.magnitude', ...])

Why it happens

.magnitude isn't in the converter's known-suffix list, so _group_by_layer falls through to a blind rsplit(".", 2). That cuts inside the module path and fuses every magnitude in a block into one bogus parent layer, which then matches no layer type.

Affected suffix tables (all on main):

  • invokeai/backend/patches/lora_conversions/wan_lora_conversion_utils.py:223-228
  • invokeai/backend/patches/lora_conversions/z_image_lora_conversion_utils.py:229-234
  • invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py:152-157
  • invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py:168-173
  • invokeai/backend/patches/lora_conversions/flux_bfl_peft_lora_conversion_utils.py:78 uses a similar fixed suffix set and is worth checking too

None of these map .lora_magnitude_vector.weight either, so unlike Krea-2 they never produced a mis-oriented DoRA layer — this is a plain "fails to load", not a silent-wrong-weights bug. That's why it's a follow-up rather than part of #9517.

Suggested fix

The machinery already exists after #9517 — this is mostly plumbing:

  1. Add .magnitude (and .lora_magnitude_vector.weight) to each converter's suffix table, routing both to the dora_magnitude value key, not dora_scale. The two conventions index opposite axes of the (out_features, in_features) weight and must not be mixed up — see the DoRALayer docstring added in fix(lora): support ai-toolkit DoRA magnitudes for Krea-2 LoRAs #9517.
  2. any_lora_layer_from_state_dict already dispatches on dora_magnitude, so no change is needed there.
  3. Watch the LoKr interaction: anima_lora_conversion_utils._make_layer_patch strips only dora_scale from LoKr layers before dispatch. If Anima learns .magnitude, that strip needs to cover dora_magnitude as well, or a DoRA+LoKr adapter will take the DoRA branch and die on a missing lora_up.weight.
  4. Per-architecture verification is worth doing rather than assuming: instantiate the target transformer on the meta device and confirm every converted key resolves to a real module with magnitude.numel() == out_features. That's the check that gave fix(lora): support ai-toolkit DoRA magnitudes for Krea-2 LoRAs #9517 its 256/256 clean result.

Notes

Non-blocking follow-up raised during the review of #9517 (#9517 (review)), finding 4. No user has reported this for a specific non-Krea-2 base yet, so priority should probably follow demand — ai-toolkit is a popular trainer, so it's likely only a matter of time.

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions