Add opt-in PiD memory optimizations - #9460
Open
JPPhoto wants to merge 13 commits into
Open
Conversation
github-actions
Bot
added
python
JPPhoto
changed the title
5 tasks
JPPhoto
force-pushed
the
pid-optimization
branch
3 times, most recently
from
August 5, 2026 17:48
47cd9fb to
42097b6
Compare
Member
PR #9460 — Add opt-in PiD memory optimizations
|
Pfannkuchensack
added a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 6, 2026
…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.
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 7, 2026
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 7, 2026 11:37
5df8ac3 to
25adf56
Compare
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 7, 2026
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 7, 2026 14:33
25adf56 to
ebe622c
Compare
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 7, 2026
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 7, 2026 19:05
ebe622c to
8dff4eb
Compare
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 8, 2026
…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.
Collaborator
Author
|
@Pfannkuchensack Thanks for the detailed review. I addressed all four findings in commit
The documentation now clarifies that activation figures exclude the temporary model-sized autocast copies. Relevant tests (CUDA and non-CUDA) pass locally. |
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 9, 2026
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 9, 2026 17:25
223ce29 to
4411e01
Compare
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 9, 2026
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 9, 2026 18:30
4411e01 to
cfa38a0
Compare
Member
Findings1. High - the bf16 autocast weight-cache term does not exist in production and inflates every PiD working-memory reservation by ~2.6 GiB
That cache never materializes, because Measured with the installed
Impact. 2048px is the normal case (PiD is a 4x SR decoder; a 512px-source latent produces exactly that). Via Note the history: the calibration figures introduced in Fix
Test changes. The current tests pin the defect and must be updated, not merely kept green:
To expose this issue, add a test that passes a model with 2. Low - the docs promise a VRAM saving that
|
Collaborator
Author
|
Ignoring finding 1 per our Discord conversation regarding thrashing on this, and continuing work on the rest. |
Collaborator
Author
|
@Pfannkuchensack This is ready for another review. |
Member
Correction to my 2026-08-08 review, and the corrected numbersMy finding 1 from 2026-08-08 was wrong, and the term added in Eight lines, no model needed: import torch
lin = torch.nn.Linear(8, 8, bias=False).cuda().float()
lin.requires_grad_(False) # what PiDDecoderLoader._load_model does
x = torch.ones(1, 8, device="cuda")
with torch.no_grad(), torch.autocast("cuda", torch.bfloat16):
a = lin(x).clone()
lin.weight.mul_(1000) # mutate the fp32 weight inside the autocast region
b = lin(x)
print(torch.allclose(a, b)) # True = cached (stale copy), False = not cached
Corrected measurementsReal (4096px unoptimized was not measured - ~16 GiB of activations on top of the 5197 MiB resident net does not fit on a 24 GiB card. That part of my earlier review stands: at 4096px the flag is what makes the decode run at all.) So the real shortfall I should have reported on 2026-08-08 was 103-134 MiB in the unoptimized path, not 1429-3662 MiB, and the term that was added to cover it over-reserves by ~2.6 GiB on every decode - including with the flag off, which makes it a regression against What I am asking for - three changes, and this is my final position on it1. Remove the weight-cache term.
To expose this, a test that applies 2. Raise 3. Keep the chunked constants exactly as they are. Verified against the real checkpoint - please do not change
|
Pfannkuchensack
approved these changes
Aug 11, 2026
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 12, 2026
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 12, 2026 18:13
d0a5bdb to
cbe8430
Compare
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 13, 2026
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 13, 2026 01:23
cbe8430 to
bebe71e
Compare
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Aug 13, 2026
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 13, 2026 02:10
bebe71e to
8e5e9b8
Compare
…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.
JPPhoto
force-pushed
the
pid-optimization
branch
from
August 13, 2026 02:38
8e5e9b8 to
194d326
Compare
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
2 participants
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)