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

Merge diffusion-auto-install: Dtype defaults to auto with disk gate

Browse source
This commit is contained in:
Daniel Han 2026-07-04 09:44:58 +00:00
commit 45fe22d8eb
6 changed files with 132 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -86,6 +86,7 @@ from .diffusion_prequant import (
)
from .diffusion_auto_policy import build_resolved_record, resolve_dense_quant_candidate
from .diffusion_transformer_quant import (
TQ_AUTO,
DEFAULT_MIN_LINEAR_FEATURES,
dense_transformer_supported,
normalize_transformer_quant,
@ -439,7 +440,12 @@ class DiffusionBackend:
unload/cancellation cannot preempt the download. Mirrors the dense-path
gates in ``load_pipeline``: quant requested and supported for this device,
and no pre-quantized checkpoint that would shortcut the dense build."""
mode = normalize_transformer_quant(kwargs.get("transformer_quant"))
raw = kwargs.get("transformer_quant")
# Unset defaults to the hardware ladder (mirrors load_pipeline's tri-state).
if raw is None or str(raw).strip().lower() in ("", "auto"):
mode = TQ_AUTO
else:
mode = normalize_transformer_quant(raw)
if mode is None:
return False
try:
@ -984,7 +990,16 @@ class DiffusionBackend:
repo_id = repo_id,
)
# Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it
# Dtype tri-state: an UNSET request (or "auto") hands the decision to
# the hardware ladder -- on a dense-capable GPU the quantised build
# (int8 minimum, fp8 on data-center silicon) beats running the GGUF
# as-is, so auto is the DEFAULT. An explicit "none"/"off" pins
# GGUF-as-is and an explicit scheme pins that scheme. The overwritten
# "auto" still records source=auto in the resolved provenance.
if transformer_quant is None or str(transformer_quant).strip().lower() in ("", "auto"):
transformer_quant = TQ_AUTO
# Default-on fast path: load the DENSE bf16 transformer and torchao-quantise it
# (int8 / fp8 / fp4 tensor cores), which beats GGUF's bf16-rate per-matmul
# dequant on both speed and quality, at the cost of a higher-memory dense
# load. Gated on CUDA + bf16 + a resident fit; ANY failure (unsupported arch

43
studio/backend/core/inference/diffusion_auto_policy.py
View file

@ -27,6 +27,7 @@ the decision logic unit-tests on CPU-only hosts.
from __future__ import annotations
import logging
import os
from dataclasses import dataclass
from typing import Any, Optional
@ -130,6 +131,25 @@ def estimate_dense_quant(
)
def _hf_cache_free_mib() -> Optional[int]:
"""Free MiB on the filesystem holding the HF model cache (None when unprobeable)."""
try:
import shutil
try:
from huggingface_hub.constants import HF_HUB_CACHE as cache_dir
except Exception: # noqa: BLE001 -- hub missing/old: probe the conventional path
cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
probe = str(cache_dir)
while probe and not os.path.isdir(probe):
parent = os.path.dirname(probe)
if parent == probe:
break
probe = parent
return int(shutil.disk_usage(probe).free // (1024 * 1024))
except Exception: # noqa: BLE001 -- disk probing must never sink the candidate
return None
def resolve_dense_quant_candidate(
*,
fam: Any,
@ -179,6 +199,29 @@ def resolve_dense_quant_candidate(
estimate.companions_mib,
prequant_available,
)
if estimate is not None:
# The dense path may DOWNLOAD the artifact (the multi-GB bf16 base
# transformer, or the prequant checkpoint) into the HF cache; with Dtype
# defaulting to auto this must never wedge a nearly-full disk. An
# already-cached model re-download is a no-op, so the gate can only
# false-positive on a disk that is already critically full -- where
# falling back to the GGUF build is the right call anyway.
needed_mib = (
estimate.steady_transformer_mib
if estimate.prequant
else estimate.transient_transformer_mib
)
free_mib = _hf_cache_free_mib()
if free_mib is not None and free_mib < needed_mib + 10 * 1024:
if logger is not None:
logger.info(
"diffusion.auto_policy: skipping dense %s (~%d MiB download, "
"only %d MiB free in the model cache)",
scheme,
needed_mib,
free_mib,
)
return None
return estimate

15
studio/backend/models/inference.py
View file

@ -1750,14 +1750,15 @@ class DiffusionLoadRequest(BaseModel):
"memory-vs-quality tradeoff (shifts fine detail), not free; "
"pairs well with balanced mode.",
)
transformer_quant: Optional[Literal["auto", "int8", "fp8", "nvfp4", "mxfp8"]] = Field(
transformer_quant: Optional[Literal["auto", "none", "off", "int8", "fp8", "nvfp4", "mxfp8"]] = Field(
None,
description = "Opt-in fast transformer: load the DENSE bf16 transformer instead "
"of the GGUF and torchao-quantise it onto the low-precision tensor "
"cores (faster than GGUF's bf16-rate dequant, at higher VRAM). auto "
"picks the best for the GPU (Blackwell nvfp4/mxfp8, Ada/Hopper fp8, "
"Ampere int8); an explicit scheme forces it. Needs CUDA + bf16 + room "
"for the dense load; falls back to GGUF otherwise.",
description = "Transformer compute dtype. UNSET or auto (the default) picks the "
"fastest precision the hardware supports: the DENSE bf16 transformer "
"is loaded instead of the GGUF and torchao-quantised onto the "
"low-precision tensor cores (data-center fp8, consumer/Ampere int8), "
"falling back to the GGUF when the device, VRAM or disk cannot take "
"it. none/off pins running the GGUF as-is; an explicit scheme forces "
"that scheme. Dense path needs CUDA + bf16.",
)
transformer_quant_fast_accum: Optional[bool] = Field(
None,

25
studio/backend/tests/test_diffusion_auto_policy.py
View file

@ -132,6 +132,31 @@ def test_candidate_none_when_no_scheme_resolves(monkeypatch):
assert resolve_dense_quant_candidate(fam = _fam(), target = object(), requested = "auto") is None
def test_candidate_disk_gate_skips_when_cache_disk_low(monkeypatch):
# The dense artifact may be a multi-GB download; a nearly-full model-cache disk
# drops the candidate (the loader then keeps the GGUF build).
import core.inference.diffusion_auto_policy as ap
_patch_selector(monkeypatch, scheme = "int8")
monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 1024)
assert (
resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto")
is None
)
def test_candidate_disk_gate_unprobeable_disk_passes(monkeypatch):
# Disk probing must never sink the candidate: unprobeable (None) passes through.
import core.inference.diffusion_auto_policy as ap
_patch_selector(monkeypatch, scheme = "int8")
monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: None)
est = resolve_dense_quant_candidate(
fam = _fam("z-image"), target = object(), requested = "auto"
)
assert isinstance(est, DenseQuantEstimate)
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)

41
studio/backend/tests/test_diffusion_backend.py
View file

@ -1786,19 +1786,44 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"):
return calls
def test_default_load_skips_dense_quant_path(fake_runtime, tmp_path, monkeypatch):
# With no transformer_quant flag the GGUF path is taken and the dense gate is
# never even consulted (short-circuit), so the default cannot regress.
def test_default_load_autos_dense_gate_and_falls_back(fake_runtime, tmp_path, monkeypatch):
# UNSET Dtype defaults to the hardware ladder: the dense gate IS consulted, and a
# device without dense support (this fake runtime) falls back to the GGUF build.
from core.inference import diffusion as dmod
consulted = {"n": 0}
def _supported(*a, **k):
consulted["n"] += 1
return False
monkeypatch.setattr(dmod, "dense_transformer_supported", _supported)
(tmp_path / "m.gguf").write_bytes(b"x")
backend = DiffusionBackend()
status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image")
assert consulted["n"] >= 1
assert status["transformer_quant"] is None
assert _FakeTransformer.last["path"] # GGUF from_single_file was used
def test_explicit_off_load_skips_dense_quant_path(fake_runtime, tmp_path, monkeypatch):
# An EXPLICIT "none" pins running the GGUF as-is: the dense gate is never even
# consulted (short-circuit), so the pinned-off contract cannot regress.
from core.inference import diffusion as dmod
monkeypatch.setattr(
dmod,
"dense_transformer_supported",
lambda *a, **k: pytest.fail("dense path must not run without the flag"),
lambda *a, **k: pytest.fail("dense path must not run with an explicit off"),
)
(tmp_path / "m.gguf").write_bytes(b"x")
backend = DiffusionBackend()
status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image")
status = backend.load_pipeline(
str(tmp_path),
gguf_filename = "m.gguf",
family_override = "z-image",
transformer_quant = "none",
)
assert status["transformer_quant"] is None
assert _FakeTransformer.last["path"] # GGUF from_single_file was used
@ -2025,8 +2050,10 @@ def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch):
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True
# No quant requested -> never widen.
assert backend._dense_quant_prefetch_needed(fam, {}) is False
# UNSET defaults to the hardware ladder (Dtype default-auto) -> widens too.
assert backend._dense_quant_prefetch_needed(fam, {}) is True
# An explicit off pins running the GGUF as-is -> never widen.
assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "none"}) is False
# A resolvable pre-quantized checkpoint shortcuts the dense download.
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object())
assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False

8
studio/frontend/src/features/images/images-page.tsx
View file

@ -1024,7 +1024,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
const [speedMode, setSpeedMode] = useState<"auto" | "off" | "eager" | "default" | "max">("auto");
const [transformerQuant, setTransformerQuant] = useState<
"none" | "auto" | "int8" | "fp8" | "nvfp4" | "mxfp8"
>("none");
>("auto");
const [attentionBackend, setAttentionBackend] = useState<"auto" | "native" | "cudnn" | "flash3" | "sage">(
"auto",
);
@ -1509,7 +1509,7 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
hf_token: hfApiToken(getHfToken()),
cpu_offload: cpuOffload,
speed_mode: speedMode === "auto" ? undefined : speedMode,
transformer_quant: transformerQuant === "none" ? undefined : transformerQuant,
transformer_quant: transformerQuant === "auto" ? undefined : transformerQuant,
attention_backend: attentionBackend === "auto" ? undefined : attentionBackend,
memory_mode: memoryMode === "auto" ? undefined : memoryMode,
transformer_cache: transformerCache === "off" ? undefined : transformerCache,
@ -1916,13 +1916,13 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
{!status?.loaded || status.model_kind === "gguf" ? (
<AdvancedSelect
label="Dtype"
hint="Optional speed-up for GGUF models. Off runs the GGUF as-is. FP8/INT8/FP4 instead load the FULL base model and quantise its transformer onto low-precision tensor cores: faster per step, but a larger download and more VRAM, and it falls back to the GGUF if it can't fit. Needs CUDA."
hint="Transformer compute dtype. Auto picks the fastest precision the hardware supports (at least INT8 on a capable GPU; FP8 on data-center cards) by loading the FULL base model and quantising its transformer onto low-precision tensor cores, and falls back to running the GGUF as-is when the device, VRAM or disk can't take it. Off always runs the GGUF as-is."
badge={<ResolvedBadge status={status} controlKey="transformer_quant" />}
value={transformerQuant}
onValueChange={(v) => setTransformerQuant(v as typeof transformerQuant)}
options={[
["none", "Off (run the GGUF)"],
["auto", "Auto (fastest for GPU)"],
["none", "Off (run the GGUF)"],
["fp8", "FP8"],
["int8", "INT8"],
["nvfp4", "NVFP4 (Blackwell)"],