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