From 8d16ef977b77d3a5e16654ab07eeb7db58368dd3 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 1 Jul 2026 18:01:12 -0300 Subject: [PATCH] Fix diffusion GGUF memory over-estimate and torch.compile crashes Three chained bugs that made Z-Image (and other GGUF DiTs) crash at generation on anything but a huge, fully-idle GPU. Verified end to end on an RTX 6000 Ada: Q2_K now plans resident and generates a real 1024x1024 PNG on both the resident and forced-group-offload paths. - Memory planner over-estimated the GGUF transformer's resident size. diffusers keeps GGUF weights PACKED (uint8 GGUFParameter) and dequantises per-matmul transiently, so resident VRAM is ~= the on-disk size, not the unpacked bf16 size (measured: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). The old per-quant expansion (x8 for Q2) over-estimated ~7.6x, so a 3.6 GB model on a 48 GB-free card was judged a "tight fit" and forced into group offload. Replace the multiplier table with estimate_gguf_resident_mib = storage * 1.05 (matches diffusers' own get_memory_footprint of a loaded GGUF model). - torch.compile with fullgraph=True crashed under CPU offload: group/model/ sequential offload installs a @torch.compiler.disable'd ModuleGroup.onload_ hook, which graph-breaks. Drop fullgraph when offloading is planned, same as the existing step-cache case (fullgraph = not (cache_active or offload_active)). This mirrors diffusers' documented compile+offload guidance. - compile_repeated_blocks compiles one graph per distinct block shape, but Z-Image's "repeated" blocks are heterogeneous (~11 variants), above dynamo's default recompile_limit of 8, so a resident load hard-errored under fullgraph. Raise the limit (diffusers' documented fix for regional-compile recompilation). Confirmed force_parameter_static_shapes=False is the wrong lever: same variant count, ~6x slower compile. Also drops the now-dead infer_gguf_quant_label / gguf_filename plumbing and adds regression tests for the estimate and the offload fullgraph drop. --- studio/backend/core/inference/diffusion.py | 17 +++--- .../core/inference/diffusion_memory.py | 53 +++++-------------- .../backend/core/inference/diffusion_speed.py | 44 ++++++++++++--- studio/backend/tests/test_diffusion_memory.py | 35 ++++-------- studio/backend/tests/test_diffusion_speed.py | 18 +++++++ 5 files changed, 85 insertions(+), 82 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 6d728df035..ecf1a1dcf9 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -39,10 +39,9 @@ from .diffusion_device import ( from .diffusion_memory import ( OFFLOAD_NONE, apply_memory_plan, - estimate_gguf_dense_mib, + estimate_gguf_resident_mib, estimate_image_runtime_mib, file_size_mib, - infer_gguf_quant_label, plan_diffusion_memory, snapshot_device_memory, ) @@ -522,7 +521,7 @@ class DiffusionBackend: # dense bf16 transformer must fit resident, so the fast path is offered only # when the plan is `none`. plan = self._plan_memory( - target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload + target, gguf_path, base, fam, memory_mode, cpu_offload ) # Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it @@ -641,6 +640,9 @@ class DiffusionBackend: family = fam, speed_mode = effective_speed, cache_active = cache_engaged is not None, + # The planned offload policy: group/model/sequential offload installs + # compiler-disabled onload hooks, so compile must drop fullgraph. + offload_active = plan.offload_policy != OFFLOAD_NONE, logger = logger, ) if transformer_quant_engaged is not None and not speed_applied.get("compiled"): @@ -797,7 +799,6 @@ class DiffusionBackend: self, target: DiffusionDeviceTarget, gguf_path: str, - gguf_filename: Optional[str], base: str, fam: DiffusionFamily, memory_mode: Optional[str], @@ -808,16 +809,14 @@ class DiffusionBackend: offload policy + VAE memory savers. Kept on the backend so the cached base repo (companion text-encoder / VAE) feeds the size estimate.""" device_memory = snapshot_device_memory(target) - transformer_dense = estimate_gguf_dense_mib( - file_size_mib(gguf_path), infer_gguf_quant_label(gguf_filename) - ) + transformer_resident = estimate_gguf_resident_mib(file_size_mib(gguf_path)) # The companion components (VAE + text encoders) load near their on-disk # size; sum whatever the prefetch already placed in the base-repo cache. companion = self._cache_bytes(base) companion_mib = int(companion // (1024 * 1024)) if companion else None model_dense_mib = None - if transformer_dense is not None: - model_dense_mib = transformer_dense + (companion_mib or 0) + if transformer_resident is not None: + model_dense_mib = transformer_resident + (companion_mib or 0) runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = fam.name) return plan_diffusion_memory( target = target, diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 9172ef7cf8..b2c1e25f11 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -218,51 +218,22 @@ def file_size_mib(path: Any) -> Optional[int]: return None -def infer_gguf_quant_label(filename: Optional[str]) -> Optional[str]: - """Pull a quant tag (Q4_K_M, Q8_0, BF16, ...) out of a GGUF filename.""" - if not filename: - return None - from pathlib import Path +def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]: + """Approximate the RESIDENT device size of a GGUF transformer loaded through + diffusers' ``GGUFQuantizationConfig``. - stem = Path(filename).name - if stem.lower().endswith(".gguf"): - stem = stem[:-5] - parts = [p.upper() for p in stem.replace("-", "_").split("_") if p] - for index, part in enumerate(parts): - if part in ("BF16", "F16", "FP16", "FP8", "Q8", "Q6", "Q5", "Q4", "Q3", "Q2"): - suffix = parts[index + 1 :] - # Quant names carry either a K-family suffix (Q4_K_M) or a legacy - # numeric one (Q8_0, Q5_1); keep up to two suffix tokens. - if suffix and suffix[0] in ("K", "M", "S", "L", "XS", "XXS", "0", "1"): - return "_".join([part] + suffix[:2]) - return part - if part.startswith("IQ") or part.startswith("UD"): - return "_".join(parts[index : index + 3]) - return None + The weights stay PACKED on the device as quantised bytes (``GGUFParameter`` / + uint8); ``GGUFLinear.forward`` dequantises each weight to the bf16 compute dtype + transiently for its matmul and frees it immediately, so the persistent footprint + is ~= the on-disk tensor size, NOT the unpacked bf16 size. Measured on + Z-Image-Turbo: Q2_K 3.64 GiB -> 3.68 GiB, Q8_0 7.22 GiB -> 7.25 GiB resident. + The transient per-op dequant is covered by the separate runtime headroom. - -def estimate_gguf_dense_mib(storage_mib: Optional[int], quant: Optional[str]) -> Optional[int]: - """Approximate the dequantised (device) size of a GGUF from its on-disk size - and quant label. The compute dtype is bf16/fp16, so a 4-bit file roughly - quadruples once unpacked; higher-bit quants expand less.""" + (The prior per-quant expansion assumed a full unpack that never happens on this + path; it over-estimated e.g. Q2 ~7.6x, forcing needless offload.)""" if storage_mib is None: return None - q = (quant or "").upper() - if any(t in q for t in ("BF16", "F16", "FP16")): - return storage_mib - if "FP8" in q or "Q8" in q: - return int(storage_mib * 2.0) - if "Q6" in q: - return int(storage_mib * 2.8) - if "Q5" in q: - return int(storage_mib * 3.3) - if "Q4" in q or "IQ4" in q or "UD" in q: - return int(storage_mib * 4.0) - if "Q3" in q or "IQ3" in q: - return int(storage_mib * 5.3) - if "Q2" in q or "Q1" in q or "IQ2" in q or "IQ1" in q: - return int(storage_mib * 8.0) - return int(storage_mib * 4.0) # unknown: assume 4-bit-ish + return int(storage_mib * 1.05) # small margin for allocator + bf16 norms/biases def estimate_image_runtime_mib( diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index fba962f3ba..7c4a9018d8 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -148,11 +148,17 @@ def apply_speed_optims( family: Any, speed_mode: str = SPEED_OFF, cache_active: bool = False, + offload_active: bool = False, logger: Any = None, ) -> dict[str, bool]: """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline, BEFORE placement / offload. Returns which optimisations actually engaged. Every - step is best-effort: a pipeline that doesn't support one is simply skipped.""" + step is best-effort: a pipeline that doesn't support one is simply skipped. + + ``offload_active`` is the planned offload policy != none: group/model/sequential + offloading installs ``@torch.compiler.disable``d onload hooks, so the compile must + drop ``fullgraph`` (same reason as an active step cache) or it crashes at the first + denoise step.""" applied = { "channels_last": False, "cudnn_benchmark": False, @@ -182,7 +188,11 @@ def apply_speed_optims( # max-autotune (longer compile, autotuned kernels). if compile_eligible(target, is_gguf = is_gguf, family = family): applied["compiled"] = _compile_repeated_blocks( - pipe, logger, max_autotune = mode == SPEED_MAX, cache_active = cache_active + pipe, + logger, + max_autotune = mode == SPEED_MAX, + cache_active = cache_active, + offload_active = offload_active, ) if mode == SPEED_MAX: @@ -213,6 +223,7 @@ def _compile_repeated_blocks( *, max_autotune: bool = False, cache_active: bool = False, + offload_active: bool = False, ) -> bool: transformer = getattr(pipe, "transformer", None) fn = getattr(transformer, "compile_repeated_blocks", None) @@ -225,14 +236,33 @@ def _compile_repeated_blocks( # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled # block because its static output buffer is overwritten across denoise steps. # - # fullgraph drops to False when a step cache is engaged: FBCache's per-step decision is - # ``@torch.compiler.disable``d, i.e. a graph break, which fullgraph=True rejects ("Skip - # inlining torch.compiler.disable()d function"). The break is cheap and the rest of the - # block still compiles. - kwargs: dict[str, Any] = {"fullgraph": not cache_active, "dynamic": not max_autotune} + # fullgraph drops to False when a step cache OR CPU offloading is engaged: both insert + # an ``@torch.compiler.disable``d function into the forward -- FBCache's per-step + # decision, and group/model/sequential offload's ``ModuleGroup.onload_`` streaming hook + # -- i.e. a graph break, which fullgraph=True rejects ("Skip inlining + # torch.compiler.disable()d function"). The break is cheap and the rest of the block + # still compiles. + kwargs: dict[str, Any] = { + "fullgraph": not (cache_active or offload_active), + "dynamic": not max_autotune, + } if max_autotune: kwargs["mode"] = "max-autotune-no-cudagraphs" try: + import torch + # Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block + # shape through compile_repeated_blocks; Z-Image needs ~11, above dynamo's default + # recompile_limit of 8. Once the limit is hit a resident load hard-errors under + # fullgraph (and an offload/cache load silently drops the overflow blocks to eager), + # so raise it well past that (64) for headroom on larger heterogeneous DiTs. This is + # diffusers' own documented fix for regional-compile recompilation (their guide bumps + # cache_size_limit). Deliberately NOT force_parameter_static_shapes=False: it doesn't + # cut the variant count here and makes each compile ~6x slower (24s -> 143s cold). + dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None) + if dynamo_cfg is not None: + for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver + if hasattr(dynamo_cfg, _limit_attr): + setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64)) fn(**kwargs) return True except Exception as exc: # noqa: BLE001 — optimisation only diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index bc09ffbd0f..cfdddb051e 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -27,9 +27,8 @@ from core.inference.diffusion_memory import ( DeviceMemory, MemoryPlan, apply_memory_plan, - estimate_gguf_dense_mib, + estimate_gguf_resident_mib, estimate_image_runtime_mib, - infer_gguf_quant_label, normalize_memory_mode, plan_diffusion_memory, snapshot_device_memory, @@ -70,29 +69,15 @@ def test_normalize_memory_mode_accepts_and_rejects(): # ── filename / size estimates ───────────────────────────────────────────────── -@pytest.mark.parametrize( - "filename,expected", - [ - ("z-image-turbo-Q4_K_M.gguf", "Q4_K_M"), - ("flux1-dev-Q8_0.gguf", "Q8_0"), - ("model-BF16.gguf", "BF16"), - ("qwen-image-IQ4_XS.gguf", "IQ4_XS"), - ("no-quant-here.gguf", None), - (None, None), - ], -) -def test_infer_gguf_quant_label(filename, expected): - assert infer_gguf_quant_label(filename) == expected - - -def test_estimate_gguf_dense_mib_expansion(): - # 4-bit roughly quadruples once dequantised to bf16; F16 is already dense. - assert estimate_gguf_dense_mib(1000, "Q4_K_M") == 4000 - assert estimate_gguf_dense_mib(1000, "Q8_0") == 2000 - assert estimate_gguf_dense_mib(1000, "BF16") == 1000 - assert estimate_gguf_dense_mib(None, "Q4_K_M") is None - # Unknown quant falls back to the conservative 4-bit-ish factor. - assert estimate_gguf_dense_mib(1000, None) == 4000 +def test_estimate_gguf_resident_mib_matches_packed_size(): + # GGUF weights stay packed (uint8) on-device; diffusers dequantises per-matmul + # transiently, so the resident footprint ~= the on-disk size regardless of quant + # level (measured on Z-Image-Turbo: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). A + # small margin covers allocator overhead. The prior per-quant expansion over- + # estimated (Q2 ~7.6x) and forced needless offload on a roomy card. + assert estimate_gguf_resident_mib(1000) == 1050 + assert estimate_gguf_resident_mib(7220) == 7581 + assert estimate_gguf_resident_mib(None) is None def test_estimate_image_runtime_scales_with_pixels_and_family(): diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index b0ca10af4c..1f64b31b7a 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -216,6 +216,24 @@ def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch): assert applied["tf32"] is False and applied["fused_qkv"] is False +def test_offload_active_drops_fullgraph(monkeypatch): + # Group/model/sequential offload installs a torch.compiler.disable'd onload hook; + # compiling with fullgraph=True then crashes at the first denoise step. Same reason + # as an active step cache -> fullgraph must drop to False when offload is planned. + _stub_torch(monkeypatch) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, + _target(), + is_gguf = True, + family = _family(), + speed_mode = SPEED_DEFAULT, + offload_active = True, + ) + assert applied["compiled"] is True + assert pipe.compile_kwargs["fullgraph"] is False + + def test_speed_default_compiles_gguf(monkeypatch): _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True)