Skip to content

Enable VAE tiling automatically after CUDA OOM - #3297

Merged
bghira merged 1 commit into
mainfrom
fix/vae-oom-auto-tiling
Oct 26, 2026
Merged

bghira merged 1 commit into
mainfrom
fix/vae-oom-auto-tiling

Conversation

@bghira

@bghira bghira commented Oct 26, 2026

Copy link
Copy Markdown
Owner

High-resolution VAE encoding or validation decoding can exhaust CUDA memory even when denoiser training fits. Enable tiling automatically on supported, untiled VAEs after a CUDA OOM and retry that call once. Keep tiling enabled for later calls; already-tiled VAEs, unsupported VAEs, other errors and failed retries still propagate their errors.

Install the encode/decode recovery through the shared model foundation VAE getter so cache creation and pipeline decoding use the same behavior. Release the failed call's traceback before clearing the CUDA allocator cache and retrying. Preserve method signatures, gradients, deepcopy and serialization. Document the behavior in OPTIONS and all five translations, and update the existing VAE tiling tooltip.

Validation:

  • Regression suite: 7,102 tests passed, 323 skipped. Final focused recovery suite: 12 tests passed.
  • Black and staged diff checks passed.
  • Real Krea2 Qwen Image VAE on L40S: independently reproduced CUDA OOM during untiled 2048px encoding and decoding with 8 GiB of additional CUDA residency representing text-encoder memory. Both wrapped calls enabled tiling exactly once and returned finite outputs; subsequent decoding retained tiling. The tests use the actual checkpoint, without a mocked allocation failure.

@bghira
bghira merged commit 498e8f5 into main Oct 26, 2026
4 checks passed
@bghira
bghira deleted the fix/vae-oom-auto-tiling branch October 26, 2026 22:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant