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

Wire hosted pre-quantized DiT checkpoints into the image families

Browse source 
Point prequant_repos for flux.1, flux.2-klein, flux.2-dev, qwen-image
(int8 only there; fp8 is family-denied), z-image and krea-2 at the
unsloth/<Model>-FP8 Hub repos carrying gate-validated int8 and fp8
transformer checkpoints, so the fast quant path loads the small
pre-quantized file instead of materialising the dense bf16 transformer
and quantising on device. Measured on FLUX.2-dev int8: build peak drops
from 60.7 GB (dense + quantize) to 30.7 GB (hosted prequant), identical
30.7 GB resident after either path since loading a checkpoint is
bit-identical to on-the-fly quantisation.

The hosted repos name files <Model>-<SCHEME>.pt, so resolve_prequant_source
now derives that model-name filename from the repo id (scheme suffix
stripped case-insensitively) and carries the legacy transformer_<scheme>.pt
as a fallback the resolver tries when the primary 404s, keeping older
repos loadable.

Wiring a repo also exposed a fallback hazard: with a prequant source
present, the dense-fit preflight used to be skipped entirely, so a failed
prequant download would fall through to the dense bf16 load the memory
plan never budgeted, OOMing after eviction. The preflight now always runs
and gates an allow_dense_fallback flag through _load_dense_quant_pipeline:
a dense misfit still skips the fast path when no prequant exists, but with
one it proceeds and a prequant failure raises to the GGUF build instead of
loading dense. The same flag is set when the auto-policy replans an
offloaded GGUF against a prequant-sized transient.

