ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix diffusion GGUF memory over-estimate and torch.compile crashes

Browse source 
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.
This commit is contained in:
oobabooga 2026-07-01 18:01:12 -03:00
commit 8d16ef977b
5 changed files with 87 additions and 84 deletions
Download patch file Download diff file Expand all files Collapse all files

17
studio/backend/core/inference/diffusion.py
View file

@ -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,

53
studio/backend/core/inference/diffusion_memory.py
View file

@ -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(

44
studio/backend/core/inference/diffusion_speed.py
View file

@ -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

35
studio/backend/tests/test_diffusion_memory.py
View file

@ -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():

18
studio/backend/tests/test_diffusion_speed.py
View file

@ -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)