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

Add the diffusion auto-policy layer: per-family footprint estimates and the dense-quant re-plan

Browse source 
The loader used to plan memory from the GGUF file size and only offer the dense
transformer-quant fast path when that plan was already resident, so on a card where
the GGUF forced offload the int8/fp8 build (roughly half the bf16 bytes, or exactly
the quantised size when a pre-quantized checkpoint exists) was never attempted.

diffusion_auto_policy.py is a pure decision layer: a bf16-resident component table
per family (transformer / text encoders / VAE, with base-repo overrides for the
multi-size families), per-scheme size factors with separate steady and transient
(build peak) numbers, and resolve_dense_quant_candidate which the loader now uses to
re-plan memory against the candidate artifact before settling for offload. The
engaged plan is adopted only when the dense build succeeds; the GGUF fallback keeps
its own plan.

Status now carries a resolved provenance record per Advanced control (value, source
auto or explicit, reason) so the UI can label backend decisions.
This commit is contained in:
Daniel Han 2026-07-04 06:49:49 +00:00
commit a8e708f66e
3 changed files with 532 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -59,6 +59,7 @@ from .diffusion_speed import (
SPEED_OFF,
apply_speed_optims,
compile_eligible,
normalize_speed_mode,
resolve_speed_mode,
restore_backend_flags,
snapshot_backend_flags,
@ -75,6 +76,7 @@ from .diffusion_prequant import (
load_prequantized_transformer,
resolve_prequant_source,
)
from .diffusion_auto_policy import build_resolved_record, resolve_dense_quant_candidate
from .diffusion_transformer_quant import (
DEFAULT_MIN_LINEAR_FEATURES,
dense_transformer_supported,
@ -283,6 +285,10 @@ class _LoadState:
compile_cache_ctx: Any = None
# Token kept so LoRA adapters selected at generate time can be fetched from the Hub.
hf_token: Optional[str] = None
# Per-control provenance from the auto-policy: {control: {value, source, reason}}.
# source is "auto" when the backend decided (request left unset / "auto") and
# "explicit" when the caller pinned the value. Surfaced via status for the UI badges.
resolved: Optional[dict] = None
@dataclass
@ -971,11 +977,48 @@ class DiffusionBackend:
# pipeline) do not have.
pipe = None
transformer_quant_engaged = None
quant_plan = None
if (
kind == "gguf"
and normalize_transformer_quant(transformer_quant) is not None
and dense_transformer_supported(target)
and plan.offload_policy == OFFLOAD_NONE
and plan.offload_policy != OFFLOAD_NONE
):
# The GGUF-size plan picked offload, but the dense-quant artifact has
# a DIFFERENT footprint: int8/fp8 weights are ~half the bf16 bytes, and
# a pre-quantized checkpoint never materialises dense bf16 at all. Ask
# the auto-policy for the candidate's estimate and re-plan against it:
# a resident quantised build beats an offloaded GGUF on speed AND
# quality, so it must be attempted before settling for offload.
candidate = resolve_dense_quant_candidate(
fam = fam,
target = target,
requested = transformer_quant,
base_repo = base,
prequant_path = transformer_prequant_path,
logger = logger,
)
if candidate is not None:
replanned = self._plan_memory(
target,
single_file_path,
base,
fam,
memory_mode,
cpu_offload,
kind = kind,
repo_id = repo_id,
transformer_resident_override_mib = (
candidate.transient_transformer_mib
),
)
if replanned.offload_policy == OFFLOAD_NONE:
quant_plan = replanned
if (
kind == "gguf"
and normalize_transformer_quant(transformer_quant) is not None
and dense_transformer_supported(target)
and (plan.offload_policy == OFFLOAD_NONE or quant_plan is not None)
):
try:
pipe, transformer_quant_engaged = self._load_dense_quant_pipeline(
@ -1004,6 +1047,10 @@ class DiffusionBackend:
# GGUF build (the OOM-fallback path this cleanup exists for).
del exc
clear_gpu_cache()
if transformer_quant_engaged is not None and quant_plan is not None:
# The re-planned resident placement is the one the engaged dense build
# actually uses; the GGUF-size plan stays in force for the fallback.
plan = quant_plan
if pipe is None:
if kind == "pipeline":
@ -1228,6 +1275,57 @@ class DiffusionBackend:
pipe, plan, device = device, logger = logger
)
# Per-control provenance for status: what engaged and who decided it
# (the caller, or this backend's auto resolution). cpu_offload=False is
# the unset default, so only True counts as an explicit request.
resolved = build_resolved_record(
{
"speed_mode": (
speed_mode,
effective_speed,
"quantized transformer requires compile"
if transformer_quant_engaged is not None
and normalize_speed_mode(speed_mode) in (None, SPEED_OFF)
else "per-kind default"
if speed_mode is None
else "requested",
),
"transformer_quant": (
transformer_quant,
transformer_quant_engaged or "off",
"not engaged (GGUF transformer loaded)"
if transformer_quant_engaged is None
else "re-planned resident for the quantised artifact"
if quant_plan is not None
else "engaged on the dense fast path",
),
"attention_backend": (
attention_backend,
attention_engaged or "native",
"cuDNN fused attention upgrade"
if attention_engaged and attention_backend is None
else "diffusers default"
if attention_engaged is None
else "requested",
),
"memory_mode": (
memory_mode,
effective_policy,
"planned from measured free VRAM vs estimated footprint",
),
"transformer_cache": (
transformer_cache,
cache_engaged or "off",
"off by default" if transformer_cache is None else "requested",
),
"cpu_offload": (
True if cpu_offload else None,
effective_policy != OFFLOAD_NONE,
"legacy flag" if cpu_offload else "from the memory plan",
),
}
)
self._state = _LoadState(
pipe = pipe,
family = fam,
@ -1250,6 +1348,7 @@ class DiffusionBackend:
eager_patched = eager_patched,
compile_cache_ctx = compile_ctx,
hf_token = hf_token,
resolved = resolved,
)
state_committed = True
finally:
@ -1383,6 +1482,7 @@ class DiffusionBackend:
*,
kind: str = "gguf",
repo_id: Optional[str] = None,
transformer_resident_override_mib: Optional[int] = None,
):
"""Build the memory plan for this load: snapshot free device memory and
estimate the model's resident footprint, then let the planner pick an
@ -1392,7 +1492,10 @@ class DiffusionBackend:
The size estimate is per-kind: diffusers keeps GGUF weights packed (per-matmul
transient dequant), so a GGUF loads near its on-disk size; a safetensors
single-file loads near its on-disk size (it carries its dtype); and a full
pipeline is one cached download (transformer + companions), already compressed."""
pipeline is one cached download (transformer + companions), already compressed.
``transformer_resident_override_mib`` replaces the file-size transformer estimate
when the loader is planning for a DIFFERENT artifact than the file on disk (the
dense transformer-quant candidate, whose footprint the auto-policy estimates)."""
device_memory = snapshot_device_memory(target)
if kind == "pipeline":
# The whole repo (transformer + companions) is one cached download; the
@ -1408,7 +1511,12 @@ class DiffusionBackend:
model_dense_mib = estimate_safetensors_dense_mib(cached_mib)
companion_mib = None
else:
if kind == "single_file":
if transformer_resident_override_mib is not None:
# Planning for a different artifact than the file on disk (the dense
# transformer-quant candidate): the auto-policy's estimate replaces the
# file-size derivation; companions below stay measured from the cache.
transformer_resident = transformer_resident_override_mib
elif kind == "single_file":
# Safetensors single-file: no dequant expansion (it carries its dtype).
transformer_resident = estimate_safetensors_dense_mib(
file_size_mib(single_file_path)
@ -2133,6 +2241,7 @@ class DiffusionBackend:
"workflows": [],
"supports_lora": False,
"supports_controlnet": False,
"resolved": None,
}
from core.inference import diffusion_controlnet, diffusion_lora
@ -2154,6 +2263,7 @@ class DiffusionBackend:
"transformer_quant": state.transformer_quant,
"attention_backend": state.attention_backend,
"transformer_cache": state.transformer_cache,
"resolved": state.resolved,
# Image-conditioned workflows the loaded family supports, so the UI can gate
# its tabs. txt2img is always available on the diffusers engine.
"workflows": _family_workflows(state.family),

206
studio/backend/core/inference/diffusion_auto_policy.py Normal file
View file

@ -0,0 +1,206 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Hardware-aware auto-policy for the diffusion loader: a pure decision layer.
The loader historically resolved each Advanced control on its own, and -- critically --
planned memory from the GGUF file size BEFORE the dense transformer-quant fast path was
considered. That ordering hid the fast path exactly where it matters: on a card where the
GGUF-size plan picks offload, a dense int8/fp8 transformer (roughly half the bf16 bytes)
would often still fit fully resident and beat the offloaded GGUF on every axis.
This module supplies the two pieces the loader needs to fix that, without moving any of
the existing per-control executors:
* a per-family bf16 component-size table (transformer / text encoders / VAE) with
per-scheme scaling, so the candidate artifact's footprint can be estimated BEFORE
anything is downloaded or materialised; and
* ``resolve_dense_quant_candidate``, which turns a request + device into a concrete
(scheme, steady, transient) estimate the loader re-plans memory against.
It also builds the ``resolved`` record surfaced through status: for every Advanced
control, the engaged value plus whether it came from the user (explicit) or the policy
(auto), with a short reason. Pure by design: no torch import at module import time, so
the decision logic unit-tests on CPU-only hosts.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Any, Optional
_MIB_PER_GB = 1000.0 ** 3 / (1024.0 * 1024.0) # component sizes below are decimal GB
# Steady-state size of a torchao-quantised transformer relative to its bf16 weights:
# int8 / fp8 store one byte per param plus per-row scales (~0.52x) with a little slack
# for non-quantised modules (norms, embeddings, proj_out stay bf16); nvfp4 packs two
# params per byte plus block scales. Measured on the live int8/fp8 loads this session.
_QUANT_STEADY_FACTOR: dict[str, float] = {
"int8": 0.55,
"fp8": 0.55,
"mxfp8": 0.58,
"nvfp4": 0.33,
}
# bf16-RESIDENT component sizes in decimal GB: (transformer, text encoders, VAE).
# These are what the components occupy on device after the loader's dtype cast, NOT the
# repo download size (Z-Image ships its transformer in fp32: 24.6 GB of shards that load
# as 12.3 GB bf16). Sourced from the HF sibling metadata of each family's base repo with
# precision duplicates removed, cross-checked against the training-side dense_bf16_gb
# table and the measured loads in this repo's GPU verification runs.
_FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = {
"flux.1": (23.8, 9.8, 0.2),
"flux.1-kontext": (23.8, 9.8, 0.2),
"flux.2-klein": (7.8, 8.0, 0.2),
"flux.2-dev": (64.5, 48.0, 0.4),
"qwen-image": (40.9, 16.6, 0.3),
"qwen-image-edit": (40.9, 16.6, 0.3),
"z-image": (12.3, 8.0, 0.2),
"krea-2": (26.3, 8.9, 0.5),
}
# Base-repo overrides for families whose picker offers multiple sizes under one family
# entry (the table above carries the family default base).
_BASE_REPO_BF16_GB: dict[str, tuple[float, float, float]] = {
"black-forest-labs/FLUX.2-klein-9B": (18.2, 16.4, 0.2),
}
def family_bf16_components_gb(
fam: Any, base_repo: Optional[str] = None
) -> Optional[tuple[float, float, float]]:
"""(transformer, text encoders, VAE) bf16-resident sizes in GB, or None when the
family is not in the table (callers must then fall back to file-size estimates)."""
if base_repo:
override = _BASE_REPO_BF16_GB.get(base_repo)
if override is not None:
return override
name = getattr(fam, "name", None)
return _FAMILY_BF16_GB.get(name) if name else None
@dataclass(frozen = True)
class DenseQuantEstimate:
"""Footprint estimate for one dense transformer-quant candidate.
``transient_transformer_mib`` is the build peak: the dense bf16 transformer when
quantising on the fly, or the quantised size itself when a pre-quantized checkpoint
is available (it loads via the meta device, so dense bf16 never lands on the GPU).
``steady_transformer_mib`` is what stays resident for generation."""
scheme: str
steady_transformer_mib: int
transient_transformer_mib: int
companions_mib: int
prequant: bool
@property
def transient_total_mib(self) -> int:
return self.transient_transformer_mib + self.companions_mib
@property
def steady_total_mib(self) -> int:
return self.steady_transformer_mib + self.companions_mib
def estimate_dense_quant(
fam: Any,
scheme: str,
*,
base_repo: Optional[str] = None,
prequant_available: bool = False,
) -> Optional[DenseQuantEstimate]:
"""Estimate the candidate's footprint from the family table, or None when the
family (or scheme factor) is unknown."""
components = family_bf16_components_gb(fam, base_repo)
factor = _QUANT_STEADY_FACTOR.get(scheme)
if components is None or factor is None:
return None
transformer_gb, text_encoders_gb, vae_gb = components
steady = int(transformer_gb * factor * _MIB_PER_GB)
transient = steady if prequant_available else int(transformer_gb * _MIB_PER_GB)
companions = int((text_encoders_gb + vae_gb) * _MIB_PER_GB)
return DenseQuantEstimate(
scheme = scheme,
steady_transformer_mib = steady,
transient_transformer_mib = transient,
companions_mib = companions,
prequant = prequant_available,
)
def resolve_dense_quant_candidate(
*,
fam: Any,
target: Any,
requested: Optional[str],
base_repo: Optional[str] = None,
prequant_path: Optional[str] = None,
logger: Optional[logging.Logger] = None,
) -> Optional[DenseQuantEstimate]:
"""The dense-quant candidate the loader should re-plan memory against, or None.
None means "no basis to re-plan": the request is off, the device cannot run the
dense path, no scheme resolves, or the family has no size entry. The loader then
keeps today's behaviour (fast path only when the GGUF-size plan is already
resident), so unlisted families see no change."""
from .diffusion_transformer_quant import (
dense_transformer_supported,
normalize_transformer_quant,
select_transformer_quant_scheme,
)
if normalize_transformer_quant(requested) is None:
return None
if not dense_transformer_supported(target):
return None
scheme = select_transformer_quant_scheme(target, requested)
if scheme is None:
return None
prequant_available = False
try:
from .diffusion_prequant import resolve_prequant_source
prequant_available = (
resolve_prequant_source(fam, scheme, path_override = prequant_path) is not None
)
except Exception: # noqa: BLE001 -- prequant probing must never sink the candidate
prequant_available = False
estimate = estimate_dense_quant(
fam, scheme, base_repo = base_repo, prequant_available = prequant_available
)
if estimate is not None and logger is not None:
logger.info(
"diffusion.auto_policy: dense %s candidate steady=%d MiB transient=%d MiB "
"companions=%d MiB prequant=%s",
scheme,
estimate.steady_transformer_mib,
estimate.transient_transformer_mib,
estimate.companions_mib,
prequant_available,
)
return estimate
# ── resolved-record (status surface) ─────────────────────────────────────────
def build_resolved_record(
controls: dict[str, tuple[Optional[Any], Any, str]],
) -> dict[str, dict[str, Any]]:
"""The per-control ``resolved`` record for status: engaged value + provenance.
``controls`` maps a control name to ``(explicit, engaged, reason)`` where
``explicit`` is the raw request value (None / "" / "auto" meaning the caller left
the decision to the backend) and ``engaged`` is what actually applied. The record
is what the frontend renders as an "Auto: X" badge next to each Advanced row."""
record: dict[str, dict[str, Any]] = {}
for name, (explicit, engaged, reason) in controls.items():
left_to_backend = explicit is None or (
isinstance(explicit, str) and explicit.strip().lower() in ("", "auto")
)
record[name] = {
"value": engaged,
"source": "auto" if left_to_backend else "explicit",
"reason": reason,
}
return record

213
studio/backend/tests/test_diffusion_auto_policy.py Normal file
View file

@ -0,0 +1,213 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""CPU-only unit tests for the diffusion auto-policy decision layer.
Covers the per-family footprint estimator (bf16-resident component sizes x per-scheme
factors, transient vs steady, base-repo overrides), the dense-quant candidate resolution
(with the quant selector / prequant probe monkeypatched, no torch), and the resolved
provenance record. The loader-side ordering fix is exercised through the planner: the
regression case is a GGUF whose file-size plan forces offload while the candidate's
estimate fits resident."""
from __future__ import annotations
from types import SimpleNamespace
import core.inference.diffusion_auto_policy as ap
from core.inference.diffusion_auto_policy import (
DenseQuantEstimate,
build_resolved_record,
estimate_dense_quant,
family_bf16_components_gb,
resolve_dense_quant_candidate,
)
from core.inference.diffusion_memory import (
OFFLOAD_NONE,
DeviceMemory,
plan_diffusion_memory,
)
def _fam(name = "z-image"):
return SimpleNamespace(name = name)
# ── the per-family table ──────────────────────────────────────────────────────
def test_family_table_covers_the_dit_families():
for name in ("flux.1", "flux.2-klein", "flux.2-dev", "qwen-image", "z-image", "krea-2"):
comps = family_bf16_components_gb(_fam(name))
assert comps is not None, f"{name} missing from the bf16 component table"
transformer, text_encoders, vae = comps
assert transformer > 1.0 and text_encoders > 0.0 and vae > 0.0
def test_family_table_unknown_family_returns_none():
assert family_bf16_components_gb(_fam("not-a-family")) is None
def test_base_repo_override_wins_over_the_family_default():
# flux.2-klein's family default is the 4B base; loading the 9B GGUF passes the 9B
# base repo, whose transformer is more than twice the size.
default = family_bf16_components_gb(_fam("flux.2-klein"))
nine_b = family_bf16_components_gb(
_fam("flux.2-klein"), base_repo = "black-forest-labs/FLUX.2-klein-9B"
)
assert nine_b is not None and default is not None
assert nine_b[0] > 2 * default[0]
# ── the estimator ─────────────────────────────────────────────────────────────
def test_estimate_int8_steady_is_roughly_half_bf16():
est = estimate_dense_quant(_fam("z-image"), "int8")
assert est is not None
bf16_mib = 12.3 * ap._MIB_PER_GB
assert 0.5 * bf16_mib < est.steady_transformer_mib < 0.6 * bf16_mib
# On-the-fly quantisation transiently materialises the dense bf16 transformer.
assert est.transient_transformer_mib == int(bf16_mib)
assert est.prequant is False
def test_estimate_prequant_transient_equals_steady():
# A pre-quantized checkpoint loads via the meta device: dense bf16 never lands on
# the GPU, so the build peak IS the quantised size.
est = estimate_dense_quant(_fam("z-image"), "int8", prequant_available = True)
assert est is not None
assert est.transient_transformer_mib == est.steady_transformer_mib
assert est.prequant is True
def test_estimate_nvfp4_is_smaller_than_int8():
int8 = estimate_dense_quant(_fam("flux.1"), "int8")
nvfp4 = estimate_dense_quant(_fam("flux.1"), "nvfp4")
assert int8 is not None and nvfp4 is not None
assert nvfp4.steady_transformer_mib < int8.steady_transformer_mib
def test_estimate_unknown_family_or_scheme_returns_none():
assert estimate_dense_quant(_fam("not-a-family"), "int8") is None
assert estimate_dense_quant(_fam("z-image"), "q4_k") is None
# ── candidate resolution (selector + prequant probe stubbed) ─────────────────
def _patch_selector(monkeypatch, *, supported = True, scheme = "int8", prequant = None):
import core.inference.diffusion_transformer_quant as tq
monkeypatch.setattr(tq, "dense_transformer_supported", lambda target: supported)
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, req: scheme)
import core.inference.diffusion_prequant as pq
monkeypatch.setattr(
pq, "resolve_prequant_source", lambda fam, s, path_override = None: prequant
)
def test_candidate_resolves_for_a_supported_request(monkeypatch):
_patch_selector(monkeypatch, scheme = "int8")
est = resolve_dense_quant_candidate(
fam = _fam("z-image"), target = object(), requested = "auto"
)
assert isinstance(est, DenseQuantEstimate)
assert est.scheme == "int8"
assert est.transient_transformer_mib > est.steady_transformer_mib
def test_candidate_none_when_request_is_off(monkeypatch):
_patch_selector(monkeypatch)
for off in (None, "", "none", "off"):
assert (
resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = off)
is None
)
def test_candidate_none_when_device_unsupported(monkeypatch):
_patch_selector(monkeypatch, supported = False)
assert (
resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = "auto")
is None
)
def test_candidate_none_when_no_scheme_resolves(monkeypatch):
_patch_selector(monkeypatch, scheme = None)
assert (
resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = "auto")
is None
)
def test_candidate_none_for_an_unlisted_family(monkeypatch):
# No size entry -> no basis to re-plan; the loader keeps today's resident-only gate.
_patch_selector(monkeypatch)
assert (
resolve_dense_quant_candidate(
fam = _fam("not-a-family"), target = object(), requested = "auto"
)
is None
)
def test_candidate_uses_prequant_transient_when_available(monkeypatch):
_patch_selector(monkeypatch, prequant = object())
est = resolve_dense_quant_candidate(
fam = _fam("z-image"), target = object(), requested = "int8"
)
assert est is not None and est.prequant is True
assert est.transient_transformer_mib == est.steady_transformer_mib
# ── the ordering-fix regression, at the planner level ─────────────────────────
def _cuda_target():
return SimpleNamespace(device = "cuda", supports_model_cpu_offload = True)
def test_quant_candidate_fits_resident_where_gguf_plan_offloads():
# The ordering-fix mechanism, on a 32 GiB consumer card (RTX 5090 class): the user
# picked a LARGE GGUF (the BF16 file), so the file-size plan forces offload -- but
# the dense-quant candidate is far smaller (int8 prequant of z-image: the transient
# IS the quantised size), and re-planning against the candidate keeps everything
# resident. Before the fix the loader never attempted the fast path here.
memory = DeviceMemory("cuda", "cuda", "discrete_vram", 30000, 32768)
z_bf16_gguf_mib = int(12.3 * ap._MIB_PER_GB * 1.05) # BF16 GGUF resident estimate
companions_mib = 2600 # fp8-quantised text encoders + VAE
gguf_plan = plan_diffusion_memory(
target = _cuda_target(),
device_memory = memory,
model_dense_mib = z_bf16_gguf_mib + companions_mib,
companion_dense_mib = companions_mib,
runtime_headroom_mib = 6963,
)
assert gguf_plan.offload_policy != OFFLOAD_NONE
est = estimate_dense_quant(_fam("z-image"), "int8", prequant_available = True)
assert est is not None
assert est.transient_transformer_mib < z_bf16_gguf_mib / 1.8
quant_plan = plan_diffusion_memory(
target = _cuda_target(),
device_memory = memory,
model_dense_mib = est.transient_transformer_mib + companions_mib,
companion_dense_mib = companions_mib,
runtime_headroom_mib = 6963,
)
assert quant_plan.offload_policy == OFFLOAD_NONE
# ── the resolved provenance record ────────────────────────────────────────────
def test_resolved_record_marks_auto_and_explicit():
record = build_resolved_record(
{
"speed_mode": (None, "default", "per-kind default"),
"transformer_quant": ("auto", "fp8", "auto ladder"),
"attention_backend": ("cudnn", "_native_cudnn", "requested"),
"memory_mode": ("", "none", "planned"),
"cpu_offload": (True, True, "legacy flag"),
}
)
assert record["speed_mode"]["source"] == "auto"
assert record["transformer_quant"]["source"] == "auto" # "auto" delegates to backend
assert record["attention_backend"]["source"] == "explicit"
assert record["memory_mode"]["source"] == "auto" # blank string delegates
assert record["cpu_offload"]["source"] == "explicit"
assert record["transformer_quant"]["value"] == "fp8"
assert all("reason" in v for v in record.values())