147d323912
fp8 auto precision defaulted to precise accumulate on any non-consumer GPU, which made fp8 2.05x slower than int8 on RTX 6000 Ada and slower than not quantising at all. NVIDIA's professional whitepapers do publish equal FP8 rates for both accumulate modes there, so the hardware premise held, but the cost is in the cuBLAS path rather than the published rate. Default to fast accumulate: measured on B200 the flag is a no-op (4096^3 _scaled_mm at 3023.8 vs 3041.8 TFLOP/s, bitwise-identical output, 1.213 s vs 1.230 s end to end), so it is a large win where it bites and free where it does not. Precise accumulate stays available via transformer_quant_fast_accum. Z-Image's DiT is a Lumina2 derivative, so unsloth/Z-Image-GGUF and unsloth/Z-Image-Turbo-GGUF both declare general.architecture = "lumina2" and the whole line was tagged image-diffusion-unsupported and hidden from the Images "On Device" list, though validate_load_request loads them. Resolve shared archs from the repo/file name like bare "wan" already does, with a test asserting the picker and the loader agree for every family. The sage attention on-demand install ran an unpinned `pip install sageattention`, but PyPI's newest wheel is 1.0.6 and diffusers refuses anything below 2.1.1: the install always "succeeded", wrote an unusable version into the running venv, and was rejected on the next line. Carry the dispatcher's floor so pip resolves nothing instead. The dense-quant disk gate sized the download from the bf16-RESIDENT table. The fp32 families download twice that (Z-Image: 23,479 MiB against a 21,970 MiB gate), leaving a window where the check passed and the download filled the disk; Ideogram 4 ships fp8 and was overcharged the other way. Size the gate by published bytes, verified against HF sibling metadata for all 12 families. Patch installs went through unsloth_zoo, which refuses to import unless UNSLOTH_IS_PRESENT is set, and that is set by unsloth itself. The server imports unsloth at boot so it never showed there, but any other process ran silently unpatched with every install returning False, which is 13 test failures on a clean environment. Import unsloth and retry once, memoised per process. Also: the GGUF+LoRA refusal pointed at the native engine without saying a GPU host only selects it under UNSLOTH_DIFFUSION_ENGINE=sd_cpp, so the suggestion was unreachable; the gallery recipe recorded loras from the generate request alone, losing a load-time bake; load-progress claimed "40.07 GB downloaded" for a fully cached load; and pickers.tsx imported three catalog-group helpers it never used. Reported by oobabooga.
284 lines
13 KiB
Python
284 lines
13 KiB
Python
# 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, family = None: scheme
|
|
)
|
|
import core.inference.diffusion_prequant as pq
|
|
|
|
monkeypatch.setattr(
|
|
pq,
|
|
"resolve_prequant_source",
|
|
lambda fam, s, path_override = None, base_repo = None: prequant,
|
|
)
|
|
# Neutralize the cache-disk gate by default so resolution tests don't depend on the runner's free
|
|
# space. The two disk-gate tests re-patch this to exercise it explicitly.
|
|
monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: None)
|
|
|
|
|
|
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_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_disk_gate_sizes_fp32_families_by_their_real_download(monkeypatch):
|
|
# The size table is bf16-RESIDENT, but the disk gate is about bytes landing in the HF cache.
|
|
# Z-Image publishes fp32 shards (measured 23,479 MiB against 11,730 MiB resident), so gating on
|
|
# the resident figure left a window where the check passed and the download then filled the
|
|
# disk. The gate must ask for the download size; the memory estimate must NOT change.
|
|
import core.inference.diffusion_auto_policy as ap
|
|
|
|
_patch_selector(monkeypatch, scheme = "int8")
|
|
est = ap.estimate_dense_quant(_fam("z-image"), "int8")
|
|
assert est.transient_transformer_mib == 11_730 # resident: unchanged
|
|
assert est.download_transformer_mib == 23_460 # download: measured 23,479 MiB
|
|
# Free space that clears the old (resident-based) bar but not the real download is refused.
|
|
monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 11_730 + 10 * 1024 + 512)
|
|
assert (
|
|
resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto")
|
|
is None
|
|
)
|
|
monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 23_460 + 10 * 1024 + 512)
|
|
assert isinstance(
|
|
resolve_dense_quant_candidate(fam = _fam("z-image"), target = object(), requested = "auto"),
|
|
DenseQuantEstimate,
|
|
)
|
|
|
|
|
|
def test_disk_gate_does_not_overcharge_a_family_published_below_bf16(monkeypatch):
|
|
# The correction runs both ways: Ideogram 4 ships fp8 (transformer + unconditional_transformer
|
|
# measured at 17,718 MiB) and doubles on the way to bf16, so charging the resident 35,477 MiB
|
|
# of disk would refuse the candidate on a disk that easily holds it.
|
|
import core.inference.diffusion_auto_policy as ap
|
|
|
|
_patch_selector(monkeypatch, scheme = "int8")
|
|
est = ap.estimate_dense_quant(_fam("ideogram-4"), "int8")
|
|
assert est.transient_transformer_mib == 35_476
|
|
assert est.download_transformer_mib == 17_738
|
|
monkeypatch.setattr(ap, "_hf_cache_free_mib", lambda: 17_738 + 10 * 1024 + 512)
|
|
assert isinstance(
|
|
resolve_dense_quant_candidate(fam = _fam("ideogram-4"), target = object(), requested = "auto"),
|
|
DenseQuantEstimate,
|
|
)
|
|
|
|
|
|
def test_disk_gate_matches_download_for_bf16_published_families(monkeypatch):
|
|
# Families that publish bf16 download what they occupy (measured 0.99-1.07x), so the two
|
|
# numbers stay equal and no factor entry is needed.
|
|
import core.inference.diffusion_auto_policy as ap
|
|
|
|
for name in ("flux.1", "flux.2-dev", "qwen-image", "krea-2", "hidream-i1"):
|
|
est = ap.estimate_dense_quant(_fam(name), "int8")
|
|
assert est.download_transformer_mib == est.transient_transformer_mib, name
|
|
|
|
|
|
def test_candidate_none_for_an_unlisted_family(monkeypatch):
|
|
# No size entry means 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):
|
|
# A hosted-repo prequant source is available without a local-path check.
|
|
_patch_selector(monkeypatch, prequant = SimpleNamespace(kind = "repo", location = "org/int8"))
|
|
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 on a 32 GiB consumer card: the user picked a LARGE (BF16) GGUF, so the
|
|
# file-size plan forces offload -- but the dense-quant candidate is far smaller, and re-planning
|
|
# against it keeps everything resident. Before the fix the loader never attempted the fast path.
|
|
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())
|