Skip to content

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

Merged
lstein merged 2 commits into
mainfrom
fix/vram-available-allocator-credit
Aug 31, 2026
Merged

fix(model cache): credit the allocator's reclaimable reserve in the VRAM budget#192
lstein merged 2 commits into
mainfrom
fix/vram-available-allocator-credit

Conversation

@lstein

@lstein lstein commented Aug 31, 2026

Copy link
Copy Markdown
Collaborator

Problem

_get_vram_available budgeted 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 an empty_cache was 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_available reported 128 MB when 1,152 MB was obtainable.

Fix

  • Credit reserved − allocated − inactive_split_bytes in the CUDA/ROCm and dedicated-XPU branches: the reserve is reusable directly, and empty_cache() returns whole unoccupied segments to the driver. Intra-segment fragmentation slack stays excluded. Side effect: _offload_unlocked_models now sees freed bytes immediately and stops evicting once the budget suffices, instead of always dumping every unlocked model.
  • Withheld entirely under expandable-segments mode (parsed from PYTORCH_ALLOC_CONF / PYTORCH_CUDA_ALLOC_CONF / PYTORCH_HIP_ALLOC_CONF, which is where InvokeAI's own pytorch_cuda_alloc_conf setting lands before torch import): there, holes are not inactive splits, empty_cache reclaims nothing, and adversarial review reproduced a hard OOM on a load the credited budget claimed would fit.
  • New first-pass WARNING when the budget stays negative after offloading, naming every model still resident in VRAM (full cache key, size, lock state) and the effective working-memory reservation — until now nothing in the default log said what was occupying the device when a load fell back to minimum weights.

Verification

  • New CUDA tests (run on an RTX 5060 Ti): the invisible-reserve reproduction (128 → 1152 MB), the offload-under-working-memory-pressure guard, the diagnostic warning naming locked residents, and the expandable-segments withhold. Full model-cache suite: 336 passed.
  • Adversarial fresh-context review ran before push; its confirmed findings (expandable-segments over-credit → hard OOM; warning printing 0 MB for a defaulted reservation; ambiguous submodel keys in the resident list) are fixed in the second commit. It also verified: no double counting across all _get_vram_available call sites, the cudaMallocAsync backend 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

lstein and others added 2 commits August 30, 2026 21:55
…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
lstein merged commit 77cbecf into main Aug 31, 2026
19 checks passed
@lstein
lstein deleted the fix/vram-available-allocator-credit branch August 31, 2026 03:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant