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

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-06-26 06:17:15 +00:00
commit 3cfb29dd6b
5 changed files with 151 additions and 72 deletions
Download patch file Download diff file Expand all files Collapse all files

135
scripts/quant_probe.py
View file

@ -43,8 +43,9 @@ def _lpips(ref_arr, arr):
try:
import torch
import lpips
if _LPIPS["fn"] is None:
_LPIPS["fn"] = lpips.LPIPS(net="alex", verbose=False).cuda().eval()
_LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval()
def t(x):
t = torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
@ -53,7 +54,7 @@ def _lpips(ref_arr, arr):
with torch.no_grad():
return float(_LPIPS["fn"](t(ref_arr), t(arr)).item())
except Exception as exc: # noqa: BLE001
print(f" (lpips unavailable: {type(exc).__name__}: {str(exc)[:80]})", flush=True)
print(f" (lpips unavailable: {type(exc).__name__}: {str(exc)[:80]})", flush = True)
return None
@ -62,9 +63,9 @@ def _load_dense():
import diffusers
t = diffusers.ZImageTransformer2DModel.from_pretrained(
BASE, subfolder="transformer", torch_dtype=torch.bfloat16
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
)
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16, transformer=t)
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
pipe.to("cuda")
return pipe
@ -76,10 +77,12 @@ def _load_gguf():
t = diffusers.ZImageTransformer2DModel.from_single_file(
hf_hub_download(REPO, GGUF),
quantization_config=diffusers.GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
torch_dtype=torch.bfloat16, config=BASE, subfolder="transformer",
quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = torch.bfloat16),
torch_dtype = torch.bfloat16,
config = BASE,
subfolder = "transformer",
)
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16, transformer=t)
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
pipe.to("cuda")
return pipe
@ -91,6 +94,7 @@ def _quant_config(name):
Int8DynamicActivationInt8WeightConfig,
Float8DynamicActivationFloat8WeightConfig,
)
if name == "int8wo":
return Int8WeightOnlyConfig()
if name == "int8dq":
@ -105,7 +109,7 @@ def _quant_config(name):
try:
import torch
return MXDynamicActivationMXWeightConfig(
activation_dtype=torch.float8_e4m3fn, weight_dtype=torch.float8_e4m3fn
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
)
except TypeError:
return MXDynamicActivationMXWeightConfig()
@ -118,7 +122,7 @@ def _make_filter_fn(min_features):
timestep/pooled projections (in_features=256) run at M=1 and crash it -- skip them."""
import torch.nn as nn
def filter_fn(module, fqn=""):
def filter_fn(module, fqn = ""):
return (
isinstance(module, nn.Linear)
and getattr(module, "in_features", 0) >= min_features
@ -131,11 +135,12 @@ def _make_filter_fn(min_features):
def _apply_quant(pipe, name, log, min_features):
import torch.nn as nn
from torchao.quantization import quantize_
cfg = _quant_config(name)
total = sum(1 for m in pipe.transformer.modules() if isinstance(m, nn.Linear))
filt = _make_filter_fn(min_features)
q = sum(1 for n, m in pipe.transformer.named_modules() if filt(m, n))
quantize_(pipe.transformer, cfg, filter_fn=filt)
quantize_(pipe.transformer, cfg, filter_fn = filt)
log(f" quantized transformer with {name} ({q}/{total} linears >= {min_features} feat)")
@ -156,11 +161,17 @@ def _compile(pipe, log):
def _gen(pipe, steps, seed, res):
import torch
g = torch.Generator(device="cuda").manual_seed(seed)
g = torch.Generator(device = "cuda").manual_seed(seed)
torch.cuda.synchronize()
t0 = time.time()
img = pipe(prompt=PROMPT, width=res, height=res, num_inference_steps=steps,
guidance_scale=0.0, generator=g).images[0]
img = pipe(
prompt = PROMPT,
width = res,
height = res,
num_inference_steps = steps,
guidance_scale = 0.0,
generator = g,
).images[0]
torch.cuda.synchronize()
return img, time.time() - t0
@ -169,23 +180,37 @@ def _median(xs):
return sorted(xs)[len(xs) // 2]
def main(argv=None) -> int:
def main(argv = None) -> int:
p = argparse.ArgumentParser()
p.add_argument("--steps", type=int, default=8)
p.add_argument("--res", type=int, default=1024)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--iters", type=int, default=3)
p.add_argument("--min-feat", type=int, default=512,
help="only quantize Linear with in&out features >= this (int8 _int_mm needs M>16)")
p.add_argument("--configs", default="bf16,bf16_c,gguf_c,int8dq_c,fp8dq_c,nvfp4_c,mxfp8_c,int8wo_c",
help="comma list; suffix _c = +compile")
p.add_argument("--steps", type = int, default = 8)
p.add_argument("--res", type = int, default = 1024)
p.add_argument("--seed", type = int, default = 42)
p.add_argument("--iters", type = int, default = 3)
p.add_argument(
"--min-feat",
type = int,
default = 512,
help = "only quantize Linear with in&out features >= this (int8 _int_mm needs M>16)",
)
p.add_argument(
"--configs",
default = "bf16,bf16_c,gguf_c,int8dq_c,fp8dq_c,nvfp4_c,mxfp8_c,int8wo_c",
help = "comma list; suffix _c = +compile",
)
args = p.parse_args(argv)
steps, res, seed, iters = args.steps, args.res, args.seed, args.iters
import torch
OUT.mkdir(parents=True, exist_ok=True)
def run(tag, *, source, quant=None, compile=False):
OUT.mkdir(parents = True, exist_ok = True)
def run(
tag,
*,
source,
quant = None,
compile = False,
):
torch.compiler.reset()
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
@ -210,61 +235,79 @@ def main(argv=None) -> int:
torch.cuda.empty_cache()
return tag, _median(dts), arr, load_peak, gen_peak
print_ = lambda s: print(s, flush=True) # noqa: E731
print_ = lambda s: print(s, flush = True) # noqa: E731
# config table: tag -> (source, quant, compile)
table = {
"bf16": ("dense", None, False),
"bf16_c": ("dense", None, True),
"gguf_c": ("gguf", None, True),
"int8wo_c": ("dense", "int8wo", True),
"int8dq_c": ("dense", "int8dq", True),
"fp8dq_c": ("dense", "fp8dq", True),
"nvfp4_c": ("dense", "nvfp4", True),
"mxfp8_c": ("dense", "mxfp8", True),
"bf16": ("dense", None, False),
"bf16_c": ("dense", None, True),
"gguf_c": ("gguf", None, True),
"int8wo_c": ("dense", "int8wo", True),
"int8dq_c": ("dense", "int8dq", True),
"fp8dq_c": ("dense", "fp8dq", True),
"nvfp4_c": ("dense", "nvfp4", True),
"mxfp8_c": ("dense", "mxfp8", True),
}
want = [c.strip() for c in args.configs.split(",") if c.strip()]
print(f"== quant probe (Z-Image-Turbo, {res}px, {steps} steps, seed {seed}) ==", flush=True)
print(f"== quant probe (Z-Image-Turbo, {res}px, {steps} steps, seed {seed}) ==", flush = True)
ref_arr = None
rows = []
for tag in want:
if tag not in table:
print(f" {tag}: unknown config, skipping", flush=True)
print(f" {tag}: unknown config, skipping", flush = True)
continue
source, quant, compile = table[tag]
print(f"-- {tag} (source={source} quant={quant} compile={compile}) --", flush=True)
print(f"-- {tag} (source={source} quant={quant} compile={compile}) --", flush = True)
try:
_, med, arr, lp, gp = run(tag, source=source, quant=quant, compile=compile)
_, med, arr, lp, gp = run(tag, source = source, quant = quant, compile = compile)
except Exception as exc: # noqa: BLE001
import traceback
print(f" {tag:10s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush=True)
print(f" {tag:10s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True)
traceback.print_exc()
rows.append((tag, None, None, None, None, None))
continue
if ref_arr is None and tag == "bf16":
ref_arr = arr
psnr = _psnr(ref_arr, arr) if ref_arr is not None else None
lpips_v = _lpips(ref_arr, arr) if (ref_arr is not None and tag != "bf16") else (0.0 if tag == "bf16" else None)
lpips_v = (
_lpips(ref_arr, arr)
if (ref_arr is not None and tag != "bf16")
else (0.0 if tag == "bf16" else None)
)
rows.append((tag, med, psnr, lpips_v, lp, gp))
ps = f"{psnr:.1f}dB" if psnr is not None else "n/a"
lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a"
print(f" {tag:10s} {med:.3f}s PSNR={ps:>7s} LPIPS={lps:>6s} loadVRAM={lp:.1f}G genVRAM={gp:.1f}G", flush=True)
print(
f" {tag:10s} {med:.3f}s PSNR={ps:>7s} LPIPS={lps:>6s} loadVRAM={lp:.1f}G genVRAM={gp:.1f}G",
flush = True,
)
base = next((r[1] for r in rows if r[0] == "bf16" and r[1]), None)
gguf = next((r[1] for r in rows if r[0] == "gguf_c" and r[1]), None)
print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush=True)
print(f"{'config':10s} {'sec':>7s} {'vs_bf16':>8s} {'vs_gguf':>8s} {'PSNR':>8s} {'LPIPS':>7s} {'loadG':>6s} {'genG':>6s}", flush=True)
print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True)
print(
f"{'config':10s} {'sec':>7s} {'vs_bf16':>8s} {'vs_gguf':>8s} {'PSNR':>8s} {'LPIPS':>7s} {'loadG':>6s} {'genG':>6s}",
flush = True,
)
for tag, med, psnr, lpips_v, lp, gp in rows:
if med is None:
print(f"{tag:10s} {'FAILED':>7s}", flush=True)
print(f"{tag:10s} {'FAILED':>7s}", flush = True)
continue
vb = f"{base/med:.2f}x" if base else "-"
vg = f"{gguf/med:.2f}x" if gguf else "-"
ps = f"{psnr:.1f}" if psnr is not None and psnr != float('inf') else ("inf" if psnr == float('inf') else "n/a")
ps = (
f"{psnr:.1f}"
if psnr is not None and psnr != float("inf")
else ("inf" if psnr == float("inf") else "n/a")
)
lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a"
print(f"{tag:10s} {med:>7.3f} {vb:>8s} {vg:>8s} {ps:>8s} {lps:>7s} {lp:>6.1f} {gp:>6.1f}", flush=True)
print("QUANT-PROBE-DONE", flush=True)
print(
f"{tag:10s} {med:>7.3f} {vb:>8s} {vg:>8s} {ps:>8s} {lps:>7s} {lp:>6.1f} {gp:>6.1f}",
flush = True,
)
print("QUANT-PROBE-DONE", flush = True)
return 0

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

@ -488,8 +488,14 @@ class DiffusionBackend:
):
try:
pipe, transformer_quant_engaged = self._load_dense_quant_pipeline(
transformer_cls, pipeline_cls, base, device, dtype, hf_token,
target, transformer_quant,
transformer_cls,
pipeline_cls,
base,
device,
dtype,
hf_token,
target,
transformer_quant,
)
except Exception as exc: # noqa: BLE001 — fall back to the GGUF build
logger.warning(

10
studio/backend/core/inference/diffusion_transformer_quant.py
View file

@ -52,8 +52,8 @@ DEFAULT_MIN_LINEAR_FEATURES = 512
# its kernels are available.
_AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = (
((10, 0), (TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8)), # Blackwell sm_100+
((8, 9), (TQ_FP8, TQ_INT8)), # Ada sm_89 / Hopper sm_90
((8, 0), (TQ_INT8,)), # Ampere sm_80 / sm_86
((8, 9), (TQ_FP8, TQ_INT8)), # Ada sm_89 / Hopper sm_90
((8, 0), (TQ_INT8,)), # Ampere sm_80 / sm_86
)
# Cache of (scheme, device) -> bool so the quantise+matmul smoke test runs once.
@ -84,7 +84,6 @@ def dense_transformer_supported(target: Any) -> bool:
return False
try:
import torch
return getattr(target, "dtype", None) is torch.bfloat16
except Exception:
return False
@ -118,7 +117,6 @@ def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Op
def _capability() -> Optional[tuple[int, int]]:
try:
import torch
major, minor = torch.cuda.get_device_capability()
return (int(major), int(minor))
except Exception:
@ -129,7 +127,6 @@ def _scheme_supported(scheme: str, device: str) -> bool:
"""CUDA + (for fp8) the fp8 dtype + a cached quantise+matmul smoke test for ``scheme``."""
try:
import torch
if not torch.cuda.is_available():
return False
if scheme == TQ_FP8 and not hasattr(torch, "float8_e4m3fn"):
@ -179,12 +176,10 @@ def _make_quant_config(scheme: str) -> Any:
return Float8DynamicActivationFloat8WeightConfig()
if scheme == TQ_NVFP4:
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig
return NVFP4DynamicActivationNVFP4WeightConfig()
if scheme == TQ_MXFP8:
import torch
from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
try:
return MXDynamicActivationMXWeightConfig(
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
@ -201,7 +196,6 @@ def make_filter_fn(min_features: int):
def filter_fn(module: Any, fqn: str = "") -> bool:
try:
import torch
if not isinstance(module, torch.nn.Linear):
return False
except Exception:

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

@ -922,7 +922,8 @@ def test_default_load_skips_dense_quant_path(fake_runtime, tmp_path, monkeypatch
from core.inference import diffusion as dmod
monkeypatch.setattr(
dmod, "dense_transformer_supported",
dmod,
"dense_transformer_supported",
lambda *a, **k: pytest.fail("dense path must not run without the flag"),
)
(tmp_path / "m.gguf").write_bytes(b"x")
@ -940,7 +941,9 @@ def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatc
calls = _stub_dense_quant(monkeypatch, scheme = "fp8")
(tmp_path / "m.gguf").write_bytes(b"x")
status = backend.load_pipeline(
str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image",
str(tmp_path),
gguf_filename = "m.gguf",
family_override = "z-image",
transformer_quant = "fp8",
)
assert status["transformer_quant"] == "fp8"
@ -971,12 +974,14 @@ def test_transformer_quant_falls_back_to_gguf_on_failure(fake_runtime, tmp_path,
monkeypatch.setattr(dmod, "quantize_transformer", lambda pipe, target, **kw: None)
(tmp_path / "m.gguf").write_bytes(b"x")
status = backend.load_pipeline(
str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image",
str(tmp_path),
gguf_filename = "m.gguf",
family_override = "z-image",
transformer_quant = "fp8",
)
assert status["loaded"] is True
assert status["transformer_quant"] is None # fell back
assert _FakeTransformer.last["path"] # GGUF from_single_file used
assert status["transformer_quant"] is None # fell back
assert _FakeTransformer.last["path"] # GGUF from_single_file used
def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, monkeypatch):
@ -996,8 +1001,11 @@ def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, mo
monkeypatch.setattr(_FakeTransformer, "from_pretrained", _fp_fail, raising = False)
(tmp_path / "m.gguf").write_bytes(b"x")
status = backend.load_pipeline(
str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image",
transformer_quant = "fp8", memory_mode = "low_vram",
str(tmp_path),
gguf_filename = "m.gguf",
family_override = "z-image",
transformer_quant = "fp8",
memory_mode = "low_vram",
)
assert status["transformer_quant"] is None
assert status["offload_policy"] == "model"

46
studio/backend/tests/test_diffusion_transformer_quant.py
View file

@ -33,7 +33,13 @@ def _target(*, device = "cuda", dtype = "bfloat16"):
return types.SimpleNamespace(device = device, dtype = dtype)
def _stub_torch(monkeypatch, *, cc = (10, 0), with_fp8 = True, cuda_available = True):
def _stub_torch(
monkeypatch,
*,
cc = (10, 0),
with_fp8 = True,
cuda_available = True,
):
torch = types.ModuleType("torch")
torch.bfloat16 = "bfloat16"
torch.float16 = "float16"
@ -155,6 +161,7 @@ def test_smoke_probe_caches_and_tolerates_failure(monkeypatch):
class _Lin:
def __init__(self, *a, **k):
pass
def to(self, **k):
return self
@ -167,8 +174,14 @@ def test_smoke_probe_caches_and_tolerates_failure(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", torch)
tqz = types.ModuleType("torchao.quantization")
def _quantize_ok(module, config, filter_fn = None):
def _quantize_ok(
module,
config,
filter_fn = None,
):
calls["n"] += 1
tqz.quantize_ = _quantize_ok
tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8cfg"
tqz.Float8DynamicActivationFloat8WeightConfig = lambda: "fp8cfg"
@ -182,8 +195,14 @@ def test_smoke_probe_caches_and_tolerates_failure(monkeypatch):
# A scheme whose quantize_ raises -> probe False (and cached).
tq._SMOKE_CACHE.clear()
def _quantize_boom(module, config, filter_fn = None):
def _quantize_boom(
module,
config,
filter_fn = None,
):
raise RuntimeError("kernel unavailable")
tqz.quantize_ = _quantize_boom
assert tq._smoke_probe(TQ_FP8, "cuda") is False
@ -195,15 +214,16 @@ def test_make_filter_fn(monkeypatch):
class _Lin:
def __init__(self, i, o):
self.in_features, self.out_features = i, o
torch = types.ModuleType("torch")
torch.nn = types.SimpleNamespace(Linear = _Lin)
monkeypatch.setitem(sys.modules, "torch", torch)
keep = make_filter_fn(512)
assert keep(_Lin(1024, 4096), "blocks.0.attn.to_q") is True
assert keep(_Lin(256, 4096), "time_proj") is False # small in_features -> skip
assert keep(_Lin(4096, 256), "out_proj") is False # small out_features -> skip
assert keep(object(), "not_linear") is False # non-Linear -> skip
assert keep(_Lin(256, 4096), "time_proj") is False # small in_features -> skip
assert keep(_Lin(4096, 256), "out_proj") is False # small out_features -> skip
assert keep(object(), "not_linear") is False # non-Linear -> skip
assert keep(types.SimpleNamespace(), "no_attrs") is False
@ -215,14 +235,16 @@ def test_quantize_transformer_applies_and_marks(monkeypatch):
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: f"{scheme}cfg")
recorder: list = []
tqz = types.ModuleType("torchao.quantization")
tqz.quantize_ = lambda module, config, filter_fn = None: recorder.append((module, config, filter_fn))
tqz.quantize_ = lambda module, config, filter_fn = None: recorder.append(
(module, config, filter_fn)
)
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
transformer = types.SimpleNamespace()
pipe = types.SimpleNamespace(transformer = transformer)
assert quantize_transformer(pipe, _target(), mode = "fp8") == TQ_FP8
assert len(recorder) == 1 and recorder[0][0] is transformer and recorder[0][1] == "fp8cfg"
assert callable(recorder[0][2]) # a filter_fn was passed
assert callable(recorder[0][2]) # a filter_fn was passed
assert transformer._unsloth_runtime_quant == TQ_FP8 # diagnostic marker set
@ -236,8 +258,14 @@ def test_quantize_transformer_tolerates_failure(monkeypatch):
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_INT8)
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: "cfg")
tqz = types.ModuleType("torchao.quantization")
def _boom(module, config, filter_fn = None):
def _boom(
module,
config,
filter_fn = None,
):
raise RuntimeError("partial quant failure")
tqz.quantize_ = _boom
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())