Repository navigation
feat(models): calibrate VAE working memory per model from the Model Manager - #9910
Pfannkuchensack wants to merge 178 commits into
Conversation
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.
…o feat/working-memory-user-calibration
|
Thanks for the review. All four blockers are addressed; this branch is now at ef7ca0f. 1. Graph export against current 2. Component identities. Fixed in f1080e2.
3. Retry lineage durability. Fixed in f1080e2. The async listener is gone.
4.
|
…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.
…ipelines in the plan
…e default VAE of SD mains
…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
left a comment
There was a problem hiding this comment.
Merge blockers:
-
invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:1040andinvokeai/app/invocations/model.py:60:ModelIdentifierFieldomits 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 a409 model_changed. -
scripts/family_e2e/__main__.py:200,scripts/family_e2e/runner.py:385,scripts/family_e2e/server.py:556:stopcan 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 beforeon_spawn()records its identity, the server escapes the captured tree and is absent from the heartbeat. Effect:stopcan 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, callstop, 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
…king-memory-user-calibration
|
@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)
2.
|
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)
Estimates
device_working_mem_gb), nothing changes.Robustness found in hardware runs
invokeai/backend/working_memory/fragmentation.py).Before/after on an RX 9060 XT (SDXL, fp32 decode)
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 checkandformat --check.make frontend-openapi/frontend-typegenequivalent);pnpm lintandpnpm testininvokeai/frontend/api.pnpm build(format, oxlint, tsc, architecture);src/features/models255;node --test scripts/mock-backend-fixtures.test.mjs28;pnpm test:performance:architecturepasses.pnpm testhere: 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:
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):Browser
pnpm dev:mock; fault scenarios throughPOST /__faults) and on both real servers.Linux NVIDIA: the NVML counter reader and this calibration round ran on 2× RTX PRO 6000 Blackwell (#9910 (comment)); with #9897's
edba0dd882it 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
Remaining limits
Compatibility / Rollout
<db_dir>/working_memory/comes from feat(working-memory): measure working memory in admin calibration runs #9897. New optional fields (attempt_id, runsoft) are left out when absent, so older files load. No database migration.GET /api/v2/models/i/{key}/working_memory/planand/summary(404 for unknown models).attempt_idon calibrate. Line fieldsmeasured_by,own_windows, and size fieldscounted,excluded_reason. Generatedopenapi.jsonandschema.tsare updated.browser-baseline.jsonrecords the measuredotherRawBytes(215114 B on every route, main 211323 B); nothing else in the baseline changed.Checklist
What's Newcopy (if doing a release after this PR)