Tests updated to the new filename convention plus new coverage for the
derivation and the legacy-name fallback; the prequant-skips-refit test now
asserts the re-check runs and forbids the dense fallback. Verified end to
end on GPU: z-image int8 resolves the hosted repo, downloads the
model-name file and renders (6.8s load, 5.9 GB peak).
This commit is contained in:
Daniel Han 2026-07-17 05:47:45 +00:00
commit fbcf070fce
5 changed files with 168 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1210,6 +1210,10 @@ class DiffusionBackend:
# The GGUF-size `plan` can mis-budget the fast path two ways, so preflight the real
# footprint BEFORE eviction; both branches need the base repo + a resolved scheme.
dense_declined = False
# False when the memory plan only holds a PREQUANT-sized build: if the prequant
# load then fails, the loader must raise to GGUF instead of materialising the
# dense bf16 transformer the plan never budgeted for.
dense_fallback_allowed = True
if (
kind == "gguf"
and normalize_transformer_quant(transformer_quant) is not None
@ -1246,6 +1250,10 @@ class DiffusionBackend:
)
if replanned.offload_policy == OFFLOAD_NONE:
quant_plan = replanned
# The GGUF plan already declined resident; a prequant-sized
# replan says nothing about the (larger) dense transformer.
if candidate.prequant:
dense_fallback_allowed = False
else:
# The GGUF fits resident, but this path first materialises the base's dense
# bf16 transformer (bigger), so re-check the fit against THAT -- a card that
@ -1267,23 +1275,29 @@ class DiffusionBackend:
if scheme is not None
else None
)
if prequant is None:
dense_mib = int(
self._dense_transformer_resident_bytes(base) // (1024 * 1024)
dense_mib = int(
self._dense_transformer_resident_bytes(base) // (1024 * 1024)
)
if dense_mib > 0:
dense_plan = self._plan_memory(
target,
single_file_path,
base,
fam,
memory_mode,
cpu_offload,
kind = kind,
repo_id = repo_id,
transformer_resident_override_mib = dense_mib,
)
if dense_mib > 0:
dense_plan = self._plan_memory(
target,
single_file_path,
base,
fam,
memory_mode,
cpu_offload,
kind = kind,
repo_id = repo_id,
transformer_resident_override_mib = dense_mib,
)
dense_declined = dense_plan.offload_policy != OFFLOAD_NONE
if dense_plan.offload_policy != OFFLOAD_NONE:
dense_fallback_allowed = False
# Without a prequant source the dense build is the ONLY path,
# so a dense misfit skips the fast path entirely (as before); with
# one, the small prequant load proceeds and only the dense
# fallback is forbidden.
if prequant is None:
dense_declined = True
if (
kind == "gguf"
and normalize_transformer_quant(transformer_quant) is not None
@ -1305,6 +1319,7 @@ class DiffusionBackend:
fam = fam,
base_local_dir = _base_local_dir,
prequant_path = transformer_prequant_path,
allow_dense_fallback = dense_fallback_allowed,
)
except Exception as exc: # noqa: BLE001 — fall back to the GGUF build
logger.warning(
@ -1713,6 +1728,7 @@ class DiffusionBackend:
fam: Optional[DiffusionFamily] = None,
prequant_path: Optional[str] = None,
base_local_dir: Optional[str] = None,
allow_dense_fallback: bool = True,
) -> tuple[Any, str]:
"""Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
@ -1761,6 +1777,12 @@ class DiffusionBackend:
return pipe, scheme
# 2. Fallback: materialise the dense bf16 transformer and quantise it on-device.
if not allow_dense_fallback:
# The memory plan only budgeted the prequant-sized build; materialising the dense
# bf16 transformer here would exceed it after eviction. Raise to the GGUF build.
raise RuntimeError(
"prequant checkpoint unavailable and the dense transformer does not fit resident"
)
transformer = transformer_cls.from_pretrained(
base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
)

26
studio/backend/core/inference/diffusion_families.py
View file

@ -112,6 +112,14 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
pipeline_class = "FluxPipeline",
transformer_class = "FluxTransformer2DModel",
base_repo = "black-forest-labs/FLUX.1-schnell",
# Hosted pre-quantized DiT checkpoints (gate-validated vs same-seed bf16). The loader
# verifies the checkpoint's baked base_model_id against the repo actually being loaded,
# so a non-default base (e.g. FLUX.1-dev under this family) safely falls back to the
# dense-quantize path instead of loading schnell weights.
prequant_repos = (
("int8", "unsloth/FLUX.1-schnell-FP8"),
("fp8", "unsloth/FLUX.1-schnell-FP8"),
),
aliases = ("flux1", "flux-1"),
# LoRA training targets FLUX.1-dev via the DiT trainer (QLoRA nf4); the dev repo is gated.
trainable = True,
@ -133,6 +141,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
pipeline_class = "Flux2KleinPipeline",
transformer_class = "Flux2Transformer2DModel",
base_repo = "black-forest-labs/FLUX.2-klein-4B",
prequant_repos = (
("int8", "unsloth/FLUX.2-klein-4B-FP8"),
("fp8", "unsloth/FLUX.2-klein-4B-FP8"),
),
aliases = ("flux2-klein",),
# Flux2KleinPipeline takes reference image(s) via `image`, so it exposes a "reference"
# workflow atop text-to-image. It has an inpaint pipeline (no img2img) -> inpaint + extend.
@ -156,6 +168,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
pipeline_class = "Flux2Pipeline",
transformer_class = "Flux2Transformer2DModel",
base_repo = "black-forest-labs/FLUX.2-dev",
prequant_repos = (
("int8", "unsloth/FLUX.2-dev-FP8"),
("fp8", "unsloth/FLUX.2-dev-FP8"),
),
aliases = ("flux2-dev", "flux2dev"),
sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"),
sd_cpp_vae_format = "flux2",
@ -201,6 +217,8 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
pipeline_class = "QwenImagePipeline",
transformer_class = "QwenImageTransformer2DModel",
base_repo = "Qwen/Qwen-Image",
# int8 only: fp8 is family-denied (_FAMILY_SCHEME_DENY) so a repo entry would be dead.
prequant_repos = (("int8", "unsloth/Qwen-Image-FP8"),),
cfg_kwarg = "true_cfg_scale",
aliases = ("qwen_image", "qwenimage"),
# LoRA training via the DiT trainer, defaulting to the prequant nf4 repo (QLoRA).
@ -229,6 +247,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
pipeline_class = "ZImagePipeline",
transformer_class = "ZImageTransformer2DModel",
base_repo = "Tongyi-MAI/Z-Image-Turbo",
prequant_repos = (
("int8", "unsloth/Z-Image-Turbo-FP8"),
("fp8", "unsloth/Z-Image-Turbo-FP8"),
),
aliases = ("zimage", "z_image"),
# LoRA training via the DiT trainer (bf16); defaults to the prequant nf4 repo for QLoRA.
trainable = True,
@ -250,6 +272,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = (
pipeline_class = "Krea2Pipeline",
transformer_class = "Krea2Transformer2DModel",
base_repo = "krea/Krea-2-Turbo",
prequant_repos = (
("int8", "unsloth/Krea-2-Turbo-FP8"),
("fp8", "unsloth/Krea-2-Turbo-FP8"),
),
aliases = ("krea2",),
# LoRA training via the DiT trainer (no prequant repo yet, so nf4 quantizes on the fly).
# Krea's guidance: train on the undistilled Raw, run adapters on Turbo, so Raw is the

39
studio/backend/core/inference/diffusion_prequant.py
View file

@ -90,18 +90,32 @@ def local_prequant_path_ready(path: str) -> bool:
@dataclass(frozen = True)
class PrequantSource:
"""Where a pre-quantized checkpoint lives. ``kind`` is "path" (a local file) or "repo"
(Hub repo id in ``location`` + ``filename``)."""
(Hub repo id in ``location`` + ``filename``; ``fallback_filename`` is tried when the
primary name is absent, covering repos still on the legacy transformer_<scheme>.pt)."""
kind: str
location: str
filename: Optional[str] = None
fallback_filename: Optional[str] = None
def prequant_filename(scheme: str) -> str:
"""The conventional checkpoint filename for ``scheme`` inside a Hub repo."""
"""The legacy checkpoint filename for ``scheme`` inside a Hub repo."""
return f"transformer_{scheme}.pt"
def prequant_repo_filename(repo_id: str, scheme: str) -> str:
"""The model-name checkpoint filename for ``scheme`` in ``repo_id``: the hosted repos are
named <Model>-FP8 (or -INT8 / -quantized) and carry <Model>-<SCHEME>.pt files, e.g.
unsloth/Z-Image-Turbo-FP8 -> Z-Image-Turbo-INT8.pt / Z-Image-Turbo-FP8.pt."""
model = repo_id.rsplit("/", 1)[-1]
for suffix in ("-fp8", "-int8", "-quantized"):
if model.lower().endswith(suffix):
model = model[: -len(suffix)]
break
return f"{model}-{scheme.upper()}.pt"
def resolve_prequant_source(
fam: Any,
scheme: str,
@ -122,7 +136,12 @@ def resolve_prequant_source(
except Exception: # noqa: BLE001 — a bad family object must not break the load
repo_id = None
if repo_id:
return PrequantSource(kind = "repo", location = repo_id, filename = prequant_filename(scheme))
return PrequantSource(
kind = "repo",
location = repo_id,
filename = prequant_repo_filename(repo_id, scheme),
fallback_filename = prequant_filename(scheme),
)
return None
@ -243,7 +262,19 @@ def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) ->
return expanded if os.path.isfile(expanded) else None
if source.kind == "repo":
from huggingface_hub import hf_hub_download
return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token)
try:
from huggingface_hub.errors import EntryNotFoundError
except Exception: # noqa: BLE001 — older hub layouts; fall back to a private marker
class EntryNotFoundError(Exception): # type: ignore[no-redef]
pass
try:
return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token)
except EntryNotFoundError:
if not source.fallback_filename or source.fallback_filename == source.filename:
raise
return hf_hub_download(
repo_id = source.location, filename = source.fallback_filename, token = hf_token
)
return None

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

