Skip to content

Handle mixed-dtype mismatches in autocast linear and conv wrappers#9006

Open
JPPhoto wants to merge 3 commits intoinvoke-ai:mainfrom
JPPhoto:fix-dtype-mismatch-error-during-invocation
Open

Handle mixed-dtype mismatches in autocast linear and conv wrappers#9006
JPPhoto wants to merge 3 commits intoinvoke-ai:mainfrom
JPPhoto:fix-dtype-mismatch-error-during-invocation

Conversation

@JPPhoto
Copy link
Copy Markdown
Collaborator

@JPPhoto JPPhoto commented Mar 29, 2026

Summary

Fixes mixed-dtype mismatch failures in autocast-wrapped Linear and Conv2d layers during invocation.

This change:

  • fixes CustomLinear so it handles dtype mismatches in the plain path, the bias-only mismatch case, and the sidecar aggregated-parameter patch path
  • updates shared patch aggregation to preserve real torch.Tensor params instead of degrading them to meta tensors - adds focused regression tests for plain mixed-dtype inference and sidecar parameter patching, with CPU dtype parametrization for portability

Practical exposure:

  • partial-load / autocast execution where activations are float16 or bfloat16 but stored weights or bias remain float32
  • bias-only dtype drift after manual casting, state-dict loading, or patch application - sidecar patching during invocation, especially non-LoRA parameter patches that produce residual tensors in the wrong dtype
  • mixed-precision inference with adapters or patch systems that introduce tensors at a different dtype than the active execution path
  • CPU fallback and non-CUDA environments too, since the underlying failure is a generic PyTorch dtype mismatch (mat1 and mat2 must have the same dtype, self and mat2 must have the same dtype)

In practice, this prevents invocation-time failures that would otherwise appear only under certain precision / patching combinations, making them easy to miss and hard to reproduce.

Related Issues / Discussions

Solves the issue of partial-load execution that I experienced.

QA Instructions

Run the focused mixed-dtype regression tests via pytest tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_all_custom_modules.py -k 'mixed_dtype_inference_without_patches or mixed_dtype_sidecar_parameter_patch or bias_only_mismatch'

Merge Plan

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)

@github-actions github-actions bot added python PRs that change python files backend PRs that change backend files python-tests PRs that change python tests labels Mar 29, 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.

1 participant