Add opt-in PiD memory optimizations - #9460
Open
JPPhoto wants to merge 7 commits into
Open
Conversation
5 tasks
JPPhoto
force-pushed
the
pid-optimization
branch
2 times, most recently
from
August 5, 2026 16:21
451d823 to
47cd9fb
Compare
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 5, 2026 17:48
47cd9fb to
42097b6
Compare
Collaborator
PR #9460 — Add opt-in PiD memory optimizations
|
…at it costs Addresses every point from the review of invoke-ai#9460. The setting freed activation memory that the cache then withheld anyway. `estimate_pid_decode_working_memory` was not flag-aware, so a decode that peaks at 1.5 GiB reserved the unoptimized 3.9 GiB; the cache takes max(working_mem_bytes, device_working_mem_gb) and subtracts that from the weight budget, so the saving never became weight residency - it only avoided a hard OOM, on precisely the low-VRAM machines this feature exists for. The estimate now takes the flag, and each node reads the setting once and feeds both the estimate and the decode from it, so the two cannot drift apart. Calibrated against measured peaks (RTX 4090, fp32 PidNet, bf16 autocast, 4 steps, B=1): 1024px 509 MiB 1536px 934 MiB 2048px 1533 MiB which is 85.3 * U + 167 MiB, not a pure multiple of the output area - chunking bounds the per-block activations to a fixed working set, so a single scaling constant would under-reserve at small sizes or over-reserve at large ones. The constants carry ~15% headroom. Below the chunk size the pixel blocks run unchunked, so the estimate clamps to the unoptimized one rather than charging for a working set that is never allocated. The documented cost was the wrong cost. "at the cost of slower decoding" is not observable - 2.78 s either way at 2048px, median of 3 with warmup - while the cost that is real went unmentioned: the option changes the decoded image. Both parts are non-bit-exact and the few-step sampler amplifies them (~43 dB PSNR end-to-end, visually indistinguishable, not reproducible against an unoptimized decode). The setting description and both docs pages now state that, with the measured VRAM numbers. The equivalence test could not fail. It asserted `assert_close` at pixel_hidden_size=4 / BL=8 / chunk=3 on the CPU, while the shipped path is CUDA under bf16 autocast with BL in the thousands. Measured at production dimensions: CPU fp32 is bit-identical (including B=2, where chunks straddle images), CUDA fp32 differs by 9.5e-07, CUDA bf16 by 1.57e-02 - systematic, not noise, both paths being internally deterministic. So chunking is exact as *mathematics*; the divergence is cuBLAS picking different kernels for 1024-row slices. `test_pid_chunked_equivalence.py` now pins both halves: exactness on the CPU, and an absolute tolerance contract on CUDA with 3x headroom (relative tolerances are meaningless here - activations cross zero, so max|rel| reaches 1e3 on elements whose absolute error is one bf16 ULP). Two review points did not survive measurement, and are documented rather than "fixed": - The fp32 `_velocity_to_x0` branch is kept. Fusing the multiply-subtract in fp64 is bit-identical to the default expression but frees nothing (288 MiB either way at 2048px), so the 192 MiB is bought entirely with precision - max|diff| 4.8e-07 per call, 8.6% of the flag's total saving. That is a fair trade for someone who opted into trading quality for VRAM, but it is now written down, in the function, in the setting description and in the docs. - The doubled `adaLN_modulation` per chunk stays. Global attention sits between the two halves, so reusing the slices means holding them for every chunk - the full-resolution tensor the path exists to avoid (536 MiB in bf16 at 2048px). Projecting only the needed slices is no better: they are interleaved per pixel position, so selecting them gathers rows of a 1536x24576 weight, ~50 MiB per call. The recompute is ~9.9 TFLOP per 2048px decode, about 4% wall clock - and the honest source of the "slower decoding" the setting advertises. Observability: a decode with the flag on now logs the resolution, the patch-token count and whether chunking actually engaged. The setting is server-wide and never enters image metadata, so this is the only record that a given decode ran optimized, and the only feedback that a yaml-only, restart-required knob took effect at all. Tests: AST sweeps assert that every module building a `PiDDecodeConfig` forwards the flag, estimates working memory for the same mode, and reads the setting exactly once - so an eighth PiD node is covered the day it lands. Plus estimate behaviour (shrinks when enabled, keeps the fixed term, never exceeds the unoptimized estimate, still returns 0 for unsupported backbones), batch > 1 across chunk boundaries, and the two paths the flag must not reach: the discriminator feature extraction returns before the pixel blocks, and context parallelism is unreachable in this codebase (its only caller lives in a vendored class InvokeAI never instantiates) - both pinned so a refactor cannot quietly put them under the flag. All new tests mutation-verified: breaking the chunked assembly, making the estimate ignore the flag, dropping the fixed term, and omitting the flag from a node's decode or estimate each fail at least one test. tests/app + tests/backend/pid: 2202 passed. The 9 failures are the pre-existing network-dependent ones in test_model_install / test_load_api / test_download_queue.
…cuous
CI caught this on macos-default py3.11; every other job in the matrix was cancelled by fail-fast.
Test-only change, no production code touched.
Two separate mistakes, both mine:
1. The CPU comparison asserted `torch.equal`. That held on x86-64 with MKL and failed on
macOS/Accelerate. Splitting a GEMM along its row dimension can select a micro-kernel with
different K-blocking, so bit-exactness there is a property of the BLAS, not of the chunking.
Only reassociation-closeness is portable.
2. Worse, and only found while investigating the first: the `batch_size=1` parametrization never
entered the chunked path at all. The dispatch guard is `BL > chunk_size`, and 512px with B=1 puts
BL at exactly 1024 - so it compared the unchunked path against itself and passed for the wrong
reason. Verified by spying on `_forward_chunked`: zero calls.
Both cases now demonstrably chunk - 768px/B=1 (BL 2304, boundaries inside one image) and 512px/B=2
(BL 2048, boundaries straddling images) - and a context manager fails the test if
`_forward_chunked` is not entered, so the comparison cannot silently empty out again.
Bit-equality is replaced by a signal-relative bound, calibrated rather than guessed. At these
dimensions the signal is ~5.7, so one fp32 ULP is ~6.8e-07:
correct code, x86-64/MKL max|diff| = 0
attention contribution off by 1e-6 max|diff| = 7.2e-07 (1.3e-07 relative, sub-ULP)
attention contribution off by 1e-4 max|diff| = 1.0e-05 (1.9e-06 relative, ~15 ULP)
1e-5 relative is ~84 ULP: above any BLAS reassociation, four orders of magnitude below a structural
break. The docstring states what that gives up - a uniform scaling error below ~2e-06 relative is
indistinguishable from legitimate reassociation and no portable test can claim it - and what it
still guards, which is the bug class that matters.
Mutation-verified against realistic breakage: an off-by-one on the last chunk, a wrong `s_cond`
slice, and skipping the chunked path each fail 5 of the 7 tests. 98 passed locally, ruff clean.
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.
Summary
Adds opt-in PiD memory optimizations for GPUs with limited VRAM.
When
pid_memory_optimization: trueis set ininvokeai.yaml, PiD uses float32 sampler intermediates and chunks full-resolution PiT activations. The option applies to every supported PiD decoder. It defaults to false, preserving existing behavior.Chunking is configured per decode call so cached PiD models cannot retain optimization state between requests.
This PR adds backend tests, generated configuration types, the complete generated settings entry, and PiD/low-VRAM documentation.
Related Issues / Discussions
QA Instructions
pid_memory_optimization: truetoinvokeai.yaml.Merge Plan
Checklist
What's Newcopy (if doing a release after this PR)