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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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