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

docs/src/generated/invocation-context.json

@@ -343,6 +343,27 @@
           "returns": "A LoadedModelWithoutConfig object.",
           "signature": "(source: str | AnyHttpUrl, loader: Optional[Callable[[Path], AnyModel]] = None) -> LoadedModelWithoutConfig"
         },
+        {
+          "description": "Offload unlocked cached models from VRAM until `vram_bytes_needed` bytes are free on this thread's\nexecution device.\nUse this before placing a model on the GPU *outside* the model cache (e.g. a BitsAndBytes-quantized model\nthat cannot be moved between devices). Loads that go through `load()` never need this - the cache makes room\nfor them itself when they are locked - but an out-of-cache load competes with the cached models for VRAM\nand would otherwise only get whatever they happened to leave free.",
+          "name": "make_room_in_vram",
+          "parameters": [
+            {
+              "default": "",
+              "description": "The VRAM footprint the caller is about to allocate.",
+              "name": "vram_bytes_needed",
+              "type": "int"
+            },
+            {
+              "default": "None",
+              "description": "Working memory to keep free on top of the model, floored at the configured default.",
+              "name": "working_mem_bytes",
+              "type": "Optional[int]"
+            }
+          ],
+          "return_type": "int",
+          "returns": "The number of VRAM bytes freed.",
+          "signature": "(vram_bytes_needed: int, working_mem_bytes: Optional[int] = None) -> int"
+        },
         {
           "description": "Move a model (and all of its submodels) from VRAM to RAM, freeing its VRAM but keeping it cached.\nUse this when an invocation is done with a model for the rest of the run - e.g. a one-shot text encoder -\nso the next, larger load does not have to compete with it for VRAM. The model stays in the RAM cache, so\na subsequent load only re-streams it back to VRAM rather than rebuilding it from disk.",
           "name": "offload_from_vram",

invokeai/app/invocations/qwen_image_text_encoder.py

@@ -1,4 +1,5 @@
-from typing import Literal
+import traceback
+from typing import Any, Literal
 
 import torch
 from PIL import Image as PILImage
@@ -14,6 +15,8 @@
 from invokeai.app.invocations.model import QwenVLEncoderField
 from invokeai.app.invocations.primitives import QwenImageConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.load.model_cache.model_cache import MODEL_LOAD_LOCK
+from invokeai.backend.model_manager.load.model_util import calc_model_size_by_fs
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     ConditioningFieldData,
     QwenImageConditioningInfo,
@@ -39,6 +42,13 @@
 )
 _GENERATE_DROP_IDX = 34
 
+# Fraction of the on-disk (bf16) encoder that stays resident after BitsAndBytes quantization. Linear weights
+# shrink to 8 or 4 bits, but embeddings, norms, biases and the (excluded) lm_head stay in bf16, and nf4 carries
+# per-block absmax scales, so the ratios sit above the pure 1/2 and 1/4. Over-estimating only offloads a little
+# more of the cached models to RAM, whereas under-estimating leaves the encoder without enough VRAM, so both
+# values are deliberately conservative.
+_QUANTIZED_SIZE_RATIO: dict[str, float] = {"int8": 0.6, "nf4": 0.4}
+
 _IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
 
 
@@ -212,43 +222,21 @@ def _encode(
                 padding=True,
                 return_tensors="pt",
             ).to(device=device)
-
-            outputs = text_encoder(
-                input_ids=model_inputs.input_ids,
-                attention_mask=model_inputs.attention_mask,
-                pixel_values=getattr(model_inputs, "pixel_values", None),
-                image_grid_thw=getattr(model_inputs, "image_grid_thw", None),
-                output_hidden_states=True,
-            )
-
-            # Use last hidden state (matching diffusers pipeline)
-            hidden_states = outputs.hidden_states[-1]
-
-            # Extract valid (non-padding) tokens using the attention mask,
-            # then drop the system prompt prefix tokens.
-            # The drop index differs between edit mode (64) and generate mode (34).
-            drop_idx = _EDIT_DROP_IDX if images else _GENERATE_DROP_IDX
-
-            attn_mask = model_inputs.attention_mask
-            bool_mask = attn_mask.bool()
-            valid_lengths = bool_mask.sum(dim=1)
-            selected = hidden_states[bool_mask]
-            split_hidden = torch.split(selected, valid_lengths.tolist(), dim=0)
-
-            # Drop system prefix tokens and build padded output
-            trimmed = [h[drop_idx:] for h in split_hidden]
-            attn_mask_list = [torch.ones(h.size(0), dtype=torch.long, device=device) for h in trimmed]
-            max_seq_len = max(h.size(0) for h in trimmed)
-
-            prompt_embeds = torch.stack(
-                [torch.cat([h, h.new_zeros(max_seq_len - h.size(0), h.size(1))]) for h in trimmed]
-            )
-            encoder_attention_mask = torch.stack(
-                [torch.cat([m, m.new_zeros(max_seq_len - m.size(0))]) for m in attn_mask_list]
-            )
-
-            prompt_embeds = prompt_embeds.to(dtype=torch.bfloat16)
+            prompt_embeds, encoder_attention_mask = self._run_encoder(text_encoder, model_inputs, bool(images))
+        except BaseException as exc:
+            # The in-flight traceback references the forward's frames, and through their locals the model, so the
+            # release below would otherwise be a no-op on this path. Clearing the finished frames drops those
+            # references while keeping the traceback's line information for the error report.
+            traceback.clear_frames(exc.__traceback__)
+            raise
         finally:
