Skip to content

feat(models): calibrate VAE working memory per model from the Model Manager - #9910

Open
Pfannkuchensack wants to merge 178 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/working-memory-user-calibration
Open

Pfannkuchensack wants to merge 178 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/working-memory-user-calibration

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Oct 18, 2026 •

Copy link
Copy Markdown
Member

Stacked on #9899 (and through it #9894, #9896, #9897, #9898). Until those merge, the diff also shows their commits. This PR's own changes are 87079a62d9, cb91aab513 and 7714be86b6; the rest are merges.

Summary

The shipped VAE working-memory estimates come from a few reference GPUs. A user's card, driver or allocator can need less, or more. Admins can now calibrate a model's VAE decode on their own GPU from the Model Manager, and estimates on that GPU then use the measurement.

Model detail, "Working memory" section (actions are admin-only)

  • Calibrate: a dialog shows the planned runs (sizes per dtype, 1 warm-up + 2 measured each). The runs then go through the queue as short generations, with progress and Stop.
  • Reset: removes this model's measurements on this GPU.
  • Per line (operation · dtype): a table with Measured, Built-in and Calibrated estimate per size. Each column header has an info popover. Sizes that did not count are muted with their reason.
  • States: never calibrated, in use, may be outdated, not used, too few sizes, running, incomplete, failed, unsupported, no GPU, more than one generation device, store unreadable, and measurement method not verified on this platform.
  • Shared lines: models whose decode prices on the same line share the calibration. The section lists who measured it.

Estimates

  • Within the measured sizes: the user line, an upper envelope over the counted measurements plus 5 % (the shipped margin).
  • Above the largest measured size: the larger of the user line and built-in.
  • No usable fit: built-in exactly as before.
  • Environment changes: a hard change (backend, allocator, torch, settings) stops using the fit. A soft change (driver) keeps it and marks it "may be outdated".
  • Below the working-memory floor (device_working_mem_gb), nothing changes.

