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
)