+            # Drop this frame's references before `cleanup` runs: the quantized encoder is only released once
+            # nothing holds it, and `cleanup` calls empty_cache() right after. With the model still alive here,
+            # that empty_cache() ran too early and ~9 GB of encoder weights stayed *reserved* by torch after the
+            # node finished. The cache budgets from allocated + driver-free VRAM, so reserved-but-unused memory
+            # looked like it was in use and the next model (the transformer) was needlessly partial-loaded. The
+            # activations live in `_run_encoder`'s frame and are gone by now for the same reason.
+            del text_encoder
             if cleanup is not None:
                 cleanup()
 
@@ -258,6 +246,51 @@ def _encode(
 
         return prompt_embeds, encoder_attention_mask
 
+    @staticmethod
+    def _run_encoder(
+        text_encoder: torch.nn.Module, model_inputs: Any, edit_mode: bool
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run the encoder and build the padded embeddings + mask.
+
+        Kept out of `_encode` on purpose: the full-vocabulary logits and per-layer hidden states the forward returns
+        are only referenced by this frame, so they are released as soon as it returns (or is cleared on error),
+        before `_encode` empties the CUDA cache.
+        """
+        device = model_inputs.input_ids.device
+        outputs = text_encoder(
+            input_ids=model_inputs.input_ids,
+            attention_mask=model_inputs.attention_mask,
+            pixel_values=getattr(model_inputs, "pixel_values", None),
+            image_grid_thw=getattr(model_inputs, "image_grid_thw", None),
+            output_hidden_states=True,
+        )
+
+        # Use last hidden state (matching diffusers pipeline)
+        hidden_states = outputs.hidden_states[-1]
+
+        # Extract valid (non-padding) tokens using the attention mask,
+        # then drop the system prompt prefix tokens.
+        # The drop index differs between edit mode (64) and generate mode (34).
+        drop_idx = _EDIT_DROP_IDX if edit_mode else _GENERATE_DROP_IDX
+
+        attn_mask = model_inputs.attention_mask
+        bool_mask = attn_mask.bool()
+        valid_lengths = bool_mask.sum(dim=1)
+        selected = hidden_states[bool_mask]
+        split_hidden = torch.split(selected, valid_lengths.tolist(), dim=0)
+
+        # Drop system prefix tokens and build padded output
+        trimmed = [h[drop_idx:] for h in split_hidden]
+        attn_mask_list = [torch.ones(h.size(0), dtype=torch.long, device=device) for h in trimmed]
+        max_seq_len = max(h.size(0) for h in trimmed)
+
+        prompt_embeds = torch.stack([torch.cat([h, h.new_zeros(max_seq_len - h.size(0), h.size(1))]) for h in trimmed])
+        encoder_attention_mask = torch.stack(
+            [torch.cat([m, m.new_zeros(max_seq_len - m.size(0))]) for m in attn_mask_list]
+        )
+
+        return prompt_embeds.to(dtype=torch.bfloat16), encoder_attention_mask
+
     def _load_cached_encoder(self, context: InvocationContext):
         """Load the text encoder through the model cache (no quantization)."""
         from transformers import Qwen2_5_VLForConditionalGeneration
@@ -277,6 +310,13 @@ def _load_quantized_encoder(self, context: InvocationContext):
         BnB-quantized models are pinned to GPU and can't be moved between devices,
         so they can't go through the standard model cache. The model is loaded fresh
         each time and freed after use via the cleanup callback.
+
+        Because the load bypasses the cache, it also bypasses the cache's usual
+        make-room-for-the-model-being-locked step, so this path has to ask the cache
+        for VRAM explicitly. Without that, whatever the resident transformer/VAE left
+        free is all the encoder gets: `device_map="auto"` then silently plans to spill
+        layers to the CPU, which BnB int8 refuses with "Some modules are dispatched on
+        the CPU or the disk" (issue #9147).
         """
         import gc
         import warnings
@@ -302,23 +342,34 @@ def _load_quantized_encoder(self, context: InvocationContext):
         else:  # int8
             bnb_config = BitsAndBytesConfig(load_in_8bit=True)
 
-        context.util.signal_progress("Loading Qwen2.5-VL encoder (quantized)")
-        with warnings.catch_warnings():
-            # BnB int8 internally casts bfloat16→float16; the warning is harmless
-            warnings.filterwarnings("ignore", message="MatMul8bitLt.*cast.*float16")
-            text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                str(encoder_path),
-                quantization_config=bnb_config,
-                device_map="auto",
-                torch_dtype=torch.bfloat16,
-                local_files_only=True,
-            )
-
-        device = next(text_encoder.parameters()).device
+        # Load onto this worker's execution device, never `device_map="auto"`: "auto" sizes its plan from whatever
+        # VRAM is free *right now* and quietly spills to the CPU when the cached models fill the card, and in
+        # multi-GPU mode it may also pick a device other than the one this worker is pinned to. With an explicit
+        # device, a genuine shortfall surfaces as a plain OOM instead of a misleading offload error.
+        device = TorchDevice.choose_torch_device()
+        quantized_bytes = int(calc_model_size_by_fs(encoder_path) * _QUANTIZED_SIZE_RATIO[self.quantization])
 
+        context.util.signal_progress("Loading Qwen2.5-VL encoder (quantized)")
+        # Both the offload and the load below assign real parameters (`register_parameter` via `load_state_dict` /
+        # `setattr`), which a concurrent cache construction on another worker would hijack onto the meta device
+        # (see MODEL_LOAD_LOCK). Hold the read lock across them, like every other VRAM move, and take it *before*
+        # the cache lock that `make_room_in_vram` acquires, per the lock-ordering contract on MODEL_LOAD_LOCK.
+        with MODEL_LOAD_LOCK.read_lock():
+            context.models.make_room_in_vram(quantized_bytes)
+            with warnings.catch_warnings():
+                # BnB int8 internally casts bfloat16→float16; the warning is harmless
+                warnings.filterwarnings("ignore", message="MatMul8bitLt.*cast.*float16")
+                text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                    str(encoder_path),
+                    quantization_config=bnb_config,
+                    device_map={"": device},
+                    torch_dtype=torch.bfloat16,
+                    local_files_only=True,
+                )
+
+        # Hand the model out without keeping a reference in this closure, so that once `_encode` drops its own the
+        # weights are actually free by the time empty_cache() runs.
         def cleanup():