Robustness found in hardware runs

  • Finished attempts only: a calibration attempt's windows join no fit until every run of that attempt has finished. Larger cases are never priced from the same attempt's smaller ones.
  • Fragmentation: allocator fragmentation in a warm server is detected per attempt and excluded. This happens on ROCm with expandable segments: reserved 1.3–1.47× allocated while allocated stays flat. The shipped-data builder uses the same rule (invokeai/backend/working_memory/fragmentation.py).
  • Full offload before the measured lock: before a measured calibration lock, other unlocked models are moved off the GPU whole. The cache's partial offload otherwise leaves the measured model partly resident ([bug]: model cache leaves the requested model partly resident after splitting the last offload candidate #9907). Normal locks are unchanged.
  • Plan capacity: the plan sizes cases for the capacity the cache will grant, including driver memory kept after earlier runs (HIP's retained pool).

Before/after on an RX 9060 XT (SDXL, fp32 decode)

1024² reservation
Built-in ~6.1 GiB
User fit without the fragmentation rule (first hardware run) 8.3 GiB, which would evict the UNet again
This PR follows the counted line up to 1152², below built-in (1152²: measured 6.95, calibrated 7.3, built-in 7.48 GB)

Related Issues / Discussions

QA Instructions

Tests

  • uv run --no-sync pytest -n 4: 4818 passed, 158 skipped, 5 xfailed. Covered: working-memory backend, calibration service, router and runner, session processor, tests/backend/model_manager/load, tests/app/invocations, the builder, tests/family_e2e.
  • uv tool run ruff@0.11.2 check and format --check.
  • API regenerated (make frontend-openapi/frontend-typegen equivalent); pnpm lint and pnpm test in invokeai/frontend/api.
  • Frontend:
    • pnpm build (format, oxlint, tsc, architecture);
    • vitest src/features/models 255;
    • browser tests for the section and the graph compiler 24;
    • node --test scripts/mock-backend-fixtures.test.mjs 28;
    • pnpm test:performance:architecture passes.
  • Full pnpm test here: 3 failures from the de-DE locale of this machine (number formatting) and 2 load timeouts that pass alone. All unrelated.

Mutation checks, each caught:

  • the margin applied to the whole reservation;
  • a user fit outside its range;
  • live-attempt windows joining a fit;
  • attempts pooled in the fragmentation rule;
  • the calibration offload disabled, or applied outside calibration;
  • splitting instead of whole moves;
  • the plan's outside-usage floor removed;
  • reset scope, stale-failure display, uncounted-row display.

Hardware, SDXL from the UI (18 runs per attempt) on an RTX 4090 (CUDA, cudaMallocAsync) and an RX 9060 XT (ROCm 10, Windows, expandable segments):

  • Final attempt: 36 of 36 windows fully resident on both cards.
  • 9060 XT: fp32 1152² counts; the earlier fragmented 1216² is excluded.
  • 4090: calibrated below built-in at every size (0.91–0.96×).

Browser

  • Every state on the mock backend (pnpm dev:mock; fault scenarios through POST /__faults) and on both real servers.
  • Keyboard flows: Calibrate, Stop and Reset confirmations; focus lands on Stop after start and returns after each dialog; popovers open with Enter and close with Escape.
  • Screenshots: to be attached.

Linux NVIDIA: the NVML counter reader and this calibration round ran on 2× RTX PRO 6000 Blackwell (#9910 (comment)); with #9897's edba0dd882 it counts as verified.

Not verified on hardware: the Linux AMD, MPS and XPU counter readers. The UI marks them as unverified.

Review

Resolved material findings

  • Shared lines: they showed as "built-in". Reset copy now says other models' measurements still apply.
  • Stale state: it never cleared after a recalibration. The header now ignores other environments and superseded estimator versions.
  • Reset scope: Reset removed less than the summary attributed to the model.
  • Fit refresh during an attempt: fits were refreshed mid-attempt and the plan was built from stale fits. Now held back per attempt.
  • Fragmentation: fragmented windows set the user line, and windows from different attempts were pooled.
  • Margin test: it could not catch a margin on analytic terms.
  • UI states: a stale failure was shown after a lost connection, and the table columns were misaligned.
  • Plan sizes: they ran partly resident on 16 GB cards.

Remaining limits

  • First calibration on ROCm: the plan reads capacity live, so its largest size may not count until the retained pool has been recorded once. Partly resident windows never count.
  • Growth rule margin: a clean attempt's first-to-largest growth was 1.095×, close to the 1.10 threshold. A false flag falls back towards built-in.
  • Plan capacity formula: it re-implements the cache's budget rule in the app layer (without the reclaimable-block credit, so it is conservative).
  • Encodes: image-to-image and inpainting encodes keep built-in estimates.

Compatibility / Rollout

  • Store: the store file per device class under <db_dir>/working_memory/ comes from feat(working-memory): measure working memory in admin calibration runs #9897. New optional fields (attempt_id, run soft) are left out when absent, so older files load. No database migration.
  • API: new GET /api/v2/models/i/{key}/working_memory/plan and /summary (404 for unknown models). attempt_id on calibrate. Line fields measured_by, own_windows, and size fields counted, excluded_reason. Generated openapi.json and schema.ts are updated.
  • Without calibration: estimates are unchanged; the resolver falls back to the shipped values.
  • Locale: +3.8 KB of strings. browser-baseline.json records the measured otherRawBytes (215114 B on every route, main 211323 B); nothing else in the baseline changed.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Meaningful regression coverage added / updated where needed; obsolete tests/code removed
  • Persisted-state and API changes include required migrations / compatibility validation
  • Relevant performance/efficiency opportunities considered; material claims have evidence
  • Material review findings resolved and relevant checks rerun
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

Runs each installed model family through a launched scratch server and compares PSNR, time and driver-level VRAM against accepted references.
Graphs come from the frontend builders via an offline graph export; the generate submission plan is shared with the workbench.
VAE estimators return a WorkingMemoryEstimate built from fitted and analytic terms; model_on_device accepts it as working_memory=.
The AutoencoderKL coefficients and their measured points move to backend/working_memory/data, checked by a data test; reservations are unchanged.
A probe records each estimated model_on_device window's device-level demand inside a calibration scope; results go to a per-device-class store under databases/working_memory.
Admin endpoints start, inspect, reset and export runs; the reserved queue origin is refused for ordinary enqueues.
VAE estimators price a measured slope and base per conv backend, allocator and dtype instead of scaling one bf16 constant by element size; CogView4 gets its own lines.
SD/SDXL nodes and their estimate share one run-dtype rule; on ROCm Windows an SDXL fp32 1024² decode reserves 6.2 instead of 14.6 GB and the UNet stays resident.
A lower second reading at a measured size no longer tilts a fitted line, and points that include cast weight growth no longer fit one; two shipped lines move by under 1 %.
…n fact

Environment ids now include windows, linux or other, so calibrations from different operating systems are kept apart; the OS build string moves to soft fact os_build.
… runs

start --calibrate submits each case through the calibration endpoint and saves the store export per server lifetime; scripts/working_memory/build_shipped.py turns exports into probe points and refits.
Probe points bound estimates from below only, and the builder refuses data that the shared bound rule rejects.
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
…anager

Adds a working-memory section to the model detail: plan, calibrate, stop and reset per GPU, with the measured, built-in and calibrated estimate per size.
User fits from finished calibration attempts price VAE decodes within their measured range (+10 % + 256 MiB on the fitted part) and never below built-in values above it.
New plan and summary routes; summary reports lines shared with other models and returns 404 for unknown models.
…calibration

Judges allocator fragmentation per calibration attempt with one rule shared with the shipped-data builder, and moves other models off the GPU before a measured calibration lock.
User fits use the shipped 5 % margin; the plan sizes cases from the capacity the cache will grant, including memory the driver keeps after earlier runs.
The table marks uncounted sizes with their reason and explains each value column.
# Conflicts:
#	invokeai/backend/util/vae_working_memory.py
#	tests/app/invocations/test_flux2_working_memory.py
# Conflicts:
#	invokeai/backend/util/vae_working_memory.py
Main added useNsfwChecker and useWatermark to the generation project settings; exported calibration and E2E graphs run without both, as the defaults do.
…ser-calibration

# Conflicts:
#	invokeai/frontend/webv2/performance/browser-baseline.json
#	invokeai/frontend/webv2/scripts/mock-backend-fixtures.test.mjs
#	invokeai/frontend/webv2/scripts/mock-backend.mjs
…e baseline

Raises otherRawBytes to the measured 215114 B on every route; the 3791 B growth are the Model Manager working-memory strings.
@github-actions github-actions Bot added api python PRs that change python files Root invocations PRs that change invocations labels Oct 18, 2026
@Pfannkuchensack

Copy link
Copy Markdown
Member Author

Thanks for the review. All four blockers are addressed; this branch is now at ef7ca0f.

1. Graph export against current main. Fixed in #9894 (e558c51) and merged up through #9899. The export now adds reference images through the Reference Images section's own typing (withAddedReferenceImages), so SDXL IP-Adapter, FLUX.2 Kontext and Krea-2 style/edit references match the UI. Cases with reference images on a model that takes none are refused. graphExport.test.ts covers each of these; pnpm lint (tsc included) and the test pass on the merge result.

2. Component identities. Fixed in f1080e2.

  • graph_models entries now record the hash as well as base, type, format and variant.
  • A later attempt step compares every model its graph shares with all earlier runs of the attempt, and a mismatch refuses the step naming the changed model. Test: test_a_later_step_whose_component_model_changed_is_refused[variant|format|hash].
  • A component changed after the graph was posted is caught at admission, as before.
  • One window stays open on purpose: a component edited between compiling the graph and posting the first step goes unseen. Models cannot be locked against editing today, and closing it would mean re-resolving every component at enqueue. The only person who can hit it is the admin running the calibration, and the damage is limited to that admin's own measurement. The docs (model-manager.mdx) state this as a known limit.

3. Retry lineage durability. Fixed in f1080e2. The async listener is gone.

  • SessionQueue.retry_items_by_id counts the retry on its calibration run (retries_in_flight, under the store's file lock) before its admit transaction, and records each retry with its lineage after the commit.
  • A count that is left over keeps the attempt out of the fit for good. That covers a stop or crash between the count and the record, a record that could not be written, and a commit that raised after its insertions and so may have gone through.
  • A rollback or a full queue counts down without recording.
  • If the store cannot count, nothing is retried and the route answers 503.
  • Tests:
    • test_a_retry_canceled_and_pruned_before_any_refresh_still_aborts_its_attempt_also_after_a_restart: your scenario, with a real restart.
    • test_a_retry_whose_commit_failed_after_its_insertions_keeps_its_count
    • test_a_retry_the_full_queue_had_no_room_for_leaves_no_count

4. start / stop / resume race. Fixed in #9894 (d492e4b, 279b062) and merged up.

  • A run.lock serializes every status change in run.json.
  • The runner registers in the heartbeat under that lock before it does anything else.
  • stop marks the run stopped under the lock before it reads the heartbeat. So either it finds the runner and kills it, or the runner finds the run stopped and exits untouched (exit 130).
  • The runner checks again before every server and every step, and never overwrites a status other than running.
  • start and resume return only once the runner has registered, and resume refuses while the run's runner is alive.
  • Tests:
    • test_a_runner_that_comes_up_after_stop_leaves_the_stopped_run_untouched: pauses before the first heartbeat, then stops and resumes.
    • test_stop_kills_the_runner_tree_but_no_process_whose_pid_was_reused[readable|refused]
    • test_a_run_stopped_during_the_graph_export_is_not_put_back_to_running

…ast unlock

The last unlock now checks every name of a tied weight and every CPU weight with a RAM copy against the tensor the cache counts, not only weights whose dtype differs.
A cast to fp32 and back left tied names (distinct Parameters over one tensor, as adoption and cache moves create) and CPU weights on their own copies, holding more memory than counted.
… a process survives

The runner writes the server launcher into heartbeat.json right after the spawn, before it waits for startup.
stop checks that every killed process exited; while one still runs it keeps the scratch roots and stores and exits with 2.
…ettings name

The plan sizes each case for the default VAE Generate applies on selecting the model, picks a GPU VAE where that one runs on the CPU, and reports SD models without such a pick as unsupported.
A standalone SD VAE is calibrated through a main model whose default VAE does not run on the CPU.
… removes

Resetting a run, a model or every calibration drops the removed runs' item outcomes and file entries, so repeated calibrate-and-reset cycles no longer grow them.
… Manager attempts

A run sent through the API without an attempt counts on its own and keeps what it measured when stopped.
…d trusts no reused server pid

Once the runner is dead, stop reads heartbeat.json again and kills and checks the server it names now.
The heartbeat records the serving process with the create time captured at its start.
…n unlock

A failed lock or paced continuation that drops the last pin now returns cast weights to their counted dtypes too, all in one release path.
Records that are detached or about to be evicted as stale are skipped, and the docstrings state how a tied weight split across devices is re-pointed.
…r the fits lock

Outcomes read from the queue are kept per run and item id, so a renumbered database cannot let one run's item decide another's attempt.
The summary, the attempt check and resets judge them under the lock removals hold, so a reset cannot race a judgement.
…lan from accepts_vae

The planner asks the shared rule the model loaders and Generate's picker read, latent channels and variant included, instead of restating it.
…ixel line

Only a spec keyed by no platform trait may carry a hand-set value; a keyed spec's value off a pixel line fails the data test.

@JPPhoto JPPhoto left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Merge blockers:

  • invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:1040 and invokeai/app/invocations/model.py:60: ModelIdentifierField omits model format and variant. If either changes without changing key, hash, base, or type before POST, _changed_models() accepts the stale graph and records the new identity. Effect: calibration can measure a different model configuration. Likelihood: deterministic when either field changes. Recovery: compare a compile-time format/variant identity before enqueue. Test: change each field before POST and assert a 409 model_changed.

  • scripts/family_e2e/__main__.py:200, scripts/family_e2e/runner.py:385, scripts/family_e2e/server.py:556: stop can snapshot the runner's children just before it spawns the server. If the server starts in its own POSIX session before the runner is killed, but before on_spawn() records its identity, the server escapes the captured tree and is absent from the heartbeat. Effect: stop can report success and remove scratch state while the server still uses it. Likelihood: narrow race during stop and server startup. Recovery: synchronize spawning and identity registration with the run lock. Test: pause between the child snapshot and spawn, call stop, and assert the server exits before scratch cleanup.

…ibration

# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/app/util/torch_cuda_allocator.py
#	invokeai/backend/working_memory/coefficients.py
#	invokeai/backend/working_memory/fit.py
#	invokeai/backend/working_memory/platform.py
#	tests/backend/model_manager/load/model_cache/test_model_cache_unlock_restores_casts.py
#	tests/backend/working_memory/test_fit.py
#	tests/backend/working_memory/test_shipped_coefficients.py
The runner checks the run is still running, spawns the server and records it in the heartbeat in one run-lock hold, so stop either sees the server or no server starts.
A server whose identity cannot be recorded is killed before the lock is released.
Points of estimators keyed by no platform trait name no platform in build_shipped.
…er model identity

Calibrate requests carry graph_models, the hash, base, type, format and variant of every model the graph names as the client compiled it; a difference from the current records, or from a keyed repeat's run, is refused with 409 model_changed.
The model manager and the family E2E tool send the identities of the records they compiled over.
# Conflicts:
#	invokeai/backend/working_memory/fit.py
#	scripts/working_memory/build_shipped.py
…ibration

# Conflicts:
#	docs/src/content/docs/development/Guides/family-e2e.mdx
#	scripts/family_e2e/__main__.py
#	scripts/family_e2e/client.py
#	scripts/family_e2e/records.py
#	scripts/family_e2e/references.py
#	scripts/family_e2e/report.py
#	scripts/family_e2e/runner.py
#	scripts/family_e2e/server.py
#	tests/family_e2e/fake_invokeai_server.py
#	tests/family_e2e/test_cli.py
#	tests/family_e2e/test_references.py
#	tests/family_e2e/test_report.py
#	tests/family_e2e/test_runner.py
@Pfannkuchensack

Pfannkuchensack commented Oct 28, 2026 •

Copy link
Copy Markdown
Member Author

@JPPhoto Thanks. Both blockers are fixed. This branch is now at db6105d, and each new test fails without its fix.

1. Format and variant before POST (fixed in f87babb)

  • Each calibrate request now has a required field, graph_models. For every model the graph names, it gives the hash, base, type, format and variant the client compiled the graph against. The type is CalibrationCompiledModel: every fact is required, and only format and variant may be null.
  • The field must name exactly the graph's models, or the server answers 422.
  • Before anything is enqueued, the server compares each entry with the current record and with the graph's own identifier. Any difference gets 409 model_changed, naming the model and the fact.
  • The run then records the identities that passed this check.
  • A keyed repeat of a run that has no items is compared with that run's recorded identities.
  • Clients:
    • The model manager sends the identities of the records it compiled over, not a fresh lookup at POST time.
    • scripts/family_e2e sends the identities of the configs its cases were exported against. It finds the graph's models the same way the server does.
    • The mock backend validates the field like the server.
  • The documented known limit is removed; model-manager.mdx now states the check.
  • Tests:
    • a change before POST, for each of hash, base, type, format and variant, gets 409;
    • a stale graph identifier with current graph_models gets 409;
    • a missing, extra or unstated entry gets 422;
    • a keyed repeat stating other identities is refused;
    • on the UI side, the posted identities are the compiled ones even when the records change during compile;
    • a compiler snapshot shows the Python and TS walks find the same models.

2. stop versus the server spawn (fixed in #9899, 66f446f, merged up here)

  • The runner now does three steps in one hold of the run lock: it checks that the run is still running, spawns the server, and records the server's identity in the heartbeat.
  • stop marks the run stopped under the same lock before it reads the heartbeat. So a server is either in that read, or it never starts.
  • Only the spawn and the record are under the lock; the readiness waits are outside it.
  • If the identity cannot be recorded, the server is killed before the lock is released.
  • Tests:
    • your sequence, with stop arriving before and during the spawn: the server is either never spawned, or recorded and dead before scratch cleanup;
    • a check that the status test and the record happen with the run lock held, in the same hold;
    • a server whose identity cannot be recorded is dead before the lock is released.

@JPPhoto JPPhoto left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good!

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

7.0.0 api backend PRs that change backend files docs PRs that change docs frontend PRs that change frontend files frontend-deps PRs that change frontend dependencies invocations PRs that change invocations python PRs that change python files python-deps PRs that change python dependencies python-tests PRs that change python tests Root services PRs that change app services

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants