fix(model cache): credit the allocator's reclaimable reserve in the VRAM budget - #192
Merged
Conversation
…RAM budget _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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01S4B5exWsbC2z2Uu2tA167A
…al review - 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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01S4B5exWsbC2z2Uu2tA167A
lstein
requested review from
JPPhoto,
Pfannkuchensack and
blessedcoolant
as code owners
August 31, 2026 02:08
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
_get_vram_availablebudgeted VRAM 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 anempty_cachewas skipped (several paths skip it deliberately) — 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 in the wild on a 48 GB W7900 running a MiniMax H3 Ref2VA job: the 20 GB transformer loaded at 0.0% VRAM residency ("Loading 1.14 MB into VRAM, but only −13113.5 MB were requested") and streamed every block from RAM per step. Reproduced deterministically at small scale on CUDA: after freeing 1 GiB without an
empty_cache,_get_vram_availablereported 128 MB when 1,152 MB was obtainable.Fix
reserved − allocated − inactive_split_bytesin the CUDA/ROCm and dedicated-XPU branches: the reserve is reusable directly, andempty_cache()returns whole unoccupied segments to the driver. Intra-segment fragmentation slack stays excluded. Side effect:_offload_unlocked_modelsnow sees freed bytes immediately and stops evicting once the budget suffices, instead of always dumping every unlocked model.PYTORCH_ALLOC_CONF/PYTORCH_CUDA_ALLOC_CONF/PYTORCH_HIP_ALLOC_CONF, which is where InvokeAI's ownpytorch_cuda_alloc_confsetting lands before torch import): there, holes are not inactive splits,empty_cachereclaims nothing, and adversarial review reproduced a hard OOM on a load the credited budget claimed would fit.Verification
_get_vram_availablecall sites, thecudaMallocAsyncbackend credits honestly, the native allocator's release-and-retry makes whole-segment credit safe, no non-default-stream usage in any load path, and the paced-lock warning fires exactly once per stream.🤖 Generated with Claude Code
https://claude.ai/code/session_01S4B5exWsbC2z2Uu2tA167A