Skip to content

fix: SDXL DoRA LoRA fails with enable_partial_loading=true#9063

Open
Pfannkuchensack wants to merge 6 commits intomainfrom
fix/sdxl-dora-partial-loading
Open

fix: SDXL DoRA LoRA fails with enable_partial_loading=true#9063
Pfannkuchensack wants to merge 6 commits intomainfrom
fix/sdxl-dora-partial-loading

Conversation

@Pfannkuchensack
Copy link
Copy Markdown
Collaborator

Summary

cast_to_device returns a plain torch.Tensor instead of torch.nn.Parameter when moving parameters between devices. This causes _aggregate_patch_parameters to replace valid weights with meta device dummies, falsely triggering DoRA's quantization guard.

Added a check for plain torch.Tensor in _aggregate_patch_parameters so that parameters cast to a different device are recognized as valid non-quantized weights.

Fixes #8624

Related Issues / Discussions

  • #8624 — SDXL LoRA (DoRA) fails when enable_partial_loading=true

QA Instructions

  1. Load an SDXL model with enable_partial_loading=true
  2. Apply a DoRA-based LoRA (e.g. Stabilizer ILNAI)
  3. Run a denoise — should complete without the "base model quantization format (likely bitsandbytes) is not compatible with DoRA patches" error

Merge Plan

Single-file change, safe to merge anytime.

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)

cast_to_device returns plain torch.Tensor instead of torch.nn.Parameter,
causing _aggregate_patch_parameters to replace valid weights with meta
device dummies, falsely triggering DoRA's quantization guard.

Fixes #8624
@github-actions github-actions Bot added python PRs that change python files backend PRs that change backend files labels Apr 17, 2026
@JPPhoto JPPhoto self-assigned this Apr 20, 2026
@Pfannkuchensack
Copy link
Copy Markdown
Collaborator Author

@JPPhoto
Copy link
Copy Markdown
Collaborator

JPPhoto commented Apr 20, 2026

  • invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_module_mixin.py: there is still no targeted regression coverage for the DoRA plus partial-loading plus CPU-to-device autocast path that motivated this PR. A useful follow-up test would exercise a wrapped patched module with base weights on CPU, input and patch on CUDA or MPS, and assert that _aggregate_patch_parameters() preserves real tensors instead of substituting meta.

Otherwise it's good - ping me if you write this test or decide it's too difficult, @Pfannkuchensack.

…ocast

Adds targeted coverage for the bug fixed in a0a8721 (#8624, PR #9063):

- test_aggregate_patch_parameters_preserves_plain_tensor_with_dora:
  CPU-only unit test that feeds a plain torch.Tensor (as handed in by
  _cast_weight_bias_for_input) into _aggregate_patch_parameters with a
  DoRA patch. Pre-fix, the tensor was replaced by a meta-device dummy,
  tripping DoRA's quantization guard.

- "single_dora" variant in the patch_under_test fixture: exercises the
  full CUDA/MPS autocast hot path via
  test_linear_sidecar_patches_with_autocast_from_cpu_to_device.
@github-actions github-actions Bot added the python-tests PRs that change python tests label Apr 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backend PRs that change backend files python PRs that change python files python-tests PRs that change python tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[bug]: SDXL Lora (DoRA) fails when enable_partial_loading=true

2 participants