-            nonlocal text_encoder
-            del text_encoder
             gc.collect()
             TorchDevice.empty_cache()
 

invokeai/app/services/shared/invocation_context.py

@@ -600,6 +600,24 @@ def offload_from_vram(self, identifier: Union[str, "ModelIdentifierField"]) -> i
         key = identifier if isinstance(identifier, str) else identifier.key
         return self._services.model_manager.load.ram_cache.offload_model_from_vram(key)
 
+    def make_room_in_vram(self, vram_bytes_needed: int, working_mem_bytes: Optional[int] = None) -> int:
+        """Offload unlocked cached models from VRAM until `vram_bytes_needed` bytes are free on this thread's
+        execution device.
+
+        Use this before placing a model on the GPU *outside* the model cache (e.g. a BitsAndBytes-quantized model
+        that cannot be moved between devices). Loads that go through `load()` never need this - the cache makes room
+        for them itself when they are locked - but an out-of-cache load competes with the cached models for VRAM
+        and would otherwise only get whatever they happened to leave free.
+
+        Args:
+            vram_bytes_needed: The VRAM footprint the caller is about to allocate.
+            working_mem_bytes: Working memory to keep free on top of the model, floored at the configured default.
+
+        Returns:
+            The number of VRAM bytes freed.
+        """
+        return self._services.model_manager.load.ram_cache.make_room_in_vram(vram_bytes_needed, working_mem_bytes)
+
     @staticmethod
     def _raise_if_external(model: AnyModelConfig) -> None:
         if model.base == BaseModelType.External or model.format == ModelFormat.ExternalApi:

invokeai/backend/model_manager/load/model_cache/model_cache.py

@@ -2749,6 +2749,29 @@ def drop_model(self, model_key: str) -> int:
             TorchDevice.empty_cache()
         return len(dropped)
 
+    @synchronized
+    def make_room_in_vram(self, vram_bytes_needed: int, working_mem_bytes: Optional[int] = None) -> int:
+        """Offload unlocked models from VRAM to RAM until `vram_bytes_needed` bytes are free on the execution device.
+
+        This is the entry point for code that has to put a model on the GPU *outside* the cache - e.g. a
+        BitsAndBytes-quantized text encoder, which is pinned to the device it was quantized on and so cannot be
+        managed by the cache. Such a load competes with cached models for VRAM, but never passes through `lock()`,
+        which is where the cache normally makes room for the model being locked. Without an explicit request, the
+        out-of-cache load only sees whatever VRAM the resident models happened to leave free.
+
+        The same policy as `lock()` is used (`_offload_unlocked_models`): unlocked models are offloaded to RAM until
+        the availability check is satisfied, and kept in the cache so a later use re-streams weights instead of
+        rebuilding from disk. Locked (in-use) models are never touched. `working_mem_bytes` is the operation's
+        working memory and is floored at the configured default, exactly as in `lock()`.
+
+        A CPU execution device has no VRAM to make room in, so the call is a no-op there (as `lock()` is).
+
+        Returns the number of VRAM bytes freed based on believed model sizes.
+        """
+        if self._execution_device.type == "cpu":
+            return 0
+        return self._offload_unlocked_models(vram_bytes_needed, working_mem_bytes)
+
     @synchronized
     def offload_model_from_vram(self, model_key: str) -> int:
         """Move a model (and its submodels) from VRAM to RAM without dropping it from the cache.