Skip to content

Add opt-in PiD memory optimizations - #9460

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

Add opt-in PiD memory optimizations#9460
JPPhoto wants to merge 5 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.

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