Skip to content

Add opt-in PiD memory optimizations - #9460

Open
JPPhoto wants to merge 7 commits into
invoke-ai:mainfrom
JPPhoto:pid-optimization
Open

Add opt-in PiD memory optimizations#9460
JPPhoto wants to merge 7 commits into
invoke-ai:mainfrom
JPPhoto:pid-optimization

Conversation

@JPPhoto

@JPPhoto JPPhoto commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds opt-in PiD memory optimizations for GPUs with limited VRAM.

When pid_memory_optimization: true is set in invokeai.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

  1. Add pid_memory_optimization: true to invokeai.yaml.
  2. Restart InvokeAI.
  3. Run a PiD decode and confirm it completes with reduced peak VRAM.
  4. Remove the setting or set it to false, restart, and confirm the original unchunked path is used.

Merge Plan

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable)
  • ❗Changes to a redux slice have a corresponding migration
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

@JPPhoto
JPPhoto requested a review from blessedcoolant as a code owner August 4, 2026 22:17
@JPPhoto JPPhoto added the 6.14.0 label Aug 4, 2026
@JPPhoto JPPhoto moved this to 6.14.x Theme: USER EXPERIENCE in Invoke - Community Roadmap Aug 4, 2026
@github-actions github-actions Bot added python PRs that change python files invocations PRs that change invocations backend PRs that change backend files services PRs that change app services frontend PRs that change frontend files python-tests PRs that change python tests docs PRs that change docs labels Aug 4, 2026
@JPPhoto JPPhoto changed the title Add opt-in PiD optimizations Add opt-in PiD memory optimizations Aug 5, 2026
@JPPhoto
JPPhoto force-pushed the pid-optimization branch 2 times, most recently from 451d823 to 47cd9fb Compare August 5, 2026 16:21
@Pfannkuchensack

Pfannkuchensack commented Aug 6, 2026

Copy link
Copy Markdown
Collaborator

PR #9460 — Add opt-in PiD memory optimizations

  • The freed VRAM stays reserved, so the feature is only half wired. estimate_pid_decode_working_memory (invokeai/backend/pid/decode.py:137) is not flag-aware — it still returns out_h * out_w * 4 * 250 = 3.9 GiB at 2048², while the measured peak with the flag on is 1.5 GiB. The cache takes max(working_mem_bytes, device_working_mem_gb) (invokeai/backend/model_manager/load/model_cache/model_cache.py:1093) and subtracts it from the weight budget. The comment above the constant says exactly what is at stake: "an over-large value ... forces PidNet to partial-load onto the CPU (slow)." On the low-VRAM systems this feature targets, the saving prevents a hard OOM but never becomes weight residency.
  • The documented cost is not the real cost. "at the cost of slower decoding" (invokeai/app/services/config/config_default.py:116 and :224, plus both docs pages) is not observable: 2.78 s both ways at 2048px, median of 3 with warmup. What is real and undocumented is that the output changes. Fix the description and both docs pages.
  • Chunked ≠ unchunked on the target hardware, and the test cannot fail. On CUDA under bf16 autocast at production dimensions, a single PiTBlock gives max|diff| = 1.59e-2 and assert_close fails; the same block on CPU fp32 is bit-identical. test_pit_block_chunked_forward_matches_unchunked_and_bounds_adaln_batch passes only because it runs at pixel_hidden_size=4 / BL=8 / chunk=3 on CPU — both use_autocast parametrizations are CPU-only. End-to-end effect: PSNR 43.2 dB, max|diff| 0.66 on [-1,1], 13.1% of pixels differ by >2 LSB. Add a CUDA + bf16 test at production dimensions with BL >= 2 * chunk_size and either fix the divergence or make the tolerance an explicit, documented contract.
  • Split the fp32 _velocity_to_x0 change out of the memory flag, or justify it. Measured in isolation at 2048²: 288 → 96 MiB, i.e. 192 MiB — against 2.18 GiB from chunking alone (one unchunked PiTBlock peaks at 2864.7 MiB vs 598.4 MiB chunked). It is a precision reduction that saves almost no peak memory but does change the image (PSNR 44.2 dB standalone). The two tests for it only assert which branch is taken — neither measures memory nor bounds the delta.
  • adaLN_modulation(s_cond) is computed twice per chunk (invokeai/backend/pid/_src/networks/pixeldit_official.py:541 and :548), each time discarding most slices. That is ~9.9 TFLOP extra per 2048px decode and the only real basis for the "slower" claim. Return the remaining four slices from _compress_activation_chunk instead.
  • Nothing enforces that the nodes forward the setting. All seven construction sites pass it today (verified by grep), but the tests cover _student_sample_loop, _velocity_to_x0 and config loading — never the invocation layer. An eighth PiD node could silently omit it.
  • The setting leaves no trace on the output. It is a server config value, never a node input, so it never enters graph metadata. The same workflow with the same seed produces different pixels on two servers with nothing recording why.
  • No observability: no log line when the flag is enabled, and _PID_ACTIVATION_CHUNK_SIZE (invokeai/backend/pid/decode.py:31) is hardcoded and unconfigurable — a yaml-only, restart-required knob with zero feedback.
  • Untested paths: batch B > 1 (chunk boundaries then straddle images), the context-parallel (_cp_group) branch, and the feature_indices / discriminator branch of PidNet.forward.

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

6.14 Nice-to-Have 6.14.0 backend PRs that change backend files docs PRs that change docs frontend PRs that change frontend files invocations PRs that change invocations python PRs that change python files python-tests PRs that change python tests services PRs that change app services

Projects

Status: 6.14.x Theme: USER EXPERIENCE

Development

Successfully merging this pull request may close these issues.

2 participants