@ -2656,10 +2656,11 @@ def test_dense_quant_skipped_when_dense_transformer_does_not_fit(
assert _FakeTransformer.last["path"] # GGUF path used
def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypatch):
# With a prequant checkpoint, the fast path loads the small quantized file, not the dense
# bf16 -- so the dense-transformer re-check must NOT run and must NOT decline the fast path,
# even when the base's dense shards are cached and large.
def test_dense_quant_prequant_proceeds_but_forbids_dense_fallback(fake_runtime, tmp_path, monkeypatch):
# With a prequant checkpoint, the fast path loads the small quantized file, so a dense
# misfit must NOT decline the fast path -- but the dense re-check still runs to gate the
# in-loader fallback: if the prequant later fails, the loader must raise to GGUF instead of
# materialising the dense bf16 the plan never budgeted (allow_dense_fallback=False).
from core.inference import diffusion as dmod
backend = DiffusionBackend()
@ -2688,13 +2689,15 @@ def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypa
):
if transformer_resident_override_mib is not None:
dense_refit_ran.append(True)
# GGUF budget fits (real plan -> none); the dense-transformer preflight does not.
return types.SimpleNamespace(offload_policy = "model")
return orig_plan(self, *a, **k)
monkeypatch.setattr(DiffusionBackend, "_plan_memory", spy_plan)
attempted = []
def fake_dense_load(self, *a, **k):
attempted.append(True)
attempted.append(k.get("allow_dense_fallback"))
return None, None # fall through to GGUF; we only assert the path was reached
monkeypatch.setattr(DiffusionBackend, "_load_dense_quant_pipeline", fake_dense_load)
@ -2705,8 +2708,8 @@ def test_dense_quant_prequant_skips_dense_refit(fake_runtime, tmp_path, monkeypa
family_override = "z-image",
transformer_quant = "fp8",
)
assert dense_refit_ran == [] # prequant -> dense re-check skipped
assert attempted == [True] # fast path still attempted (with the prequant)
assert dense_refit_ran == [True] # the re-check runs (it gates the fallback)...
assert attempted == [False] # ...fast path still attempted, dense fallback forbidden
def test_dense_quant_unusable_prequant_path_runs_dense_refit(fake_runtime, tmp_path, monkeypatch):

