From ecc84ebac31b6a98fc3d696e0de1f0454de57703 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 30 Aug 2026 21:55:19 -0400 Subject: [PATCH 1/2] fix(model cache): credit the allocator's reclaimable reserve in the VRAM budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _get_vram_available budgeted on driver-free memory alone (mem_get_info), so blocks the torch caching allocator holds but is not using — freed weights and dead activations from earlier pipeline stages, whenever an empty_cache was skipped — counted as consumed. On big jobs the budget under-reported by gigabytes and a model that would have fit was partial-loaded down to its minimum weight set, while the allocator happily reused the 'missing' memory for activations (observed: a MiniMax H3 transformer at 0% VRAM residency streaming 20GB of weights from RAM per step). Credit (reserved - allocated - inactive_split) for CUDA/ROCm and dedicated XPU: the reserve is reusable directly and empty_cache returns whole unoccupied segments to the driver; the inactive-split slack (free space inside partially-occupied segments) stays excluded since a large contiguous allocation cannot use it. As a side effect _offload_unlocked_models now sees freed bytes immediately and stops evicting once the budget suffices, instead of always dumping every unlocked model. Also emit a first-pass WARNING when the budget stays negative after offloading, naming every model still resident in VRAM and its lock state - the offload pass cannot touch locked entries, and until now nothing in the default log said what was occupying the device when a load fell back to minimum weights. Verified on CUDA (RTX 5060 Ti): the new tests reproduce the invisible 1 GiB reserve (available 128MB -> 1152MB), guard the offload-under-pressure path, and assert the diagnostic names locked residents. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01S4B5exWsbC2z2Uu2tA167A --- .../load/model_cache/model_cache.py | 59 +++++++- .../test_model_cache_vram_accounting.py | 136 ++++++++++++++++++ 2 files changed, 189 insertions(+), 6 deletions(-) create mode 100644 tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.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 e737c8b7af9..23d66d9c38e 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -1206,6 +1206,24 @@ 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.split(':')[-1]}={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 + ] + 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: {(working_mem_bytes or 0) / 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 +1378,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 +1413,29 @@ 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. + """ + 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..f92c37654e6 --- /dev/null +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_vram_accounting.py @@ -0,0 +1,136 @@ +"""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() From 46bcf19943226a167970e4f31bc34efe95a6f292 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 30 Aug 2026 22:07:56 -0400 Subject: [PATCH 2/2] fix(model cache): harden the allocator-reserve credit after adversarial review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Withhold the credit entirely under expandable-segments mode (any of PYTORCH_ALLOC_CONF / PYTORCH_CUDA_ALLOC_CONF / PYTORCH_HIP_ALLOC_CONF — InvokeAI's own pytorch_cuda_alloc_conf setting lands in the same env var before torch import). There, freed blocks inside a segment are not counted as inactive splits, empty_cache reclaims nothing, and a large allocation cannot use the holes — the review reproduced a hard OOM on a load the credited budget claimed would fit, where the uncredited budget degraded gracefully to weight streaming. - The budget-short warning now reports the EFFECTIVE working-memory reservation (callers passing None get the configured default, smaller values are clamped up) instead of printing 0MB for the very number being diagnosed, and lists residents by full cache key so two models' same-type submodels are distinguishable. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01S4B5exWsbC2z2Uu2tA167A --- .../load/model_cache/model_cache.py | 32 +++++++++++++++++-- .../test_model_cache_vram_accounting.py | 20 ++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) 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 23d66d9c38e..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 @@ -1213,14 +1228,19 @@ def _load_locked_model( # too-small budget is diagnosable from the default log. First pass only (paced # continuations would repeat it). resident = [ - f"{entry.key.split(':')[-1]}={entry.cached_model.cur_vram_bytes() / MB:.0f}MB" + 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: {(working_mem_bytes or 0) / MB:.0f}MB); the model will " + 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'}." ) @@ -1418,8 +1438,14 @@ def _get_reclaimable_allocator_bytes(self) -> int: 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. + 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) 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 index f92c37654e6..84d78826670 100644 --- 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 @@ -134,3 +134,23 @@ def test_negative_budget_warns_and_names_locked_residents(): 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