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 e737c8b7af9..daca5cfbf15 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -1,5 +1,6 @@ import gc import logging +import os import queue import threading import time @@ -37,6 +38,20 @@ from invokeai.backend.util.logging import InvokeAILogger from invokeai.backend.util.prefix_logger_adapter import PrefixedLoggerAdapter + +def _expandable_segments_enabled() -> bool: + """Whether the torch caching allocator runs in expandable-segments mode. + + The mode is configured through the allocator env vars before torch import (InvokeAI's own + `pytorch_cuda_alloc_conf` setting is plumbed into PYTORCH_CUDA_ALLOC_CONF the same way), so + parsing them is authoritative for the life of the process. + """ + for var in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF", "PYTORCH_HIP_ALLOC_CONF"): + if "expandable_segments:true" in os.environ.get(var, "").replace(" ", "").lower(): + return True + return False + + # Size of a GB in bytes. GB = 2**30 @@ -1206,6 +1221,29 @@ def _load_locked_model( f"Unloaded {vram_bytes_freed_from_own_model / MB:.2f}MB from the model being locked ({cache_entry.key})." ) + if vram_available < 0 and stream_started_at is None: + # The budget is still short after offloading everything offloadable: the model will run + # with its minimum weight set streamed from RAM. Name what is occupying the device — + # in particular anything still LOCKED, which the offload pass cannot touch — so a + # too-small budget is diagnosable from the default log. First pass only (paced + # continuations would repeat it). + resident = [ + f"{entry.key}={entry.cached_model.cur_vram_bytes() / MB:.0f}MB" + + (" [locked]" if entry.is_locked else "") + for entry in self._cached_models.values() + if entry.cached_model.cur_vram_bytes() > 0 and entry.key != cache_entry.key + ] + # The reservation _get_vram_available actually applied: callers passing None (the + # majority) get the configured default, and smaller values are clamped up to it — + # printing the raw argument would report 0MB for the very number being diagnosed. + working_mem_bytes_default = int(self._execution_device_working_mem_gb * GB) + effective_working_mem = max(working_mem_bytes or working_mem_bytes_default, working_mem_bytes_default) + self._logger.warning( + f"VRAM budget for '{cache_entry.key}' is short by {-vram_available / MB:.0f}MB even after " + f"offloading (working memory reservation: {effective_working_mem / MB:.0f}MB); the model will " + f"run with minimum weights resident. Other models still in VRAM: {', '.join(resident) or 'none'}." + ) + # Move as much of the model as possible into VRAM. # For testing, only allow 10% of the model to be loaded into VRAM. # vram_available = int(model_vram_needed * 0.1) @@ -1360,17 +1398,23 @@ def _get_vram_available(self, working_mem_bytes: Optional[int]) -> int: return vram_total_available_to_cache - self._get_vram_in_use() if self._execution_device.type == "cuda": - # TODO(ryand): It is debatable whether we should use memory_reserved() or memory_allocated() here. - # memory_reserved() includes memory reserved by the torch CUDA memory allocator that may or may not be - # re-used for future allocations. For now, we use memory_allocated() to be conservative. - # vram_reserved = torch.cuda.memory_reserved(self._execution_device) vram_allocated = torch.cuda.memory_allocated(self._execution_device) vram_free, _vram_total = torch.cuda.mem_get_info(self._execution_device) - vram_available_to_process = vram_free + vram_allocated + # Blocks the caching allocator holds but is not using are just as available to this + # process as driver-free memory: the allocator reuses them directly, and empty_cache() + # returns whole unoccupied segments to the driver. mem_get_info() alone counts them as + # consumed — so whenever weights or a previous stage's activations were freed without + # an empty_cache() (which several paths deliberately skip), the budget under-reported + # by that whole amount and a model that would have fit was partial-loaded down to its + # minimum weight set (observed: a fully-evictable multi-GB reserve left a 20 GB + # transformer at 0% residency while the allocator happily reused the "missing" memory + # for activations). Credit the reclaimable reserve, excluding intra-segment + # fragmentation slack that a large contiguous allocation could not use. + vram_available_to_process = vram_free + vram_allocated + self._get_reclaimable_allocator_bytes() elif self._execution_device.type == "xpu" and _has_dedicated_vram(self._execution_device): vram_allocated = torch.xpu.memory_allocated(self._execution_device) vram_free, _vram_total = TorchDevice.xpu_mem_get_info(self._execution_device) - vram_available_to_process = vram_free + vram_allocated + vram_available_to_process = vram_free + vram_allocated + self._get_reclaimable_allocator_bytes() elif self._execution_device.type in ("mps", "xpu"): # Shared-memory devices: MPS, and Intel integrated GPUs, whose reported "VRAM" is # system RAM. Budget against actual free system memory instead of device totals. @@ -1389,6 +1433,35 @@ 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_reclaimable_allocator_bytes(self) -> int: + """Bytes the torch caching allocator holds for this device but is not using, excluding + inactive-split slack (free space inside partially-occupied segments, which cannot serve a + large contiguous allocation and which empty_cache() cannot return to the driver). + + Best-effort: 0 when the backend does not expose allocator stats, and 0 under + expandable-segments mode — there, freed blocks inside a segment are NOT counted as + inactive splits, empty_cache() reclaims nothing, and a large allocation cannot use the + holes, so the whole (reserved - allocated) figure is untrustworthy (a measured hard OOM + on an allocation the credited budget claimed would fit). + """ + if _expandable_segments_enabled(): + return 0 + try: + if self._execution_device.type == "cuda": + reserved = torch.cuda.memory_reserved(self._execution_device) + allocated = torch.cuda.memory_allocated(self._execution_device) + stats = torch.cuda.memory_stats(self._execution_device) + elif self._execution_device.type == "xpu": + reserved = torch.xpu.memory_reserved(self._execution_device) + allocated = torch.xpu.memory_allocated(self._execution_device) + stats = torch.xpu.memory_stats(self._execution_device) + else: + return 0 + inactive_split = int(stats.get("inactive_split_bytes.all.current", 0)) + return max(0, int(reserved) - int(allocated) - inactive_split) + except Exception: + return 0 + def _get_vram_in_use(self) -> int: """Get the amount of VRAM currently in use by the cache.""" if self._execution_device.type == "cuda": diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.py b/tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.py new file mode 100644 index 00000000000..84d78826670 --- /dev/null +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.py @@ -0,0 +1,156 @@ +"""Tests for the VRAM budget arithmetic behind lock()/partial loads. + +`_get_vram_available` must count every byte this process can actually obtain: driver-free memory +PLUS the torch caching allocator's reserved-but-unused blocks (the allocator reuses those +directly, and `empty_cache()` returns whole unoccupied segments to the driver). Budgeting on +driver-free alone under-reported by whatever earlier stages freed without an `empty_cache()` +— observed in the wild as a fully-evictable multi-GB reserve pushing a 20 GB transformer down +to 0% VRAM residency while the allocator happily reused the "missing" memory for activations. +""" + +import logging +from unittest.mock import MagicMock + +import pytest +import torch + +from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache + +GB = 1024**3 +MB = 1024**2 + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA.") + + +class BigModule(torch.nn.Module): + """A module whose single parameter is `mb` MiB of fp32.""" + + def __init__(self, mb: int): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(mb * MB // 4, dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + +def _make_cache(logger: logging.Logger | MagicMock | None = None) -> ModelCache: + if logger is None: + logger = MagicMock() + logger.getEffectiveLevel.return_value = logging.INFO + return ModelCache( + execution_device_working_mem_gb=0.1, + enable_partial_loading=True, + keep_ram_copy_of_weights=True, + execution_device="cuda:0", + storage_device="cpu", + logger=logger, + shared_cpu_weights=None, + ) + + +@requires_cuda +def test_lock_offloads_unlocked_models_under_working_memory_pressure(): + """A big working-memory reservation must evict resident-but-unlocked models so the locked + model still loads fully (the freed bytes become budget via the allocator-reserve credit and + the post-offload empty_cache).""" + torch.cuda.empty_cache() + cache = _make_cache() + cache.put("A", BigModule(512)) + rec_a = cache.get("A") + cache.lock(rec_a, None) + cache.unlock(rec_a) + assert rec_a.cached_model.cur_vram_bytes() >= 512 * MB + + free, _total = torch.cuda.mem_get_info(torch.device("cuda:0")) + # Without offloading A there is only 128 MB of budget — far less than B needs. + working = free - 128 * MB + + cache.put("B", BigModule(256)) + rec_b = cache.get("B") + cache.lock(rec_b, working) + try: + assert rec_a.cached_model.cur_vram_bytes() == 0, "unlocked resident model was not offloaded" + assert rec_b.cached_model.cur_vram_bytes() == rec_b.cached_model.total_bytes() + finally: + cache.unlock(rec_b) + + +@requires_cuda +def test_get_vram_available_credits_reserved_but_free_allocator_blocks(): + """Freed-but-not-empty_cache'd allocator blocks are reclaimable and must count as available.""" + torch.cuda.empty_cache() + cache = _make_cache() + + # Simulate a previous pipeline stage's freed activations: 1 GiB allocated then dropped, with + # no empty_cache — the bytes stay in the allocator's reserve, invisible to mem_get_info. + junk = torch.empty(1 * GB, dtype=torch.uint8, device="cuda:0") + del junk + + free, _total = torch.cuda.mem_get_info(torch.device("cuda:0")) + working = free - 128 * MB + available = cache._get_vram_available(working) + + # Driver-free alone would report ~128 MB; the credited reserve must dominate. The margin + # tolerates concurrent allocations by other processes on a shared dev GPU. + assert available >= 900 * MB, f"reserved-but-free blocks not credited (available={available / MB:.0f}MB)" + + torch.cuda.empty_cache() + + +@requires_cuda +def test_negative_budget_warns_and_names_locked_residents(): + """When the budget stays short after offloading, the first-pass warning must name what is + still occupying the device — locked entries especially, since the offload cannot touch them.""" + torch.cuda.empty_cache() + logger = MagicMock() + logger.getEffectiveLevel.return_value = logging.INFO + cache = _make_cache(logger) + + cache.put("stuck", BigModule(256)) + rec_stuck = cache.get("stuck") + cache.lock(rec_stuck, None) # deliberately left locked + + free, _total = torch.cuda.mem_get_info(torch.device("cuda:0")) + impossible_working = free + 10 * GB + + cache.put("victim", BigModule(64)) + rec_victim = cache.get("victim") + cache.lock(rec_victim, impossible_working) + try: + # ModelCache wraps its logger in a PrefixedLoggerAdapter, so adapter.warning() reaches + # the underlying (mock) logger as .log(WARNING, msg). + warnings = [str(call.args[0]) for call in logger.warning.call_args_list] + warnings += [ + str(call.args[1]) + for call in logger.log.call_args_list + if call.args and call.args[0] == logging.WARNING and len(call.args) > 1 + ] + budget_warnings = [message for message in warnings if "VRAM budget for 'victim' is short by" in message] + + assert budget_warnings, f"no budget-short warning emitted; warnings: {warnings}" + assert "stuck=" in budget_warnings[0] + assert "[locked]" in budget_warnings[0] + finally: + cache.unlock(rec_victim) + cache.unlock(rec_stuck) + torch.cuda.empty_cache() + + +@requires_cuda +def test_reclaimable_credit_withheld_under_expandable_segments(monkeypatch: pytest.MonkeyPatch): + """Under expandable-segments mode the (reserved - allocated) figure counts intra-segment + holes that empty_cache cannot reclaim and a large allocation cannot use — the credit must be + withheld entirely (the env parse is authoritative: the mode is fixed before torch import).""" + cache = _make_cache() + + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + assert cache._get_reclaimable_allocator_bytes() == 0 + + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:512") + monkeypatch.setenv("PYTORCH_HIP_ALLOC_CONF", "expandable_segments: true") + assert cache._get_reclaimable_allocator_bytes() == 0 + + monkeypatch.delenv("PYTORCH_HIP_ALLOC_CONF") + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF") + # With no allocator config the credit path is active again (>= 0 by construction). + assert cache._get_reclaimable_allocator_bytes() >= 0