60
studio/backend/tests/test_diffusion_prequant.py
View file

@ -46,7 +46,18 @@ def test_resolve_family_repo_by_scheme():
fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"), ("int8", "org/hosted-int8")))
src = resolve_prequant_source(fam, "int8")
assert src.kind == "repo" and src.location == "org/hosted-int8"
assert src.filename == "transformer_int8.pt"
# Model-name convention first (repo scheme suffix stripped), legacy name as fallback.
assert src.filename == "hosted-INT8.pt"
assert src.fallback_filename == "transformer_int8.pt"
def test_prequant_repo_filename_convention():
from core.inference.diffusion_prequant import prequant_repo_filename
assert prequant_repo_filename("unsloth/Z-Image-Turbo-FP8", "int8") == "Z-Image-Turbo-INT8.pt"
assert prequant_repo_filename("unsloth/Z-Image-Turbo-FP8", "fp8") == "Z-Image-Turbo-FP8.pt"
assert prequant_repo_filename("unsloth/Qwen-Image-2512-INT8", "int8") == "Qwen-Image-2512-INT8.pt"
assert prequant_repo_filename("org/Some-Model-quantized", "fp8") == "Some-Model-FP8.pt"
assert prequant_repo_filename("org/PlainRepo", "int8") == "PlainRepo-INT8.pt"
def test_resolve_wrong_scheme_is_none():
@ -480,6 +491,53 @@ def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path):
assert result is not None
def test_load_repo_source_falls_back_to_legacy_filename(monkeypatch, tmp_path):
# A repo still carrying the legacy transformer_<scheme>.pt name serves the download after
# the model-name filename 404s; both names are requested in order.
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, _good_ckpt())
monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
downloaded = tmp_path / "transformer_fp8.pt"
downloaded.write_bytes(b"x")
class _NotFound(Exception):
pass
errors = types.ModuleType("huggingface_hub.errors")
errors.EntryNotFoundError = _NotFound
requested = []
def _dl(repo_id, filename, token = None):
requested.append(filename)
if filename != "transformer_fp8.pt":
raise _NotFound(filename)
return str(downloaded)
hub = types.ModuleType("huggingface_hub")
hub.hf_hub_download = _dl
hub.errors = errors
monkeypatch.setitem(sys.modules, "huggingface_hub", hub)
monkeypatch.setitem(sys.modules, "huggingface_hub.errors", errors)
source = PrequantSource(
kind = "repo", location = "org/Z-Image-Turbo-FP8",
filename = "Z-Image-Turbo-FP8.pt", fallback_filename = "transformer_fp8.pt",
)
result = load_prequantized_transformer(
_FakeTransformer,
"Tongyi-MAI/Z-Image-Turbo",
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
logger = None,
)
assert result is not None
assert requested == ["Z-Image-Turbo-FP8.pt", "transformer_fp8.pt"]
def test_load_local_path_outside_allowlist_refused(monkeypatch, tmp_path):
# Even with the opt-in set, a path OUTSIDE every allowlisted directory must not be
# unpickled: enabling one trusted dir is not a wildcard for arbitrary request paths.