From 855ebe2c02cd82d0ee81e95fb6f8e6aa73a06990 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 17 Oct 2026 04:57:58 +0200 Subject: [PATCH] feat(working-memory): ship VAE working-memory coefficients as data 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. --- .../contributing/new-model-integration.mdx | 4 +- .../Architecture/model-manager.mdx | 20 + .../vae/cogview4_image_to_latents.py | 2 +- .../vae/cogview4_latents_to_image.py | 2 +- .../app/invocations/vae/flux2_vae_decode.py | 2 +- .../app/invocations/vae/flux2_vae_encode.py | 2 +- .../app/invocations/vae/flux_vae_decode.py | 4 +- .../app/invocations/vae/flux_vae_encode.py | 2 +- .../vae/ideogram4_latents_to_image.py | 2 +- .../app/invocations/vae/image_to_latents.py | 2 +- .../app/invocations/vae/latents_to_image.py | 2 +- .../invocations/vae/sd3_image_to_latents.py | 2 +- .../invocations/vae/sd3_latents_to_image.py | 2 +- .../vae/z_image_image_to_latents.py | 2 +- .../vae/z_image_latents_to_image.py | 4 +- invokeai/backend/AGENTS.md | 1 + .../flux/extensions/kontext_extension.py | 2 +- invokeai/backend/flux2/ref_image_extension.py | 2 +- .../backend/model_manager/load/load_base.py | 18 +- invokeai/backend/util/vae_working_memory.py | 542 +++++++----- invokeai/backend/working_memory/__init__.py | 6 + .../backend/working_memory/coefficients.py | 191 +++++ invokeai/backend/working_memory/data/vae.json | 769 ++++++++++++++++++ invokeai/backend/working_memory/estimate.py | 52 ++ invokeai/backend/working_memory/platform.py | 45 + invokeai/backend/working_memory/spec.py | 99 +++ pyproject.toml | 2 + scripts/calibrate_flux2_working_memory.py | 37 +- scripts/calibrate_flux_vae_working_memory.py | 6 +- tests/app/invocations/conftest.py | 18 + .../invocations/test_flux2_working_memory.py | 235 ++---- tests/app/invocations/test_flux_vae_decode.py | 2 +- .../test_flux_vae_decode_tiling.py | 18 +- .../test_vae_working_memory_backend.py | 76 +- .../invocations/test_z_image_tiled_decode.py | 135 +-- .../test_z_image_working_memory.py | 12 +- .../model_manager/load/test_loaded_model.py | 64 +- .../util/test_rocm_sdpa_head_dim_guard.py | 30 +- tests/backend/working_memory/__init__.py | 0 .../test_shipped_coefficients.py | 312 +++++++ .../working_memory/test_spec_registry.py | 72 ++ 41 files changed, 2266 insertions(+), 534 deletions(-) create mode 100644 invokeai/backend/working_memory/__init__.py create mode 100644 invokeai/backend/working_memory/coefficients.py create mode 100644 invokeai/backend/working_memory/data/vae.json create mode 100644 invokeai/backend/working_memory/estimate.py create mode 100644 invokeai/backend/working_memory/platform.py create mode 100644 invokeai/backend/working_memory/spec.py create mode 100644 tests/backend/working_memory/__init__.py create mode 100644 tests/backend/working_memory/test_shipped_coefficients.py create mode 100644 tests/backend/working_memory/test_spec_registry.py diff --git a/docs/src/content/docs/contributing/new-model-integration.mdx b/docs/src/content/docs/contributing/new-model-integration.mdx index eda7821880a..04e4530b2d8 100644 --- a/docs/src/content/docs/contributing/new-model-integration.mdx +++ b/docs/src/content/docs/contributing/new-model-integration.mdx @@ -601,11 +601,11 @@ Packages are discovered automatically. A directory of node modules without an `_ ... ``` - Report progress with `context.util.sd_step_callback(state, BaseModelType.NewModel)`; it renders previews through the architecture's `LatentSpaceFacet`. The denoise node is also where these are checked: interrupts (`context.util.is_canceled()`), inpainting, img2img through `denoising_start`, and a working-memory estimate passed to `model_on_device(working_mem_bytes=...)`. + Report progress with `context.util.sd_step_callback(state, BaseModelType.NewModel)`; it renders previews through the architecture's `LatentSpaceFacet`. The denoise node is also where these are checked: interrupts (`context.util.is_canceled()`), inpainting, img2img through `denoising_start`, and a working-memory estimate passed to `model_on_device(working_mem_bytes=)`, or `working_memory=` for an estimator that declares an `EstimatorSpec`. 4. **VAE encode and decode** (`vae/new_model_image_to_latents.py`, `vae/new_model_latents_to_image.py`) - Apply the VAE's latent normalization, scaling and shift or per-channel mean and std, exactly as the reference pipeline does. Handle tiling and give a working-memory estimate. The estimate helpers live in `invokeai/backend/util/vae_working_memory.py`, with calibration scripts in `scripts/calibrate_*_working_memory.py`. Node types conventionally end in `_i2l` and `_l2i`. + Apply the VAE's latent normalization, scaling and shift or per-channel mean and std, exactly as the reference pipeline does. Handle tiling and give a working-memory estimate. The estimate helpers live in `invokeai/backend/util/vae_working_memory.py`, with calibration scripts in `scripts/calibrate_*_working_memory.py`. A new estimator declares an `EstimatorSpec`, returns a `WorkingMemoryEstimate` and takes its fitted coefficients from `invokeai/backend/working_memory/data/`, next to the measured points they were fitted on: the estimate must be at least each point and, where a fitted term decides it, at most 1.3 times it (or the point's `max_ratio`). The [model manager](/development/architecture/model-manager/) lists what such an estimator needs. `working_mem_bytes=` remains for a plain reservation. Node types conventionally end in `_i2l` and `_l2i`. 5. **LoRA loaders** (if LoRA is supported, see [Optional features](#12-optional-features)) diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index 846883395b4..2dfa5d93010 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -972,6 +972,26 @@ with loaded_model.model_on_device() as (state_dict, vae): Since not all models have state dicts, the `state_dict` return value can be None. +An operation that needs more device memory than `device_working_mem_gb` reserves it when it locks the model, so the cache offloads other models instead of running out of memory. Pass the operation's working-memory estimate: + +```py +estimate = estimate_vae_working_memory_flux2( + operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device +) +with vae_info.model_on_device(working_memory=estimate) as (_, vae): + image = vae.decode(latents)[0] +``` + +An estimator that declares an `EstimatorSpec` returns a `WorkingMemoryEstimate` (`invokeai/backend/working_memory/estimate.py`): the reservation, the terms it adds up to, the features it was evaluated from, and the source of its coefficients, which `model_on_device()` logs at debug level. A plain byte count goes in as `working_mem_bytes=` instead; never both. The cache never reserves less than `device_working_mem_gb`. + +To write such an estimator (`invokeai/backend/util/vae_working_memory.py` has five): + +- Register an `EstimatorSpec` at module import, and list the module in `ESTIMATOR_MODULES` (`working_memory/spec.py`); the data test fails on a module that declares a spec without being listed. +- Split the estimator in two. The impure part reads the model and the device and records every runtime decision (tile, element size, score-matrix bytes) as a feature. The pure `terms(features, values)` turns features and coefficient values into terms, building fitted terms with `WorkingMemoryTerm.fitted(...)`. The public function returns `SPEC.estimate(features, device)`. +- Bump `version` whenever `terms` changes the formula. +- Ship one coefficient entry per measured platform value in `invokeai/backend/working_memory/data/`, either with its own `values` or with `values_from` another estimator's entry. Unmeasured platform values (MPS, XPU, CPU) are priced as `UNMEASURED_PRICED_AS` in `working_memory/platform.py` says. +- Ship the measured points next to them: peak reserved bytes (`reserved_peak`), one fresh process per point. `tests/backend/working_memory/` holds each estimate to at least its point, and where a fitted term decides it, to at most 1.3 times the point, or the point's tighter `max_ratio`. The `device_working_mem_gb` floor counts only where an analytic term decides. Every value an entry carries must be decided by at least one point. + ### Emitting model loading events When the `context` argument is passed to `load_model_*()`, it will retrieve the invocation event bus from the passed `InvocationContext` object to emit events on the invocation bus. The two events are "model_load_started" and "model_load_completed". Both carry the following payload: diff --git a/invokeai/app/invocations/vae/cogview4_image_to_latents.py b/invokeai/app/invocations/vae/cogview4_image_to_latents.py index 09cde601d32..dd4a3862d98 100644 --- a/invokeai/app/invocations/vae/cogview4_image_to_latents.py +++ b/invokeai/app/invocations/vae/cogview4_image_to_latents.py @@ -43,7 +43,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso estimated_working_memory = estimate_vae_working_memory_cogview4( operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae.disable_tiling() diff --git a/invokeai/app/invocations/vae/cogview4_latents_to_image.py b/invokeai/app/invocations/vae/cogview4_latents_to_image.py index 99fb2f21261..e8b655b0ce8 100644 --- a/invokeai/app/invocations/vae/cogview4_latents_to_image.py +++ b/invokeai/app/invocations/vae/cogview4_latents_to_image.py @@ -50,7 +50,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), - vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae), + vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae), ): context.util.signal_progress("Running VAE") assert isinstance(vae, (AutoencoderKL)) diff --git a/invokeai/app/invocations/vae/flux2_vae_decode.py b/invokeai/app/invocations/vae/flux2_vae_decode.py index f0852f1880f..cb5ed93c10a 100644 --- a/invokeai/app/invocations/vae/flux2_vae_decode.py +++ b/invokeai/app/invocations/vae/flux2_vae_decode.py @@ -57,7 +57,7 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would diff --git a/invokeai/app/invocations/vae/flux2_vae_encode.py b/invokeai/app/invocations/vae/flux2_vae_encode.py index 926f3a7c043..4e024158f1c 100644 --- a/invokeai/app/invocations/vae/flux2_vae_encode.py +++ b/invokeai/app/invocations/vae/flux2_vae_encode.py @@ -60,7 +60,7 @@ def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torc device=device, ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype image_tensor = image_tensor.to(device=device, dtype=vae_dtype) diff --git a/invokeai/app/invocations/vae/flux_vae_decode.py b/invokeai/app/invocations/vae/flux_vae_decode.py index 170162e8077..f50c2a3272c 100644 --- a/invokeai/app/invocations/vae/flux_vae_decode.py +++ b/invokeai/app/invocations/vae/flux_vae_decode.py @@ -80,7 +80,7 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents if ( not use_tiling and config.auto_tiled_decode - and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory.reservation_bytes) ): use_tiling = True tile_size = self.tile_size @@ -92,7 +92,7 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents device=vae_info.compute_device, ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae_dtype = next(iter(vae.parameters())).dtype # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from diff --git a/invokeai/app/invocations/vae/flux_vae_encode.py b/invokeai/app/invocations/vae/flux_vae_encode.py index 5984fa737ed..a5a10900608 100644 --- a/invokeai/app/invocations/vae/flux_vae_encode.py +++ b/invokeai/app/invocations/vae/flux_vae_encode.py @@ -90,7 +90,7 @@ def vae_encode( # generator held outside would already have been drawn from by the attempt that failed, so # the retry would produce a different latent than a directly tiled run. generator = torch.Generator(device=vae_info.compute_device).manual_seed(0) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae_dtype = next(iter(vae.parameters())).dtype # The VAE's own device, as in the decode nodes: a cpu_only VAE (or another worker's GPU) diff --git a/invokeai/app/invocations/vae/ideogram4_latents_to_image.py b/invokeai/app/invocations/vae/ideogram4_latents_to_image.py index f2e72d6e94c..0016a13ccec 100644 --- a/invokeai/app/invocations/vae/ideogram4_latents_to_image.py +++ b/invokeai/app/invocations/vae/ideogram4_latents_to_image.py @@ -81,7 +81,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: operation="decode", image_tensor=z, vae=vae_info.model, device=device ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): context.util.signal_progress("Running VAE") vae_dtype = next(vae.parameters()).dtype TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/vae/image_to_latents.py b/invokeai/app/invocations/vae/image_to_latents.py index 71d3951d79c..d599fadd862 100644 --- a/invokeai/app/invocations/vae/image_to_latents.py +++ b/invokeai/app/invocations/vae/image_to_latents.py @@ -83,7 +83,7 @@ def vae_encode( fp32=upcast, device=vae_info.compute_device, ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)), "VAE must be of type SD-1.5 or SDXL" orig_dtype = vae.dtype if upcast: diff --git a/invokeai/app/invocations/vae/latents_to_image.py b/invokeai/app/invocations/vae/latents_to_image.py index 4fbc87b8659..29c6a1b9db8 100644 --- a/invokeai/app/invocations/vae/latents_to_image.py +++ b/invokeai/app/invocations/vae/latents_to_image.py @@ -65,7 +65,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), - vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae), + vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae), ): context.util.signal_progress("Running VAE decoder") assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)) diff --git a/invokeai/app/invocations/vae/sd3_image_to_latents.py b/invokeai/app/invocations/vae/sd3_image_to_latents.py index 5d558eeed52..fe1cb2faaa5 100644 --- a/invokeai/app/invocations/vae/sd3_image_to_latents.py +++ b/invokeai/app/invocations/vae/sd3_image_to_latents.py @@ -39,7 +39,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso estimated_working_memory = estimate_vae_working_memory_sd3( operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae.disable_tiling() diff --git a/invokeai/app/invocations/vae/sd3_latents_to_image.py b/invokeai/app/invocations/vae/sd3_latents_to_image.py index 442ea0a83b6..87f67d73d8f 100644 --- a/invokeai/app/invocations/vae/sd3_latents_to_image.py +++ b/invokeai/app/invocations/vae/sd3_latents_to_image.py @@ -52,7 +52,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), - vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae), + vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae), ): context.util.signal_progress("Running VAE") assert isinstance(vae, (AutoencoderKL)) diff --git a/invokeai/app/invocations/vae/z_image_image_to_latents.py b/invokeai/app/invocations/vae/z_image_image_to_latents.py index de71e54ea0b..149116ae9a0 100644 --- a/invokeai/app/invocations/vae/z_image_image_to_latents.py +++ b/invokeai/app/invocations/vae/z_image_image_to_latents.py @@ -72,7 +72,7 @@ def vae_encode( device=vae_info.compute_device, ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): if not isinstance(vae, AutoencoderKL): raise TypeError( f"Expected AutoencoderKL, got {type(vae).__name__}. VAE model type changed unexpectedly after loading." diff --git a/invokeai/app/invocations/vae/z_image_latents_to_image.py b/invokeai/app/invocations/vae/z_image_latents_to_image.py index 2c5055402f3..af4c4bf07f6 100644 --- a/invokeai/app/invocations/vae/z_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/z_image_latents_to_image.py @@ -74,7 +74,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: if ( not use_tiling and config.auto_tiled_decode - and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory.reservation_bytes) ): use_tiling = True estimated_working_memory = estimate_vae_working_memory_flux( @@ -90,7 +90,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: # FLUX VAE now honours the axes the UI has always offered for it. seamless_context = SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes) - with seamless_context, vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with seamless_context, vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): context.util.signal_progress("Running VAE") if not isinstance(vae, AutoencoderKL): raise TypeError( diff --git a/invokeai/backend/AGENTS.md b/invokeai/backend/AGENTS.md index 3b30a1c98c3..2f12dc1df2e 100644 --- a/invokeai/backend/AGENTS.md +++ b/invokeai/backend/AGENTS.md @@ -6,6 +6,7 @@ - Use existing device/model-management abstractions, `util/devices.py` included. Keep CPU/MPS/ROCm/XPU/multi-device support. No hardcoded CUDA, GPU zero, or dtype assumptions. - Respect session/thread device ownership. Restore temporary model patches, device state, hooks, allocations on success, cancellation, and exceptions. No request-specific mutable globals. - Define model/cache lifetimes and working-memory estimates. Bound queues, caches, downloads, batches, buffers. Consider peak memory and throughput. +- New or refitted working-memory coefficients go in `working_memory/data/` with the measured points that bound them; their estimators declare an `EstimatorSpec` and return a `WorkingMemoryEstimate`. - Keep model validation/loading safeguards, vendored exclusions, upstream provenance. No unrelated vendored rewrites or formatting. ## Performance diff --git a/invokeai/backend/flux/extensions/kontext_extension.py b/invokeai/backend/flux/extensions/kontext_extension.py index 7e801c54c04..1f31859909f 100644 --- a/invokeai/backend/flux/extensions/kontext_extension.py +++ b/invokeai/backend/flux/extensions/kontext_extension.py @@ -138,7 +138,7 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]: operation="encode", image_tensor=image_tensor, vae=vae_info.model ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae_dtype = next(iter(vae.parameters())).dtype image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) diff --git a/invokeai/backend/flux2/ref_image_extension.py b/invokeai/backend/flux2/ref_image_extension.py index 9184b15c1d2..c0e5d36b404 100644 --- a/invokeai/backend/flux2/ref_image_extension.py +++ b/invokeai/backend/flux2/ref_image_extension.py @@ -218,7 +218,7 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]: device=TorchDevice.choose_torch_device(), ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 84a8d539516..fddf9a368dc 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -25,6 +25,8 @@ ModelCache, ) from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType +from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate class LoadedModelWithoutConfig: @@ -194,13 +196,27 @@ def __exit__(self, *args: Any, **kwargs: Any) -> None: @contextmanager def model_on_device( - self, working_mem_bytes: Optional[int] = None + self, working_mem_bytes: Optional[int] = None, *, working_memory: Optional[WorkingMemoryEstimate] = None ) -> Generator[Tuple[Optional[Dict[str, torch.Tensor]], AnyModel], None, None]: """Return a tuple consisting of the model's state dict (if it exists) and the locked model on execution device. :param working_mem_bytes: The amount of working memory to keep available on the compute device when loading the model. + :param working_memory: A working-memory estimate whose `reservation_bytes` is kept available instead; see + `invokeai.backend.working_memory`. Pass this or `working_mem_bytes`, not both. """ + if isinstance(working_mem_bytes, WorkingMemoryEstimate): + raise TypeError( + "model_on_device() got a WorkingMemoryEstimate as working_mem_bytes; pass it as working_memory=." + ) + if working_memory is not None: + if working_mem_bytes is not None: + raise ValueError("Pass either working_mem_bytes or working_memory to model_on_device(), not both.") + working_mem_bytes = working_memory.reservation_bytes + InvokeAILogger.get_logger(__name__).debug( + f"Working memory for {self._cache_record.key}: {working_memory.estimator} reserves " + f"{working_mem_bytes / 2**20:.0f} MB ({working_memory.source})" + ) self._lock_paced(working_mem_bytes) self._end_first_use_window() try: diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 4c92894f1ee..6462dc1f41a 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -1,4 +1,6 @@ -from typing import Literal +from collections.abc import Mapping +from dataclasses import replace +from typing import Literal, cast import torch from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL @@ -13,6 +15,8 @@ from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger from invokeai.backend.util.vae_tiling_scope import resolve_tile_size +from invokeai.backend.working_memory.estimate import FeatureValue, WorkingMemoryEstimate, WorkingMemoryTerm +from invokeai.backend.working_memory.spec import EstimatorSpec, register # Share of the device's memory an untiled decode may claim before the decode nodes tile it up front (see # `should_pretile_vae_decode`). High enough that common resolutions on common cards stay untiled -- tiled output is @@ -110,6 +114,96 @@ def _vae_mid_block_score_matrix_bytes( ) +# The SD1/SDXL, SD3, CogView4, FLUX.1 and FLUX.2 VAEs run the same diffusers AutoencoderKL convolution stack, and pay +# the same for it: peak reserved bytes per output pixel per element byte, one fitted coefficient per operation +# (`k_decode`, `k_encode`) and convolution backend. The values ship in `backend/working_memory/data/vae.json`, with the +# points they were fitted on and the per-card table behind them; `scripts/calibrate_flux2_working_memory.py --only vae` +# reproduces that table on any build. +# +# MIOpen's convolution workspaces are simply larger than cuDNN's -- ~1.6x for a decode, ~2.4x for an encode. This is +# not the attention term: it shows up identically on the fused path, and the implied coefficients are flat across +# resolution on both backends, so the linear model itself holds. Nor is "encoding costs half of decoding" architectural: +# it is a cuDNN property (0.49 there, 0.78 on MIOpen), which is why each operation carries its own coefficient instead +# of a ratio. Shipping the MIOpen values everywhere would add ~60% to every cuDNN decode for nothing, so the +# coefficients follow the backend of the device the VAE runs on (`platform_for`). +# +# Caveat for AMD users: `MIOPEN_FIND_MODE=2` (FAST) selects convolution algorithms heuristically instead of by +# benchmark, and measured a uniform 1.28x more memory at every resolution when the MIOpen values were first fitted. +# Invoke has since made FAST its default (`run_app.py`); the RX 9060 XT measurement was taken under it and did not show +# the 1.28x. If you see decodes exceed their reservation on another card, raise `device_working_mem_gb`. +_AUTOENCODER_KL_COEFFICIENTS = frozenset({"k_decode", "k_encode"}) +_CONV_BACKEND_KEYS = ("conv_backend",) + + +def _peak_of_two_phases( + convolution_phase: tuple[WorkingMemoryTerm, ...], score_matrix_bytes: int +) -> tuple[WorkingMemoryTerm, ...]: + """The convolution phase and the mid-block score matrix, of which only the larger is reserved. + + max, not sum: the two peak in different phases of the same forward. The mid-block sits at the 8x-downsampled + bottleneck -- first in the decoder, last in the encoder -- so the full-resolution convolution feature maps that drive + the linear term are not live while the score matrix is, and peak *reserved* is a high-water mark, not a running + total. Measured on the FLUX.2 VAE (the points and their method are in `data/vae.json`): on gfx1100 and gfx1201 + forcing `math` moves the measured peak by nothing at all up to 1024px, and the totals stay flat-linear in area + either way. On CUDA the score matrix only pokes above the convolution peak at 1536px, and then by 2.6GB against the + 21.5GB this term prices standalone, because the attention phase reuses blocks the allocator is already holding. + Summing them reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7GB; taking the max reserves 6.9GB. + + Where a max model is weakest is the crossover, where the two terms are near-equal and whatever overlap exists is + no longer hidden. There is exactly one measured point like that: a 768px encode with cuDNN's linear coefficient and + a materializing kernel measures 1.80GB against a 1.35GB max. It is not reachable as a shortfall, because the cache + floors every reservation at `device_working_mem_gb` (3GB by default, see `ModelCache._get_vram_available`) and the + whole crossover region sits under that floor. On the builds that really do materialize -- gfx1100 and gfx1201 -- + the MIOpen coefficient keeps the linear term above the score term across the measured range, so the crossover does + not arise there at all. + + On a tie the convolutions are the ones reserved. + """ + convolution_bytes: float = 0 + for term in convolution_phase: + convolution_bytes += term.bytes + scores_win = score_matrix_bytes > convolution_bytes + return ( + *(replace(term, included=not scores_win) for term in convolution_phase), + WorkingMemoryTerm("mid_block_scores", score_matrix_bytes, included=scores_win), + ) + + +def _sd_terms(features: Mapping[str, FeatureValue], values: Mapping[str, float]) -> tuple[WorkingMemoryTerm, ...]: + coefficient = f"k_{features['operation']}" + element_size = cast(int, features["element_size"]) + tile = cast(int | None, features["tile"]) + if tile is not None: + area_bytes = tile * tile * element_size + # We add 25% to the working memory estimate when tiling is enabled to account for factors like tile overlap + # and number of tiles. We could make this more precise in the future, but this should be good enough for + # most use cases. + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes * 1.25, values) + else: + area_bytes = cast(int, features["height"]) * cast(int, features["width"]) * element_size + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes, values) + + terms: tuple[WorkingMemoryTerm, ...] = (convolutions,) + # The tiny autoencoder has no mid-block attention, so it is priced without a score matrix. + if features["vae_class"] == "AutoencoderKL": + terms = _peak_of_two_phases(terms, cast(int, features["score_matrix_bytes"])) + if features["fp32"]: + # If we are running in FP32, then we should account for the likely increase in model size (~250MB). + terms += (WorkingMemoryTerm("fp32_weights", 250 * 2**20, True),) + return terms + + +VAE_SD_SPEC = register( + EstimatorSpec( + id="vae.sd", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_sd_terms, + ) +) + + def estimate_vae_working_memory_sd15_sdxl( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, @@ -117,47 +211,75 @@ def estimate_vae_working_memory_sd15_sdxl( tile_size: int | None, fp32: bool, device: torch.device | None = None, -) -> int: - """Estimate the working memory required to encode or decode the given tensor.""" - # It was found experimentally that the peak working memory scales linearly with the number of pixels and the - # element size (precision). This estimate is accurate for both SD1 and SDXL. - element_size = 4 if fp32 else 2 - - # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per - # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where - # the VAE will run; it defaults to the session device. - scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) +) -> WorkingMemoryEstimate: + """Estimate the working memory required to encode or decode the given tensor. + It was found experimentally that the peak working memory scales linearly with the number of pixels and the element + size (precision). This estimate is accurate for both SD1 and SDXL. `device` is where the VAE will run; it picks the + convolution backend's coefficients and defaults to the session device. + """ + element_size = 4 if fp32 else 2 + device = device if device is not None else TorchDevice.choose_torch_device() latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + height = latent_scale_factor_for_operation * image_tensor.shape[-2] + width = latent_scale_factor_for_operation * image_tensor.shape[-1] + + if tile_size == 0: + tile_size = vae.tile_sample_min_size + assert isinstance(tile_size, int) + + features: dict[str, FeatureValue] = { + "operation": operation, + "height": height, + "width": width, + "tile": tile_size, + "element_size": element_size, + "fp32": fp32, + "vae_class": "AutoencoderKL" if isinstance(vae, AutoencoderKL) else "AutoencoderTiny", + } + if isinstance(vae, AutoencoderKL): + score_dtype = torch.float32 if fp32 else next(vae.parameters()).dtype + score_h, score_w = (tile_size, tile_size) if tile_size is not None else (height, width) + features["score_matrix_bytes"] = _vae_mid_block_score_matrix_bytes(score_h, score_w, score_dtype, device=device) - if tile_size is not None: - if tile_size == 0: - tile_size = vae.tile_sample_min_size - assert isinstance(tile_size, int) - h = tile_size - w = tile_size - working_memory = h * w * element_size * scaling_constant + return VAE_SD_SPEC.estimate(features, device) - # We add 25% to the working memory estimate when tiling is enabled to account for factors like tile overlap - # and number of tiles. We could make this more precise in the future, but this should be good enough for - # most use cases. - working_memory = working_memory * 1.25 - else: - h = latent_scale_factor_for_operation * image_tensor.shape[-2] - w = latent_scale_factor_for_operation * image_tensor.shape[-1] - working_memory = h * w * element_size * scaling_constant - if isinstance(vae, AutoencoderKL): - # max, not sum: the mid-block attention and the full-resolution convolutions peak in - # different phases of the forward (see the FLUX.2 estimator for the measurements). - score_dtype = torch.float32 if fp32 else next(vae.parameters()).dtype - working_memory = max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, score_dtype, device=device)) +def _single_pass_terms( + features: Mapping[str, FeatureValue], values: Mapping[str, float] +) -> tuple[WorkingMemoryTerm, ...]: + """An untiled AutoencoderKL pass: the convolutions over the whole image, or the mid-block score matrix.""" + coefficient = f"k_{features['operation']}" + area_bytes = cast(int, features["height"]) * cast(int, features["width"]) * cast(int, features["element_size"]) + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes, values) + return _peak_of_two_phases((convolutions,), cast(int, features["score_matrix_bytes"])) - if fp32: - # If we are running in FP32, then we should account for the likely increase in model size (~250MB). - working_memory += 250 * 2**20 - return int(working_memory) +def _single_pass_features( + operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL, device: torch.device +) -> dict[str, FeatureValue]: + latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + height = latent_scale_factor_for_operation * image_tensor.shape[-2] + width = latent_scale_factor_for_operation * image_tensor.shape[-1] + param = next(vae.parameters()) + return { + "operation": operation, + "height": height, + "width": width, + "element_size": param.element_size(), + "score_matrix_bytes": _vae_mid_block_score_matrix_bytes(height, width, param.dtype, device=device), + } + + +VAE_COGVIEW4_SPEC = register( + EstimatorSpec( + id="vae.cogview4", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_single_pass_terms, + ) +) def estimate_vae_working_memory_cogview4( @@ -165,25 +287,14 @@ def estimate_vae_working_memory_cogview4( image_tensor: torch.Tensor, vae: AutoencoderKL, device: torch.device | None = None, -) -> int: - """Estimate the working memory required by the invocation in bytes.""" - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 - - h = latent_scale_factor_for_operation * image_tensor.shape[-2] - w = latent_scale_factor_for_operation * image_tensor.shape[-1] - element_size = next(vae.parameters()).element_size() - - # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per - # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where - # the VAE will run; it defaults to the session device. - scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) - working_memory = h * w * element_size * scaling_constant +) -> WorkingMemoryEstimate: + """Estimate the working memory required by the invocation in bytes. - # max, not sum: the mid-block attention and the full-resolution convolutions peak in different - # phases of the forward (see the FLUX.2 estimator for the measurements). - return int( - max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, next(vae.parameters()).dtype, device=device)) - ) + `device` is where the VAE will run; it picks the convolution backend's coefficients and defaults to the session + device. + """ + device = device if device is not None else TorchDevice.choose_torch_device() + return VAE_COGVIEW4_SPEC.estimate(_single_pass_features(operation, image_tensor, vae, device), device) # What a tiled decode does *not* bound: the assembled image, several times over, per output pixel. @@ -217,13 +328,80 @@ def estimate_vae_working_memory_cogview4( _VAE_MOMENT_CHANNELS_PER_LATENT = 2 +def _flux_terms(features: Mapping[str, FeatureValue], values: Mapping[str, float]) -> tuple[WorkingMemoryTerm, ...]: + operation = features["operation"] + coefficient = f"k_{operation}" + out_h = cast(int, features["height"]) + out_w = cast(int, features["width"]) + element_size = cast(int, features["element_size"]) + tile = cast(int | None, features["tile"]) + + if tile is None: + area_bytes = out_h * out_w * element_size + convolution_phase: tuple[WorkingMemoryTerm, ...] = ( + WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes, values), + ) + return _peak_of_two_phases(convolution_phase, cast(int, features["score_matrix_bytes"])) + + # `calc_tiles_min_overlap` clamps the tile to the image *per axis*, so the largest tile the + # decode builds is this, not `tile` squared. Without the clamp an 8192px tile on a + # 1024x8200 image -- tiled, because one axis exceeds it -- prices 8192x8192. + area_bytes = min(tile, out_h) * min(tile, out_w) * element_size + # A 25% margin for tile overlap and the number of tiles, mirroring the SD1/SDXL estimator. + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes * 1.25, values) + if operation == "decode": + convolution_phase = ( + convolutions, + WorkingMemoryTerm( + "assembled_image", + out_h * out_w * _IMAGE_CHANNELS * (_FLUX_VAE_TILED_IMAGE_COPIES * element_size + 1), + True, + ), + ) + else: + # A tiled encode leaves two things un-bounded: the full-resolution image it slices tiles + # from, and the moments it assembles. + # + # The moments are *not* one copy. `AutoencoderKL._tiled_encode` appends every encoded + # tile to `rows`, then builds `result_rows` from them -- blending against `rows[i-1][j]`, + # so `rows` stays live throughout -- and finally concatenates `enc` on top of both. The + # tiles overlap, so their total area is `1 / (1 - overlap_factor)**2` times the output; + # `result_rows` and `enc` are one output each. Hence the factor below, which is 3.78 at + # the 0.25 every `AutoencoderKL` ships. + live_moment_copies = 1.0 / (1.0 - cast(float, features["tile_overlap_factor"])) ** 2 + 2.0 + moments = ( + (out_h // LATENT_SCALE_FACTOR) + * (out_w // LATENT_SCALE_FACTOR) + * _VAE_MOMENT_CHANNELS_PER_LATENT + * cast(int, features["latent_channels"]) + * element_size + ) + convolution_phase = ( + convolutions, + WorkingMemoryTerm("input_image", out_h * out_w * _IMAGE_CHANNELS * element_size, True), + WorkingMemoryTerm("assembled_moments", moments * live_moment_copies, True), + ) + return _peak_of_two_phases(convolution_phase, cast(int, features["score_matrix_bytes"])) + + +VAE_FLUX_SPEC = register( + EstimatorSpec( + id="vae.flux", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_flux_terms, + ) +) + + def estimate_vae_working_memory_flux( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL, tile_size: int | None = None, device: torch.device | None = None, -) -> int: +) -> WorkingMemoryEstimate: """Estimate the working memory required by the invocation in bytes. `tile_size` is in output pixels and defaults to None, i.e. a single-pass decode -- the existing @@ -237,16 +415,11 @@ def estimate_vae_working_memory_flux( for the decode side. `device` is where the VAE will run (a `cpu_only` VAE runs on the CPU); it picks the convolution - backend's constant. Defaults to the session device. + backend's coefficients. Defaults to the session device. """ - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 - element_size = next(vae.parameters()).element_size() - - # The FLUX.1 autoencoder shares the FLUX.2 VAE's convolution stack and measures the same constants on both - # backends; see `_FLUX_VAE_SCALING_CONSTANTS`. + param = next(vae.parameters()) device = device if device is not None else TorchDevice.choose_torch_device() - scaling_constant = _flux_vae_scaling_constant(operation, device) out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] @@ -261,51 +434,27 @@ def estimate_vae_working_memory_flux( # field carries no upper bound: a 2048px tile on a 1024px image would otherwise reserve ~23GB # for a decode that needs ~4.6GB, evicting the transformer from the cache for nothing. if tile is not None and (tile < out_h or tile < out_w): - # `calc_tiles_min_overlap` clamps the tile to the image *per axis*, so the largest tile the - # decode builds is this, not `tile` squared. Without the clamp an 8192px tile on a - # 1024x8200 image -- tiled, because one axis exceeds it -- prices 8192x8192. - tile_h = min(tile, out_h) - tile_w = min(tile, out_w) - # A 25% margin for tile overlap and the number of tiles, mirroring the SD1/SDXL estimator. - working_memory = tile_h * tile_w * element_size * scaling_constant * 1.25 - if operation == "decode": - working_memory += out_h * out_w * _IMAGE_CHANNELS * (_FLUX_VAE_TILED_IMAGE_COPIES * element_size + 1) - else: - # A tiled encode leaves two things un-bounded: the full-resolution image it slices tiles - # from, and the moments it assembles. - # - # The moments are *not* one copy. `AutoencoderKL._tiled_encode` appends every encoded - # tile to `rows`, then builds `result_rows` from them -- blending against `rows[i-1][j]`, - # so `rows` stays live throughout -- and finally concatenates `enc` on top of both. The - # tiles overlap, so their total area is `1 / (1 - overlap_factor)**2` times the output; - # `result_rows` and `enc` are one output each. Hence the factor below, which is 3.78 at - # the 0.25 every `AutoencoderKL` ships. - # - # Read from the VAE rather than hard-coded: an overlap factor of 0.5 would make the - # tiles alone four times the output, and a fixed 3.78 would then under-reserve. An - # estimate consumed as `free >= estimate` has to err upwards. - config = getattr(vae, "config", None) - latent_channels = getattr(config, "latent_channels", 16) - overlap_factor = getattr(vae, "tile_overlap_factor", 0.25) - live_moment_copies = 1.0 / (1.0 - overlap_factor) ** 2 + 2.0 - moments = ( - (out_h // LATENT_SCALE_FACTOR) - * (out_w // LATENT_SCALE_FACTOR) - * _VAE_MOMENT_CHANNELS_PER_LATENT - * latent_channels - * element_size - ) - working_memory += out_h * out_w * _IMAGE_CHANNELS * element_size - working_memory += moments * live_moment_copies - score_h, score_w = tile_h, tile_w + score_h, score_w = min(tile, out_h), min(tile, out_w) else: - working_memory = out_h * out_w * element_size * scaling_constant + tile = None score_h, score_w = out_h, out_w - # max, not sum: the mid-block attention and the full-resolution convolutions peak in different - # phases of the forward (see the FLUX.2 estimator for the measurements). - dtype = next(vae.parameters()).dtype - return int(max(working_memory, _vae_mid_block_score_matrix_bytes(score_h, score_w, dtype, device=device))) + features: dict[str, FeatureValue] = { + "operation": operation, + "height": out_h, + "width": out_w, + "element_size": param.element_size(), + "tile": tile, + "score_matrix_bytes": _vae_mid_block_score_matrix_bytes(score_h, score_w, param.dtype, device=device), + } + if tile is not None and operation != "decode": + # Read from the VAE rather than hard-coded: an overlap factor of 0.5 would make the tiles + # alone four times the output, and a fixed 3.78 would then under-reserve. An estimate + # consumed as `free >= estimate` has to err upwards. + features["latent_channels"] = getattr(getattr(vae, "config", None), "latent_channels", 16) + features["tile_overlap_factor"] = getattr(vae, "tile_overlap_factor", 0.25) + + return VAE_FLUX_SPEC.estimate(features, device) # The FLUX.2 VAE runs one attention block at the bottom of the encoder and one at the top of the @@ -317,51 +466,32 @@ def estimate_vae_working_memory_flux( _FLUX2_VAE_MID_BLOCK_HEAD_DIM = 512 _FLUX2_VAE_SPATIAL_COMPRESSION = 8 -# Peak reserved bytes per output pixel per element byte, per conv backend and operation, for the FLUX.2 VAE and the -# FLUX.1 autoencoder alike. Fitted with `scripts/calibrate_flux2_working_memory.py --only vae`, which reproduces -# this table on any build: -# -# decode encode encode/decode -# cuDNN RTX 4090, torch 2.7.1 2185 1072 0.49 -# MIOpen RX 9070 XT, torch 2.10 3453 2688 0.78 -# MIOpen PRO W7900, torch 2.10 3525 2688 0.76 -# MIOpen RX 9060 XT, torch 2.12, Windows, FLUX.1 VAE 3451 2688 0.78 -# -# Two AMD generations (RDNA3 and RDNA4), the same numbers: the encode column agrees to the byte and -# the decode column to 2%. So this is MIOpen, not a per-card quirk, and the column below is fitted -# to the larger with ~2% headroom. The FLUX.1 row (bf16, 768-1536px, Invoke's ROCm SDPA guard installed, -# `MIOPEN_FIND_MODE=FAST`) lands on the same numbers: the two VAEs share their convolution stack. The -# cuDNN column is the FLUX.1 constant this estimator shipped with before. -# -# MIOpen's convolution workspaces are simply larger than cuDNN's. This is not the attention term -- -# it shows up identically on the fused path, and the implied constants are flat across resolution on -# both backends, so the linear model itself holds. And the "encoding costs half of decoding" ratio -# the other estimators in this module use turns out to be a cuDNN property rather than an -# architectural one, which is why the two operations carry their own numbers here instead of a ratio. -# -# Shipping the MIOpen numbers everywhere would add ~60% to every cuDNN decode for nothing, so the -# constant follows the backend. -# -# Caveat for AMD users: `MIOPEN_FIND_MODE=2` (FAST) selects convolution algorithms heuristically instead of -# by benchmark, and measured a uniform 1.28x more memory at every resolution when these constants were -# first fitted. Invoke has since made FAST its default (`run_app.py`); the RX 9060 XT row above was taken -# under it and did not show the 1.28x. If you see decodes exceed their reservation on another card, raise -# `device_working_mem_gb`. -_FLUX_VAE_SCALING_CONSTANTS: dict[str, dict[str, int]] = { - "cudnn": {"decode": 2200, "encode": 1100}, - "miopen": {"decode": 3600, "encode": 2750}, -} - -def _flux_vae_scaling_constant(operation: Literal["encode", "decode"], device: torch.device) -> int: - """Pick the pixel-area constant for the convolution backend this device will actually use. - - A HIP build reports ``device.type == "cuda"``, so the torch build -- not the device string -- is - what separates MIOpen from cuDNN. MPS and CPU are unmeasured and take the cuDNN column; on MPS - that pairs with a score-matrix term the probe always charges there, so the total is not thin. - """ - is_rocm = device.type == "cuda" and torch.version.hip is not None - return _FLUX_VAE_SCALING_CONSTANTS["miopen" if is_rocm else "cudnn"][operation] +def _flux2_terms(features: Mapping[str, FeatureValue], values: Mapping[str, float]) -> tuple[WorkingMemoryTerm, ...]: + coefficient = f"k_{features['operation']}" + element_size = cast(int, features["element_size"]) + batch_size = cast(int, features["batch_size"]) + tile = cast(int | None, features["tile"]) + if tile is not None: + area_bytes = tile * tile * element_size + # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes * 1.25 * batch_size, values) + else: + area_bytes = cast(int, features["height"]) * cast(int, features["width"]) * element_size + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes * batch_size, values) + # Both terms are per sample; the larger is reserved (see `_peak_of_two_phases` for why, and the measurements). + return _peak_of_two_phases((convolutions,), cast(int, features["score_matrix_bytes"])) + + +VAE_FLUX2_SPEC = register( + EstimatorSpec( + id="vae.flux2", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_flux2_terms, + ) +) def estimate_vae_working_memory_flux2( @@ -370,15 +500,15 @@ def estimate_vae_working_memory_flux2( vae: AutoencoderKLFlux2, tile_size: int | None = None, device: torch.device | None = None, -) -> int: +) -> WorkingMemoryEstimate: """Estimate the working memory required to encode or decode with the FLUX.2 (32-channel) VAE. Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE, and - the implied constant is flat across 512-1536px on every backend measured. What is *not* constant - is the constant itself: MIOpen's convolution workspaces cost ~1.6x cuDNN's for a decode and ~2.4x - for an encode, so it is looked up per backend -- see ``_FLUX_VAE_SCALING_CONSTANTS`` for the - fitted table and the caveats. Peak *reserved* memory is what is measured throughout, the - conservative quantity that includes allocator overhead. + the implied coefficient is flat across 512-1536px on every backend measured. What is *not* constant + is the coefficient itself: MIOpen's convolution workspaces cost ~1.6x cuDNN's for a decode and ~2.4x + for an encode, so it is shipped per backend -- see `_AUTOENCODER_KL_COEFFICIENTS` for the caveats + and `backend/working_memory/data/vae.json` for the fitted values and their points. Peak *reserved* + memory is what is measured throughout, the conservative quantity that includes allocator overhead. For reference, decoding 1024x1024 peaks at ~4.3GB on cuDNN and ~6.6GB on MIOpen, and 1536x1536 at ~9.6GB -- far above the default ``device_working_mem_gb``, which is why this estimate must be @@ -394,10 +524,10 @@ def estimate_vae_working_memory_flux2( for everything else, so the estimate is right on all of them. The two terms are independent: a build can have a fused kernel and still need the larger - convolution constant (gfx1201 did, before the ROCm rule above). They also do not add -- see the - ``max`` at the end of this function for why, and for the measurements behind it. (Unlike the transformer, this attention does not go through - diffusers' attention dispatcher -- ``AttnProcessor2_0`` calls ``F.scaled_dot_product_attention`` - itself -- so torch's own answer is the whole answer here.) + convolution coefficient (gfx1201 did, before the ROCm rule above). They also do not add -- see + `_peak_of_two_phases` for why, and for the measurements behind it. (Unlike the transformer, this + attention does not go through diffusers' attention dispatcher -- ``AttnProcessor2_0`` calls + ``F.scaled_dot_product_attention`` itself -- so torch's own answer is the whole answer here.) When tiling is enabled the peak is bounded by a single tile instead of the full image (measured ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap), and the score @@ -411,53 +541,34 @@ def estimate_vae_working_memory_flux2( call sites all pass batch 1 today; the shared estimator does not assume it. """ param = next(vae.parameters()) - element_size = param.element_size() - device = device if device is not None else TorchDevice.choose_torch_device() - scaling_constant = _flux_vae_scaling_constant(operation, device) batch_size = image_tensor.shape[0] if image_tensor.dim() >= 4 else 1 + latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] + out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] if tile_size is not None: - # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. - working_memory = tile_size * tile_size * element_size * scaling_constant * 1.25 mid_block_seq_len = (tile_size // _FLUX2_VAE_SPATIAL_COMPRESSION) ** 2 else: - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 - out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] - out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] - working_memory = out_h * out_w * element_size * scaling_constant mid_block_seq_len = (out_h // _FLUX2_VAE_SPATIAL_COMPRESSION) * (out_w // _FLUX2_VAE_SPATIAL_COMPRESSION) - working_memory *= batch_size - score_matrix_bytes = sdpa_score_matrix_bytes( - device=device, - dtype=param.dtype, - # The score matrix is (batch, heads, S, S); one head per sample prices the whole batch. - num_heads=_FLUX2_VAE_MID_BLOCK_HEADS * batch_size, - head_dim=_FLUX2_VAE_MID_BLOCK_HEAD_DIM, - seq_len=mid_block_seq_len, - ) - - # max, not sum: the two terms peak in different phases of the same forward. The mid-block sits at - # the 8x-downsampled bottleneck -- first in the decoder, last in the encoder -- so the full- - # resolution convolution feature maps that drive the linear term are not live while the score - # matrix is, and peak *reserved* is a high-water mark, not a running total. Measured (see the - # constants table above for the method): on gfx1100 and gfx1201 forcing `math` moves the measured - # peak by nothing at all up to 1024px, and the totals stay flat-linear in area either way. On - # CUDA the score matrix only pokes above the convolution peak at 1536px, and then by 2.6GB - # against the 21.5GB this term prices standalone, because the attention phase reuses blocks the - # allocator is already holding. Summing them reserved 11.1GB for a 1024px gfx1100 decode that - # measures 6.7GB; taking the max reserves 6.9GB. - # - # Where a max model is weakest is the crossover, where the two terms are near-equal and whatever - # overlap exists is no longer hidden. There is exactly one measured point like that: a 768px - # encode with cuDNN's linear constant and a materializing kernel measures 1.80GB against a - # 1.35GB max. It is not reachable as a shortfall, because the cache floors every reservation at - # `device_working_mem_gb` (3GB by default, see `ModelCache._get_vram_available`) and the whole - # crossover region sits under that floor. On the builds that really do materialize -- gfx1100 and - # gfx1201 -- the MIOpen constant keeps the linear term above the score term across the measured - # range, so the crossover does not arise there at all. - return int(max(working_memory, score_matrix_bytes)) + features: dict[str, FeatureValue] = { + "operation": operation, + "height": out_h, + "width": out_w, + "tile": tile_size, + "batch_size": batch_size, + "element_size": param.element_size(), + "score_matrix_bytes": sdpa_score_matrix_bytes( + device=device, + dtype=param.dtype, + # The score matrix is (batch, heads, S, S); one head per sample prices the whole batch. + num_heads=_FLUX2_VAE_MID_BLOCK_HEADS * batch_size, + head_dim=_FLUX2_VAE_MID_BLOCK_HEAD_DIM, + seq_len=mid_block_seq_len, + ), + } + return VAE_FLUX2_SPEC.estimate(features, device) def estimate_vae_working_memory_anima( @@ -954,30 +1065,27 @@ def estimate_vae_working_memory_qwen_image( return int(working_memory) +VAE_SD3_SPEC = register( + EstimatorSpec( + id="vae.sd3", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_single_pass_terms, + ) +) + + def estimate_vae_working_memory_sd3( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL, device: torch.device | None = None, -) -> int: - """Estimate the working memory required by the invocation in bytes.""" - # Encode operations use approximately 50% of the memory required for decode operations - - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 - - h = latent_scale_factor_for_operation * image_tensor.shape[-2] - w = latent_scale_factor_for_operation * image_tensor.shape[-1] - element_size = next(vae.parameters()).element_size() - - # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per - # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where - # the VAE will run; it defaults to the session device. - scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) - - working_memory = h * w * element_size * scaling_constant +) -> WorkingMemoryEstimate: + """Estimate the working memory required by the invocation in bytes. - # max, not sum: the mid-block attention and the full-resolution convolutions peak in different - # phases of the forward (see the FLUX.2 estimator for the measurements). - return int( - max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, next(vae.parameters()).dtype, device=device)) - ) + `device` is where the VAE will run; it picks the convolution backend's coefficients and defaults to the session + device. + """ + device = device if device is not None else TorchDevice.choose_torch_device() + return VAE_SD3_SPEC.estimate(_single_pass_features(operation, image_tensor, vae, device), device) diff --git a/invokeai/backend/working_memory/__init__.py b/invokeai/backend/working_memory/__init__.py new file mode 100644 index 00000000000..dd71f347f51 --- /dev/null +++ b/invokeai/backend/working_memory/__init__.py @@ -0,0 +1,6 @@ +"""Working-memory estimates as data: what an operation needs on the compute device, and how that was priced. + +An estimator turns an operation's runtime facts (`features`) and a set of fitted coefficients into a +`WorkingMemoryEstimate` that `LoadedModel.model_on_device(working_memory=...)` reserves. The coefficients ship as data +under `data/`, next to the measured points that bound them; `tests/backend/working_memory/` enforces both. +""" diff --git a/invokeai/backend/working_memory/coefficients.py b/invokeai/backend/working_memory/coefficients.py new file mode 100644 index 00000000000..d8e07b13854 --- /dev/null +++ b/invokeai/backend/working_memory/coefficients.py @@ -0,0 +1,191 @@ +"""The coefficients estimators are evaluated with, and the measured points that bound them. + +Both ship as JSON under `data/`, read through `importlib.resources` so a wheel install finds them. A coefficient entry +belongs to one estimator version on one measured platform value, and either carries its values or borrows another +entry's (`values_from`); a point records the features of an operation that was measured and its measured peak. +`tests/backend/working_memory/test_shipped_coefficients.py` resolves every entry against its `EstimatorSpec` and holds +every point to the estimate the shipped coefficients give it. +""" + +import functools +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from importlib.resources import files +from importlib.resources.abc import Traversable +from types import MappingProxyType +from typing import TYPE_CHECKING, Literal, Self + +from pydantic import BaseModel, ConfigDict, ValidationError, model_validator + +from invokeai.backend.working_memory.estimate import FeatureValue +from invokeai.backend.working_memory.platform import UNMEASURED_PRICED_AS, Platform + +if TYPE_CHECKING: + from invokeai.backend.working_memory.spec import EstimatorSpec + + +class _DataModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class ValuesFrom(_DataModel): + """The entry whose values another entry takes: same platform, this estimator id and version.""" + + estimator: str + version: int + + +class CoefficientEntry(_DataModel): + """The coefficient values of one estimator version on one platform: its own, or another entry's.""" + + estimator: str + version: int + platform: dict[str, str] + values: dict[str, float] | None = None + values_from: ValuesFrom | None = None + """Borrow the values of another estimator's entry on the same platform, so that its measured points bound this + estimator too. One hop: the entry borrowed from carries its own values.""" + provenance: str + """Card, torch build, method and script the values come from, and which points they were fitted on.""" + + @model_validator(mode="after") + def _values_or_values_from(self) -> Self: + if (self.values is None) == (self.values_from is None): + raise ValueError("a coefficient entry carries exactly one of `values` and `values_from`") + return self + + +class MeasuredPoint(_DataModel): + """One measured operation. The estimate for its features must be at least `measured_bytes`, and where a fitted + term decides it, at most 1.3 times that (or `max_ratio`).""" + + estimator: str + version: int + platform: dict[str, str] + features: dict[str, FeatureValue] + measured_bytes: int + method: Literal["reserved_peak"] + """``reserved_peak``: the caching allocator's peak reserved bytes over the operation, one fresh process per + point. Reserved rather than allocated, because reserved is what the model cache's reservation has to cover.""" + device: str + provenance: str + max_ratio: float | None = None + """A tighter upper bound than the default 1.3, as a multiple of `measured_bytes`, for a point the estimate has to + stay closer to. It applies whichever term decides the estimate.""" + + +class WorkingMemoryData(_DataModel): + """One file under `data/`.""" + + schema_version: Literal[1] + coefficients: list[CoefficientEntry] + points: list[MeasuredPoint] + + +@dataclass(frozen=True) +class ResolvedCoefficients: + """The coefficient values an estimate is evaluated with, and where they came from.""" + + values: Mapping[str, float] + source: str + + +def _parse(entry: Traversable) -> WorkingMemoryData: + try: + return WorkingMemoryData.model_validate_json(entry.read_text(encoding="utf-8")) + except ValidationError as e: + raise ValueError(f"Malformed working-memory data file {entry}: {e}") from e + + +def load_data(directory: Traversable) -> Mapping[str, WorkingMemoryData]: + """Every JSON data file in `directory`, parsed, by file name. + + A directory with none raises, naming it: every estimator resolves its coefficients from there, so an install that + lost the files would otherwise fail later as a missing entry for whichever estimator ran first. + """ + entries = ( + sorted((entry for entry in directory.iterdir() if entry.name.endswith(".json")), key=lambda entry: entry.name) + if directory.is_dir() + else [] + ) + if not entries: + raise FileNotFoundError(f"No working-memory data files in {directory}; every EstimatorSpec reads them there.") + return MappingProxyType({entry.name: _parse(entry) for entry in entries}) + + +@functools.cache +def shipped_data() -> Mapping[str, WorkingMemoryData]: + """The data files shipped under `data/`. Read once per process.""" + return load_data(files("invokeai.backend.working_memory") / "data") + + +def _where(keys: Mapping[str, str]) -> str: + return ",".join(f"{key}={value}" for key, value in keys.items()) + + +def _entry( + entries: list[CoefficientEntry], estimator: str, version: int, keys: Mapping[str, str], where: str +) -> CoefficientEntry: + candidates = [entry for entry in entries if entry.estimator == estimator and entry.platform == keys] + matches = [entry for entry in candidates if entry.version == version] + if len(matches) > 1: + raise LookupError( + f"Working-memory coefficients for estimator '{estimator}' version {version} on {where} are declared " + f"{len(matches)} times." + ) + if not matches: + found = f" (found only version {sorted({entry.version for entry in candidates})})" if candidates else "" + raise LookupError( + f"No working-memory coefficients for estimator '{estimator}' version {version} on {where}{found}." + ) + return matches[0] + + +def resolve_coefficients( + data: Iterable[WorkingMemoryData], spec: "EstimatorSpec", platform: Platform, origin: str +) -> ResolvedCoefficients: + """The values `spec` takes on `platform`, and a `source` that says how they were found. + + A platform value nothing is measured for is priced as the measured value `UNMEASURED_PRICED_AS` names. The entry + for the spec's id, version and those platform values either carries the values or borrows them, one hop, from + another estimator's entry on the same platform. No other fallback: an estimate priced with coefficients fitted for + something else is the shortfall this data exists to prevent, so a missing, ambiguous, chained or mismatched entry + raises, naming what is wrong. `origin` prefixes the `source`. + """ + requested = {key: getattr(platform, key) for key in spec.platform_keys} + keys = {key: UNMEASURED_PRICED_AS.get(key, {}).get(value, value) for key, value in requested.items()} + where = _where(keys) if keys == requested else f"{_where(keys)} (priced for {_where(requested)})" + entries = [entry for file in data for entry in file.coefficients] + + entry = _entry(entries, spec.id, spec.version, keys, where) + values, lender = entry.values, None + if entry.values_from is not None: + lender = entry.values_from.estimator + lent = _entry(entries, lender, entry.values_from.version, keys, where) + if lent.values is None: + raise LookupError( + f"Working-memory coefficients for estimator '{spec.id}' on {where} borrow from '{lender}', whose " + "entry borrows in turn; borrow from the entry that carries the values." + ) + values = lent.values + assert values is not None, "a validated entry carries values or values_from" + if set(values) != spec.coefficients: + raise LookupError( + f"Working-memory coefficients for estimator '{spec.id}' on {where} provide {sorted(values)}" + f"{f' (from {lender})' if lender else ''}; the estimator reads {sorted(spec.coefficients)}." + ) + + priced = ",".join(value for key, value in keys.items() if value != requested[key]) + if priced and lender: + note = f" ({priced} values of {lender})" + elif priced: + note = f" ({priced} values)" + else: + note = f" (values of {lender})" if lender else "" + return ResolvedCoefficients(values=MappingProxyType(dict(values)), source=f"{origin}:{_where(requested)}{note}") + + +@functools.cache +def shipped_coefficients(spec: "EstimatorSpec", platform: Platform) -> ResolvedCoefficients: + """The shipped values `spec` takes on `platform`. Resolved once per spec and platform.""" + return resolve_coefficients(shipped_data().values(), spec, platform, origin="shipped") diff --git a/invokeai/backend/working_memory/data/vae.json b/invokeai/backend/working_memory/data/vae.json new file mode 100644 index 00000000000..bae9081384f --- /dev/null +++ b/invokeai/backend/working_memory/data/vae.json @@ -0,0 +1,769 @@ +{ + "schema_version": 1, + "coefficients": [ + { + "estimator": "vae.sd", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the SD1/SDXL VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.sd", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the SD1/SDXL VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.sd3", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the SD3 VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.sd3", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the SD3 VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.cogview4", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the CogView4 VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.cogview4", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the CogView4 VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.flux", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the FLUX.1 VAE runs the same diffusers AutoencoderKL convolution stack. The FLUX.1 point below falls inside them." + }, + { + "estimator": "vae.flux", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the FLUX.1 VAE runs the same diffusers AutoencoderKL convolution stack. Confirmed on the FLUX.1 VAE: RX 9060 XT (gfx1200), torch 2.12, Windows, bf16, 768-1536px, Invoke's ROCm SDPA guard installed, MIOPEN_FIND_MODE=FAST: implied 3451 decode / 2688 encode." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values": { + "k_decode": 2200, + "k_encode": 1100 + }, + "provenance": "Peak reserved bytes per output pixel per element byte of the diffusers AutoencoderKL convolution stack. Fitted on the FLUX.2 VAE (the vae.flux2 cudnn points) with scripts/calibrate_flux2_working_memory.py --only vae: RTX 4090, torch 2.7.1, bf16, untiled, fused SDPA; implied 2180/2185/2165 decode and 1072/1063/1061 encode at 512/768/1024px, flat across resolution. Rounded up to 2200 / 1100, the pair the FLUX.1 estimator already shipped." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values": { + "k_decode": 3600, + "k_encode": 2750 + }, + "provenance": "Peak reserved bytes per output pixel per element byte of the diffusers AutoencoderKL convolution stack on MIOpen. Fitted on the FLUX.2 VAE (the vae.flux2 miopen points) with scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled: RX 9070 XT (gfx1201) implied 3453/3369/3368 decode, PRO W7900 (gfx1100) 3525 at 512px; encode 2687/2688/2688 on both cards, to the byte. Two AMD generations agreeing that closely make this MIOpen rather than a per-card quirk, so the values are the larger card's plus ~2% headroom: 3600 / 2750. MIOpen's convolution workspaces are larger than cuDNN's (~1.6x decode, ~2.4x encode), identically with fused and math attention, and its encode/decode ratio is 0.78 against cuDNN's 0.49." + } + ], + "points": [ + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 1142461301, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 2576980378, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 4540854174, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 561566974, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 1254130450, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 2224793059, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1280, + "width": 1280, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 7075958620, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, bf16, untiled, fused SDPA: the fused column of the max-versus-sum measurement (#9519). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 1138753536, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 2531262464, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 4466933760, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1328, + "width": 1328, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 7493124096, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1536, + "width": 1536, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 10043260928, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 1328, + "width": 1328, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 3168796672, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 285212672 + }, + "measured_bytes": 561566974, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 1443889152 + }, + "measured_bytes": 1933809025, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 4372276707, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 3787087413, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1280, + "width": 1280, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 11141120000 + }, + "measured_bytes": 6360846565, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1536, + "width": 1536, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 23102226432 + }, + "measured_bytes": 12847320924, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 285212672 + }, + "measured_bytes": 1810328715, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 1443889152 + }, + "measured_bytes": 3973918491, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 7063073718, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 285212672 + }, + "measured_bytes": 1847909679, + "method": "reserved_peak", + "device": "PRO W7900 (gfx1100)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals.", + "max_ratio": 1.15 + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 1443889152 + }, + "measured_bytes": 4059817837, + "method": "reserved_peak", + "device": "PRO W7900 (gfx1100)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals.", + "max_ratio": 1.15 + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 7197291446, + "method": "reserved_peak", + "device": "PRO W7900 (gfx1100)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals.", + "max_ratio": 1.15 + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "encode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 285212672 + }, + "measured_bytes": 1408749273, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201) and PRO W7900 (gfx1100), identical to the byte", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "encode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 1443889152 + }, + "measured_bytes": 3170759606, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201) and PRO W7900 (gfx1100), identical to the byte", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "encode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 5637144576, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201) and PRO W7900 (gfx1100), identical to the byte", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 2, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 8541616210, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "#9519: 1024px decode by latent batch, CUDA, bf16, fused SDPA, one fresh process per point: 4.23 / 7.96 / 11.89 GiB at batch 1 / 2 / 3. The test table wrote them as GiB x 1000 (4229, 7955, 11889); batch 1 is the 4.229 GiB 1024px decode above." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 3, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 12765716546, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "#9519: 1024px decode by latent batch, CUDA, bf16, fused SDPA, one fresh process per point: 4.23 / 7.96 / 11.89 GiB at batch 1 / 2 / 3. The test table wrote them as GiB x 1000 (4229, 7955, 11889); batch 1 is the 4.229 GiB 1024px decode above." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 2024, + "width": 2024, + "tile": 512, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 585105408, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "#9519: reference-image encode at the 2024px reference cap with the forced 512px tile, CUDA, bf16, fused SDPA; flat at ~0.55 GB from 1024px up. Recorded in MiB." + }, + { + "estimator": "vae.flux", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 3072, + "width": 3072, + "element_size": 2, + "tile": null, + "score_matrix_bytes": 0 + }, + "measured_bytes": 20229295964, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "Commit 326bdf1bd0: the FLUX.1 VAE as a diffusers AutoencoderKL, bf16, fused SDPA, untiled 3072px encode: 18.84 GiB. Re-measured 18.57 GiB on torch 2.13+cu130, 2026-10-17, with scripts/calibrate_flux_vae_working_memory.py." + } + ] +} diff --git a/invokeai/backend/working_memory/estimate.py b/invokeai/backend/working_memory/estimate.py new file mode 100644 index 00000000000..1f388687557 --- /dev/null +++ b/invokeai/backend/working_memory/estimate.py @@ -0,0 +1,52 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field + +FeatureValue = int | float | str | bool | None +"""One runtime fact an estimator is evaluated on: a size, an element width, a tile, a byte count it measured.""" + + +@dataclass(frozen=True) +class WorkingMemoryTerm: + """One part of a working-memory estimate. + + A fitted term is proportional to a data coefficient; `fitted` builds it, so its `bytes` are exactly the + coefficient's value times `scale`. An analytic term (a score matrix, an image buffer) carries neither. `included` is + False for a term the estimate priced but does not reserve, such as the smaller of two phases that peak apart. + """ + + name: str + bytes: float + included: bool + coefficient: str | None = None + scale: float | None = None + + @classmethod + def fitted(cls, name: str, coefficient: str, scale: float, values: Mapping[str, float]) -> "WorkingMemoryTerm": + """A term of `scale` units at `values[coefficient]` bytes per unit.""" + return cls(name, values[coefficient] * scale, True, coefficient, scale) + + +@dataclass(frozen=True) +class WorkingMemoryEstimate: + """What an operation is expected to need on its compute device, and the terms that add up to it.""" + + estimator: str + """The id of the `EstimatorSpec` that produced it, e.g. ``vae.flux2``.""" + version: int + features: Mapping[str, FeatureValue] + """Everything the estimate was evaluated from, runtime decisions such as the tile and the score-matrix size + included, so it can be evaluated again from these alone.""" + terms: tuple[WorkingMemoryTerm, ...] + source: str + """Where the coefficients came from, e.g. ``shipped:conv_backend=miopen (values of vae.flux2)``.""" + reservation_bytes: int = field(init=False) + """What the model cache keeps free: the included terms' bytes, truncated. Derived from the terms, never given.""" + + def __post_init__(self) -> None: + # One term at a time in term order rather than `sum()`, which compensates float rounding since Python 3.12, + # so the reservation is exactly the running total the terms describe. + total: float = 0 + for term in self.terms: + if term.included: + total += term.bytes + object.__setattr__(self, "reservation_bytes", int(total)) diff --git a/invokeai/backend/working_memory/platform.py b/invokeai/backend/working_memory/platform.py new file mode 100644 index 00000000000..5dfc2a00997 --- /dev/null +++ b/invokeai/backend/working_memory/platform.py @@ -0,0 +1,45 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal + +import torch + +ConvBackend = Literal["cudnn", "miopen", "mps", "xpu", "cpu"] + + +@dataclass(frozen=True) +class Platform: + """The traits of the machine an operation runs on, as coefficients are keyed. Each field is a platform key an + `EstimatorSpec` can name, and holds the true value: what was measured, and what was not, is `UNMEASURED_PRICED_AS`'s + business.""" + + conv_backend: ConvBackend + + +UNMEASURED_PRICED_AS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( + { + # MPS, XPU and CPU convolutions are unmeasured and take the cuDNN coefficients. On MPS that pairs with a + # score-matrix term the SDPA probe always charges there, so the total is not thin. + "conv_backend": MappingProxyType({"mps": "cudnn", "xpu": "cudnn", "cpu": "cudnn"}), + } +) +"""Per platform key, the values no coefficients are measured for, and the measured value each is priced as. The one +place that borrowing is declared: the resolver applies it and names it in an estimate's `source`, and shipped data +carries entries for measured values only.""" + + +def platform_for(device: torch.device) -> Platform: + """The platform an operation on `device` runs on. + + A HIP build reports ``device.type == "cuda"``, so the torch build -- not the device string -- is what separates + MIOpen from cuDNN, and a `cpu_only` model on a ROCm build is not on MIOpen. Any device type other than CUDA, MPS and + XPU runs torch's CPU kernels, or none at all, and is priced as the CPU. + """ + if device.type == "cuda": + return Platform(conv_backend="miopen" if torch.version.hip is not None else "cudnn") + if device.type == "mps": + return Platform(conv_backend="mps") + if device.type == "xpu": + return Platform(conv_backend="xpu") + return Platform(conv_backend="cpu") diff --git a/invokeai/backend/working_memory/spec.py b/invokeai/backend/working_memory/spec.py new file mode 100644 index 00000000000..3f0c116a2eb --- /dev/null +++ b/invokeai/backend/working_memory/spec.py @@ -0,0 +1,99 @@ +import importlib +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Final + +import torch + +from invokeai.backend.working_memory.coefficients import ResolvedCoefficients, shipped_coefficients +from invokeai.backend.working_memory.estimate import FeatureValue, WorkingMemoryEstimate, WorkingMemoryTerm +from invokeai.backend.working_memory.platform import platform_for + +TermsFunction = Callable[[Mapping[str, FeatureValue], Mapping[str, float]], tuple[WorkingMemoryTerm, ...]] +"""An estimator's pure part: its features and its coefficient values in, its terms out.""" + +ESTIMATOR_MODULES: Final = ("invokeai.backend.util.vae_working_memory",) +"""Every module that registers an `EstimatorSpec`. The registry holds only the specs of modules imported so far, so a +caller that needs all of them imports these (`import_estimator_modules`); `tests/backend/working_memory/` fails on a +module that declares a spec without being listed here.""" + + +@dataclass(frozen=True) +class EstimatorSpec: + """What a working-memory estimator declares: its id and version, the coefficients its fitted terms read, and the + platform keys those coefficients are shipped per. + + Registered at import of the module that defines the estimator. `terms` is that estimator's pure part, so an + estimate can be evaluated again from recorded features alone -- which is how the measured points under `data/` are + checked. It is behaviour rather than declaration and stays out of equality and hashing. + """ + + id: str + version: int + """Bumped whenever `terms` changes the formula, so coefficients fitted to the old one no longer resolve.""" + coefficients: frozenset[str] + platform_keys: tuple[str, ...] + terms: TermsFunction = field(compare=False, repr=False) + + def estimate(self, features: Mapping[str, FeatureValue], device: torch.device) -> WorkingMemoryEstimate: + """The estimate for an operation on `device`, priced with the shipped coefficients of the platform it runs on.""" + return self.evaluate(features, shipped_coefficients(self, platform_for(device))) + + def evaluate( + self, features: Mapping[str, FeatureValue], coefficients: ResolvedCoefficients + ) -> WorkingMemoryEstimate: + """The estimate for `features` with the given coefficients.""" + return WorkingMemoryEstimate( + estimator=self.id, + version=self.version, + features=features, + terms=self.terms(features, coefficients.values), + source=coefficients.source, + ) + + +_SPECS: Final[dict[str, EstimatorSpec]] = {} + + +def _defined_at(terms: TermsFunction) -> str: + return f"{terms.__module__}.{getattr(terms, '__qualname__', repr(terms))}" + + +def register(spec: EstimatorSpec) -> EstimatorSpec: + """Declare an estimator, once per id, from the module that defines it. + + `importlib.reload` of that module declares the same spec again, from a new `terms` function of the same module and + name: a no-op that adopts the re-executed function. Anything else under an id already taken raises -- differing + declared data, the same data from another function, or the very same function object again (two specs of one module + sharing a `terms` function and, by a copy-paste slip, an id). Each would otherwise silently replace one estimator's + formula or coefficients with another's. + + jurigged (`--dev_reload`) patches function bodies in place and re-runs only the top-level statements that changed: + an edited `terms` body takes effect with nothing re-registered, but an edited `register(...)` statement raises even + when its data is unchanged (a reformat, a comment), because it hands over the same function object. Restart the + server then. + """ + declared = _SPECS.get(spec.id) + if declared is not None: + if declared != spec: + raise ValueError(f"Working-memory estimator '{spec.id}' is already registered as {declared}.") + if declared.terms is spec.terms or _defined_at(declared.terms) != _defined_at(spec.terms): + raise ValueError( + f"Working-memory estimator '{spec.id}' is already registered by {_defined_at(declared.terms)}; " + "ids are unique per estimator." + ) + _SPECS[spec.id] = spec + return spec + + +def registered_specs() -> Mapping[str, EstimatorSpec]: + """Every spec registered so far, by id: those of the estimator modules imported in this process.""" + return MappingProxyType(_SPECS) + + +def import_estimator_modules() -> Mapping[str, EstimatorSpec]: + """Every spec there is: imports `ESTIMATOR_MODULES` first, so the answer does not depend on import order.""" + for name in ESTIMATOR_MODULES: + importlib.import_module(name) + return registered_specs() diff --git a/pyproject.toml b/pyproject.toml index 786a71e56d0..e57bc3dd037 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -548,6 +548,8 @@ version = { attr = "invokeai.version.__version__" } # absent from a wheel install and that loader raises FileNotFoundError. "invokeai.backend.minimax_h3" = ["*.json"] "invokeai.backend.t5" = ["tokenizer/*.json"] +# Fitted working-memory coefficients, read by every estimator that declares an EstimatorSpec. +"invokeai.backend.working_memory" = ["data/*.json"] "invokeai.backend.clip" = ["*.json", "tokenizer/*.json", "tokenizer/*.gz"] "invokeai.app.services.image_index" = ["*.txt"] "invokeai.app.services.workflow_records.default_workflows" = ["*.json"] diff --git a/scripts/calibrate_flux2_working_memory.py b/scripts/calibrate_flux2_working_memory.py index d7d6ca34349..c429c5b9a49 100644 --- a/scripts/calibrate_flux2_working_memory.py +++ b/scripts/calibrate_flux2_working_memory.py @@ -7,8 +7,8 @@ 1. ``SDPA_MATH_BYTES_PER_SCORE_ELEMENT`` (``backend/util/attention.py``) -- bytes per element of the score matrix, charged only where SDPA has no fused kernel and materializes it. -2. ``estimate_vae_working_memory_flux2``'s 2200 / 1100 bytes per pixel per element byte - (``backend/util/vae_working_memory.py``). +2. ``estimate_vae_working_memory_flux2``'s ``k_decode`` / ``k_encode`` bytes per pixel per element byte, + shipped per convolution backend in ``backend/working_memory/data/vae.json``. 3. ``Flux2DenoiseInvocation._estimate_working_memory``'s 0.4 MB per token, defined at the Klein 9B width (4096) and scaled linearly by the loaded variant's width. 4. The dispatch question underneath all of it: does this build's SDPA fuse these shapes, or does it @@ -67,14 +67,19 @@ from invokeai.backend.model_manager.configs.flux2_variant import flux2_hidden_size from invokeai.backend.model_manager.taxonomy import Flux2VariantType from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT, sdpa_score_matrix_bytes -from invokeai.backend.util.vae_working_memory import ( - _flux_vae_scaling_constant, - estimate_vae_working_memory_flux2, -) +from invokeai.backend.util.vae_working_memory import VAE_FLUX2_SPEC, estimate_vae_working_memory_flux2 +from invokeai.backend.working_memory.coefficients import shipped_coefficients +from invokeai.backend.working_memory.platform import platform_for GIB = 1024**3 MIB = 1024**2 + +def _shipped_vae_coefficient(operation: str) -> float: + """The VAE coefficient this build's estimate uses for `operation`: the column its convolution backend selects.""" + return shipped_coefficients(VAE_FLUX2_SPEC, platform_for(torch.device("cuda"))).values[f"k_{operation}"] + + # The FLUX.2 VAE attends on the 8x-downsampled grid with a single 512-wide head. Mirrors # `_FLUX2_VAE_*` in vae_working_memory.py. VAE_SPATIAL_COMPRESSION = 8 @@ -264,7 +269,9 @@ def run() -> None: if peak is None: return row - estimate = estimate_vae_working_memory_flux2(operation=operation, image_tensor=x, vae=vae, device=device) + estimate = estimate_vae_working_memory_flux2( + operation=operation, image_tensor=x, vae=vae, device=device + ).reservation_bytes seq_len = (px // VAE_SPATIAL_COMPRESSION) ** 2 score_bytes = sdpa_score_matrix_bytes( device=device, dtype=param.dtype, num_heads=1, head_dim=VAE_MID_BLOCK_HEAD_DIM, seq_len=seq_len @@ -278,8 +285,8 @@ def run() -> None: "reserved_delta": peak, "estimate": estimate, "score_term": score_bytes, - # The whole measured peak over the pixel area, which is what the 2200/1100-style literals - # name. The score matrix is deliberately NOT backed out: the two terms peak in different + # The whole measured peak over the pixel area, which is what the `k_decode`/`k_encode` + # coefficients name. The score matrix is deliberately NOT backed out: the two terms peak in different # phases of the forward, so the estimate takes the larger of them rather than the sum, and # the convolution phase is what sets the high-water mark until the quadratic term overtakes # it past ~1280px. Where `score_term` exceeds the measured peak, this column is not the @@ -457,10 +464,10 @@ def report_sdpa(shapes: list[tuple[int, int, int]], dtype_name: str) -> list[dic def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[dict]: - decode_k = _flux_vae_scaling_constant("decode", torch.device("cuda")) - encode_k = _flux_vae_scaling_constant("encode", torch.device("cuda")) - print(f"\n=== 3. VAE linear constants (this build selects {decode_k} decode / {encode_k} encode) ===") - print("`implied_k` is the measured peak over pixel area, directly comparable to those literals;") + decode_k = _shipped_vae_coefficient("decode") + encode_k = _shipped_vae_coefficient("encode") + print(f"\n=== 3. VAE linear coefficients (this build selects {decode_k:.0f} decode / {encode_k:.0f} encode) ===") + print("`implied_k` is the measured peak over pixel area, directly comparable to those coefficients;") print("fit on the rows whose `math` column matches what this build really does (section 1).") print("`covered` is the question that matters: is the shipped estimate an upper bound here?") print("Caveat: forcing math on a build that HAS a fused kernel is not equivalent to a build that") @@ -495,7 +502,7 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di _flag_non_monotonic(rows) for operation in ("decode", "encode"): # Compare against the column this build actually selects, not against a hard-coded number. - shipped = _flux_vae_scaling_constant(operation, torch.device("cuda")) + shipped = _shipped_vae_coefficient(operation) for force_math in (False, True): ks = [ r["implied_linear_constant"] @@ -511,7 +518,7 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di continue mode = "math" if force_math else "fused" verdict = "OK" if shipped >= max(ks) else "SHORT" - print(f" {operation} ({mode}): implied_k max = {max(ks):.0f}, shipped = {shipped} -> {verdict}") + print(f" {operation} ({mode}): implied_k max = {max(ks):.0f}, shipped = {shipped:.0f} -> {verdict}") short = [r for r in rows if not r.get("oom") and not r["covered"]] if short: print("\n Points the shipped estimate does NOT cover:") diff --git a/scripts/calibrate_flux_vae_working_memory.py b/scripts/calibrate_flux_vae_working_memory.py index c737005939e..9c62db1c2ef 100644 --- a/scripts/calibrate_flux_vae_working_memory.py +++ b/scripts/calibrate_flux_vae_working_memory.py @@ -93,7 +93,7 @@ def _worker(vae_path: str, operation: str, h: int, w: int, dtype_name: str, tile else: x = torch.randn(1, 3, h, w, device="cuda", dtype=dtype) - estimate = estimate_vae_working_memory_flux(operation, x, vae, tile_size=tile) + estimate = estimate_vae_working_memory_flux(operation, x, vae, tile_size=tile).reservation_bytes torch.cuda.synchronize() torch.cuda.empty_cache() @@ -218,8 +218,8 @@ def main() -> None: print(f" {r['operation']} {r['h']}x{r['w']} tile={r['tile']} headroom {r['headroom']:.2f}x") print() print(" The cache under-reserves here, so it keeps other models resident and the VAE can") - print(" OOM in the last step of a generation. The constants in `vae_working_memory.py` need") - print(" a backend branch, the way `estimate_vae_working_memory_qwen_image` already has one.") + print(" OOM in the last step of a generation. The vae.flux coefficients for this convolution") + print(" backend in `invokeai/backend/working_memory/data/vae.json` need refitting.") else: worst = min(measured, key=lambda r: r["headroom"]) if measured else None if worst: diff --git a/tests/app/invocations/conftest.py b/tests/app/invocations/conftest.py index 6aed50a4190..0151a2d484f 100644 --- a/tests/app/invocations/conftest.py +++ b/tests/app/invocations/conftest.py @@ -5,6 +5,8 @@ import pytest from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate, WorkingMemoryTerm + def _build_flux_shaped_vae() -> AutoencoderKL: """A real `AutoencoderKL` with the FLUX.1 VAE's *shape* and almost none of its weight. @@ -47,3 +49,19 @@ def _build_flux_shaped_vae() -> AutoencoderKL: def flux_shaped_vae() -> Callable[[], AutoencoderKL]: """Factory, not an instance: a test that needs two independent VAEs should not share one.""" return _build_flux_shaped_vae + + +def _fixed_working_memory(reservation_bytes: int) -> WorkingMemoryEstimate: + return WorkingMemoryEstimate( + estimator="test.fixed", + version=1, + features={}, + terms=(WorkingMemoryTerm("fixed", reservation_bytes, included=True),), + source="test", + ) + + +@pytest.fixture +def fixed_working_memory() -> Callable[[int], WorkingMemoryEstimate]: + """An estimate that reserves exactly the given bytes, for a node whose estimator a test patches out.""" + return _fixed_working_memory diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 286b6b44511..6eeb2180bb1 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -7,7 +7,10 @@ a 1024x1024 generation. See https://github.com/invoke-ai/InvokeAI/issues/9500. The `MEASURED_*` tables below are peak *reserved* memory measured on CUDA in bf16 (the conservative -quantity, including allocator overhead). Every estimate must stay an upper bound on them. +quantity, including allocator overhead). Every estimate must stay an upper bound on them. The VAE's +measured points ship with its coefficients in `invokeai/backend/working_memory/data/vae.json`, where +`tests/backend/working_memory/test_shipped_coefficients.py` holds the estimate to at least each point +and, where a fitted coefficient decides it, at most 1.3x (or the point's `max_ratio`). """ from contextlib import contextmanager @@ -35,14 +38,19 @@ rocm_sdpa_uses_math_kernel, sdpa_score_matrix_bytes, ) -from invokeai.backend.util.vae_working_memory import ( - _FLUX_VAE_SCALING_CONSTANTS, - estimate_vae_working_memory_flux2, -) +from invokeai.backend.util.vae_working_memory import VAE_FLUX2_SPEC, estimate_vae_working_memory_flux2 +from invokeai.backend.working_memory.coefficients import shipped_coefficients +from invokeai.backend.working_memory.platform import Platform MB = 1024**2 GB = 1024**3 + +def _vae_k(conv_backend, operation): + """The shipped FLUX.2 VAE coefficient the estimator resolves for one convolution backend and operation.""" + return shipped_coefficients(VAE_FLUX2_SPEC, Platform(conv_backend)).values[f"k_{operation}"] + + # The estimator's default width, and the head count that follows from it. The suite's pinned # `MEASURED_DENOISE` table was taken on Klein 9B, which is that default, so those rows are unmoved # by the width scaling. @@ -272,18 +280,6 @@ def test_the_lora_margin_scales_with_the_batch(self): class TestFlux2VaeWorkingMemoryEstimate: - # (operation, pixel size, measured peak reserved MB), bf16, untiled. - MEASURED_VAE = [ - ("decode", 512, 1086), - ("decode", 768, 2414), - ("decode", 1024, 4260), - ("decode", 1328, 7146), - ("decode", 1536, 9578), - ("encode", 512, 536), - ("encode", 1024, 2122), - ("encode", 1328, 3022), - ] - def _mock_bf16_vae(self): vae = MagicMock(spec=AutoencoderKLFlux2) vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) # element_size == 2 @@ -293,19 +289,12 @@ def _tensor_for(self, operation, px): # decode receives 32-channel latents at pixels/8; encode receives a pixel image. return torch.zeros(1, 32, px // 8, px // 8) if operation == "decode" else torch.zeros(1, 3, px, px) - @pytest.mark.parametrize("operation, px, measured_mb", MEASURED_VAE) - def test_estimate_is_an_upper_bound_on_measured_peak(self, operation, px, measured_mb): - estimate = estimate_vae_working_memory_flux2( - operation=operation, image_tensor=self._tensor_for(operation, px), vae=self._mock_bf16_vae(), device=FUSED - ) - assert estimate >= measured_mb * MB - - @pytest.mark.parametrize("operation, expected_constant", [("decode", 2200), ("encode", 1100)]) - def test_constant_scales_pixel_area_and_element_size(self, operation, expected_constant): + @pytest.mark.parametrize("operation", ["decode", "encode"]) + def test_constant_scales_pixel_area_and_element_size(self, operation): estimate = estimate_vae_working_memory_flux2( operation=operation, image_tensor=self._tensor_for(operation, 1024), vae=self._mock_bf16_vae(), device=FUSED ) - assert estimate == 1024 * 1024 * 2 * expected_constant + assert estimate.reservation_bytes == 1024 * 1024 * 2 * _vae_k("cudnn", operation) def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): """Reference-image encoding forces 512px tiling precisely so the peak stops following the @@ -317,79 +306,26 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): vae=self._mock_bf16_vae(), tile_size=512, device=FUSED, - ) + ).reservation_bytes for px in (1024, 1328, 2024) ] assert len(set(estimates)) == 1 - assert estimates[0] == int(512 * 512 * 2 * 1100 * 1.25) - assert estimates[0] >= 558 * MB # measured tiled peak at 2024px + assert estimates[0] == int(512 * 512 * 2 * _vae_k("cudnn", "encode") * 1.25) # The whole point of tiling: it must shrink the reservation, not just bound the VAE. untiled = estimate_vae_working_memory_flux2( operation="encode", image_tensor=torch.zeros(1, 3, 2024, 2024), vae=self._mock_bf16_vae(), device=FUSED - ) + ).reservation_bytes assert estimates[0] < untiled / 4 class TestVaeConstantsFollowTheConvBackend: - """The pixel-area constant is not one number. MIOpen's convolution workspaces cost far more than - cuDNN's for the same decode, and the gap is not the attention term -- it is identical on the - fused path. Fitted with `scripts/calibrate_flux2_working_memory.py --only vae`, peak reserved, - each point in a fresh process: - - decode encode - cuDNN RTX 4090 2180 2185 2165 1072 1063 1061 - MIOpen RX 9070 XT 3453 3369 3368 2687 2688 2688 - - Flat across 512/768/1024px on both, so the linear model holds; only the coefficient moves. Note - the encode column: cuDNN's is half its decode, MIOpen's is four fifths, so "encoding costs half - of decoding" -- the ratio the other estimators in the module use -- is a cuDNN property, not an - architectural one. + """The pixel-area coefficient is not one number. MIOpen's convolution workspaces cost far more + than cuDNN's for the same decode, and the gap is not the attention term -- it is identical on the + fused path. Both columns, the cards they were fitted on and the points that bound them are in + `invokeai/backend/working_memory/data/vae.json`; what is pinned here is that a device selects the + column of the backend it really runs on. """ - # (backend, operation, px, measured GiB) - MEASURED = [ - ("cudnn", "decode", 512, 1.064), - ("cudnn", "decode", 768, 2.400), - ("cudnn", "decode", 1024, 4.229), - ("cudnn", "encode", 512, 0.523), - ("cudnn", "encode", 768, 1.168), - ("cudnn", "encode", 1024, 2.072), - ("miopen", "decode", 512, 1.686), - ("miopen", "decode", 768, 3.701), - ("miopen", "decode", 1024, 6.578), - # PRO W7900 (gfx1100). Costs slightly more per pixel than gfx1201 and sets the column. - ("miopen", "decode", 512, 1.721), - ("miopen", "decode", 768, 3.781), - ("miopen", "decode", 1024, 6.703), - ("miopen", "encode", 512, 1.312), - ("miopen", "encode", 768, 2.953), - ("miopen", "encode", 1024, 5.250), - ] - - @pytest.mark.parametrize("backend, operation, px, measured_gib", MEASURED) - def test_the_constant_upper_bounds_its_own_backend(self, backend, operation, px, measured_gib): - """Each column has to bound the hardware it was fitted on. Before this the cuDNN column was - used everywhere, leaving a 2.3 GiB shortfall on a 1024px MIOpen decode.""" - constant = _FLUX_VAE_SCALING_CONSTANTS[backend][operation] - estimate = px * px * 2 * constant # bf16 element size - assert estimate >= measured_gib * GB - assert estimate <= 1.3 * measured_gib * GB - - def test_the_cudnn_constant_would_not_cover_miopen(self): - """The regression this guards: a 1024px MIOpen decode needs 6.58 GiB and the cuDNN constant - reserves 4.30. That is the same class of shortfall #9500 reports, one backend over.""" - cudnn = _FLUX_VAE_SCALING_CONSTANTS["cudnn"]["decode"] - assert 1024 * 1024 * 2 * cudnn < 6.578 * GB - - def test_the_encode_ratio_is_not_architectural(self): - """cuDNN's encode is half its decode; MIOpen's is four fifths. A shared ratio cannot express - both, which is why the table carries the two operations separately.""" - ratios = { - backend: consts["encode"] / consts["decode"] for backend, consts in _FLUX_VAE_SCALING_CONSTANTS.items() - } - assert ratios["cudnn"] == pytest.approx(0.50, abs=0.02) - assert ratios["miopen"] == pytest.approx(0.76, abs=0.03) - def test_a_rocm_build_selects_the_miopen_column(self): """A HIP build reports `device.type == "cuda"`, so the torch build is the discriminator.""" vae = MagicMock(spec=AutoencoderKLFlux2) @@ -401,7 +337,7 @@ def estimate(): v.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) return estimate_vae_working_memory_flux2( operation="decode", image_tensor=latents, vae=v, device=torch.device("cuda") - ) + ).reservation_bytes # `_IS_ROCM` is pinned because it is bound at import: on a real HIP build the head-dim # guard short-circuits ahead of the probe patched here and charges a score matrix to both @@ -421,15 +357,15 @@ def estimate(): cuda = estimate() assert rocm > cuda - assert rocm / cuda == pytest.approx(3600 / 2200, rel=1e-3) + assert rocm / cuda == pytest.approx(_vae_k("miopen", "decode") / _vae_k("cudnn", "decode"), rel=1e-3) def test_a_rocm_build_composes_the_miopen_column_with_the_guard_score_matrix(self): """The state every ROCm rig is actually in, driven through the estimator. A HIP build takes the MIOpen column *and* the head-dim guard's score matrix -- the guard is unconditional for this VAE's 512-wide head -- so `max` has to pick between two terms that - are both larger than the cuDNN pair. `TestVaeTermsDoNotAdd` checks the same two terms - against the constants table, which cannot see an estimator that stops combining them: a + are both larger than the cuDNN pair. The shipped points check the same two terms against the + measurements, which cannot see an estimator that stops combining them: a plausible "the score matrix covers it anyway" refactor that drops back to the cuDNN column whenever SDPA materializes under-reserves a 1024px decode by 1.6x and passes every other test in this module. @@ -443,7 +379,7 @@ def estimate(px): image_tensor=torch.zeros(1, 32, px // 8, px // 8), vae=vae, device=torch.device("cuda"), - ) + ).reservation_bytes with ( patch("torch.version.hip", "7.1.25424"), @@ -456,9 +392,8 @@ def estimate(px): at_1024 = estimate(1024) at_1536 = estimate(1536) - miopen = _FLUX_VAE_SCALING_CONSTANTS["miopen"]["decode"] + miopen = _vae_k("miopen", "decode") assert at_1024 == 1024 * 1024 * 2 * miopen, "at 1024px the MIOpen linear term is the larger" - assert at_1024 >= 6.703 * GB, "and it still covers the measured W7900 peak it was fitted to" assert at_1536 == 36864 * 36864 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT, "at 1536px the score matrix is" assert at_1536 > 1536 * 1536 * 2 * miopen @@ -470,13 +405,10 @@ class TestFlux2VaeBatchIsBudgeted: the OOM happens anyway. Measured at 1024px on CUDA/bf16, peak reserved, each point in a fresh process: 4.23GB at batch 1, - 7.96GB at batch 2, 11.89GB at batch 3. Linear, and slightly sub-linear per sample, so scaling the - single-sample estimate is an upper bound rather than a fit. + 7.96GB at batch 2, 11.89GB at batch 3 (shipped as `vae.flux2` points). Linear, and slightly + sub-linear per sample, so scaling the single-sample estimate is an upper bound rather than a fit. """ - # (batch, measured peak reserved MB) for a 1024px decode. - MEASURED_DECODE_BATCH = [(1, 4229), (2, 7955), (3, 11889)] - def _decode_estimate(self, batch, px=1024, tile_size=None, device=FUSED): vae = MagicMock(spec=AutoencoderKLFlux2) vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) @@ -486,13 +418,7 @@ def _decode_estimate(self, batch, px=1024, tile_size=None, device=FUSED): vae=vae, tile_size=tile_size, device=device, - ) - - @pytest.mark.parametrize("batch, measured_mb", MEASURED_DECODE_BATCH) - def test_estimate_is_an_upper_bound_on_the_measured_batch_peak(self, batch, measured_mb): - estimate = self._decode_estimate(batch) - assert estimate >= measured_mb * MB - assert estimate <= 2 * measured_mb * MB + ).reservation_bytes def test_estimate_scales_with_the_batch(self): """The regression in one assertion: before this, all three of these were equal.""" @@ -506,7 +432,7 @@ def test_a_three_dimensional_tensor_is_one_sample(self): vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) unbatched = estimate_vae_working_memory_flux2( operation="decode", image_tensor=torch.zeros(32, 128, 128), vae=vae, device=FUSED - ) + ).reservation_bytes assert unbatched == self._decode_estimate(1) def test_tiling_bounds_the_tile_not_the_batch(self): @@ -549,13 +475,13 @@ def _mock_vae_info(self): vae_info.model_on_device = MagicMock(return_value=cm) return vae_info - def test_decode_requests_working_memory(self): + def test_decode_requests_working_memory(self, fixed_working_memory): vae_info = self._mock_vae_info() context = MagicMock() context.models.load.return_value = vae_info context.tensors.load.return_value = torch.zeros(1, 32, 128, 128) - expected = 10 * GB + expected = fixed_working_memory(10 * GB) with patch( "invokeai.app.invocations.vae.flux2_vae_decode.estimate_vae_working_memory_flux2", return_value=expected ) as estimate: @@ -571,14 +497,14 @@ def test_decode_requests_working_memory(self): estimate.assert_called_once() assert estimate.call_args.kwargs["operation"] == "decode" - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected) + vae_info.model_on_device.assert_called_once_with(working_memory=expected) - def test_encode_requests_working_memory(self): + def test_encode_requests_working_memory(self, fixed_working_memory): vae_info = self._mock_vae_info() context = MagicMock() context.models.load.return_value = vae_info - expected = 4 * GB + expected = fixed_working_memory(4 * GB) with ( patch( "invokeai.app.invocations.vae.flux2_vae_encode.estimate_vae_working_memory_flux2", return_value=expected @@ -598,7 +524,7 @@ def test_encode_requests_working_memory(self): estimate.assert_called_once() assert estimate.call_args.kwargs["operation"] == "encode" - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected) + vae_info.model_on_device.assert_called_once_with(working_memory=expected) class _StopBeforeLoad(Exception): @@ -782,9 +708,9 @@ def _pinned_build(): score matrix here", but on a HIP build that same token carries two more decisions with it, and both would make this module's verdict depend on the runner: - * `_flux2_vae_scaling_constant` selects the MIOpen convolution column, 1.6-2.5x the cuDNN one - every expectation guarded by this helper is written against -- so the VAE assertions would - pass on a CUDA runner and fail on a ROCm one. + * `platform_for` selects the MIOpen convolution coefficients, 1.6-2.5x the cuDNN ones every + expectation guarded by this helper is written against -- so the VAE assertions would pass on + a CUDA runner and fail on a ROCm one. * `rocm_sdpa_uses_math_kernel` short-circuits ahead of the probe below, on a threshold a user can move with `INVOKE_ROCM_FUSED_SDPA_MAX_HEAD_DIM`. At the default it happens to agree with `_materializing_probe`; at a lower one it charges a score matrix for the denoise's 128-wide @@ -845,7 +771,7 @@ def _vae_estimate(self, operation, px, device, tile_size=None): vae=self._mock_bf16_vae(), tile_size=tile_size, device=device, - ) + ).reservation_bytes # (operation, pixel size, mid-block tokens). The VAE attends on the 8x-downsampled grid. @pytest.mark.parametrize( @@ -907,20 +833,17 @@ class TestVaeTermsDoNotAdd: decoder, last in the encoder -- so the full-resolution convolution feature maps that drive the linear term are not live while the score matrix is, and peak *reserved* is a high-water mark. - Measured, decode, peak reserved, fresh process per point: - - px cuDNN fused cuDNN forced-math linear(2200) score(17) - 1024 4.229 3.527 4.30 4.25 - 1280 6.590 5.924 6.71 10.38 - 1536 9.353 11.965 9.67 21.52 - - Forcing math measures *below* the fused path until 1536px, and even there it exceeds it by - 2.6GB against the 21.5GB the score term prices standalone -- the attention phase reuses blocks - the allocator is already holding. On gfx1100 and gfx1201, forcing math moves the measured peak - by nothing at all up to 1024px, and the totals stay flat-linear in area either way. + Measured, decode, peak reserved, fresh process per point (the cuDNN fused and forced-math points + in `invokeai/backend/working_memory/data/vae.json`): forcing math measures *below* the fused path + until 1536px, and even there it exceeds it by 2.6GB against the 21.5GB the score term prices + standalone -- the attention phase reuses blocks the allocator is already holding. On gfx1100 and + gfx1201, forcing math moves the measured peak by nothing at all up to 1024px, and the totals stay + flat-linear in area either way. Summing the terms reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7GB. Taking the - max reserves 6.8GB. + max reserves 6.8GB. Where a max model is not an upper bound on its own -- the crossover, a 768px + encode with cuDNN's coefficient and a materializing kernel -- the `device_working_mem_gb` floor + is, and the shipped point there is checked against the floor. """ def _estimate(self, px, device, materializing): @@ -929,64 +852,20 @@ def _estimate(self, px, device, materializing): with _materializing() if materializing else _pinned_build(): return estimate_vae_working_memory_flux2( operation="decode", image_tensor=torch.zeros(1, 32, px // 8, px // 8), vae=vae, device=device - ) - - # (pixel size, measured GiB) for a decode on the PRO W7900. Its 512px point is what sets the - # MIOpen decode constant: implied 3525, where gfx1201 asks for only 3453. - MEASURED_W7900_DECODE = [(1024, 6.703), (768, 3.781), (512, 1.721)] - - @pytest.mark.parametrize("px, measured_gib", MEASURED_W7900_DECODE) - def test_the_max_model_bounds_the_measured_miopen_peak(self, px, measured_gib): - linear = px * px * 2 * _FLUX_VAE_SCALING_CONSTANTS["miopen"]["decode"] - score = (px // 8) ** 4 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - assert max(linear, score) >= measured_gib * GB - assert max(linear, score) <= 1.15 * measured_gib * GB - - def test_the_sum_model_would_have_over_reserved_by_two_thirds(self): - """At 1024px the sum reserves 11.1GiB for a decode that measures 6.7 -- on a 16GB card that - is the difference between the transformer staying resident and being evicted.""" - linear = 1024 * 1024 * 2 * _FLUX_VAE_SCALING_CONSTANTS["miopen"]["decode"] - score = 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - assert (linear + score) / (6.703 * GB) > 1.6 - assert max(linear, score) / (6.703 * GB) < 1.15 + ).reservation_bytes def test_the_estimate_is_the_larger_term_not_the_sum(self): """At 1024px the two are within 2% of each other on the cuDNN column, which makes this the sharpest place to tell the models apart.""" - linear = 1024 * 1024 * 2 * _FLUX_VAE_SCALING_CONSTANTS["cudnn"]["decode"] + linear = 1024 * 1024 * 2 * _vae_k("cudnn", "decode") score = 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT assert self._estimate(1024, MATERIALIZING, materializing=True) == max(linear, score) assert self._estimate(1024, MATERIALIZING, materializing=True) < linear + score - # (operation, px, measured GiB) with cuDNN's linear constant and a materializing kernel -- the - # combination MPS lands in, and the only place the max model is not an upper bound on its own. - MEASURED_CUDNN_FORCED_MATH = [ - ("encode", 512, 0.523), - ("encode", 768, 1.801), - ("encode", 1024, 4.072), - ("decode", 1024, 3.527), - ("decode", 1280, 5.924), - ("decode", 1536, 11.965), - ] - - @pytest.mark.parametrize("operation, px, measured_gib", MEASURED_CUDNN_FORCED_MATH) - def test_the_working_memory_floor_covers_the_crossover(self, operation, px, measured_gib): - """A max model is weakest where the two terms are near-equal, and one measured point shows - it: a 768px encode wants 1.80GiB against a 1.35GiB max. It is not reachable as a shortfall, - because the cache floors every reservation at `device_working_mem_gb` and the whole crossover - region sits below that floor. Reproducible to three decimals across runs, so this is the - model's real shape, not noise -- which is why it is pinned rather than rounded away.""" - constant = _FLUX_VAE_SCALING_CONSTANTS["cudnn"][operation] - area = px * px if operation == "decode" else px * px - linear = area * 2 * constant - score = (px // 8) ** 4 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - floored = max(max(linear, score), 3 * GB) # what ModelCache._get_vram_available reserves - assert floored >= measured_gib * GB - def test_a_fused_build_is_unaffected(self): """The max only ever removes reservation, never adds it: with no score matrix the estimate is the linear term, exactly as before.""" - linear = 1024 * 1024 * 2 * _FLUX_VAE_SCALING_CONSTANTS["cudnn"]["decode"] + linear = 1024 * 1024 * 2 * _vae_k("cudnn", "decode") assert self._estimate(1024, FUSED, materializing=False) == linear def test_regional_prompting_adds_the_score_matrix(self): @@ -1283,7 +1162,7 @@ def estimate(): v.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) return estimate_vae_working_memory_flux2( operation="decode", image_tensor=torch.zeros(1, 32, 128, 128), vae=v, device=FUSED - ) + ).reservation_bytes with _diffusers_backend("_native_math"): forced_math = estimate() diff --git a/tests/app/invocations/test_flux_vae_decode.py b/tests/app/invocations/test_flux_vae_decode.py index 4e206cf62b2..8b93e064d3e 100644 --- a/tests/app/invocations/test_flux_vae_decode.py +++ b/tests/app/invocations/test_flux_vae_decode.py @@ -32,7 +32,7 @@ def _loaded_vae(shift_factor: float | None, scaling_factor: float = 0.3611) -> M vae_info.compute_device = torch.device("cpu") @contextmanager - def _on_device(working_mem_bytes=None): + def _on_device(working_memory=None): yield (None, vae) vae_info.model_on_device = _on_device diff --git a/tests/app/invocations/test_flux_vae_decode_tiling.py b/tests/app/invocations/test_flux_vae_decode_tiling.py index 2b6f998067e..1fca3e20687 100644 --- a/tests/app/invocations/test_flux_vae_decode_tiling.py +++ b/tests/app/invocations/test_flux_vae_decode_tiling.py @@ -107,21 +107,25 @@ def test_the_tiling_state_is_restored_afterwards(self, flux_shaped_vae): [(False, False, None), (True, False, 256), (False, True, 0)], ) def test_the_reservation_matches_the_decode_that_will_run( - self, flux_shaped_vae, tiled, force_tiled_decode, expected_tile_size + self, flux_shaped_vae, fixed_working_memory, tiled, force_tiled_decode, expected_tile_size ): """A tiled decode reserved for a single pass is the OOM the retry exists to avoid, and an untiled decode reserved for a tile evicts models it did not need to.""" path = "invokeai.app.invocations.vae.flux_vae_decode.estimate_vae_working_memory_flux" vae = flux_shaped_vae() - _, context, _ = _build_decode_mocks( + vae_info, context, _ = _build_decode_mocks( vae, torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512), force_tiled_decode=force_tiled_decode ) - with patch(path, return_value=1024) as estimate: + estimated = fixed_working_memory(1024) + with patch(path, return_value=estimated) as estimate: _build_invocation(tiled=tiled, tile_size=256 if tiled else 0).invoke(context) assert estimate.call_args.kwargs["tile_size"] == expected_tile_size + vae_info.model_on_device.assert_called_once_with(working_memory=estimated) @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) - def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(self, auto, flux_shaped_vae): + def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off( + self, auto, flux_shaped_vae, fixed_working_memory + ): """On Windows an allocation that does not fit is paged rather than refused, so the OOM retry never fires and the decision has to be made from the estimate alone.""" module = "invokeai.app.invocations.vae.flux_vae_decode" @@ -129,8 +133,9 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512) ) context.config.get.return_value.auto_tiled_decode = auto + untiled, tiled = fixed_working_memory(20 * 2**30), fixed_working_memory(2 * 2**30) with ( - patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[20 * 2**30, 2 * 2**30]) as estimate, + patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[untiled, tiled]) as estimate, patch(f"{module}.should_pretile_vae_decode", return_value=True) as pretile, ): _build_invocation().invoke(context) @@ -138,11 +143,12 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se if auto: pretile.assert_called_once_with(vae_info.compute_device, 20 * 2**30) assert estimate.call_args.kwargs["tile_size"] == 0 - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + vae_info.model_on_device.assert_called_once_with(working_memory=tiled) assert tiling_state[0]["use_tiling"] is True assert tiling_state[0]["tile_sample_min_size"] == DEFAULT_TILE_SAMPLE_MIN_SIZE else: pretile.assert_not_called() + vae_info.model_on_device.assert_called_once_with(working_memory=untiled) assert tiling_state[0]["use_tiling"] is False diff --git a/tests/app/invocations/test_vae_working_memory_backend.py b/tests/app/invocations/test_vae_working_memory_backend.py index 5829d061afe..9dfc4fade93 100644 --- a/tests/app/invocations/test_vae_working_memory_backend.py +++ b/tests/app/invocations/test_vae_working_memory_backend.py @@ -1,7 +1,7 @@ -"""Which convolution-backend constant the diffusers-AutoencoderKL estimators use. +"""Which convolution-backend coefficients the diffusers-AutoencoderKL estimators use, and what the SD estimator adds up. SD1/SDXL, SD3 and CogView4 run the same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so they take the -same measured per-backend constants: MIOpen needs ~1.64x what cuDNN does. The change is large enough to matter -- an +same measured per-backend coefficients: MIOpen needs ~1.64x what cuDNN does. The change is large enough to matter -- an SD1.5 1024px fp32 decode reserves 9.49 GB on cuDNN and 15.36 GB on MIOpen -- and the branch is invisible on a CUDA machine, so it is pinned here rather than left to whichever card the suite runs on. """ @@ -13,10 +13,15 @@ from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from invokeai.backend.util.vae_working_memory import ( + VAE_COGVIEW4_SPEC, + VAE_SD3_SPEC, + VAE_SD_SPEC, estimate_vae_working_memory_cogview4, estimate_vae_working_memory_sd3, estimate_vae_working_memory_sd15_sdxl, ) +from invokeai.backend.working_memory.coefficients import shipped_coefficients +from invokeai.backend.working_memory.platform import Platform def _vae(dtype: torch.dtype = torch.float16) -> MagicMock: @@ -25,39 +30,74 @@ def _vae(dtype: torch.dtype = torch.float16) -> MagicMock: return vae -def _estimate(estimator, device: torch.device, hip: str | None): +def _estimate(estimator, device: torch.device, hip: str | None) -> int: latents = torch.zeros(1, 4, 64, 64) # 512px output with ( patch("torch.version.hip", hip), # The mid-block score matrix is the other term in these estimates and has its own tests; hold it at zero so - # the assertions below are about the convolution constant alone. + # the assertions below are about the convolution coefficient alone. patch("invokeai.backend.util.vae_working_memory._vae_mid_block_score_matrix_bytes", return_value=0), ): if estimator is estimate_vae_working_memory_sd15_sdxl: - return estimator( + estimate = estimator( operation="decode", image_tensor=latents, vae=_vae(), tile_size=None, fp32=False, device=device ) - return estimator(operation="decode", image_tensor=latents, vae=_vae(), device=device) + else: + estimate = estimator(operation="decode", image_tensor=latents, vae=_vae(), device=device) + return estimate.reservation_bytes + + +def _decode_bytes(spec, conv_backend: str) -> float: + """A 512px fp16 decode at the shipped coefficient of that backend.""" + return 512 * 512 * 2 * shipped_coefficients(spec, Platform(conv_backend)).values["k_decode"] ESTIMATORS = [ - pytest.param(estimate_vae_working_memory_sd15_sdxl, id="sd15-sdxl"), - pytest.param(estimate_vae_working_memory_sd3, id="sd3"), - pytest.param(estimate_vae_working_memory_cogview4, id="cogview4"), + pytest.param(estimate_vae_working_memory_sd15_sdxl, VAE_SD_SPEC, id="sd15-sdxl"), + pytest.param(estimate_vae_working_memory_sd3, VAE_SD3_SPEC, id="sd3"), + pytest.param(estimate_vae_working_memory_cogview4, VAE_COGVIEW4_SPEC, id="cogview4"), ] -@pytest.mark.parametrize("estimator", ESTIMATORS) -def test_a_rocm_vae_reserves_the_miopen_constant(estimator): - assert _estimate(estimator, torch.device("cuda", 0), hip="7.2.0") == 512 * 512 * 2 * 3600 +@pytest.mark.parametrize("estimator, spec", ESTIMATORS) +def test_a_rocm_vae_reserves_the_miopen_coefficient(estimator, spec): + assert _estimate(estimator, torch.device("cuda", 0), hip="7.2.0") == _decode_bytes(spec, "miopen") -@pytest.mark.parametrize("estimator", ESTIMATORS) -def test_a_cuda_vae_reserves_the_cudnn_constant(estimator): - assert _estimate(estimator, torch.device("cuda", 0), hip=None) == 512 * 512 * 2 * 2200 +@pytest.mark.parametrize("estimator, spec", ESTIMATORS) +def test_a_cuda_vae_reserves_the_cudnn_coefficient(estimator, spec): + assert _estimate(estimator, torch.device("cuda", 0), hip=None) == _decode_bytes(spec, "cudnn") -@pytest.mark.parametrize("estimator", ESTIMATORS) -def test_a_cpu_only_vae_reserves_the_cudnn_constant_even_on_a_rocm_build(estimator): +@pytest.mark.parametrize("estimator, spec", ESTIMATORS) +def test_a_cpu_only_vae_reserves_the_cudnn_coefficient_even_on_a_rocm_build(estimator, spec): """A `cpu_only` VAE does not run on MIOpen, so naming its device has to change the answer.""" - assert _estimate(estimator, torch.device("cpu"), hip="7.2.0") == 512 * 512 * 2 * 2200 + assert _estimate(estimator, torch.device("cpu"), hip="7.2.0") == _decode_bytes(spec, "cudnn") + + +def _sd_fp32_decode(tile_size: int | None) -> int: + latents = torch.zeros(1, 4, 96, 128) # a 768x1024 output + with patch("invokeai.backend.util.vae_working_memory._vae_mid_block_score_matrix_bytes", return_value=0): + estimate = estimate_vae_working_memory_sd15_sdxl( + operation="decode", + image_tensor=latents, + vae=_vae(), + tile_size=tile_size, + fp32=True, + device=torch.device("cpu"), + ) + return estimate.reservation_bytes + + +FP32_WEIGHTS = 250 * 2**20 +"""What an fp32 decode adds for the upcast weights.""" + + +def test_a_tiled_fp32_sd_decode_is_one_tile_with_its_margin_plus_the_fp32_weights(): + k = shipped_coefficients(VAE_SD_SPEC, Platform("cudnn")).values["k_decode"] + assert _sd_fp32_decode(512) == int(512 * 512 * 4 * k * 1.25 + FP32_WEIGHTS) + + +def test_an_untiled_fp32_sd_decode_is_the_whole_image_plus_the_fp32_weights(): + k = shipped_coefficients(VAE_SD_SPEC, Platform("cudnn")).values["k_decode"] + assert _sd_fp32_decode(None) == int(768 * 1024 * 4 * k + FP32_WEIGHTS) diff --git a/tests/app/invocations/test_z_image_tiled_decode.py b/tests/app/invocations/test_z_image_tiled_decode.py index 2f8a46a894d..eb9d3a3a9b6 100644 --- a/tests/app/invocations/test_z_image_tiled_decode.py +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -10,7 +10,19 @@ from invokeai.app.invocations.vae.z_image_latents_to_image import ZImageLatentsToImageInvocation from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_tiling_scope import DEFAULT_TILE_SAMPLE_MIN_SIZE, MIN_TILE_SAMPLE_SIZE -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux +from invokeai.backend.util.vae_working_memory import VAE_FLUX_SPEC, estimate_vae_working_memory_flux +from invokeai.backend.working_memory.coefficients import shipped_coefficients +from invokeai.backend.working_memory.platform import Platform + + +def _flux_reservation(*args, **kwargs) -> int: + """What `estimate_vae_working_memory_flux` reserves; the estimator cells are about its size.""" + return estimate_vae_working_memory_flux(*args, **kwargs).reservation_bytes + + +def _k(conv_backend: str, operation: str) -> float: + """The shipped FLUX.1 VAE coefficient the estimator resolves for one convolution backend and operation.""" + return shipped_coefficients(VAE_FLUX_SPEC, Platform(conv_backend)).values[f"k_{operation}"] def _mock_flux_vae(element_size_bytes: int = 2) -> MagicMock: @@ -36,7 +48,7 @@ def _fused_non_rocm_build(self, monkeypatch): score matrix for it. On a ROCm rig that resolves to `cuda`, where the head-dim guard charges 4.25GiB for the 1024px case below -- which the expected value here clears by 1%, so these exact-value assertions turn build-dependent the moment either constant moves. Pin the fused, non-HIP regime they are written for, which - is also the cuDNN column of `_FLUX_VAE_SCALING_CONSTANTS`; `TestClassicVaeEstimators` in + also selects the cuDNN coefficients; `TestClassicVaeEstimators` in `tests/backend/util/test_rocm_sdpa_head_dim_guard.py` owns the ROCm side of this estimator. """ import invokeai.backend.util.attention as attention @@ -46,11 +58,18 @@ def _fused_non_rocm_build(self, monkeypatch): monkeypatch.setattr(vwm.TorchDevice, "choose_torch_device", classmethod(lambda cls: torch.device("cpu"))) @pytest.mark.parametrize( - ("hip", "operation", "constant"), - [(None, "decode", 2200), (None, "encode", 1100), ("7.14.0", "decode", 3600), ("7.14.0", "encode", 2750)], + ("hip", "operation", "conv_backend"), + [ + (None, "decode", "cudnn"), + (None, "encode", "cudnn"), + ("7.14.0", "decode", "miopen"), + ("7.14.0", "encode", "miopen"), + ], ids=["cudnn-decode", "cudnn-encode", "miopen-decode", "miopen-encode"], ) - def test_the_convolution_backend_of_the_vaes_device_picks_the_constant(self, monkeypatch, hip, operation, constant): + def test_the_convolution_backend_of_the_vaes_device_picks_the_constant( + self, monkeypatch, hip, operation, conv_backend + ): """MIOpen's workspaces are larger than cuDNN's: measured 3451 decode / 2688 encode on an RX 9060 XT, the same as the FLUX.2 VAE. A HIP build reports its GPU as "cuda", so the torch build decides, for the VAE's device.""" import invokeai.backend.util.vae_working_memory as vwm @@ -59,22 +78,22 @@ def test_the_convolution_backend_of_the_vaes_device_picks_the_constant(self, mon monkeypatch.setattr(vwm, "_vae_mid_block_score_matrix_bytes", lambda *args, **kwargs: 0) tensor = torch.zeros(1, 16, 64, 64) if operation == "decode" else torch.zeros(1, 3, 512, 512) - estimate = estimate_vae_working_memory_flux( + estimate = _flux_reservation( operation=operation, image_tensor=tensor, vae=_mock_flux_vae(), device=torch.device("cuda", 0) ) - on_cpu = estimate_vae_working_memory_flux( + on_cpu = _flux_reservation( operation=operation, image_tensor=tensor, vae=_mock_flux_vae(), device=torch.device("cpu") ) - assert estimate == 512 * 512 * 2 * constant - assert on_cpu == 512 * 512 * 2 * (2200 if operation == "decode" else 1100), "a cpu_only VAE runs on cuDNN terms" + assert estimate == 512 * 512 * 2 * _k(conv_backend, operation) + assert on_cpu == 512 * 512 * 2 * _k("cudnn", operation), "a cpu_only VAE runs on cuDNN terms" def test_the_default_reproduces_the_untiled_estimate(self): """Regression guard for the six call sites that pass no tile_size at all.""" latents = torch.zeros(1, 16, 128, 128) - # 1024x1024 output px * 2 bytes * 2200 - expected = 1024 * 1024 * 2 * 2200 - actual = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + # 1024x1024 output px * 2 bytes * the cuDNN decode coefficient + expected = 1024 * 1024 * 2 * _k("cudnn", "decode") + actual = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) assert actual == expected def test_the_tiled_estimate_is_a_tile_bounded_decode_plus_the_assembled_image(self): @@ -86,13 +105,13 @@ def test_the_tiled_estimate_is_a_tile_bounded_decode_plus_the_assembled_image(se under-reservation that no other assertion here can see. """ # 512px tiles, fp16, over a 1536x1536 output. - decode_term = 512 * 512 * 2 * 2200 * 1.25 + decode_term = 512 * 512 * 2 * _k("cudnn", "decode") * 1.25 # 4 RGB buffers at the VAE's element size, plus the 8-bit one the node hands to PIL. Four # is measured, not assumed: the assembled image grows the device peak by ~1.6 copies on a # diffusers VAE and ~0.1 on the host-merging BFL one, and the node adds its in-window # clamp/scale copies on top. See the constant's comment for the measurements. image_term = 1536 * 1536 * 3 * (4 * 2 + 1) - estimate = estimate_vae_working_memory_flux( + estimate = _flux_reservation( operation="decode", image_tensor=torch.zeros(1, 16, 192, 192), vae=_mock_flux_vae(), tile_size=512 ) assert estimate == int(decode_term + image_term) @@ -101,10 +120,10 @@ def test_both_tiled_terms_scale_with_the_vaes_element_size(self): """An fp32 VAE costs twice the bytes per pixel, in the tile and in the assembled image alike. Only the 8-bit buffer is fixed, which is why this is a little under 2x.""" latents = torch.zeros(1, 16, 192, 192) - fp16 = estimate_vae_working_memory_flux( + fp16 = _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(element_size_bytes=2), tile_size=512 ) - fp32 = estimate_vae_working_memory_flux( + fp32 = _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(element_size_bytes=4), tile_size=512 ) assert 1.95 < fp32 / fp16 < 2.0 @@ -120,10 +139,8 @@ def test_a_tiled_estimate_stays_a_fraction_of_the_untiled_one(self, latent_hw): the clamped tile plus the same image term, which can land above the single-pass estimate.) """ latents = torch.zeros(1, 16, *latent_hw) - untiled = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) - tiled = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=512 - ) + untiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + tiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=512) assert tiled < untiled / 2 def test_a_tile_larger_than_the_image_is_not_estimated_as_tiled(self): @@ -134,10 +151,8 @@ def test_a_tile_larger_than_the_image_is_not_estimated_as_tiled(self): decode needing ~4.6GB, evicting the transformer from the cache to make room for nothing. """ latents = torch.zeros(1, 16, 128, 128) # 1024px - untiled = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) - oversized = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=4096 - ) + untiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + oversized = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=4096) assert oversized == untiled def test_a_tile_wider_than_one_axis_is_priced_at_the_clamped_tile(self): @@ -148,10 +163,8 @@ def test_a_tile_wider_than_one_axis_is_priced_at_the_clamped_tile(self): single-pass decode that tiling is there to avoid -- ~92GB against ~37GB. """ latents = torch.zeros(1, 16, 128, 1024) # 1024 x 8192 px: only the long axis exceeds the tile - tiled = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=4096 - ) - untiled = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + tiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=4096) + untiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) assert tiled < untiled def test_a_tiled_estimate_still_covers_the_assembled_image(self): @@ -163,10 +176,10 @@ def test_a_tiled_estimate_still_covers_the_assembled_image(self): second resolution the image term is checked at, so a term that did not scale with area would fail here even if it matched at 1536px. """ - estimate = estimate_vae_working_memory_flux( + estimate = _flux_reservation( operation="decode", image_tensor=torch.zeros(1, 16, 1024, 1024), vae=_mock_flux_vae(), tile_size=512 ) - tile_term = int(512 * 512 * 2 * 2200 * 1.25) + tile_term = int(512 * 512 * 2 * _k("cudnn", "decode") * 1.25) one_full_resolution_buffer = 8192 * 8192 * 3 * 2 # The assembled output plus at least one of the post-processing copies. assert estimate >= tile_term + 2 * one_full_resolution_buffer @@ -178,8 +191,8 @@ def test_the_sentinel_does_not_read_the_size_off_the_vae(self): latents = torch.zeros(1, 16, 192, 192) vae = _mock_flux_vae() vae.tile_sample_min_size = 384 # as if a previous run had set it - estimate = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=vae, tile_size=0) - assert estimate == estimate_vae_working_memory_flux( + estimate = _flux_reservation(operation="decode", image_tensor=latents, vae=vae, tile_size=0) + assert estimate == _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=DEFAULT_TILE_SAMPLE_MIN_SIZE ) @@ -190,17 +203,15 @@ def test_the_sentinel_works_on_a_vae_without_that_attribute_at_all(self): vae = MagicMock(spec=AutoencoderKL) vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float16)]) del vae.tile_sample_min_size - estimate = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=vae, tile_size=0) - assert estimate == estimate_vae_working_memory_flux( + estimate = _flux_reservation(operation="decode", image_tensor=latents, vae=vae, tile_size=0) + assert estimate == _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=DEFAULT_TILE_SAMPLE_MIN_SIZE ) def test_a_tile_below_the_cost_floor_is_estimated_at_the_floor(self): latents = torch.zeros(1, 16, 192, 192) - estimate = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=8 - ) - assert estimate == estimate_vae_working_memory_flux( + estimate = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=8) + assert estimate == _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=MIN_TILE_SAMPLE_SIZE ) @@ -221,23 +232,21 @@ def _fused_non_rocm_build(self, monkeypatch): monkeypatch.setattr("invokeai.backend.util.vae_working_memory.sdpa_score_matrix_bytes", lambda *a, **k: 0) monkeypatch.setattr(torch.version, "hip", None, raising=False) - def test_the_untiled_estimate_is_unchanged(self): - """Four call sites pass no tile and depend on this number staying what it was.""" + def test_an_untiled_encode_is_the_whole_image_at_the_encode_coefficient(self): + """What the call sites that pass no tile reserve: the full image, priced at the cuDNN encode coefficient.""" image = torch.zeros(1, 3, 1024, 1024) - assert estimate_vae_working_memory_flux("encode", image, _mock_flux_vae()) == 1024 * 1024 * 2 * 1100 + assert _flux_reservation("encode", image, _mock_flux_vae()) == 1024 * 1024 * 2 * _k("cudnn", "encode") def test_a_tiled_estimate_is_flat_in_the_image_size_where_the_untiled_one_is_not(self): vae = _mock_flux_vae() estimates = [ - estimate_vae_working_memory_flux("encode", torch.zeros(1, 3, edge, edge), vae, tile_size=512) + _flux_reservation("encode", torch.zeros(1, 3, edge, edge), vae, tile_size=512) for edge in (1024, 2048, 3072) ] # Not exactly equal: the input image and the assembled latent are not bounded by tiling. # But they are a rounding error beside the tile term, which is what "flat" has to mean. assert max(estimates) < min(estimates) * 1.2 - untiled = [ - estimate_vae_working_memory_flux("encode", torch.zeros(1, 3, edge, edge), vae) for edge in (1024, 3072) - ] + untiled = [_flux_reservation("encode", torch.zeros(1, 3, edge, edge), vae) for edge in (1024, 3072)] assert untiled[1] == untiled[0] * 9 def test_the_tiled_estimate_is_one_tile_plus_the_input_and_every_live_copy_of_the_moments(self): @@ -247,12 +256,10 @@ def test_the_tiled_estimate_is_one_tile_plus_the_input_and_every_live_copy_of_th vae = _mock_flux_vae() edge, tile = 2048, 512 live_moment_copies = 1.0 / (1.0 - 0.25) ** 2 + 2.0 - expected = tile * tile * 2 * 1100 * 1.25 + expected = tile * tile * 2 * _k("cudnn", "encode") * 1.25 expected += edge * edge * 3 * 2 expected += (edge // 8) * (edge // 8) * 2 * 16 * 2 * live_moment_copies - assert estimate_vae_working_memory_flux("encode", torch.zeros(1, 3, edge, edge), vae, tile_size=tile) == int( - expected - ) + assert _flux_reservation("encode", torch.zeros(1, 3, edge, edge), vae, tile_size=tile) == int(expected) def test_it_does_not_inherit_the_decode_residual(self): """The decode adds four element-size copies of the assembled image plus a byte buffer, for @@ -262,8 +269,8 @@ def test_it_does_not_inherit_the_decode_residual(self): image = torch.zeros(1, 3, 2048, 2048) latents = torch.zeros(1, 16, 2048 // 8, 2048 // 8) - encode = estimate_vae_working_memory_flux("encode", image, vae, tile_size=512) - decode = estimate_vae_working_memory_flux("decode", latents, vae, tile_size=512) + encode = _flux_reservation("encode", image, vae, tile_size=512) + decode = _flux_reservation("decode", latents, vae, tile_size=512) # Same tile, same output area, so the whole difference is the residual term. assert decode - encode > 2048 * 2048 * 3 * 2, "the encode is paying for the decode's image copies" @@ -273,9 +280,7 @@ def test_a_tile_larger_than_the_image_is_priced_as_the_untiled_encode(self): oversized tile would otherwise reserve for a pass that never runs.""" vae = _mock_flux_vae() image = torch.zeros(1, 3, 512, 512) - assert estimate_vae_working_memory_flux( - "encode", image, vae, tile_size=1024 - ) == estimate_vae_working_memory_flux("encode", image, vae) + assert _flux_reservation("encode", image, vae, tile_size=1024) == _flux_reservation("encode", image, vae) def _build_decode_mocks(vae, latents: torch.Tensor, decoded: torch.Tensor, force_tiled_decode: bool = False): @@ -364,15 +369,18 @@ def test_the_tiling_state_is_restored_afterwards(self, flux_shaped_vae): assert (vae.use_tiling, vae.tile_sample_min_size, vae.tile_latent_min_size) == before @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) - def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(self, auto, flux_shaped_vae): + def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off( + self, auto, flux_shaped_vae, fixed_working_memory + ): """Where the driver pages instead of raising an OOM, the retry below never fires: the estimate decides.""" module = "invokeai.app.invocations.vae.z_image_latents_to_image" vae_info, context, state = _build_decode_mocks( flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512) ) context.config.get.return_value.auto_tiled_decode = auto + untiled, tiled = fixed_working_memory(20 * 2**30), fixed_working_memory(2 * 2**30) with ( - patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[20 * 2**30, 2 * 2**30]) as estimate, + patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[untiled, tiled]) as estimate, patch(f"{module}.should_pretile_vae_decode", return_value=True) as pretile, ): _build_invocation().invoke(context) @@ -380,18 +388,21 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se if auto: pretile.assert_called_once_with(vae_info.compute_device, 20 * 2**30) assert estimate.call_args.kwargs["tile_size"] == 0 - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + vae_info.model_on_device.assert_called_once_with(working_memory=tiled) assert state[0]["use_tiling"] is True else: pretile.assert_not_called() assert estimate.call_count == 1 + vae_info.model_on_device.assert_called_once_with(working_memory=untiled) assert state[0]["use_tiling"] is False - def test_requested_tiling_skips_the_untiled_estimate(self, flux_shaped_vae): + def test_requested_tiling_skips_the_untiled_estimate(self, flux_shaped_vae, fixed_working_memory): module = "invokeai.app.invocations.vae.z_image_latents_to_image" _, context, _ = _build_decode_mocks(flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) with ( - patch(f"{module}.estimate_vae_working_memory_flux", return_value=2 * 2**30) as estimate, + patch( + f"{module}.estimate_vae_working_memory_flux", return_value=fixed_working_memory(2 * 2**30) + ) as estimate, patch(f"{module}.should_pretile_vae_decode") as pretile, ): _build_invocation(tiled=True).invoke(context) @@ -399,11 +410,13 @@ def test_requested_tiling_skips_the_untiled_estimate(self, flux_shaped_vae): pretile.assert_not_called() @pytest.mark.parametrize("tiled,expected_tile_size", [(False, None), (True, 0)]) - def test_the_estimate_is_tile_bounded_only_when_tiling(self, flux_shaped_vae, tiled, expected_tile_size): + def test_the_estimate_is_tile_bounded_only_when_tiling( + self, flux_shaped_vae, fixed_working_memory, tiled, expected_tile_size + ): path = "invokeai.app.invocations.vae.z_image_latents_to_image.estimate_vae_working_memory_flux" vae = flux_shaped_vae() _, context, _ = _build_decode_mocks(vae, torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) - with patch(path, return_value=1024) as estimate: + with patch(path, return_value=fixed_working_memory(1024)) as estimate: _build_invocation(tiled=tiled).invoke(context) assert estimate.call_args.kwargs["tile_size"] == expected_tile_size diff --git a/tests/app/invocations/test_z_image_working_memory.py b/tests/app/invocations/test_z_image_working_memory.py index c04e44e2b0d..655714a8c9a 100644 --- a/tests/app/invocations/test_z_image_working_memory.py +++ b/tests/app/invocations/test_z_image_working_memory.py @@ -42,7 +42,7 @@ def _vae_info(vae) -> MagicMock: class TestTheEstimateReachesTheCache: - def test_the_decode_node_reserves_what_it_estimated(self, flux_shaped_vae): + def test_the_decode_node_reserves_what_it_estimated(self, flux_shaped_vae, fixed_working_memory): vae = flux_shaped_vae() vae_info = _vae_info(vae) @@ -53,7 +53,7 @@ def test_the_decode_node_reserves_what_it_estimated(self, flux_shaped_vae): # `ImageOutput.build` validates these, so they cannot be bare MagicMocks. context.images.save.return_value = MagicMock(image_name="test.png", width=512, height=512) - expected_memory = 500 * 1024 * 1024 + expected_memory = fixed_working_memory(500 * 1024 * 1024) path = "invokeai.app.invocations.vae.z_image_latents_to_image.estimate_vae_working_memory_flux" with patch(path, return_value=expected_memory) as estimate: ZImageLatentsToImageInvocation.model_construct( @@ -64,18 +64,18 @@ def test_the_decode_node_reserves_what_it_estimated(self, flux_shaped_vae): ).invoke(context) estimate.assert_called_once() - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected_memory) + vae_info.model_on_device.assert_called_once_with(working_memory=expected_memory) - def test_the_encode_node_reserves_what_it_estimated(self, flux_shaped_vae): + def test_the_encode_node_reserves_what_it_estimated(self, flux_shaped_vae, fixed_working_memory): vae = flux_shaped_vae() vae_info = _vae_info(vae) - expected_memory = 250 * 1024 * 1024 + expected_memory = fixed_working_memory(250 * 1024 * 1024) path = "invokeai.app.invocations.vae.z_image_image_to_latents.estimate_vae_working_memory_flux" with patch(path, return_value=expected_memory) as estimate: latents = ZImageImageToLatentsInvocation.vae_encode(vae_info, torch.zeros(1, 3, 512, 512)) estimate.assert_called_once() - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected_memory) + vae_info.model_on_device.assert_called_once_with(working_memory=expected_memory) # The encode really ran: a 512px image through an 8x, 16-channel VAE is a 64px latent. assert latents.shape == (1, 16, 64, 64) diff --git a/tests/backend/model_manager/load/test_loaded_model.py b/tests/backend/model_manager/load/test_loaded_model.py index 04368325bde..9793e2b89e3 100644 --- a/tests/backend/model_manager/load/test_loaded_model.py +++ b/tests/backend/model_manager/load/test_loaded_model.py @@ -1,3 +1,5 @@ +import logging + import pytest import torch @@ -12,6 +14,8 @@ from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import ( apply_custom_layers_to_model, ) +from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate, WorkingMemoryTerm class ModelWithRequiredScale(torch.nn.Module): @@ -25,10 +29,12 @@ class FakeCache: def __init__(self): self.lock_calls = 0 self.unlock_calls = 0 + self.reserved: list[int | None] = [] def lock(self, cache_record: CacheRecord, working_mem_bytes: int | None, max_move_bytes: int | None = None) -> bool: - del cache_record, working_mem_bytes, max_move_bytes + del cache_record, max_move_bytes self.lock_calls += 1 + self.reserved.append(working_mem_bytes) return True def unlock(self, cache_record: CacheRecord) -> None: @@ -63,6 +69,62 @@ def test_model_on_device_leaves_full_load_models_unchanged(): assert all(param.device.type == "cpu" for param in model.parameters()) +def _full_load_model(cache: FakeCache) -> LoadedModelWithoutConfig: + cached_model = CachedModelOnlyFullLoad( + model=torch.nn.Linear(4, 4), compute_device=torch.device("meta"), total_bytes=1, keep_ram_copy=False + ) + return LoadedModelWithoutConfig(cache_record=CacheRecord(key="vae-key", cached_model=cached_model), cache=cache) + + +_ESTIMATE = WorkingMemoryEstimate( + estimator="vae.flux2", + version=1, + features={"operation": "decode"}, + terms=(WorkingMemoryTerm.fitted("convolutions", "k_decode", 3 * 2**29 + 0.25, {"k_decode": 2.0}),), + source="shipped:conv_backend=cudnn", +) + + +def test_a_working_memory_estimate_reserves_its_bytes_and_logs_where_they_came_from(caplog): + cache = FakeCache() + logger = "invokeai.backend.model_manager.load.load_base" + # Created before the level is lowered: creating it sets the configured level, which would undo that. + InvokeAILogger.get_logger(logger) + + with ( + caplog.at_level(logging.DEBUG, logger=logger), + _full_load_model(cache).model_on_device(working_memory=_ESTIMATE), + ): + pass + + assert cache.reserved == [3 * 2**30] + assert cache.unlock_calls == 1 + (record,) = [r for r in caplog.records if r.name == logger] + assert "vae-key" in record.message + assert "vae.flux2 reserves 3072 MB (shipped:conv_backend=cudnn)" in record.message + + +def test_an_estimate_and_a_byte_count_together_are_refused_before_locking(): + cache = FakeCache() + + with pytest.raises(ValueError, match="not both"): + with _full_load_model(cache).model_on_device(1024, working_memory=_ESTIMATE): + pass + + assert cache.lock_calls == 0 + + +def test_an_estimate_passed_as_a_byte_count_says_where_it_goes(): + """The positional slot is `working_mem_bytes`; an estimate there would fail deep inside the cache instead.""" + cache = FakeCache() + + with pytest.raises(TypeError, match="working_memory="): + with _full_load_model(cache).model_on_device(_ESTIMATE): + pass + + assert cache.lock_calls == 0 + + def test_supports_partial_loading_reflects_cached_model_type(): partial_model = CachedModelWithPartialLoad( model=torch.nn.Linear(4, 4), compute_device=torch.device("meta"), keep_ram_copy=False diff --git a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py index e82d1ccead1..6710b3afbf3 100644 --- a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py +++ b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py @@ -488,7 +488,7 @@ class TestClassicVaeEstimators: def _rocm_cuda(self, monkeypatch): """The guard's ROCm switch and the convolution backend are separate questions, and this class is about the first: the assertions below carry the cuDNN linear term, so pin `torch.version.hip` (which selects MIOpen's - larger constants) away from it while the attention side stays on ROCm.""" + larger coefficients) away from it while the attention side stays on ROCm.""" import invokeai.backend.util.vae_working_memory as vwm monkeypatch.setattr(attention, "_IS_ROCM", True) @@ -508,16 +508,25 @@ def _score_bytes(px: int) -> int: tokens = (px // 8) ** 2 return tokens * tokens * attention.SDPA_MATH_BYTES_PER_SCORE_ELEMENT + @staticmethod + def _cudnn_decode_bytes(px: int) -> float: + """A bf16 decode's convolution term at the shipped cuDNN coefficient.""" + from invokeai.backend.util.vae_working_memory import VAE_SD_SPEC + from invokeai.backend.working_memory.coefficients import shipped_coefficients + from invokeai.backend.working_memory.platform import Platform + + return px * px * 2 * shipped_coefficients(VAE_SD_SPEC, Platform("cudnn")).values["k_decode"] + def test_sdxl_untiled_decode_reserves_the_score_matrix_when_it_dominates(self, monkeypatch): from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl latents = torch.zeros(1, 4, 2048 // 8, 2048 // 8) - linear = 2048 * 2048 * 2 * 2200 + linear = self._cudnn_decode_bytes(2048) with_guard = estimate_vae_working_memory_sd15_sdxl( "decode", latents, self._vae(AutoencoderKL), None, fp32=False - ) + ).reservation_bytes assert with_guard == self._score_bytes(2048) > linear monkeypatch.setattr(attention, "_IS_ROCM", False) @@ -527,7 +536,9 @@ def test_sdxl_untiled_decode_reserves_the_score_matrix_when_it_dominates(self, m real_empty = torch.empty monkeypatch.setattr(torch, "empty", lambda *a, device=None, **kw: real_empty(*a, **kw)) assert ( - estimate_vae_working_memory_sd15_sdxl("decode", latents, self._vae(AutoencoderKL), None, fp32=False) + estimate_vae_working_memory_sd15_sdxl( + "decode", latents, self._vae(AutoencoderKL), None, fp32=False + ).reservation_bytes == linear ) @@ -538,7 +549,8 @@ def test_the_other_autoencoder_kl_estimators_price_it_too(self, estimator): import invokeai.backend.util.vae_working_memory as vwm latents = torch.zeros(1, 16, 2048 // 8, 2048 // 8) - assert getattr(vwm, estimator)("decode", latents, self._vae(AutoencoderKL)) == self._score_bytes(2048) + estimate = getattr(vwm, estimator)("decode", latents, self._vae(AutoencoderKL)) + assert estimate.reservation_bytes == self._score_bytes(2048) def test_a_tiny_autoencoder_has_no_attention_to_price(self): from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny @@ -548,7 +560,7 @@ def test_a_tiny_autoencoder_has_no_attention_to_price(self): latents = torch.zeros(1, 4, 2048 // 8, 2048 // 8) assert estimate_vae_working_memory_sd15_sdxl( "decode", latents, self._vae(AutoencoderTiny), None, fp32=False - ) == (2048 * 2048 * 2 * 2200) + ).reservation_bytes == self._cudnn_decode_bytes(2048) def test_flux1_decode_prices_the_score_matrix_by_tile_when_tiled(self): from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL @@ -556,9 +568,11 @@ def test_flux1_decode_prices_the_score_matrix_by_tile_when_tiled(self): from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux latents = torch.zeros(1, 16, 2048 // 8, 2048 // 8) - untiled = estimate_vae_working_memory_flux("decode", latents, self._vae(AutoencoderKL)) + untiled = estimate_vae_working_memory_flux("decode", latents, self._vae(AutoencoderKL)).reservation_bytes assert untiled == self._score_bytes(2048) - tiled = estimate_vae_working_memory_flux("decode", latents, self._vae(AutoencoderKL), tile_size=1024) + tiled = estimate_vae_working_memory_flux( + "decode", latents, self._vae(AutoencoderKL), tile_size=1024 + ).reservation_bytes assert tiled < self._score_bytes(2048) assert tiled >= self._score_bytes(1024), "a 1024px tile still materializes its own 16384-token score matrix" diff --git a/tests/backend/working_memory/__init__.py b/tests/backend/working_memory/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/backend/working_memory/test_shipped_coefficients.py b/tests/backend/working_memory/test_shipped_coefficients.py new file mode 100644 index 00000000000..d369d03df90 --- /dev/null +++ b/tests/backend/working_memory/test_shipped_coefficients.py @@ -0,0 +1,312 @@ +"""The shipped working-memory data, checked against the estimators that read it. + +Every coefficient entry resolves for the spec it names on a measured platform value, and every value an entry carries +is decided by at least one measured point. A point is peak *reserved* bytes, one fresh process per point (see its +`provenance`), and bounds its estimate from both sides. The model cache consumes an estimate as `free >= estimate`, so +the estimate must be at least the point. Where a fitted term decides it, it must also be at most 1.3 times the point, +or the point's tighter `max_ratio`: a reservation the operation never uses is model weights evicted for nothing. Where +an analytic term (a materialized score matrix) decides it, that term bounds rather than describes, so it gets no upper +bound, and the cache's `device_working_mem_gb` floor counts toward the lower one. +""" + +import itertools +import re +from dataclasses import asdict +from pathlib import Path +from typing import get_args, get_type_hints + +import pytest +import torch +from pydantic import ValidationError + +import invokeai +from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.backend.working_memory.coefficients import ( + CoefficientEntry, + MeasuredPoint, + ResolvedCoefficients, + ValuesFrom, + WorkingMemoryData, + load_data, + resolve_coefficients, + shipped_coefficients, + shipped_data, +) +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate +from invokeai.backend.working_memory.platform import UNMEASURED_PRICED_AS, Platform, platform_for +from invokeai.backend.working_memory.spec import ESTIMATOR_MODULES, EstimatorSpec, import_estimator_modules + +MAX_OVER_RESERVATION = 1.3 +"""How far above a measured point a reservation may sit where a fitted coefficient decides it, unless the point +carries a tighter `max_ratio`.""" + +WORKING_MEMORY_FLOOR = int(InvokeAIAppConfig.model_fields["device_working_mem_gb"].default * 2**30) +"""What `ModelCache._working_mem_reserve` floors every reservation at. The field default, not a config instance an +environment variable could move.""" + +# Imported deliberately rather than relying on whichever modules ran first: a spec missing from the registry is a spec +# whose data nobody checks. +SPECS = import_estimator_modules() + + +def _platforms(*, measured: bool) -> list[Platform]: + """Every platform built from measured values only, or every platform with at least one unmeasured value.""" + every = {key: get_args(hint) for key, hint in get_type_hints(Platform).items()} + measured_values = { + key: [v for v in values if v not in UNMEASURED_PRICED_AS.get(key, {})] for key, values in every.items() + } + chosen = measured_values if measured else every + platforms = [Platform(**dict(zip(chosen, values, strict=True))) for values in itertools.product(*chosen.values())] + if measured: + return platforms + return [p for p in platforms if any(value in UNMEASURED_PRICED_AS.get(key, {}) for key, value in asdict(p).items())] + + +def _platform_key(platform: dict[str, str]) -> tuple[tuple[str, str], ...]: + return tuple(sorted(platform.items())) + + +def _point_id(point: MeasuredPoint) -> str: + features = point.features + shape = f"{features.get('height')}x{features.get('width')}" + tile = f"-tile{features['tile']}" if features.get("tile") is not None else "" + batch = f"-b{features['batch_size']}" if features.get("batch_size", 1) != 1 else "" + scores = "-scores" if features.get("score_matrix_bytes") else "" + platform = ",".join(point.platform.values()) + card = point.device.split(" (")[0].replace(" ", "") + return f"{point.estimator}-{platform}-{features.get('operation')}-{shape}{tile}{batch}{scores}-{card}" + + +def _evaluate(point: MeasuredPoint) -> WorkingMemoryEstimate: + spec = SPECS[point.estimator] + assert point.version == spec.version + return spec.evaluate(point.features, shipped_coefficients(spec, Platform(**point.platform))) + + +def _fitted_decides(estimate: WorkingMemoryEstimate) -> bool: + return any(term.included and term.coefficient is not None for term in estimate.terms) + + +POINTS = [pytest.param(point, id=_point_id(point)) for data in shipped_data().values() for point in data.points] + + +def test_every_data_file_parses_and_declares_coefficients() -> None: + for name, data in shipped_data().items(): + assert data.coefficients, f"{name} declares no coefficients" + + +def test_every_coefficient_entry_resolves_for_its_spec_on_a_measured_platform_value() -> None: + """Resolving checks what an entry provides against what its spec reads, the borrowed values included. An entry for + an unmeasured value would never be read: the resolver prices that value as `UNMEASURED_PRICED_AS` says.""" + seen: set[tuple[str, int, tuple[tuple[str, str], ...]]] = set() + for name, data in shipped_data().items(): + for entry in data.coefficients: + where = f"{name}: {entry.estimator} on {entry.platform}" + spec = SPECS.get(entry.estimator) + assert spec is not None, f"{where}: no estimator registers this id" + assert entry.version == spec.version, f"{where}: version {entry.version}, the estimator is {spec.version}" + assert set(entry.platform) == set(spec.platform_keys), f"{where}: the spec keys {spec.platform_keys}" + for key, value in entry.platform.items(): + assert value not in UNMEASURED_PRICED_AS.get(key, {}), f"{where}: {key}={value} is not measured" + key = (entry.estimator, entry.version, _platform_key(entry.platform)) + assert key not in seen, f"{where}: declared twice" + seen.add(key) + shipped_coefficients(spec, Platform(**entry.platform)) + + +@pytest.mark.parametrize("spec", sorted(SPECS.values(), key=lambda spec: spec.id), ids=lambda spec: spec.id) +def test_every_estimator_has_coefficients_on_every_measured_platform_value(spec: EstimatorSpec) -> None: + for platform in _platforms(measured=True): + assert set(shipped_coefficients(spec, platform).values) == spec.coefficients + + +@pytest.mark.parametrize("spec", sorted(SPECS.values(), key=lambda spec: spec.id), ids=lambda spec: spec.id) +def test_an_unmeasured_platform_value_is_priced_as_the_measured_one_and_says_so(spec: EstimatorSpec) -> None: + for platform in _platforms(measured=False): + priced_as = Platform( + **{key: UNMEASURED_PRICED_AS.get(key, {}).get(value, value) for key, value in asdict(platform).items()} + ) + resolved = shipped_coefficients(spec, platform) + assert resolved.values == shipped_coefficients(spec, priced_as).values + assert resolved.source.startswith( + f"shipped:conv_backend={platform.conv_backend} ({priced_as.conv_backend} values" + ) + + +def test_every_carried_value_is_decided_by_a_measured_point() -> None: + """A value no point decides is unguarded: any number would pass. An estimator without points of its own takes a + measured one's values with `values_from`, and is then held to that one's points.""" + entries = { + (entry.estimator, entry.version, _platform_key(entry.platform)): entry + for data in shipped_data().values() + for entry in data.coefficients + } + decided: set[tuple[tuple[str, int, tuple[tuple[str, str], ...]], str]] = set() + for data in shipped_data().values(): + for point in data.points: + key = (point.estimator, point.version, _platform_key(point.platform)) + lender = entries[key].values_from + owner = (lender.estimator, lender.version, key[2]) if lender is not None else key + estimate = _evaluate(point) + decided |= {(owner, term.coefficient) for term in estimate.terms if term.included and term.coefficient} + + undecided = [ + f"{entry.estimator} {entry.platform} {name}" + for key, entry in entries.items() + for name in entry.values or {} + if (key, name) not in decided + ] + assert not undecided, f"no measured point decides these shipped values: {undecided}" + + +@pytest.mark.parametrize("point", POINTS) +def test_every_measured_point_is_bounded_by_its_estimate(point: MeasuredPoint) -> None: + estimate = _evaluate(point) + fitted = _fitted_decides(estimate) + + reserved = estimate.reservation_bytes if fitted else max(estimate.reservation_bytes, WORKING_MEMORY_FLOOR) + assert reserved >= point.measured_bytes, ( + f"reserves {reserved / 2**30:.3f} GiB for a measured {point.measured_bytes / 2**30:.3f} GiB" + ) + if point.max_ratio is not None or fitted: + max_ratio = point.max_ratio if point.max_ratio is not None else MAX_OVER_RESERVATION + assert estimate.reservation_bytes <= max_ratio * point.measured_bytes, ( + f"reserves {estimate.reservation_bytes / point.measured_bytes:.2f}x the measured peak, over {max_ratio}x" + ) + + +def test_a_point_bound_is_never_looser_than_the_default() -> None: + """`max_ratio` tightens the bound for a point the estimate has to stay close to; it never excuses one.""" + for name, data in shipped_data().items(): + for point in data.points: + if point.max_ratio is not None: + assert 1.0 <= point.max_ratio < MAX_OVER_RESERVATION, f"{name}: {_point_id(point)}" + + +@pytest.mark.parametrize( + "device, hip, conv_backend", + [ + (torch.device("cuda", 0), "7.2.0", "miopen"), + (torch.device("cuda", 0), None, "cudnn"), + # A HIP build reports its GPU as "cuda"; anything else on that build is not on MIOpen. + (torch.device("cpu"), "7.2.0", "cpu"), + (torch.device("mps"), "7.2.0", "mps"), + (torch.device("xpu"), None, "xpu"), + (torch.device("meta"), None, "cpu"), + ], +) +def test_the_platform_is_the_convolution_backend_the_device_really_runs(device, hip, conv_backend, monkeypatch) -> None: + monkeypatch.setattr(torch.version, "hip", hip) + assert platform_for(device) == Platform(conv_backend=conv_backend) + + +def _no_terms(features, values): + return () + + +def _spec(spec_id: str, version: int = 1) -> EstimatorSpec: + return EstimatorSpec( + id=spec_id, version=version, coefficients=frozenset({"k"}), platform_keys=("conv_backend",), terms=_no_terms + ) + + +def _entry(estimator: str, *, version: int = 1, values=None, values_from: str | None = None) -> CoefficientEntry: + return CoefficientEntry( + estimator=estimator, + version=version, + platform={"conv_backend": "cudnn"}, + values=values, + values_from=ValuesFrom(estimator=values_from, version=1) if values_from else None, + provenance="-", + ) + + +def _data(*entries: CoefficientEntry) -> list[WorkingMemoryData]: + return [WorkingMemoryData(schema_version=1, coefficients=list(entries), points=[])] + + +def test_a_missing_or_stale_entry_names_the_estimator_and_platform() -> None: + """No fallback to another platform or version: coefficients fitted for something else are the shortfall the data + exists to prevent.""" + spec = _spec("test.estimator", version=2) + + assert resolve_coefficients( + _data(_entry(spec.id, version=2, values={"k": 7.0})), spec, Platform("cudnn"), origin="test" + ) == ResolvedCoefficients(values={"k": 7.0}, source="test:conv_backend=cudnn") + with pytest.raises(LookupError, match=r"'test\.estimator' version 2 on conv_backend=miopen\."): + resolve_coefficients(_data(_entry(spec.id, version=2, values={"k": 7.0})), spec, Platform("miopen"), "test") + with pytest.raises(LookupError, match=r"'test\.estimator' version 2 on conv_backend=cudnn \(found only version"): + resolve_coefficients(_data(_entry(spec.id, version=1, values={"k": 7.0})), spec, Platform("cudnn"), "test") + with pytest.raises(LookupError, match=r"declared 2 times"): + resolve_coefficients( + _data(_entry(spec.id, version=2, values={"k": 7.0}), _entry(spec.id, version=2, values={"k": 8.0})), + spec, + Platform("cudnn"), + "test", + ) + + +def test_borrowed_values_resolve_one_hop_and_the_source_names_both_borrowings() -> None: + data = _data(_entry("test.lender", values={"k": 7.0}), _entry("test.borrower", values_from="test.lender")) + + on_cudnn = resolve_coefficients(data, _spec("test.borrower"), Platform("cudnn"), origin="test") + assert on_cudnn == ResolvedCoefficients(values={"k": 7.0}, source="test:conv_backend=cudnn (values of test.lender)") + on_xpu = resolve_coefficients(data, _spec("test.borrower"), Platform("xpu"), origin="test") + assert on_xpu == ResolvedCoefficients( + values={"k": 7.0}, source="test:conv_backend=xpu (cudnn values of test.lender)" + ) + + +@pytest.mark.parametrize( + "entries, message", + [ + pytest.param( + (_entry("a", values_from="b"), _entry("b", values_from="c"), _entry("c", values={"k": 1.0})), + "borrows in turn", + id="chain", + ), + pytest.param((_entry("a", values_from="b"), _entry("b", values_from="a")), "borrows in turn", id="cycle"), + pytest.param( + (_entry("a", values_from="b"),), r"No working-memory coefficients for estimator 'b'", id="no-lender" + ), + pytest.param( + (_entry("a", values_from="b"), _entry("b", values={"other": 1.0})), r"provide \['other'\]", id="mismatch" + ), + ], +) +def test_a_borrowing_that_is_not_one_hop_to_matching_values_is_refused(entries, message) -> None: + with pytest.raises(LookupError, match=message): + resolve_coefficients(_data(*entries), _spec("a"), Platform("cudnn"), origin="test") + + +@pytest.mark.parametrize("values, values_from", [({"k": 1.0}, "b"), (None, None)], ids=["both", "neither"]) +def test_an_entry_carries_exactly_one_of_values_and_values_from(values, values_from) -> None: + with pytest.raises(ValidationError, match="exactly one of"): + _entry("a", values=values, values_from=values_from) + + +def test_a_data_directory_without_data_files_is_named(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match=re.escape(str(tmp_path))): + load_data(tmp_path) + with pytest.raises(FileNotFoundError, match=re.escape(str(tmp_path / "missing"))): + load_data(tmp_path / "missing") + (tmp_path / "broken.json").write_text('{"schema_version": 1}', encoding="utf-8") + with pytest.raises(ValueError, match="broken.json"): + load_data(tmp_path) + + +@pytest.mark.platform_independent +def test_every_module_that_declares_an_estimator_is_listed() -> None: + """The registry fills as modules import, so a spec declared outside `ESTIMATOR_MODULES` is missing from every + check above until something else happens to import it -- and a ROCm user is the first to find its gap.""" + package = Path(invokeai.__file__).parent + declaring = sorted( + ".".join(path.relative_to(package.parent).with_suffix("").parts) + for owner in ("app", "backend") + for path in (package / owner).rglob("*.py") + if re.search(r"\bEstimatorSpec\(", path.read_text(encoding="utf-8")) + ) + assert declaring, "found no estimator modules -- the scan is broken" + unlisted = sorted(set(declaring) - set(ESTIMATOR_MODULES)) + stale = sorted(set(ESTIMATOR_MODULES) - set(declaring)) + assert not unlisted and not stale, f"add to ESTIMATOR_MODULES: {unlisted}; declares no spec: {stale}" diff --git a/tests/backend/working_memory/test_spec_registry.py b/tests/backend/working_memory/test_spec_registry.py new file mode 100644 index 00000000000..0c323aa179a --- /dev/null +++ b/tests/backend/working_memory/test_spec_registry.py @@ -0,0 +1,72 @@ +"""Registry mechanics, against throwaway specs rather than the shipped ones.""" + +import pytest + +from invokeai.backend.working_memory import spec as spec_module +from invokeai.backend.working_memory.spec import EstimatorSpec, register, registered_specs + + +@pytest.fixture(autouse=True) +def isolated_registry(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(spec_module, "_SPECS", {}) + + +def _module_terms(): + """A fresh `terms` function on every call, with one module and qualified name: what `importlib.reload` of its + defining module produces.""" + + def terms(features, values): + return () + + return terms + + +def _another_estimators_terms(features, values): + return () + + +def _spec(terms, version: int = 1) -> EstimatorSpec: + return EstimatorSpec( + id="test.estimator", + version=version, + coefficients=frozenset({"k"}), + platform_keys=("conv_backend",), + terms=terms, + ) + + +def test_re_executing_the_defining_module_adopts_its_new_terms() -> None: + """`importlib.reload`: the same declaration from the re-executed function.""" + before, after = _module_terms(), _module_terms() + register(_spec(before)) + register(_spec(after)) + + assert registered_specs()["test.estimator"].terms is after + + +def test_a_differing_declaration_under_a_taken_id_raises() -> None: + register(_spec(_module_terms())) + with pytest.raises(ValueError, match="already registered"): + register(_spec(_module_terms(), version=2)) + + assert registered_specs()["test.estimator"].version == 1, "the first one stands" + + +def test_the_same_declaration_from_another_function_raises() -> None: + """`terms` is not part of a spec's equality, so this is the case equality alone would let through.""" + first = _module_terms() + register(_spec(first)) + with pytest.raises(ValueError, match="ids are unique"): + register(_spec(_another_estimators_terms)) + + assert registered_specs()["test.estimator"].terms is first + + +def test_the_same_function_registered_twice_raises() -> None: + """Two specs of one module sharing a `terms` function -- as SD3 and CogView4 do -- under one id. `importlib.reload` + makes new function objects, so it never hands over the same one twice; jurigged re-running an edited `register` + statement does, and raises too (see `register`).""" + terms = _module_terms() + register(_spec(terms)) + with pytest.raises(ValueError, match="ids are unique"): + register(_spec(terms))