fix(model cache): credit the allocator's reclaimable reserve in the VRAM budget

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":

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