fix(qwen): make room in VRAM for the quantized Qwen2.5-VL encoder and release it afterwards (#9147)

[lstein]

Summary

Fixes #9147 (the part that was still failing after #9305): with Encoder quantization set to int8 or nf4, the Qwen Image prompt node failed on every prompt change after the first generation with

Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the quantized model...

Root cause. The BitsAndBytes path loads the Qwen2.5-VL encoder outside the model cache (BnB models are pinned to the device they were quantized on), so it also skipped the make-room step that lock() performs for every cached load. With the transformer and VAE still resident, from_pretrained(device_map="auto") planned the encoder onto the CPU and BnB int8 refused. Nothing in the path ever asked the cache for VRAM.

A second defect in the same path explains the "transformer swings between 57% and 100% in VRAM" symptom from the issue: _encode still held the encoder when its cleanup callback ran, so the callback's del freed nothing and empty_cache() ran while the weights were alive. The weights were released only when the frame exited, after which ~9 GB stayed reserved by torch. The cache budgets from memory_allocated + driver-free VRAM, so that reserved-but-unused memory looked like it was in use and the next model (the transformer) was needlessly partial-loaded.

Changes

• ModelCache.make_room_in_vram(bytes, working_mem_bytes=None) – exposes the offload policy lock() already uses (_offload_unlocked_models) for loads that live outside the cache. Locked models are never touched; offloaded models stay cached and re-stream later. No-op on a CPU execution device (where _get_vram_available would raise).
• context.models.make_room_in_vram(...) – invocation-context wrapper, resolving to the calling worker's per-device cache in multi-GPU mode.
• Quantized encoder path (qwen_image_text_encoder.py):
  • estimates the post-quantization footprint from the bf16 checkpoint size (0.6× int8, 0.4× nf4; measured 0.56× on the 7B encoder) and asks the cache for it;
  • loads with an explicit device_map={"": <execution device>} instead of "auto", so a real shortfall is a plain OOM rather than a misleading offload error, and in multi-GPU mode the load is pinned to the worker's own device;
  • offload + load run under the MODEL_LOAD_LOCK read lock, like every other VRAM move, so a concurrent cache construction cannot hijack the encoder's parameter assignment onto the meta device;
  • _encode drops its reference to the encoder before the cleanup callback runs, so empty_cache() actually returns the weights to the driver. The forward and post-processing live in a helper so the logits/hidden states die with its frame, and on an exception inside the forward the finished traceback frames are cleared first (the in-flight traceback otherwise keeps the model alive exactly when VRAM is scarcest).

Verification on real hardware (W7900, ROCm, bf16 Qwen2.5-VL-7B → int8)

A 37.5 GB filler model in the cache, then the prompt node's quantized path wired to the real cache:

	allocated	reserved	driver-free	cache _get_vram_available
card filled	37.5 GB	37.5 GB	6.7 GB	–
after encode, before this PR's release fix	31.9 GB	40.8 GB	3.2 GB	0.2 GB
after encode, with this PR	31.9 GB	31.9 GB	12.1 GB	9.1 GB

make_room_in_vram was called with 9.27 GB, partially offloaded the filler by 5.6 GB, the encoder landed on cuda:0, and a forward pass produced finite bf16 embeddings. (The original CPU-dispatch error itself cannot be reproduced on AMD: the HIP driver spills oversubscribed allocations into system memory instead of refusing, so device_map="auto" "succeeds" there. It reproduces on NVIDIA.)

Tests

• tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py – offload loop contract with fake VRAM accounting (stops when satisfied, skips locked, forwards working memory, holds the cache lock, CPU no-op) plus a real-accelerator test (skipped without CUDA/ROCm) that checks weights actually move for both cached-model flavours.
• tests/app/invocations/test_qwen_image_text_encoder.py – make_room_in_vram is called with the estimated size before from_pretrained, both under the model-load lock; device_map is explicit; single-file checkpoints still fall back to the cache; the encoder is released before cleanup runs, on the normal path and when the forward raises.
• tests/app/services/shared/test_invocation_context_make_room_in_vram.py – context wrapper delegation.

Each test was checked against a mutant of the production line it pins (revert the del, skip the traceback clearing, use device_map="auto", drop the make-room call, move it outside the lock, drop @synchronized, drop the CPU guard, drop the working-memory argument); every mutant fails its test.

Manual test plan

Set log_level: debug to see the [MODEL CACHE] lines. The encoder must be a folder (diffusers-layout) install with a text_encoder/ directory – BnB cannot quantize a single-file checkpoint, and those fall back to the cached path as before.

NVIDIA 16 GB (this rig reproduces the original failure)

1. Qwen Image Edit 2511 GGUF Q8_0 (or Q4_K_M) + bf16 Qwen2.5-VL encoder, prompt node Encoder quantization = int8. Generate once, change the prompt, generate again. Before: the second run fails with the CPU-dispatch error. After: it succeeds, and the debug log shows Offloading unlocked models with goal of making room for ~9500 MB followed by Unloaded <transformer> from VRAM ... right before the encoder loads.
2. Same with nf4 – expect ~6300 MB requested.
3. Run a batch of four with different prompts (dynamic prompts), i.e. encoder → denoise → decode → encoder → …. Expect no OOM at any step, and the transformer's Loaded model ... VRAM: N% line to be the same on every iteration (before, it drifted down on later iterations because of the reserved-memory leak).
4. Regression: Encoder quantization = none, same batch – unchanged behaviour (encoder goes through the cache).
5. Watch nvidia-smi between nodes: after the prompt node finishes, used memory should drop back to transformer-only.

Dual W7900 (48 GB each, ROCm)

1. Multi-GPU pinning: leave generation_devices: auto, enqueue two Qwen sessions with int8 quantization at once. Expect one session per card, each encoder loaded on its own device (rocm-smi --showmemuse shows both cards climbing by ~9 GB during the prompt nodes). Before, device_map="auto" could place both encoders on GPU 0.
2. Oversubscription: the bf16 diffusers "Qwen Image 2512" transformer (~40 GB) + int8 encoder cannot co-reside on one card. Expect the make-room line, a partial offload of the transformer (~6 GB), and a successful encode; subsequent denoise re-streams the offloaded slice.
3. BnB-on-ROCm sanity: same seed with quantization none vs int8 – images should be near-identical.
4. Run the accelerator-marked tests here: pytest tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py -k gpu.

Follow-ups (out of scope)

• _offload_unlocked_models cannot see freed memory until its trailing empty_cache(), because memory_allocated and driver-free move in opposite directions during the loop. In practice any shortfall offloads every unlocked model. Pre-existing in lock(); this PR only reuses and documents it.
• The quantized encoder is re-read from the 16 GB bf16 checkpoint and re-quantized on every prompt change. Routing it through the cache the way FLUX's int8 T5 is (BnbQuantizedLlmInt8bCheckpointModel) would make it cacheable and partially loadable, but needs a model format and an install-time or first-load quantization step.

🤖 Generated with Claude Code

https://claude.ai/code/session_01C7hM9XHDWGNa57jY5pat98

[Pfannkuchensack]

Findings

1. High - invokeai/app/invocations/qwen_image_text_encoder.py:357-368: the process-global MODEL_LOAD_LOCK read lock is now held across the entire BnB checkpoint read + quantization, stalling every other worker on every prompt change.

Chain:

1. MODEL_LOAD_LOCK is a single process-global instance (invokeai/backend/model_manager/load/model_cache/model_cache.py:365), shared by all session workers regardless of device.
2. The with MODEL_LOAD_LOCK.read_lock(): block at line 357 now spans make_room_in_vram() and the full Qwen2_5_VLForConditionalGeneration.from_pretrained(...) - a 16 GB disk read plus BnB quantization.
3. write_lock() (model_cache.py:349-356) waits for self._readers == 0, so any cold model construction on any GPU blocks for that whole duration (invokeai/backend/model_manager/load/load_default.py:325, invokeai/app/services/model_load/model_load_default.py:164).
4. The lock is write-preferring: read_lock() (model_cache.py:334-339) also blocks while _writers_waiting > 0, so every subsequent VRAM move on every other worker queues behind that waiting writer too.
5. The PR's own Follow-ups state the quantized encoder is re-read and re-quantized on every prompt change, so this is not a one-off cost.

Before this PR from_pretrained took no lock at all, so the regressed behavior is concurrency: on a dual-GPU or multiuser install, a Qwen prompt node on GPU 0 now freezes GPU 1's model loads and offloads for the length of the re-quantization. Note the lock itself is necessary - a concurrent diffusers construction under the write lock does install accelerate.init_empty_weights' global register_parameter patch, so removing it is not the fix; the scope of what it covers is the problem. Also note the in-code justification at lines 353-356 ("Both the offload and the load below assign real parameters ... like every other VRAM move") mis-describes from_pretrained as a VRAM move; it is a model construction, which invokeai/backend/model_manager/load/load_default.py:315-325 documents as a write-lock operation. With the pinned transformers>=5.5,<5.6 this happens to be safe in the outbound direction (transformers 5 builds under a thread-local torch.device("meta") context, not accelerate's global patch - verified by executing a two-thread probe), but the comment records the wrong rule for a future dependency bump.

To expose this issue, add a test that holds a fake long-running from_pretrained on one thread inside _load_quantized_encoder and asserts a second thread's MODEL_LOAD_LOCK.write_lock() acquisition is blocked for the duration - the test that exists only asserts the lock is held, never that the held region is bounded.

2. Medium - invokeai/app/invocations/qwen_image_text_encoder.py:357-368: the reserved-VRAM leak the PR fixes for _encode is left wide open on the load-failure path, which this same PR makes more likely.

_load_quantized_encoder has no try/except around the load. When from_pretrained raises (OOM is now the expected failure mode - see the PR's own "a genuine shortfall surfaces as a plain OOM"), the partially materialized model is kept alive by the in-flight traceback all the way up to the session processor's handler, and nothing calls gc.collect() or TorchDevice.empty_cache() on that path - cleanup is never constructed, and there is no empty_cache anywhere in invokeai/app/services/session_processor/session_processor_default.py. That is exactly the "weights stay reserved by torch, so the cache under-budgets the next load" defect the PR fixes for the success path at line 239. _offload_unlocked_models only calls empty_cache() after it has already read _get_vram_available() at the top of its loop (invokeai/backend/model_manager/load/model_cache/model_cache.py:2318-2320), so the next queue item measures against the stale reserved figure and partial-loads the transformer - the original user-visible symptom.

Secondary consequence on the same path: make_room_in_vram has already offloaded cached models to RAM before the load fails, and nothing restores them, so a failed encode leaves the GPU strictly worse off than before the node ran.

To expose this issue, add a test that makes Qwen2_5_VLForConditionalGeneration.from_pretrained raise torch.OutOfMemoryError inside _load_quantized_encoder and asserts the partially built model is unreachable and TorchDevice.empty_cache() was called before the exception leaves the invocation.

3. Medium - invokeai/app/invocations/qwen_image_text_encoder.py:365: replacing device_map="auto" with device_map={"": device} removes multi-GPU sharding, converting a previously working load into a hard OOM for small-card multi-GPU installs.

device_map="auto" routes through accelerate's infer_auto_device_map/get_balanced_memory, which distributes layers across all visible accelerators; BnB int8/nf4 multi-GPU placement is supported. {"": device} forces the whole encoder onto one card. Trigger: a user with two GPUs where the encoder (~9 GB int8 on the 7B model) does not fit on the session's own card even after make_room_in_vram has offloaded everything unlocked, but did fit when spread across both. Note this is not exotic in the default single-GPU configuration either: TorchDevice.choose_torch_device() resolves to cuda:0 (invokeai/backend/util/devices.py:102-121), so a second, completely idle GPU that "auto" used to recruit is now unreachable. The PR body presents the change purely as multi-GPU pinning and does not acknowledge the loss of sharding, and there is no fallback path.

To expose this issue, add a test that asserts the invocation surfaces an actionable error (rather than a bare OOM) when make_room_in_vram reports it freed less than the requested estimate.

4. Low - invokeai/backend/model_manager/load/model_cache/model_cache.py:2753 and invokeai/app/invocations/qwen_image_text_encoder.py:358: make_room_in_vram promises less than lock() delivers, and its only production caller discards the answer.

_load_locked_model (model_cache.py:1978-1994) re-reads _get_vram_available() after _offload_unlocked_models returns and takes a further corrective step when the result is still negative. make_room_in_vram does neither: it returns _offload_unlocked_models(...) directly (line 2773), a number "based on believed model sizes", and the caller at line 358 ignores it. So a shortfall caused by a peer session holding a locked transformer is invisible to the caller, which proceeds straight into the load. Related: working_mem_bytes has no production caller at all - the encoder path always passes None and inherits the configured default, so the encoder's own vision-tower activations are budgeted only by device_working_mem_gb. The parameter is exercised solely by the new cache test, i.e. covered code with no user.

5. Low - invokeai/app/invocations/qwen_image_text_encoder.py:350: the size estimate silently degrades to zero, turning the fix into a no-op with no log line.

calc_model_size_by_fs (invokeai/backend/model_manager/load/model_util.py:110-188) returns 0 for a directory that does not exist or whose weights match none of its known suffixes. int(0 * ratio) == 0 is then passed to make_room_in_vram, _offload_unlocked_models computes vram_bytes_to_free = 0 - vram_available <= 0 and breaks on the first iteration, and issue #9147 reappears with nothing in the log to distinguish it from the fixed case. The debug line the manual test plan tells users to look for (Offloading unlocked models with goal of making room for ~9500 MB) would read 0.00MB, which is easy to miss.

To expose this issue, add a test that points _load_quantized_encoder at a text_encoder/ directory calc_model_size_by_fs cannot size and asserts the invocation warns rather than silently requesting zero bytes.

6. Low - invokeai/app/invocations/qwen_image_text_encoder.py:226-231: traceback.clear_frames is applied unconditionally, including on the non-quantized cached path where nothing needs releasing.

The except BaseException handler sits outside any self.quantization check, so an encode error on the quantization="none" path - the default, and by far the most common - now has every frame local in its traceback destroyed. The cached encoder is owned by the model cache, not by _encode, so this buys nothing there and only costs debuggability for the common failure. The PR's mutant-testing note covers "skip the traceback clearing" but not "clear only when the encoder is out-of-cache".