From 7d07921b88200976b5924db075c00fe2cb8313e7 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 7 Sep 2026 12:22:33 -0400 Subject: [PATCH 1/7] fix(qwen): make room in VRAM before loading the quantized Qwen2.5-VL encoder The int8/nf4 encoder path bypasses the model cache (BitsAndBytes models are pinned to the device they were quantized on), so it also bypassed the make-room-for-the-model step that `lock()` performs for cached loads. With the transformer and VAE still resident, `device_map="auto"` planned the encoder onto the CPU and BnB int8 refused with "Some modules are dispatched on the CPU or the disk" as soon as the prompt changed after a generation (#9147). - `ModelCache.make_room_in_vram(bytes, working_mem)` exposes the same smallest-first offload policy `lock()` uses, for loads that live outside the cache. Surfaced on the invocation context as `context.models.make_room_in_vram`. - The quantized encoder estimates its post-quantization footprint from the bf16 checkpoint size, asks the cache for that much VRAM, and loads with an explicit `device_map={"": }` so a genuine shortfall is a plain OOM rather than a misleading offload error. In multi-GPU mode this also pins the load to the worker's own device instead of letting "auto" pick one. - Offload and 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. Fixes #9147. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01C7hM9XHDWGNa57jY5pat98 --- .../invocations/qwen_image_text_encoder.py | 51 +++-- .../app/services/shared/invocation_context.py | 18 ++ .../load/model_cache/model_cache.py | 19 ++ .../test_qwen_image_text_encoder.py | 99 ++++++++++ .../test_model_cache_make_room_in_vram.py | 174 ++++++++++++++++++ 5 files changed, 349 insertions(+), 12 deletions(-) create mode 100644 tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py diff --git a/invokeai/app/invocations/qwen_image_text_encoder.py b/invokeai/app/invocations/qwen_image_text_encoder.py index 66b94e7ec17..38dd2646bb7 100644 --- a/invokeai/app/invocations/qwen_image_text_encoder.py +++ b/invokeai/app/invocations/qwen_image_text_encoder.py @@ -14,6 +14,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 +41,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|>" @@ -277,6 +286,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,19 +318,30 @@ 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, - ) + # 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]) - device = next(text_encoder.parameters()).device + 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, + ) def cleanup(): nonlocal text_encoder diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index e261fc89393..320cc0ac538 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/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: diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index d223a899646..cb0dadc50ec 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -2749,6 +2749,25 @@ 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`): models are (partially) offloaded to RAM, + smallest first, 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()`. + + Returns the number of VRAM bytes freed based on believed model sizes. + """ + 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. diff --git a/tests/app/invocations/test_qwen_image_text_encoder.py b/tests/app/invocations/test_qwen_image_text_encoder.py index ab3beabae7f..0d8309abe7e 100644 --- a/tests/app/invocations/test_qwen_image_text_encoder.py +++ b/tests/app/invocations/test_qwen_image_text_encoder.py @@ -1,11 +1,22 @@ """Tests for the Qwen Image text encoder prompt building and image resizing.""" +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import torch from PIL import Image +from transformers import Qwen2_5_VLForConditionalGeneration +from invokeai.app.invocations.model import ModelIdentifierField, QwenVLEncoderField from invokeai.app.invocations.qwen_image_text_encoder import ( QwenImageTextEncoderInvocation, _build_prompt, ) +from invokeai.backend.model_manager.load.model_cache.model_cache import MODEL_LOAD_LOCK +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType +from invokeai.backend.util.devices import TorchDevice class TestBuildPrompt: @@ -122,3 +133,91 @@ def test_landscape_image(self): resized = QwenImageTextEncoderInvocation._resize_for_vl_encoder(img, target_pixels=512 * 512) w, h = resized.size assert w > h # should remain landscape + + +class TestQuantizedEncoderLoad: + """The BitsAndBytes path bypasses the model cache, so it must ask the cache for VRAM itself and load onto the + worker's execution device explicitly (issue #9147: `device_map="auto"` spilled to the CPU because the cached + transformer and VAE still filled the card, and BnB int8 refused to run that way). + """ + + TOTAL_SIZE = 16 * 2**30 # bf16 Qwen2.5-VL-7B on disk + + @staticmethod + def _make_invocation(quantization: str) -> QwenImageTextEncoderInvocation: + encoder = ModelIdentifierField( + key="enc", hash="h", name="qwen-vl", base=BaseModelType.QwenImage, type=ModelType.QwenVLEncoder + ) + return QwenImageTextEncoderInvocation( + prompt="a cat", + qwen_vl_encoder=QwenVLEncoderField(tokenizer=encoder, text_encoder=encoder), + quantization=quantization, + ) + + def _make_context(self, tmp_path: Path, events: list[str], single_file: bool = False) -> MagicMock: + if single_file: + model_root = tmp_path / "encoder.safetensors" + model_root.write_bytes(b"") + else: + model_root = tmp_path / "qwen-vl" + text_encoder_dir = model_root / "text_encoder" + text_encoder_dir.mkdir(parents=True) + index = {"metadata": {"total_size": self.TOTAL_SIZE}, "weight_map": {}} + (text_encoder_dir / "model.safetensors.index.json").write_text(json.dumps(index)) + + context = MagicMock() + context.models.get_absolute_path.return_value = model_root + context.models.make_room_in_vram.side_effect = lambda *a, **k: events.append("make_room") or 0 + return context + + @pytest.mark.parametrize(("quantization", "ratio"), [("int8", 0.6), ("nf4", 0.4)]) + def test_makes_room_in_vram_before_loading_onto_the_execution_device( + self, tmp_path: Path, quantization: str, ratio: float + ): + events: list[str] = [] + context = self._make_context(tmp_path, events) + fake_model = MagicMock() + device = torch.device("cuda:1") + seen: dict = {} + + def fake_from_pretrained(path, **kwargs): + events.append("from_pretrained") + seen.update(kwargs) + # The load must run under the model-load lock so a concurrent cache construction on another worker + # cannot hijack its parameter assignment onto the meta device. + seen["locked"] = MODEL_LOAD_LOCK._readers > 0 or MODEL_LOAD_LOCK._writer_active + return fake_model + + with ( + patch.object(Qwen2_5_VLForConditionalGeneration, "from_pretrained", side_effect=fake_from_pretrained), + patch.object(TorchDevice, "choose_torch_device", return_value=device), + ): + text_encoder, returned_device, cleanup = self._make_invocation(quantization)._load_quantized_encoder( + context + ) + + assert events == ["make_room", "from_pretrained"] + context.models.make_room_in_vram.assert_called_once_with(int(self.TOTAL_SIZE * ratio)) + assert seen["device_map"] == {"": device}, "must never use device_map='auto'" + assert seen["locked"] + assert text_encoder is fake_model + assert returned_device == device + cleanup() + + def test_single_file_checkpoint_falls_back_to_the_cache_without_making_room(self, tmp_path: Path): + """A single-file encoder cannot be BnB-quantized; it goes through the cache, which makes its own room.""" + events: list[str] = [] + context = self._make_context(tmp_path, events, single_file=True) + invocation = self._make_invocation("int8") + sentinel = (MagicMock(), torch.device("cuda"), None) + + with ( + patch.object(invocation, "_load_cached_encoder", return_value=sentinel) as cached, + patch.object(Qwen2_5_VLForConditionalGeneration, "from_pretrained") as from_pretrained, + ): + result = invocation._load_quantized_encoder(context) + + assert result is sentinel + cached.assert_called_once_with(context) + from_pretrained.assert_not_called() + context.models.make_room_in_vram.assert_not_called() diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py b/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py new file mode 100644 index 00000000000..d1337e03cf9 --- /dev/null +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py @@ -0,0 +1,174 @@ +"""Tests for `ModelCache.make_room_in_vram` - the entry point for loads that put a model on the GPU *outside* +the cache (e.g. a BitsAndBytes-quantized text encoder that cannot be moved between devices). + +Such a load never passes through `lock()`, which is where the cache normally makes room for the model being +locked, so without an explicit request it only gets whatever VRAM the resident models happened to leave free. +Issue #9147: the quantized Qwen2.5-VL encoder was planned onto the CPU by `device_map="auto"` because the cached +transformer and VAE still filled the card, and BitsAndBytes int8 refused to run that way. +""" + +import logging +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache +from tests.backend.model_manager.load.model_cache.cached_model.utils import DummyModule + +MB = 2**20 + + +@pytest.fixture +def mock_logger(): + logger = MagicMock() + logger.getEffectiveLevel.return_value = logging.INFO + return logger + + +@pytest.fixture +def cache(mock_logger): + cache = ModelCache( + execution_device_working_mem_gb=1.0, + enable_partial_loading=False, + keep_ram_copy_of_weights=True, + execution_device="cpu", + storage_device="cpu", + logger=mock_logger, + ) + yield cache + cache.shutdown() + + +class _FakeVram: + """Simulates the execution device's free VRAM as the cache offloads models. + + `ModelCache._get_vram_available` needs a real accelerator, so the CPU-only tests below replace it with this + bookkeeping: `available` grows by whatever `_move_model_to_ram` reports freed. + """ + + def __init__(self, available: int): + self.available = available + self.working_mem_seen: list[int | None] = [] + self.moved: list[tuple[str, int]] = [] + + def get_vram_available(self, working_mem_bytes): + self.working_mem_seen.append(working_mem_bytes) + return self.available + + def move_model_to_ram(self, cache_entry, vram_bytes_to_free, keep_required_weights_in_vram=None): + freed = cache_entry.cached_model.total_bytes() + self.moved.append((cache_entry.key, vram_bytes_to_free)) + self.available += freed + return freed + + +def _put(cache: ModelCache, key: str, size_bytes: int) -> None: + """Cache a module whose weights occupy exactly `size_bytes` (a bare tensor is sized at 0 by the cache).""" + assert size_bytes % 4 == 0 + module = torch.nn.Linear(size_bytes // 4, 1, bias=False) # fp32: 4 bytes per weight + assert sum(p.numel() * p.element_size() for p in module.parameters()) == size_bytes + cache.put(key, module) + + +def test_offloads_unlocked_models_until_the_request_is_satisfied(cache: ModelCache): + """Models are offloaded smallest-first, and the loop stops as soon as enough VRAM is free - the larger + model that is not needed stays resident.""" + _put(cache, "small", 10 * MB) + _put(cache, "medium", 20 * MB) + _put(cache, "large", 40 * MB) + vram = _FakeVram(available=5 * MB) + + with ( + patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), + patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), + ): + freed = cache.make_room_in_vram(30 * MB) + + assert [key for key, _ in vram.moved] == ["small", "medium"] + assert freed == 30 * MB + assert vram.available == 35 * MB + + +def test_locked_models_are_never_offloaded(cache: ModelCache): + """A locked model is in use by another invocation; it must be skipped even when the request cannot otherwise + be satisfied.""" + _put(cache, "in_use", 40 * MB) + _put(cache, "idle", 10 * MB) + cache._cached_models["in_use"].lock() + vram = _FakeVram(available=0) + + with ( + patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), + patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), + ): + freed = cache.make_room_in_vram(100 * MB) + + assert [key for key, _ in vram.moved] == ["idle"] + assert freed == 10 * MB + + +def test_no_op_when_enough_vram_is_already_free(cache: ModelCache): + _put(cache, "resident", 40 * MB) + vram = _FakeVram(available=50 * MB) + + with ( + patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), + patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), + ): + freed = cache.make_room_in_vram(30 * MB) + + assert vram.moved == [] + assert freed == 0 + + +def test_working_memory_is_forwarded_to_the_availability_check(cache: ModelCache): + """The caller's working memory must reach `_get_vram_available`, where it is floored at the configured default + exactly as in `lock()`; otherwise the encoder's activations would have to fit in whatever is left over.""" + _put(cache, "resident", 40 * MB) + vram = _FakeVram(available=0) + + with ( + patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), + patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), + ): + cache.make_room_in_vram(10 * MB, working_mem_bytes=7 * MB) + + assert vram.working_mem_seen and all(w == 7 * MB for w in vram.working_mem_seen) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.") +@pytest.mark.parametrize("partial", [False, True]) +def test_gpu_make_room_in_vram_actually_moves_weights_off_the_device(partial: bool): + """Real-accelerator check: after `make_room_in_vram`, the unlocked models' weights are on the CPU and the + locked model's weights are still on the GPU, with both cached-model flavours.""" + logger = MagicMock() + logger.getEffectiveLevel.return_value = logging.INFO + cache = ModelCache( + execution_device_working_mem_gb=0.0, + enable_partial_loading=partial, + keep_ram_copy_of_weights=True, + execution_device="cuda", + storage_device="cpu", + logger=logger, + ) + try: + idle, in_use = DummyModule(), DummyModule() + cache.put("idle", idle) + cache.put("in_use", in_use) + for key in ("idle", "in_use"): + cache._cached_models[key].cached_model.full_load_to_vram() + cache._cached_models["in_use"].lock() + assert all(p.device.type == "cuda" for p in idle.parameters()) + + # Ask for more than the whole card so every unlocked model has to go. + _, total = torch.cuda.mem_get_info() + freed = cache.make_room_in_vram(2 * total) + + assert freed == cache._cached_models["idle"].cached_model.total_bytes() + assert all(p.device.type == "cpu" for p in idle.parameters()) + assert all(p.device.type == "cuda" for p in in_use.parameters()) + # Same policy as lock(): the entry is offloaded, not evicted, so the next use re-streams weights. + assert "idle" in cache._cached_models + finally: + cache.shutdown() From aeb0b53b8ebf5ac8f1c65f45e0d3c14e51eae0de Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 7 Sep 2026 12:30:47 -0400 Subject: [PATCH 2/7] fix(qwen): release the quantized encoder before emptying the CUDA cache `_encode` still held the BitsAndBytes encoder when its cleanup callback ran, so the callback's `del` freed nothing and `empty_cache()` ran while ~9 GB of encoder weights were alive. They were released only when the frame exited, after which they stayed *reserved* by torch. The model cache budgets from 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 - the "transformer 57% in VRAM" swings reported in #9147. Measured on a W7900 with the bf16 Qwen2.5-VL-7B encoder quantized to int8: cleanup with the frame's references live left 40.7 GB allocated; dropping them first brought allocated to 31.9 GB, and empty_cache() then returned 8.9 GB to the driver. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01C7hM9XHDWGNa57jY5pat98 --- .../invocations/qwen_image_text_encoder.py | 10 +++- .../test_qwen_image_text_encoder.py | 60 +++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/invokeai/app/invocations/qwen_image_text_encoder.py b/invokeai/app/invocations/qwen_image_text_encoder.py index 38dd2646bb7..1db6f653135 100644 --- a/invokeai/app/invocations/qwen_image_text_encoder.py +++ b/invokeai/app/invocations/qwen_image_text_encoder.py @@ -258,6 +258,12 @@ def _encode( prompt_embeds = prompt_embeds.to(dtype=torch.bfloat16) 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 its own `del`. 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. + del text_encoder if cleanup is not None: cleanup() @@ -343,9 +349,9 @@ def _load_quantized_encoder(self, context: InvocationContext): 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() diff --git a/tests/app/invocations/test_qwen_image_text_encoder.py b/tests/app/invocations/test_qwen_image_text_encoder.py index 0d8309abe7e..6250aa65543 100644 --- a/tests/app/invocations/test_qwen_image_text_encoder.py +++ b/tests/app/invocations/test_qwen_image_text_encoder.py @@ -1,6 +1,7 @@ """Tests for the Qwen Image text encoder prompt building and image resizing.""" import json +import weakref from pathlib import Path from unittest.mock import MagicMock, patch @@ -11,6 +12,7 @@ from invokeai.app.invocations.model import ModelIdentifierField, QwenVLEncoderField from invokeai.app.invocations.qwen_image_text_encoder import ( + _GENERATE_DROP_IDX, QwenImageTextEncoderInvocation, _build_prompt, ) @@ -221,3 +223,61 @@ def test_single_file_checkpoint_falls_back_to_the_cache_without_making_room(self cached.assert_called_once_with(context) from_pretrained.assert_not_called() context.models.make_room_in_vram.assert_not_called() + + +class TestQuantizedEncoderRelease: + """The quantized encoder lives outside the cache, so `_encode` is its only owner. Its cleanup callback empties + the CUDA cache, which only returns the ~9 GB of encoder weights to the driver if nothing still references the + model at that point - otherwise they stay reserved by torch and the cache under-budgets the next load. + """ + + HIDDEN = 3584 + + class _FakeEncoder(torch.nn.Module): + def __init__(self, hidden: int): + super().__init__() + self.hidden = hidden + + def forward(self, input_ids, attention_mask, **_): + hidden_states = torch.zeros(input_ids.shape[0], input_ids.shape[1], self.hidden) + return MagicMock(hidden_states=[hidden_states]) + + def test_encoder_is_released_before_cleanup_runs(self, tmp_path: Path): + model_root = tmp_path / "qwen-vl" + (model_root / "tokenizer").mkdir(parents=True) + context = MagicMock() + context.models.get_absolute_path.return_value = model_root + + encoder = self._FakeEncoder(self.HIDDEN) + encoder_ref = weakref.ref(encoder) + alive_at_cleanup: list[bool] = [] + + def cleanup(): + alive_at_cleanup.append(encoder_ref() is not None) + + seq_len = _GENERATE_DROP_IDX + 5 + model_inputs = MagicMock() + model_inputs.input_ids = torch.zeros(1, seq_len, dtype=torch.long) + model_inputs.attention_mask = torch.ones(1, seq_len, dtype=torch.long) + model_inputs.to.return_value = model_inputs + processor = MagicMock(return_value=model_inputs) + + # Hand the encoder over through a one-shot side effect: a `return_value` tuple would keep the mock holding a + # strong reference of its own and mask the ownership being tested. + handoff = [encoder] + invocation = TestQuantizedEncoderLoad._make_invocation("int8") + with ( + patch.object( + invocation, + "_load_quantized_encoder", + side_effect=lambda _ctx: (handoff.pop(), torch.device("cpu"), cleanup), + ), + patch("transformers.AutoTokenizer.from_pretrained", return_value=MagicMock()), + patch("transformers.Qwen2_5_VLProcessor", return_value=processor), + ): + del encoder # `_encode` now holds the only strong reference + prompt_embeds, mask = invocation._encode(context, images=[]) + + assert alive_at_cleanup == [False], "cleanup ran while the encoder was still referenced" + assert prompt_embeds.shape == (1, 5, self.HIDDEN) + assert mask is None From a781393937ba9c462af023d39a437fd95c7ab6f2 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 7 Sep 2026 12:36:07 -0400 Subject: [PATCH 3/7] fix(qwen): make_room_in_vram is a no-op on a CPU execution device; tighten tests Findings from the adversarial review of the first commit: - On a CPU-only install `_get_vram_available` raises for the cpu device. `lock()` never reaches it there, but `make_room_in_vram` did, so the second quantized-encoder run (once anything was cached) failed with "Unsupported execution device: cpu" where it used to load. There is no VRAM to make room in, so short-circuit like `lock()` does. - The invocation test only sampled the model-load lock inside `from_pretrained`; moving `make_room_in_vram` outside the read lock survived it. The offload is a VRAM move like any other, so it now has to be under the lock as well. - Nothing proved `make_room_in_vram` owns the cache lock, and nothing covered the invocation-context delegation (dropping `working_mem_bytes` there survived). Both are pinned now. - The docstring claimed a selective smallest-first offload. That is the loop's contract, but on hardware the driver only sees freed memory after the trailing `empty_cache()`, so `lock()`'s policy (which this reuses) tends to offload every unlocked model. Reworded; the policy itself is out of scope. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01C7hM9XHDWGNa57jY5pat98 --- .../load/model_cache/model_cache.py | 12 ++-- .../test_qwen_image_text_encoder.py | 14 ++++- ...st_invocation_context_make_room_in_vram.py | 16 +++++ .../test_model_cache_make_room_in_vram.py | 63 +++++++++++++++++-- 4 files changed, 93 insertions(+), 12 deletions(-) create mode 100644 tests/app/services/shared/test_invocation_context_make_room_in_vram.py diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index cb0dadc50ec..f321c804ac9 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -2759,13 +2759,17 @@ def make_room_in_vram(self, vram_bytes_needed: int, working_mem_bytes: Optional[ 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`): models are (partially) offloaded to RAM, - smallest first, 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()`. + 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 diff --git a/tests/app/invocations/test_qwen_image_text_encoder.py b/tests/app/invocations/test_qwen_image_text_encoder.py index 6250aa65543..c61bd0c9316 100644 --- a/tests/app/invocations/test_qwen_image_text_encoder.py +++ b/tests/app/invocations/test_qwen_image_text_encoder.py @@ -137,6 +137,10 @@ def test_landscape_image(self): assert w > h # should remain landscape +def _model_load_lock_held() -> bool: + return MODEL_LOAD_LOCK._readers > 0 or MODEL_LOAD_LOCK._writer_active + + class TestQuantizedEncoderLoad: """The BitsAndBytes path bypasses the model cache, so it must ask the cache for VRAM itself and load onto the worker's execution device explicitly (issue #9147: `device_map="auto"` spilled to the CPU because the cached @@ -169,7 +173,13 @@ def _make_context(self, tmp_path: Path, events: list[str], single_file: bool = F context = MagicMock() context.models.get_absolute_path.return_value = model_root - context.models.make_room_in_vram.side_effect = lambda *a, **k: events.append("make_room") or 0 + + def make_room(*_args, **_kwargs): + # The offload is a VRAM move like any other, so it too must run under the model-load lock. + events.append("make_room" if _model_load_lock_held() else "make_room(unlocked)") + return 0 + + context.models.make_room_in_vram.side_effect = make_room return context @pytest.mark.parametrize(("quantization", "ratio"), [("int8", 0.6), ("nf4", 0.4)]) @@ -187,7 +197,7 @@ def fake_from_pretrained(path, **kwargs): seen.update(kwargs) # The load must run under the model-load lock so a concurrent cache construction on another worker # cannot hijack its parameter assignment onto the meta device. - seen["locked"] = MODEL_LOAD_LOCK._readers > 0 or MODEL_LOAD_LOCK._writer_active + seen["locked"] = _model_load_lock_held() return fake_model with ( diff --git a/tests/app/services/shared/test_invocation_context_make_room_in_vram.py b/tests/app/services/shared/test_invocation_context_make_room_in_vram.py new file mode 100644 index 00000000000..830fb0ef26e --- /dev/null +++ b/tests/app/services/shared/test_invocation_context_make_room_in_vram.py @@ -0,0 +1,16 @@ +"""`context.models.make_room_in_vram` must reach the calling thread's device cache with both arguments intact: +the invocation context is the only route an invocation has to the cache, and a dropped `working_mem_bytes` would +silently fall back to the configured default.""" + +from unittest.mock import MagicMock + +from invokeai.app.services.shared.invocation_context import ModelsInterface + + +def test_make_room_in_vram_delegates_to_the_threads_ram_cache(): + services = MagicMock() + services.model_manager.load.ram_cache.make_room_in_vram.return_value = 123 + models = ModelsInterface(services=services, data=MagicMock(), util=MagicMock()) + + assert models.make_room_in_vram(10, working_mem_bytes=7) == 123 + services.model_manager.load.ram_cache.make_room_in_vram.assert_called_once_with(10, 7) diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py b/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py index d1337e03cf9..1e622efe717 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py @@ -71,9 +71,57 @@ def _put(cache: ModelCache, key: str, size_bytes: int) -> None: cache.put(key, module) -def test_offloads_unlocked_models_until_the_request_is_satisfied(cache: ModelCache): - """Models are offloaded smallest-first, and the loop stops as soon as enough VRAM is free - the larger - model that is not needed stays resident.""" +@pytest.fixture +def gpu_accounting_cache(mock_logger): + """A cache whose execution device is a GPU as far as the policy is concerned, without touching a real one. + + `_get_vram_available` needs an accelerator; the tests replace it (and the VRAM moves) with `_FakeVram`. Note that + `_FakeVram` reports freed memory immediately, which is the loop's *contract*; on real hardware the driver only + sees it after the trailing `empty_cache()`, so the loop there tends to offload every unlocked model. + """ + cache = ModelCache( + execution_device_working_mem_gb=1.0, + enable_partial_loading=False, + keep_ram_copy_of_weights=True, + execution_device="cpu", + storage_device="cpu", + logger=mock_logger, + ) + cache._execution_device = torch.device("cuda") # policy only; every VRAM touch is patched out below + yield cache + cache._execution_device = torch.device("cpu") + cache.shutdown() + + +def test_cpu_execution_device_is_a_no_op(cache: ModelCache): + """A CPU-only install has no VRAM to make room in. `_get_vram_available` raises for a cpu device, and the + quantized encoder path calls this on every run once anything is cached, so it must short-circuit.""" + _put(cache, "resident", 40 * MB) + + assert cache.make_room_in_vram(30 * MB) == 0 + assert "resident" in cache._cached_models + + +def test_offload_runs_under_the_cache_lock(gpu_accounting_cache: ModelCache): + """Out-of-cache callers race the session workers' own lock()/unlock(); the offload must own the cache lock.""" + cache = gpu_accounting_cache + _put(cache, "resident", 40 * MB) + owned: list[bool] = [] + + def offload(vram_bytes_required, working_mem_bytes=None): + owned.append(cache._lock._is_owned()) + return 0 + + with patch.object(cache, "_offload_unlocked_models", side_effect=offload): + cache.make_room_in_vram(30 * MB) + + assert owned == [True] + + +def test_offloads_unlocked_models_until_the_request_is_satisfied(gpu_accounting_cache: ModelCache): + """Models are offloaded smallest-first, and the loop stops once the availability check reports enough free + VRAM - the larger model that is not needed stays resident.""" + cache = gpu_accounting_cache _put(cache, "small", 10 * MB) _put(cache, "medium", 20 * MB) _put(cache, "large", 40 * MB) @@ -90,7 +138,8 @@ def test_offloads_unlocked_models_until_the_request_is_satisfied(cache: ModelCac assert vram.available == 35 * MB -def test_locked_models_are_never_offloaded(cache: ModelCache): +def test_locked_models_are_never_offloaded(gpu_accounting_cache: ModelCache): + cache = gpu_accounting_cache """A locked model is in use by another invocation; it must be skipped even when the request cannot otherwise be satisfied.""" _put(cache, "in_use", 40 * MB) @@ -108,7 +157,8 @@ def test_locked_models_are_never_offloaded(cache: ModelCache): assert freed == 10 * MB -def test_no_op_when_enough_vram_is_already_free(cache: ModelCache): +def test_no_op_when_enough_vram_is_already_free(gpu_accounting_cache: ModelCache): + cache = gpu_accounting_cache _put(cache, "resident", 40 * MB) vram = _FakeVram(available=50 * MB) @@ -122,7 +172,8 @@ def test_no_op_when_enough_vram_is_already_free(cache: ModelCache): assert freed == 0 -def test_working_memory_is_forwarded_to_the_availability_check(cache: ModelCache): +def test_working_memory_is_forwarded_to_the_availability_check(gpu_accounting_cache: ModelCache): + cache = gpu_accounting_cache """The caller's working memory must reach `_get_vram_available`, where it is floored at the configured default exactly as in `lock()`; otherwise the encoder's activations would have to fit in whatever is left over.""" _put(cache, "resident", 40 * MB) From 42eed4b011b8e2a6897763e2f7d9a3e0783bdfd6 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 7 Sep 2026 12:47:59 -0400 Subject: [PATCH 4/7] fix(qwen): release the encoder on the error path and drop activations before empty_cache Two residuals of the release mechanism, from the adversarial review: - When the forward raises (an OOM is the likeliest failure here), the in-flight traceback references the forward's frames, and through their locals the model, so `del text_encoder` freed nothing and the ~9 GB stayed reserved into the next generation - exactly when VRAM is scarcest. Clear the finished traceback frames before re-raising; the raising line survives for the error report. - The full-vocabulary logits and per-layer hidden states (~0.5 GB in edit mode with three reference images) were still held by `_encode`'s locals during empty_cache(). Moved the forward and post-processing into a helper so they die with its frame on both the normal and the error path. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01C7hM9XHDWGNa57jY5pat98 --- .../invocations/qwen_image_text_encoder.py | 100 +++++++++++------- .../test_qwen_image_text_encoder.py | 34 ++++-- 2 files changed, 87 insertions(+), 47 deletions(-) diff --git a/invokeai/app/invocations/qwen_image_text_encoder.py b/invokeai/app/invocations/qwen_image_text_encoder.py index 1db6f653135..bfcffa9ca38 100644 --- a/invokeai/app/invocations/qwen_image_text_encoder.py +++ b/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 @@ -221,48 +222,20 @@ 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 its own `del`. 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. + # 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() @@ -273,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 diff --git a/tests/app/invocations/test_qwen_image_text_encoder.py b/tests/app/invocations/test_qwen_image_text_encoder.py index c61bd0c9316..ba97ad0761e 100644 --- a/tests/app/invocations/test_qwen_image_text_encoder.py +++ b/tests/app/invocations/test_qwen_image_text_encoder.py @@ -1,6 +1,8 @@ """Tests for the Qwen Image text encoder prompt building and image resizing.""" +import gc import json +import traceback import weakref from pathlib import Path from unittest.mock import MagicMock, patch @@ -252,17 +254,22 @@ def forward(self, input_ids, attention_mask, **_): hidden_states = torch.zeros(input_ids.shape[0], input_ids.shape[1], self.hidden) return MagicMock(hidden_states=[hidden_states]) - def test_encoder_is_released_before_cleanup_runs(self, tmp_path: Path): + class _ExplodingEncoder(torch.nn.Module): + def forward(self, input_ids, attention_mask, **_): + raise RuntimeError("CUDA out of memory (simulated)") + + def _run_encode(self, tmp_path: Path, encoder: torch.nn.Module, alive_at_cleanup: list[bool]): + """Run `_encode` as the sole owner of `encoder`, recording into `alive_at_cleanup` whether it was still alive + when cleanup ran (recorded through the argument so the record survives an `_encode` that raises).""" model_root = tmp_path / "qwen-vl" - (model_root / "tokenizer").mkdir(parents=True) + (model_root / "tokenizer").mkdir(parents=True, exist_ok=True) context = MagicMock() context.models.get_absolute_path.return_value = model_root - encoder = self._FakeEncoder(self.HIDDEN) encoder_ref = weakref.ref(encoder) - alive_at_cleanup: list[bool] = [] def cleanup(): + gc.collect() # as production does; the frame local under test is a strong root gc cannot clear alive_at_cleanup.append(encoder_ref() is not None) seq_len = _GENERATE_DROP_IDX + 5 @@ -275,6 +282,7 @@ def cleanup(): # Hand the encoder over through a one-shot side effect: a `return_value` tuple would keep the mock holding a # strong reference of its own and mask the ownership being tested. handoff = [encoder] + del encoder invocation = TestQuantizedEncoderLoad._make_invocation("int8") with ( patch.object( @@ -285,9 +293,23 @@ def cleanup(): patch("transformers.AutoTokenizer.from_pretrained", return_value=MagicMock()), patch("transformers.Qwen2_5_VLProcessor", return_value=processor), ): - del encoder # `_encode` now holds the only strong reference - prompt_embeds, mask = invocation._encode(context, images=[]) + return invocation._encode(context, images=[]) + + def test_encoder_is_released_before_cleanup_runs(self, tmp_path: Path): + alive_at_cleanup: list[bool] = [] + prompt_embeds, mask = self._run_encode(tmp_path, self._FakeEncoder(self.HIDDEN), alive_at_cleanup) assert alive_at_cleanup == [False], "cleanup ran while the encoder was still referenced" assert prompt_embeds.shape == (1, 5, self.HIDDEN) assert mask is None + + def test_encoder_is_released_before_cleanup_runs_when_the_forward_raises(self, tmp_path: Path): + """An OOM inside the forward is the likeliest failure here. The in-flight traceback holds the forward's + frames, whose locals reference the model, so without clearing them the release is a no-op exactly when + VRAM is scarcest - and the next generation starts from a mis-budgeted cache.""" + alive_at_cleanup: list[bool] = [] + with pytest.raises(RuntimeError, match="simulated") as excinfo: + self._run_encode(tmp_path, self._ExplodingEncoder(), alive_at_cleanup) + assert alive_at_cleanup == [False], "cleanup ran while the traceback still referenced the encoder" + # The traceback is still useful for the error report: the raising line is intact. + assert "simulated" in "".join(traceback.format_tb(excinfo.tb)) From 84c598d21b1b05eccb3c6c94683b8386cc1f2caf Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 7 Sep 2026 13:06:48 -0400 Subject: [PATCH 5/7] docs: regenerate invocation-context data for make_room_in_vram Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01C7hM9XHDWGNa57jY5pat98 --- docs/src/generated/invocation-context.json | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/docs/src/generated/invocation-context.json b/docs/src/generated/invocation-context.json index e3cbbcadede..6d347f00554 100644 --- a/docs/src/generated/invocation-context.json +++ b/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", From 486636194532839af1a35b050878376084113a22 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 14 Sep 2026 21:51:32 -0400 Subject: [PATCH 6/7] fix(qwen): scope the model-load lock per step, release a failed quantized load, report shortfalls Review round 2 (Pfannkuchensack): - The quantized encoder load is now three steps, each under the lock its kind of work needs: the offload under the MODEL_LOAD_LOCK read lock (a VRAM move), the checkpoint read under no lock (`_read_checkpoint`: mmap-backed safetensors tensors, pre-faulted so the disk I/O happens there; ordered after the offload so the offload's RAM cannot evict the pages), and the construction under the write lock via `from_pretrained(None, config=..., state_dict=...)`. The write lock is required: transformers builds under process-global save/restore patches (default dtype, tie_weights, linspace) and assigns every weight through register_parameter, which a concurrent init_empty_weights patch strands on meta. `TestTransformersLoadPathAssumptions` pins both facts. - A failed load releases the partially built model (frames cleared, state dict dropped, gc + empty_cache) before the error escapes, and an OOM is re-raised with the estimate, the availability after offloading and what to do about it. - `make_room_in_vram` returns the VRAM available after offloading (re-measured past the offload's empty_cache) instead of believed bytes freed, and drops the unused `working_mem_bytes` parameter; the encoder warns on a shortfall. - An unsizeable checkpoint logs a warning instead of requesting 0 bytes. - Forward-frame clearing on error is kept on both paths: the frames hold the activations, and the error report only formats the traceback. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ShTpiRm2EAfXZxboU6Hjf --- docs/src/generated/invocation-context.json | 12 +- .../invocations/qwen_image_text_encoder.py | 168 ++++++-- .../app/services/shared/invocation_context.py | 12 +- .../load/model_cache/model_cache.py | 18 +- .../test_qwen_image_text_encoder.py | 369 ++++++++++++++++-- ...st_invocation_context_make_room_in_vram.py | 10 +- .../test_model_cache_make_room_in_vram.py | 45 ++- 7 files changed, 541 insertions(+), 93 deletions(-) diff --git a/docs/src/generated/invocation-context.json b/docs/src/generated/invocation-context.json index 6d347f00554..f4330abe01f 100644 --- a/docs/src/generated/invocation-context.json +++ b/docs/src/generated/invocation-context.json @@ -344,7 +344,7 @@ "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.", + "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. The configured working-memory reserve\nis kept free on top of the request.", "name": "make_room_in_vram", "parameters": [ { @@ -352,17 +352,11 @@ "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" + "returns": "The VRAM available after offloading, less the working-memory reserve (so it may be negative). Locked (in-use) models are never offloaded, so the request is not guaranteed: compare the result with `vram_bytes_needed` before allocating.", + "signature": "(vram_bytes_needed: int) -> 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.", diff --git a/invokeai/app/invocations/qwen_image_text_encoder.py b/invokeai/app/invocations/qwen_image_text_encoder.py index bfcffa9ca38..e11d0263568 100644 --- a/invokeai/app/invocations/qwen_image_text_encoder.py +++ b/invokeai/app/invocations/qwen_image_text_encoder.py @@ -1,5 +1,9 @@ +import gc +import json +import mmap import traceback -from typing import Any, Literal +from pathlib import Path +from typing import Any, Callable, Literal import torch from PIL import Image as PILImage @@ -15,7 +19,7 @@ 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_cache.model_cache import MB, 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, @@ -51,6 +55,42 @@ _IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>" +# The encoder as handed from a loader to `_encode`: the model, the device to run it on, and a callback that releases +# it when `_encode` is done. +_LoadedEncoder = tuple[torch.nn.Module, torch.device, Callable[[], None]] + + +def _read_checkpoint(encoder_path: Path) -> dict[str, torch.Tensor] | None: + """Map the encoder's safetensors shards and pull their pages into memory. + + Returns `None` when the checkpoint has no safetensors shards (e.g. a `.bin` one); the caller then lets transformers + read it. The tensors are zero-copy views onto the mapped files, so holding all of them costs page cache rather + than RAM. Touching one byte per page is what makes the disk read happen *here*, outside the model-load lock, + instead of lazily inside `from_pretrained`, which has to run under it: with the pages already cached the touch is + nearly free (~60 ms for the 16 GB Qwen2.5-VL-7B checkpoint), and cold it runs at the disk's sequential speed. + Call it *after* offloading cached models to RAM, not before: the offload's anonymous memory is what would evict + freshly read pages on a RAM-constrained host, forcing a second read under the lock. + """ + from safetensors.torch import load_file + + index_path = encoder_path / "model.safetensors.index.json" + if index_path.is_file(): + with open(index_path) as index_file: + shard_names = sorted(set(json.load(index_file)["weight_map"].values())) + elif (encoder_path / "model.safetensors").is_file(): + shard_names = ["model.safetensors"] + else: + return None + + state_dict: dict[str, torch.Tensor] = {} + for shard_name in shard_names: + state_dict.update(load_file(encoder_path / shard_name)) + for tensor in state_dict.values(): + # A strided read of one byte per page faults the whole tensor in through the very mapping transformers will + # read from; the sum itself is discarded. + tensor.reshape(-1).view(torch.uint8)[:: mmap.PAGESIZE].sum() + return state_dict + def _build_prompt(user_prompt: str, num_images: int) -> str: """Build the full prompt with the appropriate template based on whether reference images are provided.""" @@ -224,9 +264,13 @@ def _encode( ).to(device=device) 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. + # The in-flight traceback references the forward's frames, and through their locals the activations + # (the full-vocabulary logits and every layer's hidden states, hundreds of MB at prompt length) and, + # on the quantized path, the model itself - so the release below would otherwise be a no-op exactly + # when VRAM is scarcest (an OOM inside the forward is the likeliest error here). Clearing the finished + # frames drops those references while keeping the traceback's line information, which is all the + # error report uses. This applies to the cache-owned encoder too: the cache keeps the model, but the + # activations are the forward's alone. traceback.clear_frames(exc.__traceback__) raise finally: @@ -237,8 +281,7 @@ def _encode( # 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() + cleanup() # If all tokens are valid (no padding), mask is not needed if encoder_attention_mask.all(): @@ -291,8 +334,8 @@ def _run_encoder( 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).""" + def _load_cached_encoder(self, context: InvocationContext) -> _LoadedEncoder: + """Load the text encoder through the model cache (no quantization). The cache stays the model's owner.""" from transformers import Qwen2_5_VLForConditionalGeneration text_encoder_info = context.models.load(self.qwen_vl_encoder.text_encoder) @@ -302,9 +345,13 @@ def _load_cached_encoder(self, context: InvocationContext): # temporarily offloaded all weights to RAM, which would wrongly run the whole encode on the CPU. device = text_encoder_info.compute_device assert isinstance(text_encoder, Qwen2_5_VLForConditionalGeneration) - return text_encoder, device, lambda: ctx.__exit__(None, None, None) - def _load_quantized_encoder(self, context: InvocationContext): + def release() -> None: + ctx.__exit__(None, None, None) + + return text_encoder, device, release + + def _load_quantized_encoder(self, context: InvocationContext) -> _LoadedEncoder: """Load the text encoder with BitsAndBytes quantization, bypassing the model cache. BnB-quantized models are pinned to GPU and can't be moved between devices, @@ -318,10 +365,9 @@ def _load_quantized_encoder(self, context: InvocationContext): 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 - from transformers import BitsAndBytesConfig, Qwen2_5_VLForConditionalGeneration + from transformers import BitsAndBytesConfig, Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration encoder_config = context.models.get_config(self.qwen_vl_encoder.text_encoder) model_root = context.models.get_absolute_path(encoder_config) @@ -343,34 +389,104 @@ def _load_quantized_encoder(self, context: InvocationContext): bnb_config = BitsAndBytesConfig(load_in_8bit=True) # 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. + # VRAM is free *right now* and quietly spills to the CPU when the cached models fill the card, and it shards + # across every visible GPU - in multi-GPU mode onto cards that belong to other workers, whose caches never + # learn of the intrusion. The encoder therefore lives on this session's device only, and a genuine shortfall + # surfaces as an OOM carrying the numbers (below) 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]) + if quantized_bytes == 0: + context.logger.warning( + f"Could not determine the size of the Qwen2.5-VL encoder weights in {encoder_path}, so no room will " + "be made for the quantized encoder in VRAM; the load may fail if the cached models fill the device." + ) 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(): + + # Three steps, each under the lock its kind of work requires (see MODEL_LOAD_LOCK), so that the process-global + # lock is held no longer than necessary - while a writer holds it, or waits for it, every VRAM move on every + # worker stalls, and every reader delays every cold model construction: + # 1. the offload is a VRAM move like any other (`load_state_dict(assign=True)` -> `register_parameter`), so it + # takes the read lock, acquired *before* the cache lock `make_room_in_vram` takes, per the lock-ordering + # contract; + # 2. the checkpoint read takes no lock: it creates no parameters and touches no global state, and it is the + # slow part on anything but NVMe. It runs after the offload so the offload's RAM cannot evict its pages; + # 3. the construction takes the WRITE lock, like every construction. It cannot run alongside another + # construction: `from_pretrained` builds under process-global save/restore patches (`torch.set_default_dtype`, + # `PreTrainedModel.tie_weights`, `torch.linspace`), which two overlapping builds would restore in the wrong + # order and leave installed. It cannot run alongside VRAM moves either, in either direction: a cache + # construction's `accelerate.init_empty_weights` patch would strand the weights it assigns (through + # `setattr` -> `register_parameter`) on the meta device, and its own `torch.set_default_dtype` would change + # what a concurrent move allocates. TestTransformersLoadPathAssumptions pins both facts against the + # installed transformers. + vram_available: int | None = None + state_dict: dict[str, torch.Tensor] | None = None + try: + if quantized_bytes > 0: + with MODEL_LOAD_LOCK.read_lock(): + vram_available = context.models.make_room_in_vram(quantized_bytes) + if vram_available < quantized_bytes: + context.logger.warning( + f"Only {max(vram_available, 0) / MB:.0f} MB of VRAM could be made available on {device} for " + f"the {self.quantization}-quantized Qwen2.5-VL encoder (~{quantized_bytes / MB:.0f} MB); " + "locked (in-use) models cannot be offloaded. The load may run out of memory." + ) + + model_config = Qwen2_5_VLConfig.from_pretrained(str(encoder_path), local_files_only=True) + state_dict = _read_checkpoint(encoder_path) + + with MODEL_LOAD_LOCK.write_lock(), 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), + None if state_dict is not None else str(encoder_path), + config=model_config, + state_dict=state_dict, quantization_config=bnb_config, device_map={"": device}, - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, local_files_only=True, ) + except BaseException as exc: + # A failed load (an OOM is the expected way) leaves the partially built model referenced only by the + # in-flight traceback's frames. Without releasing it here its weights stay allocated until the session + # processor drops the exception, and *reserved* by torch after that, so the next load under-budgets - + # the very leak `_encode` guards against on the success path. This frame is still executing and is + # skipped by `clear_frames`, hence the explicit drop of the checkpoint. + traceback.clear_frames(exc.__traceback__) + state_dict = None + gc.collect() + try: + TorchDevice.empty_cache() + except Exception as cache_exc: + # A sick device context fails here too; the original error is the one worth reporting. + context.logger.warning(f"Could not empty the device cache after the failed load: {cache_exc}") + if isinstance(exc, torch.OutOfMemoryError): + raise torch.OutOfMemoryError(self._oom_message(device, quantized_bytes, vram_available)) from exc + raise # 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(): + def cleanup() -> None: gc.collect() TorchDevice.empty_cache() return text_encoder, device, cleanup + + def _oom_message(self, device: torch.device, quantized_bytes: int, vram_available: int | None) -> str: + if quantized_bytes == 0: + sizing = "its size could not be estimated, so no room was made for it" + elif vram_available is None: + sizing = ( + f"it needs about {quantized_bytes / MB:.0f} MB, and offloading the cached models to make room failed" + ) + else: + sizing = ( + f"it needs about {quantized_bytes / MB:.0f} MB and {vram_available / MB:.0f} MB was available after " + "offloading the cached models (locked, in-use models cannot be offloaded)" + ) + return ( + f"Not enough VRAM on {device} for the {self.quantization}-quantized Qwen2.5-VL encoder: {sizing}. Try " + "'nf4' quantization, or 'none', which loads the encoder through the model cache and can partially " + "offload it. In a multi-GPU setup the encoder is loaded onto this session's device only." + ) diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index 320cc0ac538..175309fe4c3 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/invokeai/app/services/shared/invocation_context.py @@ -600,23 +600,25 @@ 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: + def make_room_in_vram(self, vram_bytes_needed: int) -> 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. + and would otherwise only get whatever they happened to leave free. The configured working-memory reserve + is kept free on top of the request. 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. + The VRAM available after offloading, less the working-memory reserve (so it may be negative). Locked + (in-use) models are never offloaded, so the request is not guaranteed: compare the result with + `vram_bytes_needed` before allocating. """ - return self._services.model_manager.load.ram_cache.make_room_in_vram(vram_bytes_needed, working_mem_bytes) + return self._services.model_manager.load.ram_cache.make_room_in_vram(vram_bytes_needed) @staticmethod def _raise_if_external(model: AnyModelConfig) -> None: diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index f321c804ac9..3082d803984 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -2750,7 +2750,7 @@ def drop_model(self, model_key: str) -> int: return len(dropped) @synchronized - def make_room_in_vram(self, vram_bytes_needed: int, working_mem_bytes: Optional[int] = None) -> int: + def make_room_in_vram(self, vram_bytes_needed: int) -> 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 @@ -2761,16 +2761,22 @@ def make_room_in_vram(self, vram_bytes_needed: int, working_mem_bytes: Optional[ 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()`. + rebuilding from disk. Locked (in-use) models are never touched. The configured working-memory reserve is + kept free on top of the request, 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). + A CPU execution device has no VRAM to make room in, so the call is a no-op there (as `lock()` is) and + reports 0. - Returns the number of VRAM bytes freed based on believed model sizes. + Returns the VRAM available to the caller *after* offloading, re-measured the way `lock()` re-measures it + (free VRAM less the working-memory reserve, so it can be negative) rather than the believed sizes of the + offloaded models: a locked model can leave the request unmet, and the driver only sees freed memory once + the offload's trailing `empty_cache()` has run. Callers compare it with `vram_bytes_needed` before they + allocate. """ if self._execution_device.type == "cpu": return 0 - return self._offload_unlocked_models(vram_bytes_needed, working_mem_bytes) + self._offload_unlocked_models(vram_bytes_needed) + return self._get_vram_available(None) @synchronized def offload_model_from_vram(self, model_key: str) -> int: diff --git a/tests/app/invocations/test_qwen_image_text_encoder.py b/tests/app/invocations/test_qwen_image_text_encoder.py index ba97ad0761e..82372dfc00e 100644 --- a/tests/app/invocations/test_qwen_image_text_encoder.py +++ b/tests/app/invocations/test_qwen_image_text_encoder.py @@ -2,6 +2,7 @@ import gc import json +import threading import traceback import weakref from pathlib import Path @@ -10,15 +11,17 @@ import pytest import torch from PIL import Image -from transformers import Qwen2_5_VLForConditionalGeneration +from safetensors.torch import save_file +from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration from invokeai.app.invocations.model import ModelIdentifierField, QwenVLEncoderField from invokeai.app.invocations.qwen_image_text_encoder import ( _GENERATE_DROP_IDX, QwenImageTextEncoderInvocation, _build_prompt, + _read_checkpoint, ) -from invokeai.backend.model_manager.load.model_cache.model_cache import MODEL_LOAD_LOCK +from invokeai.backend.model_manager.load.model_cache.model_cache import MB, MODEL_LOAD_LOCK from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType from invokeai.backend.util.devices import TorchDevice @@ -139,8 +142,86 @@ def test_landscape_image(self): assert w > h # should remain landscape -def _model_load_lock_held() -> bool: - return MODEL_LOAD_LOCK._readers > 0 or MODEL_LOAD_LOCK._writer_active +def _read_lock_held() -> bool: + return MODEL_LOAD_LOCK._readers > 0 and not MODEL_LOAD_LOCK._writer_active + + +def _write_lock_held() -> bool: + return MODEL_LOAD_LOCK._writer_active + + +def _tiny_config() -> Qwen2_5_VLConfig: + """A Qwen2.5-VL config small enough to build and load on the CPU in well under a second.""" + return Qwen2_5_VLConfig( + text_config={ + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 512, + "max_position_embeddings": 256, + "tie_word_embeddings": False, + "rope_scaling": {"type": "mrope", "mrope_section": [2, 3, 3]}, + }, + vision_config={ + "depth": 2, + "hidden_size": 32, + "intermediate_size": 64, + "num_heads": 4, + "out_hidden_size": 64, + "patch_size": 14, + "spatial_merge_size": 2, + "temporal_patch_size": 2, + "in_channels": 3, + "window_size": 112, + "fullatt_block_indexes": [1], + }, + ) + + +def _write_sharded_checkpoint(text_encoder_dir: Path, total_size: int) -> dict[str, torch.Tensor]: + """A two-shard safetensors checkpoint small enough to be read for real, with a config so the loader can build + the model config from it. `total_size` is what `calc_model_size_by_fs` reports for it.""" + text_encoder_dir.mkdir(parents=True) + _tiny_config().save_pretrained(text_encoder_dir) + tensors = { + "a.weight": torch.arange(6, dtype=torch.bfloat16).reshape(2, 3), + "b.bias": torch.ones(4, dtype=torch.bfloat16), + } + save_file({"a.weight": tensors["a.weight"]}, text_encoder_dir / "model-00001-of-00002.safetensors") + save_file({"b.bias": tensors["b.bias"]}, text_encoder_dir / "model-00002-of-00002.safetensors") + index = { + "metadata": {"total_size": total_size}, + "weight_map": {"a.weight": "model-00001-of-00002.safetensors", "b.bias": "model-00002-of-00002.safetensors"}, + } + (text_encoder_dir / "model.safetensors.index.json").write_text(json.dumps(index)) + return tensors + + +class TestReadCheckpoint: + def test_reads_every_shard_named_by_the_index(self, tmp_path: Path): + tensors = _write_sharded_checkpoint(tmp_path / "text_encoder", total_size=1) + + state_dict = _read_checkpoint(tmp_path / "text_encoder") + + assert state_dict is not None and set(state_dict) == set(tensors) + assert all(torch.equal(state_dict[k], tensors[k]) for k in tensors) + + def test_reads_an_unsharded_checkpoint(self, tmp_path: Path): + (tmp_path / "text_encoder").mkdir() + save_file({"w": torch.zeros(3)}, tmp_path / "text_encoder" / "model.safetensors") + + state_dict = _read_checkpoint(tmp_path / "text_encoder") + + assert state_dict is not None and list(state_dict) == ["w"] + + def test_returns_none_without_safetensors_shards(self, tmp_path: Path): + """A `.bin` checkpoint is left to transformers to read (under the lock, as before).""" + (tmp_path / "text_encoder").mkdir() + (tmp_path / "text_encoder" / "pytorch_model.bin").write_bytes(b"") + + assert _read_checkpoint(tmp_path / "text_encoder") is None class TestQuantizedEncoderLoad: @@ -162,34 +243,35 @@ def _make_invocation(quantization: str) -> QwenImageTextEncoderInvocation: quantization=quantization, ) - def _make_context(self, tmp_path: Path, events: list[str], single_file: bool = False) -> MagicMock: + def _make_context( + self, tmp_path: Path, events: list[str], single_file: bool = False, vram_available: int | None = None + ) -> tuple[MagicMock, dict[str, torch.Tensor]]: + """`vram_available` is what `make_room_in_vram` reports after offloading; by default the request is met.""" + tensors: dict[str, torch.Tensor] = {} if single_file: model_root = tmp_path / "encoder.safetensors" model_root.write_bytes(b"") else: model_root = tmp_path / "qwen-vl" - text_encoder_dir = model_root / "text_encoder" - text_encoder_dir.mkdir(parents=True) - index = {"metadata": {"total_size": self.TOTAL_SIZE}, "weight_map": {}} - (text_encoder_dir / "model.safetensors.index.json").write_text(json.dumps(index)) + tensors = _write_sharded_checkpoint(model_root / "text_encoder", total_size=self.TOTAL_SIZE) context = MagicMock() context.models.get_absolute_path.return_value = model_root - def make_room(*_args, **_kwargs): - # The offload is a VRAM move like any other, so it too must run under the model-load lock. - events.append("make_room" if _model_load_lock_held() else "make_room(unlocked)") - return 0 + def make_room(vram_bytes_needed: int) -> int: + # The offload is a VRAM move like any other, so it runs under the model-load *read* lock. + events.append("make_room" if _read_lock_held() else "make_room(wrong lock)") + return vram_bytes_needed if vram_available is None else vram_available context.models.make_room_in_vram.side_effect = make_room - return context + return context, tensors @pytest.mark.parametrize(("quantization", "ratio"), [("int8", 0.6), ("nf4", 0.4)]) def test_makes_room_in_vram_before_loading_onto_the_execution_device( self, tmp_path: Path, quantization: str, ratio: float ): events: list[str] = [] - context = self._make_context(tmp_path, events) + context, tensors = self._make_context(tmp_path, events) fake_model = MagicMock() device = torch.device("cuda:1") seen: dict = {} @@ -197,9 +279,11 @@ def test_makes_room_in_vram_before_loading_onto_the_execution_device( def fake_from_pretrained(path, **kwargs): events.append("from_pretrained") seen.update(kwargs) - # The load must run under the model-load lock so a concurrent cache construction on another worker - # cannot hijack its parameter assignment onto the meta device. - seen["locked"] = _model_load_lock_held() + seen["path"] = path + # The construction must run under the model-load *write* lock: it installs process-global patches no + # other construction may overlap, and a concurrent cache construction would hijack its parameter + # assignment onto the meta device (see TestTransformersLoadPathAssumptions). + seen["write_locked"] = _write_lock_held() return fake_model with ( @@ -213,17 +297,159 @@ def fake_from_pretrained(path, **kwargs): assert events == ["make_room", "from_pretrained"] context.models.make_room_in_vram.assert_called_once_with(int(self.TOTAL_SIZE * ratio)) assert seen["device_map"] == {"": device}, "must never use device_map='auto'" - assert seen["locked"] + assert seen["write_locked"] + # The checkpoint is read by the invocation (outside the lock) and handed over, not read by transformers. + assert seen["path"] is None + assert set(seen["state_dict"]) == set(tensors) + assert all(torch.equal(seen["state_dict"][k], tensors[k]) for k in tensors) + assert isinstance(seen["config"], Qwen2_5_VLConfig) + assert seen["quantization_config"].load_in_8bit == (quantization == "int8") + assert seen["quantization_config"].load_in_4bit == (quantization == "nf4") assert text_encoder is fake_model assert returned_device == device + context.logger.warning.assert_not_called() cleanup() + def test_checkpoint_is_read_outside_the_model_load_lock(self, tmp_path: Path): + """The disk read is the slow part of the load on anything but NVMe, and while MODEL_LOAD_LOCK is held (or + wanted by a writer) other workers' VRAM moves and cold constructions wait. Only the offload and the + construction may run under it; a construction elsewhere must be able to take the write lock while the + checkpoint is read. The read also comes *after* the offload, so the offload's RAM cannot evict its pages.""" + events: list[str] = [] + context, _ = self._make_context(tmp_path, events) + writer_acquired = threading.Event() + acquired_during_read: list[bool] = [] + + def fake_read(encoder_path: Path): + events.append("read" if not (_read_lock_held() or _write_lock_held()) else "read(locked)") + + def construction_on_another_worker(): + with MODEL_LOAD_LOCK.write_lock(): + writer_acquired.set() + + threading.Thread(target=construction_on_another_worker).start() + acquired_during_read.append(writer_acquired.wait(timeout=5)) + return {} + + def fake_from_pretrained(path, **kwargs): + events.append("from_pretrained" if _write_lock_held() else "from_pretrained(wrong lock)") + return MagicMock() + + with ( + patch("invokeai.app.invocations.qwen_image_text_encoder._read_checkpoint", side_effect=fake_read), + patch.object(Qwen2_5_VLForConditionalGeneration, "from_pretrained", side_effect=fake_from_pretrained), + patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cuda")), + ): + self._make_invocation("int8")._load_quantized_encoder(context) + + assert acquired_during_read == [True], "a construction was blocked for the duration of the checkpoint read" + assert events == ["make_room", "read", "from_pretrained"] + + def test_warns_when_less_than_the_estimate_could_be_made_available(self, tmp_path: Path): + """`make_room_in_vram` cannot offload locked models, so it reports what is actually available after + offloading; a shortfall must be visible in the log rather than surfacing only as a later OOM.""" + events: list[str] = [] + context, _ = self._make_context(tmp_path, events, vram_available=1000 * MB) + + with ( + patch.object(Qwen2_5_VLForConditionalGeneration, "from_pretrained", return_value=MagicMock()), + patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cuda:0")), + ): + self._make_invocation("int8")._load_quantized_encoder(context) + + context.logger.warning.assert_called_once() + message = context.logger.warning.call_args.args[0] + assert "1000 MB" in message + assert f"{int(self.TOTAL_SIZE * 0.6) / MB:.0f} MB" in message + assert "cuda:0" in message + + @pytest.mark.parametrize( + ("error_type", "error_message"), [(torch.OutOfMemoryError, "CUDA out of memory"), (ValueError, "bad shard")] + ) + def test_failed_load_releases_the_partial_model_before_the_error_escapes( + self, tmp_path: Path, error_type: type[Exception], error_message: str + ): + """When `from_pretrained` raises, the partially built model is referenced only by the traceback's frames. + Left alone, its weights stay allocated until the session processor drops the exception and *reserved* by + torch after that, so the next load under-budgets - the same leak `_encode` guards against on the success + path. An OOM additionally comes back with the numbers the user needs to act on.""" + events: list[str] = [] + context, _ = self._make_context(tmp_path, events, vram_available=4000 * MB) + partial_ref: list[weakref.ref] = [] + alive_at_empty_cache: list[bool] = [] + + def fake_from_pretrained(path, **kwargs): + partial_model = torch.nn.Linear(2, 2) # referenced only by this frame once we raise + partial_model.cycle = partial_model # real models carry cycles (hooks, recorders): refcounting won't do + partial_ref.append(weakref.ref(partial_model)) + raise error_type(error_message) + + def empty_cache(): + alive_at_empty_cache.append(partial_ref[0]() is not None) + + with ( + patch.object(Qwen2_5_VLForConditionalGeneration, "from_pretrained", side_effect=fake_from_pretrained), + patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cuda:0")), + patch.object(TorchDevice, "empty_cache", side_effect=empty_cache), + pytest.raises(error_type) as excinfo, + ): + self._make_invocation("int8")._load_quantized_encoder(context) + + assert alive_at_empty_cache == [False], "empty_cache() ran while the traceback still referenced the model" + assert partial_ref[0]() is None + # The loader's own frame outlives the failure (it is in the traceback), so it must have dropped the mapped + # checkpoint itself - `clear_frames` cannot clear an executing frame. + loader_locals = next( + f.f_locals for f, _ in traceback.walk_tb(excinfo.tb) if f.f_code.co_name == "_load_quantized_encoder" + ) + assert loader_locals["state_dict"] is None + if error_type is torch.OutOfMemoryError: + assert isinstance(excinfo.value.__cause__, torch.OutOfMemoryError) + assert error_message in str(excinfo.value.__cause__) + message = str(excinfo.value) + assert "cuda:0" in message + assert f"{int(self.TOTAL_SIZE * 0.6) / MB:.0f} MB" in message + assert "4000 MB" in message + assert "nf4" in message + else: + assert str(excinfo.value) == error_message, "only an OOM is re-described" + + def test_unsizeable_checkpoint_warns_instead_of_silently_requesting_zero_bytes(self, tmp_path: Path): + """`calc_model_size_by_fs` reports 0 for weights it cannot size, and a request for 0 bytes is a silent no-op + that brings issue #9147 straight back. Such a checkpoint has no safetensors shards either, so transformers + reads it itself.""" + model_root = tmp_path / "qwen-vl" + text_encoder_dir = model_root / "text_encoder" + text_encoder_dir.mkdir(parents=True) + _tiny_config().save_pretrained(text_encoder_dir) + (text_encoder_dir / "pytorch_model.pth").write_bytes(b"\0" * 16) + context = MagicMock() + context.models.get_absolute_path.return_value = model_root + seen: dict = {} + + def fake_from_pretrained(path, **kwargs): + seen.update(kwargs) + seen["path"] = path + return MagicMock() + + with ( + patch.object(Qwen2_5_VLForConditionalGeneration, "from_pretrained", side_effect=fake_from_pretrained), + patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cuda:0")), + ): + self._make_invocation("int8")._load_quantized_encoder(context) + + context.models.make_room_in_vram.assert_not_called() + context.logger.warning.assert_called_once() + assert str(text_encoder_dir) in context.logger.warning.call_args.args[0] + assert seen["path"] == str(text_encoder_dir) + assert seen["state_dict"] is None + def test_single_file_checkpoint_falls_back_to_the_cache_without_making_room(self, tmp_path: Path): """A single-file encoder cannot be BnB-quantized; it goes through the cache, which makes its own room.""" events: list[str] = [] - context = self._make_context(tmp_path, events, single_file=True) + context, _ = self._make_context(tmp_path, events, single_file=True) invocation = self._make_invocation("int8") - sentinel = (MagicMock(), torch.device("cuda"), None) + sentinel = (MagicMock(), torch.device("cuda"), lambda: None) with ( patch.object(invocation, "_load_cached_encoder", return_value=sentinel) as cached, @@ -237,6 +463,78 @@ def test_single_file_checkpoint_falls_back_to_the_cache_without_making_room(self context.models.make_room_in_vram.assert_not_called() +@pytest.fixture(scope="module") +def tiny_qwen_checkpoint(tmp_path_factory: pytest.TempPathFactory) -> Path: + path = tmp_path_factory.mktemp("tiny-qwen") + torch.manual_seed(0) + Qwen2_5_VLForConditionalGeneration(_tiny_config()).to(torch.bfloat16).save_pretrained(path, safe_serialization=True) + return path + + +class TestTransformersLoadPathAssumptions: + """`_load_quantized_encoder` runs `from_pretrained` under MODEL_LOAD_LOCK's *write* lock and reads the checkpoint + under no lock. The two facts below are why; if either test fails after a transformers bump, the lock scoping in + `_load_quantized_encoder` has to be revisited rather than the test. + """ + + @staticmethod + def _load(path: Path) -> Qwen2_5_VLForConditionalGeneration: + # The exact call shape production uses (minus the device map and quantization, which need a GPU). + return Qwen2_5_VLForConditionalGeneration.from_pretrained( + None, + config=Qwen2_5_VLConfig.from_pretrained(str(path), local_files_only=True), + state_dict=_read_checkpoint(path), + dtype=torch.bfloat16, + local_files_only=True, + ) + + def test_the_lock_is_load_bearing_a_concurrent_construction_strands_the_weights_on_meta( + self, tiny_qwen_checkpoint: Path + ): + """Why the construction must exclude the cache's constructions: transformers assigns every loaded weight + through `setattr` -> `register_parameter`, which a cache construction on another worker patches process-wide + (`accelerate.init_empty_weights`) to route new parameters to the meta device.""" + from accelerate import init_empty_weights + + entered, release = threading.Event(), threading.Event() + + def construction_on_another_worker(): + with init_empty_weights(): + entered.set() + release.wait() + + worker = threading.Thread(target=construction_on_another_worker) + worker.start() + entered.wait() + try: + model = self._load(tiny_qwen_checkpoint) + finally: + release.set() + worker.join() + + assert any(p.device.type == "meta" for p in model.parameters()), ( + "the loaded weights no longer pass through register_parameter; revisit which lock the construction needs" + ) + + def test_the_construction_mutates_process_global_state_so_it_needs_the_write_lock(self, tiny_qwen_checkpoint: Path): + """Why the construction takes the *write* lock rather than running as one more reader: `from_pretrained` + builds under process-global save/restore patches - the default dtype among them - so two overlapping + constructions would restore each other's state in the wrong order and leave a patch installed for the life + of the process, and a concurrent VRAM move would allocate under the changed default.""" + original = torch.set_default_dtype + set_during_load: list[torch.dtype] = [] + + def spy(dtype: torch.dtype) -> None: + set_during_load.append(dtype) + original(dtype) + + with patch.object(torch, "set_default_dtype", spy): + self._load(tiny_qwen_checkpoint) + + assert torch.bfloat16 in set_during_load, "the load no longer changes the process default dtype" + assert torch.get_default_dtype() == torch.float32 + + class TestQuantizedEncoderRelease: """The quantized encoder lives outside the cache, so `_encode` is its only owner. Its cleanup callback empties the CUDA cache, which only returns the ~9 GB of encoder weights to the driver if nothing still references the @@ -258,9 +556,10 @@ class _ExplodingEncoder(torch.nn.Module): def forward(self, input_ids, attention_mask, **_): raise RuntimeError("CUDA out of memory (simulated)") - def _run_encode(self, tmp_path: Path, encoder: torch.nn.Module, alive_at_cleanup: list[bool]): - """Run `_encode` as the sole owner of `encoder`, recording into `alive_at_cleanup` whether it was still alive - when cleanup ran (recorded through the argument so the record survives an `_encode` that raises).""" + def _run_encode(self, tmp_path: Path, encoder: torch.nn.Module, alive_at_cleanup: list[bool], cached: bool = False): + """Run `_encode` as the sole owner of `encoder` (or, with `cached=True`, as a borrower of a cache-owned one), + recording into `alive_at_cleanup` whether it was still alive when cleanup ran (recorded through the argument + so the record survives an `_encode` that raises).""" model_root = tmp_path / "qwen-vl" (model_root / "tokenizer").mkdir(parents=True, exist_ok=True) context = MagicMock() @@ -283,11 +582,12 @@ def cleanup(): # strong reference of its own and mask the ownership being tested. handoff = [encoder] del encoder - invocation = TestQuantizedEncoderLoad._make_invocation("int8") + invocation = TestQuantizedEncoderLoad._make_invocation("none" if cached else "int8") + loader = "_load_cached_encoder" if cached else "_load_quantized_encoder" with ( patch.object( invocation, - "_load_quantized_encoder", + loader, side_effect=lambda _ctx: (handoff.pop(), torch.device("cpu"), cleanup), ), patch("transformers.AutoTokenizer.from_pretrained", return_value=MagicMock()), @@ -295,6 +595,10 @@ def cleanup(): ): return invocation._encode(context, images=[]) + @staticmethod + def _run_encoder_frame_locals(tb) -> dict: + return next(frame.f_locals for frame, _ in traceback.walk_tb(tb) if frame.f_code.co_name == "_run_encoder") + def test_encoder_is_released_before_cleanup_runs(self, tmp_path: Path): alive_at_cleanup: list[bool] = [] prompt_embeds, mask = self._run_encode(tmp_path, self._FakeEncoder(self.HIDDEN), alive_at_cleanup) @@ -311,5 +615,16 @@ def test_encoder_is_released_before_cleanup_runs_when_the_forward_raises(self, t with pytest.raises(RuntimeError, match="simulated") as excinfo: self._run_encode(tmp_path, self._ExplodingEncoder(), alive_at_cleanup) assert alive_at_cleanup == [False], "cleanup ran while the traceback still referenced the encoder" - # The traceback is still useful for the error report: the raising line is intact. + # The traceback is still useful for the error report: the raising line is intact, only the locals are gone. + assert "simulated" in "".join(traceback.format_tb(excinfo.tb)) + assert "text_encoder" not in self._run_encoder_frame_locals(excinfo.tb) + + def test_forward_frames_are_cleared_on_the_cached_path_too(self, tmp_path: Path): + """On the (default) unquantized path the cache keeps the encoder, but the forward's frames still hold its + activations (the full-vocabulary logits and every layer's hidden states) on the GPU for as long as the + exception lives. The error report only formats the traceback, which clearing preserves.""" + alive_at_cleanup: list[bool] = [] + with pytest.raises(RuntimeError, match="simulated") as excinfo: + self._run_encode(tmp_path, self._ExplodingEncoder(), alive_at_cleanup, cached=True) assert "simulated" in "".join(traceback.format_tb(excinfo.tb)) + assert "model_inputs" not in self._run_encoder_frame_locals(excinfo.tb) diff --git a/tests/app/services/shared/test_invocation_context_make_room_in_vram.py b/tests/app/services/shared/test_invocation_context_make_room_in_vram.py index 830fb0ef26e..866cd211d0d 100644 --- a/tests/app/services/shared/test_invocation_context_make_room_in_vram.py +++ b/tests/app/services/shared/test_invocation_context_make_room_in_vram.py @@ -1,6 +1,6 @@ -"""`context.models.make_room_in_vram` must reach the calling thread's device cache with both arguments intact: -the invocation context is the only route an invocation has to the cache, and a dropped `working_mem_bytes` would -silently fall back to the configured default.""" +"""`context.models.make_room_in_vram` must reach the calling thread's device cache and hand back its answer: the +invocation context is the only route an invocation has to the cache, and the post-offload availability is what the +caller compares its request against.""" from unittest.mock import MagicMock @@ -12,5 +12,5 @@ def test_make_room_in_vram_delegates_to_the_threads_ram_cache(): services.model_manager.load.ram_cache.make_room_in_vram.return_value = 123 models = ModelsInterface(services=services, data=MagicMock(), util=MagicMock()) - assert models.make_room_in_vram(10, working_mem_bytes=7) == 123 - services.model_manager.load.ram_cache.make_room_in_vram.assert_called_once_with(10, 7) + assert models.make_room_in_vram(10) == 123 + services.model_manager.load.ram_cache.make_room_in_vram.assert_called_once_with(10) diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py b/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py index 1e622efe717..2279f4d4104 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py @@ -14,6 +14,7 @@ import torch from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache +from invokeai.backend.util.devices import TorchDevice from tests.backend.model_manager.load.model_cache.cached_model.utils import DummyModule MB = 2**20 @@ -131,11 +132,10 @@ def test_offloads_unlocked_models_until_the_request_is_satisfied(gpu_accounting_ patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), ): - freed = cache.make_room_in_vram(30 * MB) + available = cache.make_room_in_vram(30 * MB) assert [key for key, _ in vram.moved] == ["small", "medium"] - assert freed == 30 * MB - assert vram.available == 35 * MB + assert available == 35 * MB def test_locked_models_are_never_offloaded(gpu_accounting_cache: ModelCache): @@ -151,10 +151,10 @@ def test_locked_models_are_never_offloaded(gpu_accounting_cache: ModelCache): patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), ): - freed = cache.make_room_in_vram(100 * MB) + available = cache.make_room_in_vram(100 * MB) assert [key for key, _ in vram.moved] == ["idle"] - assert freed == 10 * MB + assert available == 10 * MB, "the caller must be able to see that the request was not met" def test_no_op_when_enough_vram_is_already_free(gpu_accounting_cache: ModelCache): @@ -166,26 +166,39 @@ def test_no_op_when_enough_vram_is_already_free(gpu_accounting_cache: ModelCache patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), ): - freed = cache.make_room_in_vram(30 * MB) + available = cache.make_room_in_vram(30 * MB) assert vram.moved == [] - assert freed == 0 + assert available == 50 * MB -def test_working_memory_is_forwarded_to_the_availability_check(gpu_accounting_cache: ModelCache): +def test_reports_the_availability_re_measured_after_the_offloads_empty_cache(gpu_accounting_cache: ModelCache): + """The result is what the driver sees *after* the offload, not the believed sizes of the offloaded models: the + loop's own availability checks run before its trailing `empty_cache()`, so freed weights only show up in a + measurement taken after it (the same re-measurement `lock()` does before it decides how much to load).""" cache = gpu_accounting_cache - """The caller's working memory must reach `_get_vram_available`, where it is floored at the configured default - exactly as in `lock()`; otherwise the encoder's activations would have to fit in whatever is left over.""" _put(cache, "resident", 40 * MB) vram = _FakeVram(available=0) + def move_model_to_ram(cache_entry, vram_bytes_to_free, keep_required_weights_in_vram=None): + # Believed size is freed, but the driver does not see it yet. + vram.moved.append((cache_entry.key, vram_bytes_to_free)) + return cache_entry.cached_model.total_bytes() + + def empty_cache(): + vram.available += 30 * MB # now the driver sees (some of) it + with ( patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), - patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), + patch.object(cache, "_move_model_to_ram", side_effect=move_model_to_ram), + patch.object(TorchDevice, "empty_cache", side_effect=empty_cache), ): - cache.make_room_in_vram(10 * MB, working_mem_bytes=7 * MB) + available = cache.make_room_in_vram(10 * MB) - assert vram.working_mem_seen and all(w == 7 * MB for w in vram.working_mem_seen) + assert [key for key, _ in vram.moved] == ["resident"] + assert available == 30 * MB + # The configured working-memory reserve applies through the default (None) floor, exactly as in `lock()`. + assert vram.working_mem_seen and all(w is None for w in vram.working_mem_seen) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.") @@ -214,9 +227,11 @@ def test_gpu_make_room_in_vram_actually_moves_weights_off_the_device(partial: bo # Ask for more than the whole card so every unlocked model has to go. _, total = torch.cuda.mem_get_info() - freed = cache.make_room_in_vram(2 * total) + available = cache.make_room_in_vram(2 * total) - assert freed == cache._cached_models["idle"].cached_model.total_bytes() + # The answer is the re-measured availability (which cannot meet a request of twice the card). + assert available < 2 * total + assert available == pytest.approx(cache._get_vram_available(None), abs=256 * MB) assert all(p.device.type == "cpu" for p in idle.parameters()) assert all(p.device.type == "cuda" for p in in_use.parameters()) # Same policy as lock(): the entry is offloaded, not evicted, so the next use re-streams weights. From 3960c4f51d3e9cd90ebe60b8be7c338a35cf5e1b Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 14 Sep 2026 22:46:22 -0400 Subject: [PATCH 7/7] fix(model cache): judge out-of-cache loads by physical VRAM, not the cache cap Review round 3 (Pfannkuchensack): - `make_room_in_vram` drove its offload loop and its return value through `_get_vram_available`, which under `max_vram_cache_size_gb` reports the cache's own budget rather than free VRAM. An out-of-cache model is not subject to the cap, so that offloaded every unlocked model and logged a false shortfall on every run. New `_get_physical_vram_available` (driver free + the allocator's reserved-but-unallocated pool - working-memory reserve) is used for this caller only; `_offload_unlocked_models` takes an optional availability predicate and `lock()` is unchanged. - The quantized encoder skips the make-room step on a CPU execution device instead of logging a 0 MB shortfall. - Tests: the cap is ignored; CPU makes no room and warns nothing; the transformers-4 key layout is converted on the state-dict path; and the cache-lock test no longer touches real CUDA (the CI failure on GPU-less runners). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ShTpiRm2EAfXZxboU6Hjf --- .../invocations/qwen_image_text_encoder.py | 3 +- .../load/model_cache/model_cache.py | 64 +++++++++++--- .../test_qwen_image_text_encoder.py | 47 +++++++++++ .../test_model_cache_make_room_in_vram.py | 84 ++++++++++++------- 4 files changed, 159 insertions(+), 39 deletions(-) diff --git a/invokeai/app/invocations/qwen_image_text_encoder.py b/invokeai/app/invocations/qwen_image_text_encoder.py index e11d0263568..b02bbfed161 100644 --- a/invokeai/app/invocations/qwen_image_text_encoder.py +++ b/invokeai/app/invocations/qwen_image_text_encoder.py @@ -422,7 +422,8 @@ def _load_quantized_encoder(self, context: InvocationContext) -> _LoadedEncoder: vram_available: int | None = None state_dict: dict[str, torch.Tensor] | None = None try: - if quantized_bytes > 0: + # A CPU execution device has no VRAM to make room in (recent bitsandbytes can quantize on the CPU). + if quantized_bytes > 0 and device.type != "cpu": with MODEL_LOAD_LOCK.read_lock(): vram_available = context.models.make_room_in_vram(quantized_bytes) if vram_available < quantized_bytes: diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 3082d803984..19926f1d13f 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -2152,6 +2152,39 @@ def _get_vram_available(self, working_mem_bytes: Optional[int]) -> int: vram_cur_available_to_cache = vram_total_available_to_cache - self._get_vram_in_use() return vram_cur_available_to_cache + def _get_physical_vram_available(self) -> int: + """VRAM a load *outside* the cache can still take on the execution device, less the configured working-memory + reserve. + + Unlike `_get_vram_available`, this ignores `max_vram_cache_size_gb`: the cap limits how much of the card the + cache may occupy, but an out-of-cache model (e.g. the BitsAndBytes-quantized Qwen encoder) is not subject to + it - its only limit is what the device physically has free. Measuring the cap here would report a shortfall + on every run and offload every unlocked model regardless of how much room the card has. + + Memory the torch allocator holds reserved-but-unallocated counts as available too: that is where offloaded + weights go until the offload's trailing `empty_cache()`, and the allocator reuses it for the next allocation, + so the offload loop can see its progress. + """ + working_mem_bytes = int(self._execution_device_working_mem_gb * GB) + device = self._execution_device + if device.type == "cuda": + vram_free, _vram_total = torch.cuda.mem_get_info(device) + reusable = torch.cuda.memory_reserved(device) - torch.cuda.memory_allocated(device) + elif device.type == "xpu" and _has_dedicated_vram(device): + vram_free, _vram_total = TorchDevice.xpu_mem_get_info(device) + reusable = torch.xpu.memory_reserved(device) - torch.xpu.memory_allocated(device) + elif device.type in ("mps", "xpu"): + # Shared-memory devices: "VRAM" is system RAM, and the device allocator's cached-but-unused memory is + # still the pool an offloaded model's weights land in (as in `_get_vram_available`). + vram_free = psutil.virtual_memory().available + if device.type == "mps": + reusable = torch.mps.driver_allocated_memory() - torch.mps.current_allocated_memory() + else: + reusable = torch.xpu.memory_reserved(device) - torch.xpu.memory_allocated(device) + else: + raise ValueError(f"Unsupported execution device: {device.type}") + return vram_free + reusable - working_mem_bytes + def _get_vram_in_use(self) -> int: """Get the amount of VRAM currently in use by the cache.""" if self._execution_device.type == "cuda": @@ -2302,13 +2335,24 @@ def _get_vram_state_str(self, model_cur_vram_bytes: int, model_total_bytes: int, + f"vram_available={(vram_available / MB):.0f} MB, " ) - def _offload_unlocked_models(self, vram_bytes_required: int, working_mem_bytes: Optional[int] = None) -> int: + def _offload_unlocked_models( + self, + vram_bytes_required: int, + working_mem_bytes: Optional[int] = None, + vram_available_fn: Optional[Callable[[], int]] = None, + ) -> int: """Offload models from the execution_device until vram_bytes_required bytes are available, or all models are offloaded. Of course, locked models are not offloaded. + `vram_available_fn` is the availability check the loop satisfies; it defaults to the cache's own budget + (`_get_vram_available`, which honours `max_vram_cache_size_gb`). An out-of-cache load passes + `_get_physical_vram_available` instead, because the cap does not apply to it. + Returns: int: The number of bytes freed based on believed model sizes. The actual change in VRAM may be different. """ + if vram_available_fn is None: + vram_available_fn = lambda: self._get_vram_available(working_mem_bytes) # noqa: E731 self._logger.debug( f"Offloading unlocked models with goal of making room for {vram_bytes_required / MB:.2f}MB of VRAM." ) @@ -2317,7 +2361,7 @@ def _offload_unlocked_models(self, vram_bytes_required: int, working_mem_bytes: cache_entries_increasing_size = sorted(self._cached_models.values(), key=lambda x: x.cached_model.total_bytes()) for cache_entry in cache_entries_increasing_size: # We do not fully trust the count of bytes freed, so we check again on each iteration. - vram_available = self._get_vram_available(working_mem_bytes) + vram_available = vram_available_fn() vram_bytes_to_free = vram_bytes_required - vram_available if vram_bytes_to_free <= 0: break @@ -2762,21 +2806,21 @@ def make_room_in_vram(self, vram_bytes_needed: int) -> int: 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. The configured working-memory reserve is - kept free on top of the request, exactly as in `lock()`. + kept free on top of the request, exactly as in `lock()`. The availability check is the device's *physical* + free memory (`_get_physical_vram_available`), not the cache's own budget: `max_vram_cache_size_gb` caps what + the cache may occupy, not what an out-of-cache model may. A CPU execution device has no VRAM to make room in, so the call is a no-op there (as `lock()` is) and reports 0. - Returns the VRAM available to the caller *after* offloading, re-measured the way `lock()` re-measures it - (free VRAM less the working-memory reserve, so it can be negative) rather than the believed sizes of the - offloaded models: a locked model can leave the request unmet, and the driver only sees freed memory once - the offload's trailing `empty_cache()` has run. Callers compare it with `vram_bytes_needed` before they - allocate. + Returns the VRAM available to the caller *after* offloading, re-measured (physically free VRAM less the + working-memory reserve, so it can be negative) rather than the believed sizes of the offloaded models: a + locked model can leave the request unmet, and the caller has to be able to tell before it allocates. """ if self._execution_device.type == "cpu": return 0 - self._offload_unlocked_models(vram_bytes_needed) - return self._get_vram_available(None) + self._offload_unlocked_models(vram_bytes_needed, vram_available_fn=self._get_physical_vram_available) + return self._get_physical_vram_available() @synchronized def offload_model_from_vram(self, model_key: str) -> int: diff --git a/tests/app/invocations/test_qwen_image_text_encoder.py b/tests/app/invocations/test_qwen_image_text_encoder.py index 82372dfc00e..c4a786d02aa 100644 --- a/tests/app/invocations/test_qwen_image_text_encoder.py +++ b/tests/app/invocations/test_qwen_image_text_encoder.py @@ -414,6 +414,27 @@ def empty_cache(): else: assert str(excinfo.value) == error_message, "only an OOM is re-described" + def test_cpu_execution_device_makes_no_room_and_logs_no_shortfall(self, tmp_path: Path): + """A CPU execution device has no VRAM to make room in (the cache's `make_room_in_vram` reports 0 there), and + recent bitsandbytes can quantize on the CPU, so the load proceeds without a spurious shortfall warning.""" + events: list[str] = [] + context, _ = self._make_context(tmp_path, events) + seen: dict = {} + + def fake_from_pretrained(path, **kwargs): + seen.update(kwargs) + return MagicMock() + + with ( + patch.object(Qwen2_5_VLForConditionalGeneration, "from_pretrained", side_effect=fake_from_pretrained), + patch.object(TorchDevice, "choose_torch_device", return_value=torch.device("cpu")), + ): + self._make_invocation("int8")._load_quantized_encoder(context) + + context.models.make_room_in_vram.assert_not_called() + context.logger.warning.assert_not_called() + assert seen["device_map"] == {"": torch.device("cpu")} + def test_unsizeable_checkpoint_warns_instead_of_silently_requesting_zero_bytes(self, tmp_path: Path): """`calc_model_size_by_fs` reports 0 for weights it cannot size, and a request for 0 bytes is a silent no-op that brings issue #9147 straight back. Such a checkpoint has no safetensors shards either, so transformers @@ -471,6 +492,14 @@ def tiny_qwen_checkpoint(tmp_path_factory: pytest.TempPathFactory) -> Path: return path +def _transformers_4_key(name: str) -> str: + """The key layout the Qwen-Image `text_encoder/` folders on disk use (they predate transformers 5).""" + for v5_prefix, v4_prefix in (("model.language_model.", "model."), ("model.visual.", "visual.")): + if name.startswith(v5_prefix): + return v4_prefix + name[len(v5_prefix) :] + return name + + class TestTransformersLoadPathAssumptions: """`_load_quantized_encoder` runs `from_pretrained` under MODEL_LOAD_LOCK's *write* lock and reads the checkpoint under no lock. The two facts below are why; if either test fails after a transformers bump, the lock scoping in @@ -534,6 +563,24 @@ def spy(dtype: torch.dtype) -> None: assert torch.bfloat16 in set_during_load, "the load no longer changes the process default dtype" assert torch.get_default_dtype() == torch.float32 + def test_transformers_4_key_layout_is_converted_on_the_state_dict_path(self, tmp_path: Path): + """The checkpoints in production folders carry the transformers-4 names (`model.*`, `visual.*`); the + state-dict entry point must run the same key conversion the path-based load does, or unmatched weights + would be freshly initialised with only a load warning and the encoder would silently produce garbage.""" + torch.manual_seed(0) + source = Qwen2_5_VLForConditionalGeneration(_tiny_config()).to(torch.bfloat16).state_dict() + v4_names = {_transformers_4_key(name) for name in source} + assert v4_names != set(source) and not any(name.startswith("model.language_model.") for name in v4_names) + _tiny_config().save_pretrained(tmp_path) + save_file( + {_transformers_4_key(name): tensor for name, tensor in source.items()}, tmp_path / "model.safetensors" + ) + + loaded = self._load(tmp_path).state_dict() + + assert set(loaded) == set(source) + assert all(torch.equal(loaded[name], source[name]) for name in source) + class TestQuantizedEncoderRelease: """The quantized encoder lives outside the cache, so `_encode` is its only owner. Its cleanup callback empties diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py b/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py index 2279f4d4104..d72883c7712 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_make_room_in_vram.py @@ -42,19 +42,17 @@ def cache(mock_logger): class _FakeVram: - """Simulates the execution device's free VRAM as the cache offloads models. + """Simulates the execution device's physically free VRAM as the cache offloads models. - `ModelCache._get_vram_available` needs a real accelerator, so the CPU-only tests below replace it with this - bookkeeping: `available` grows by whatever `_move_model_to_ram` reports freed. + `ModelCache._get_physical_vram_available` needs a real accelerator, so the CPU-only tests below replace it with + this bookkeeping: `available` grows by whatever `_move_model_to_ram` reports freed. """ def __init__(self, available: int): self.available = available - self.working_mem_seen: list[int | None] = [] self.moved: list[tuple[str, int]] = [] - def get_vram_available(self, working_mem_bytes): - self.working_mem_seen.append(working_mem_bytes) + def get_physical_vram_available(self): return self.available def move_model_to_ram(self, cache_entry, vram_bytes_to_free, keep_required_weights_in_vram=None): @@ -72,14 +70,7 @@ def _put(cache: ModelCache, key: str, size_bytes: int) -> None: cache.put(key, module) -@pytest.fixture -def gpu_accounting_cache(mock_logger): - """A cache whose execution device is a GPU as far as the policy is concerned, without touching a real one. - - `_get_vram_available` needs an accelerator; the tests replace it (and the VRAM moves) with `_FakeVram`. Note that - `_FakeVram` reports freed memory immediately, which is the loop's *contract*; on real hardware the driver only - sees it after the trailing `empty_cache()`, so the loop there tends to offload every unlocked model. - """ +def _make_gpu_accounting_cache(mock_logger, **kwargs) -> ModelCache: cache = ModelCache( execution_device_working_mem_gb=1.0, enable_partial_loading=False, @@ -87,8 +78,21 @@ def gpu_accounting_cache(mock_logger): execution_device="cpu", storage_device="cpu", logger=mock_logger, + **kwargs, ) - cache._execution_device = torch.device("cuda") # policy only; every VRAM touch is patched out below + cache._execution_device = torch.device("cuda") # policy only; every VRAM touch is patched out in the tests + return cache + + +@pytest.fixture +def gpu_accounting_cache(mock_logger): + """A cache whose execution device is a GPU as far as the policy is concerned, without touching a real one. + + `_get_physical_vram_available` needs an accelerator; the tests replace it (and the VRAM moves) with `_FakeVram`. + `_FakeVram` reports freed memory immediately, as the real measurement does through the allocator's + reserved-but-unallocated pool. + """ + cache = _make_gpu_accounting_cache(mock_logger) yield cache cache._execution_device = torch.device("cpu") cache.shutdown() @@ -109,11 +113,14 @@ def test_offload_runs_under_the_cache_lock(gpu_accounting_cache: ModelCache): _put(cache, "resident", 40 * MB) owned: list[bool] = [] - def offload(vram_bytes_required, working_mem_bytes=None): + def offload(vram_bytes_required, working_mem_bytes=None, vram_available_fn=None): owned.append(cache._lock._is_owned()) return 0 - with patch.object(cache, "_offload_unlocked_models", side_effect=offload): + with ( + patch.object(cache, "_offload_unlocked_models", side_effect=offload), + patch.object(cache, "_get_physical_vram_available", return_value=0), + ): cache.make_room_in_vram(30 * MB) assert owned == [True] @@ -129,7 +136,7 @@ def test_offloads_unlocked_models_until_the_request_is_satisfied(gpu_accounting_ vram = _FakeVram(available=5 * MB) with ( - patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), + patch.object(cache, "_get_physical_vram_available", side_effect=vram.get_physical_vram_available), patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), ): available = cache.make_room_in_vram(30 * MB) @@ -148,7 +155,7 @@ def test_locked_models_are_never_offloaded(gpu_accounting_cache: ModelCache): vram = _FakeVram(available=0) with ( - patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), + patch.object(cache, "_get_physical_vram_available", side_effect=vram.get_physical_vram_available), patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), ): available = cache.make_room_in_vram(100 * MB) @@ -163,7 +170,7 @@ def test_no_op_when_enough_vram_is_already_free(gpu_accounting_cache: ModelCache vram = _FakeVram(available=50 * MB) with ( - patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), + patch.object(cache, "_get_physical_vram_available", side_effect=vram.get_physical_vram_available), patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), ): available = cache.make_room_in_vram(30 * MB) @@ -173,9 +180,9 @@ def test_no_op_when_enough_vram_is_already_free(gpu_accounting_cache: ModelCache def test_reports_the_availability_re_measured_after_the_offloads_empty_cache(gpu_accounting_cache: ModelCache): - """The result is what the driver sees *after* the offload, not the believed sizes of the offloaded models: the - loop's own availability checks run before its trailing `empty_cache()`, so freed weights only show up in a - measurement taken after it (the same re-measurement `lock()` does before it decides how much to load).""" + """The result is a fresh measurement taken *after* the offload and its trailing `empty_cache()`, not the believed + sizes of the offloaded models nor the loop's last reading: whatever the driver reports at that point (here, + memory that only became visible at `empty_cache()`) is what the caller gets.""" cache = gpu_accounting_cache _put(cache, "resident", 40 * MB) vram = _FakeVram(available=0) @@ -189,7 +196,7 @@ def empty_cache(): vram.available += 30 * MB # now the driver sees (some of) it with ( - patch.object(cache, "_get_vram_available", side_effect=vram.get_vram_available), + patch.object(cache, "_get_physical_vram_available", side_effect=vram.get_physical_vram_available), patch.object(cache, "_move_model_to_ram", side_effect=move_model_to_ram), patch.object(TorchDevice, "empty_cache", side_effect=empty_cache), ): @@ -197,8 +204,29 @@ def empty_cache(): assert [key for key, _ in vram.moved] == ["resident"] assert available == 30 * MB - # The configured working-memory reserve applies through the default (None) floor, exactly as in `lock()`. - assert vram.working_mem_seen and all(w is None for w in vram.working_mem_seen) + + +def test_the_cache_vram_cap_does_not_apply_to_an_out_of_cache_load(mock_logger): + """`max_vram_cache_size_gb` caps what the *cache* may occupy. An out-of-cache model is not subject to it, so + the request is judged against physically free VRAM: with plenty of room on the card nothing is offloaded, and + the (negative, here) capped budget is not what gets reported back.""" + cache = _make_gpu_accounting_cache(mock_logger, max_vram_cache_size_gb=0.5) + try: + _put(cache, "resident", 40 * MB) + vram = _FakeVram(available=50 * MB) + with ( + patch.object(cache, "_get_physical_vram_available", side_effect=vram.get_physical_vram_available), + patch.object(cache, "_move_model_to_ram", side_effect=vram.move_model_to_ram), + patch.object(cache, "_get_vram_in_use", return_value=40 * MB), + ): + assert cache._get_vram_available(None) < 0 # the capped budget: cap (0.5 GB) - working mem (1 GB) - use + available = cache.make_room_in_vram(30 * MB) + + assert vram.moved == [] + assert available == 50 * MB + finally: + cache._execution_device = torch.device("cpu") + cache.shutdown() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.") @@ -229,9 +257,9 @@ def test_gpu_make_room_in_vram_actually_moves_weights_off_the_device(partial: bo _, total = torch.cuda.mem_get_info() available = cache.make_room_in_vram(2 * total) - # The answer is the re-measured availability (which cannot meet a request of twice the card). + # The answer is the re-measured physical availability (which cannot meet a request of twice the card). assert available < 2 * total - assert available == pytest.approx(cache._get_vram_available(None), abs=256 * MB) + assert available == pytest.approx(cache._get_physical_vram_available(), abs=256 * MB) assert all(p.device.type == "cpu" for p in idle.parameters()) assert all(p.device.type == "cuda" for p in in_use.parameters()) # Same policy as lock(): the entry is offloaded, not evicted, so the next use re-streams weights.