diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index 89f59a688a..43a92ac8ee 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -218,6 +218,10 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: memory_mode = args.memory_mode, speed_mode = args.speed_mode, text_encoder_quant = args.text_encoder_quant, + transformer_quant = args.transformer_quant, + transformer_quant_fast_accum = {"auto": None, "on": True, "off": False}[ + args.fp8_fast_accum + ], ) _wait_for_load(backend) _cuda_sync() @@ -299,6 +303,8 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: "speed_mode": args.speed_mode, "cpu_offload": args.cpu_offload, "text_encoder_quant": args.text_encoder_quant, + "transformer_quant": args.transformer_quant, + "fp8_fast_accum": args.fp8_fast_accum, }, } @@ -465,6 +471,21 @@ def _build_parser() -> argparse.ArgumentParser: choices = ["fp8", "nvfp4"], help = "quantise the companion text encoder (fp8 or nvfp4)", ) + p.add_argument( + "--transformer-quant", + default = None, + choices = ["auto", "int8", "fp8", "nvfp4", "mxfp8"], + help = "opt-in fast transformer: load the DENSE bf16 transformer and torchao-" + "quantise it onto the low-precision tensor cores (faster than GGUF, higher " + "VRAM). auto picks per GPU; falls back to GGUF if unsupported / no VRAM", + ) + p.add_argument( + "--fp8-fast-accum", + default = "auto", + choices = ["auto", "on", "off"], + help = "fp8 accumulate: auto picks by GPU class (fast on consumer, precise on " + "data-center); on/off force it", + ) p.add_argument( "--cpu-offload", action = "store_true", help = "legacy: force whole-module CPU offload" ) diff --git a/scripts/fp8_overflow_check.py b/scripts/fp8_overflow_check.py new file mode 100644 index 0000000000..157b9cce30 --- /dev/null +++ b/scripts/fp8_overflow_check.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Empirically check that fp8 dynamic quant with fast accumulation does not overflow. + +Hooks every quantised Linear's output during a real Z-Image generation and reports the +global max |output| and any non-finite (Inf/NaN) count, for use_fast_accum True vs False. +The concern fast_accum raises is accumulation *precision*, not overflow (the accumulator +stays FP32-range and torchao's dynamic per-row scale keeps FP8 inputs <= 448); this proves +it on the real model, including Z-Image's large (~9e5) activation peaks. Run on one CUDA GPU. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" + + +def _load_dense(): + import torch, diffusers + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _run(fast_accum, steps, res, seed, mf): + import torch + import torch.nn as nn + from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig + from torchao.float8 import Float8MMConfig + + pipe = _load_dense() + + def filt(mod, fqn = ""): + return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf + + quantize_( + pipe.transformer, + Float8DynamicActivationFloat8WeightConfig( + mm_config = Float8MMConfig(use_fast_accum = fast_accum) + ), + filter_fn = filt, + ) + + stats = {"max_abs": 0.0, "nonfinite": 0, "hooked": 0} + + def hook(mod, inp, out): + t = out[0] if isinstance(out, tuple) else out + if not torch.is_tensor(t): + return + finite = torch.isfinite(t) + nf = int((~finite).sum().item()) + stats["nonfinite"] += nf + m = float(t[finite].abs().max().item()) if finite.any() else float("inf") + if m > stats["max_abs"]: + stats["max_abs"] = m + + # Hook the quantised linears (where an fp8-accumulation overflow would surface). + # Run EAGER: forward hooks don't trace through torch.compile, and the fp8 fast-accum + # accumulation is identical compiled or eager -- compile only changes scheduling. + for m in pipe.transformer.modules(): + if isinstance(m, nn.Linear): + m.register_forward_hook(hook) + stats["hooked"] += 1 + + g = torch.Generator(device = "cuda").manual_seed(seed) + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + import numpy as np + + arr = np.array(img) + img_finite = bool(np.isfinite(arr).all()) + del pipe + torch.cuda.empty_cache() + return stats, img_finite + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 4) + p.add_argument("--res", type = int, default = 512) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--min-feat", type = int, default = 512) + args = p.parse_args(argv) + + print(f"== fp8 overflow check (Z-Image dense, {args.res}px, {args.steps} steps) ==", flush = True) + for fast in (True, False): + stats, img_finite = _run(fast, args.steps, args.res, args.seed, args.min_feat) + print( + f" fast_accum={str(fast):5s} hooked_linears={stats['hooked']:3d} " + f"max|linear_out|={stats['max_abs']:.1f} nonfinite_elems={stats['nonfinite']} " + f"image_all_finite={img_finite}", + flush = True, + ) + print("FP8-OVERFLOW-CHECK-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/nvfp4_probe.py b/scripts/nvfp4_probe.py new file mode 100644 index 0000000000..902cf8a2d2 --- /dev/null +++ b/scripts/nvfp4_probe.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Probe NVFP4 via torchao with use_triton_kernel=False (no MSLK) on the real dense +Z-Image transformer: is it a genuine FP4-tensor-core speedup over fp8, and is quality +in-bar? Reference for LPIPS is dense bf16 eager. Run on one CUDA (Blackwell) GPU.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_images" + + +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import torch, lpips + + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _load_dense(): + import torch, diffusers + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + 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] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +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) + p.add_argument("--out-dir", default = None, help = "image output dir (default: repo outputs/)") + args = p.parse_args(argv) + steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat + import torch + import torch.nn as nn + + global OUT + if args.out_dir: + OUT = Path(args.out_dir).expanduser() + OUT.mkdir(parents = True, exist_ok = True) + + def filt(mod, fqn = ""): + return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf + + def run( + tag, + *, + cfg = None, + compile = True, + ): + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load_dense() + if cfg is not None: + from torchao.quantization import quantize_ + quantize_(pipe.transformer, cfg, filter_fn = filt) + if compile: + try: + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + except Exception as exc: # noqa: BLE001 + print( + f" [{tag}] compile failed: {type(exc).__name__}: {str(exc)[:90]}", flush = True + ) + _gen(pipe, steps, seed, res) # warmup / compile + dts, img = [], None + for _ in range(args.iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + gp = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, gp + + from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8 + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV + + print(f"== nvfp4 probe (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", flush = True) + bref, ref, _ = run("bf16_eager", cfg = None, compile = False) + print(f" bf16 eager ref: {bref:.3f}s", flush = True) + rows = [("bf16_eager", bref, float("inf"), 0.0, None)] + + specs = [ + ("bf16_compile", None, True), + ("fp8_compile", FP8(), True), + ("nvfp4_notriton_compile", NV(use_triton_kernel = False), True), + ("nvfp4_notriton_eager", NV(use_triton_kernel = False), False), + ] + for tag, cfg, comp in specs: + try: + med, arr, gp = run(tag, cfg = cfg, compile = comp) + ps, lp = _psnr(ref, arr), _lpips(ref, arr) + rows.append((tag, med, ps, lp, gp)) + print( + f" {tag:24s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" {tag:24s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True) + rows.append((tag, None, None, None, None)) + + fp8 = next((r[1] for r in rows if r[0] == "fp8_compile" and r[1]), None) + print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True) + for tag, med, ps, lp, gp in rows: + if med is None: + print(f" {tag:24s} FAILED") + continue + vs_fp8 = f"{fp8/med:.2f}x" if fp8 else "-" + psv = "inf" if ps == float("inf") else f"{ps:.1f}" + lpv = ( + "ref" + if (lp == 0.0 and tag == "bf16_eager") + else (f"{lp:.3f}" if lp is not None else "n/a") + ) + print( + f" {tag:24s} {med:.3f}s vs_fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}", + flush = True, + ) + print("NVFP4-PROBE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/nvfp4_t211_probe.py b/scripts/nvfp4_t211_probe.py new file mode 100644 index 0000000000..e69e01b2de --- /dev/null +++ b/scripts/nvfp4_t211_probe.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""torch>=2.11 NVFP4 probe. Three parts: + + A. diagnostics -- torch/torchao versions, cpp-extension load state, device. + B. GEMM micro -- isolated per-linear forward latency (bf16 / fp8 / nvfp4-cutlass / + nvfp4-triton) at Z-Image-like shapes, to measure raw FP4 + tensor-core throughput free of pipeline overhead. + C. end-to-end -- real dense Z-Image transformer, latency + LPIPS + PSNR + VRAM, + reference = dense bf16 eager. + +Run on one CUDA (Blackwell) GPU. This is the experiment that decides whether NVFP4 +becomes a genuine speedup once torch>=2.11 + torchao's CUTLASS FP4 GEMM is present.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_t211_images") + + +# ----------------------------------------------------------------------------- diag +def diagnostics() -> None: + import torch + import torchao + + print("== A. diagnostics ==", flush = True) + print(f" torch {torch.__version__}", flush = True) + print(f" torchao {torchao.__version__}", flush = True) + print(f" cuda {torch.version.cuda}", flush = True) + if torch.cuda.is_available(): + print( + f" device {torch.cuda.get_device_name(0)} sm{torch.cuda.get_device_capability(0)}", + flush = True, + ) + print(f" torch.ops.torchao present: {hasattr(torch.ops, 'torchao')}", flush = True) + print( + f" fp4 primitives: e2m1={hasattr(torch, 'float4_e2m1fn_x2')} " + f"e8m0={hasattr(torch, 'float8_e8m0fnu')} _scaled_mm={hasattr(torch, '_scaled_mm')}", + flush = True, + ) + # torchao prints "Skipping import of cpp extensions ..." to stderr at import on torch<2.11. + # On 2.11 that line is absent -> the CUTLASS FP4 GEMM extension is live. + print( + " (no 'Skipping import of cpp extensions' line above => cpp/CUTLASS ext loaded)", + flush = True, + ) + + +# ----------------------------------------------------------------------------- micro +def _configs(): + from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8 + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV + return { + "bf16": None, + "fp8": FP8(), + "nvfp4_cutlass": NV(use_triton_kernel = False), + "nvfp4_triton": NV(use_triton_kernel = True), + } + + +def _bench_linear(K, N, M, cfg, iters, compile_): + import torch + import torch.nn as nn + from torchao.quantization import quantize_ + + torch.compiler.reset() + torch.cuda.empty_cache() + m = nn.Sequential(nn.Linear(K, N, bias = False)).cuda().to(torch.bfloat16) + if cfg is not None: + quantize_(m, cfg) + fn = torch.compile(m, fullgraph = True, dynamic = False) if compile_ else m + x = torch.randn(M, K, device = "cuda", dtype = torch.bfloat16) + with torch.no_grad(): + for _ in range(3): # warmup / compile + fn(x) + torch.cuda.synchronize() + dts = [] + for _ in range(iters): + t0 = time.perf_counter() + fn(x) + torch.cuda.synchronize() + dts.append(time.perf_counter() - t0) + del m, fn, x + torch.cuda.empty_cache() + med = sorted(dts)[len(dts) // 2] + tflops = 2.0 * M * K * N / med / 1e12 + return med, tflops + + +def micro(M, iters, compile_): + print(f"\n== B. GEMM micro (M={M}, compile={compile_}, iters={iters}) ==", flush = True) + # (K, N): qkv-ish, mlp-up, mlp-down for a ~3072-dim DiT + shapes = [(3072, 3072), (3072, 12288), (12288, 3072)] + cfgs = _configs() + for K, N in shapes: + print(f" shape K={K} N={N}:", flush = True) + base_ms = None + fp8_ms = None + for name, cfg in cfgs.items(): + try: + med, tfl = _bench_linear(K, N, M, cfg, iters, compile_) + ms = med * 1e3 + if name == "bf16": + base_ms = ms + if name == "fp8": + fp8_ms = ms + vs_bf16 = f"{base_ms/ms:.2f}x" if base_ms else "-" + vs_fp8 = f"{fp8_ms/ms:.2f}x" if fp8_ms else "-" + print( + f" {name:16s} {ms:7.3f} ms {tfl:7.1f} TFLOPS vs_bf16={vs_bf16:>6s} vs_fp8={vs_fp8:>6s}", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + print(f" {name:16s} FAILED: {type(exc).__name__}: {str(exc)[:120]}", flush = True) + + +# ----------------------------------------------------------------------------- e2e +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import lpips + import torch + + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _load_dense(): + import diffusers + import torch + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + 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] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def e2e(steps, res, seed, iters, mf): + import torch + import torch.nn as nn + + OUT.mkdir(parents = True, exist_ok = True) + + def filt(mod, fqn = ""): + return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf + + def run( + tag, + *, + cfg = None, + compile = True, + ): + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load_dense() + if cfg is not None: + from torchao.quantization import quantize_ + quantize_(pipe.transformer, cfg, filter_fn = filt) + if compile: + try: + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + except Exception as exc: # noqa: BLE001 + print( + f" [{tag}] compile failed: {type(exc).__name__}: {str(exc)[:90]}", flush = True + ) + _gen(pipe, steps, seed, res) # warmup / compile + dts, img = [], None + for _ in range(iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + gp = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, gp + + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV + from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8 + + print( + f"\n== C. end-to-end (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", flush = True + ) + bref, ref, _ = run("bf16_eager", cfg = None, compile = False) + print(f" bf16 eager ref: {bref:.3f}s", flush = True) + rows = [("bf16_eager", bref, float("inf"), 0.0, None)] + + specs = [ + ("bf16_compile", None, True), + ("fp8_compile", FP8(), True), + ("nvfp4_cutlass_compile", NV(use_triton_kernel = False), True), + ("nvfp4_triton_compile", NV(use_triton_kernel = True), True), + ] + for tag, cfg, comp in specs: + try: + med, arr, gp = run(tag, cfg = cfg, compile = comp) + ps, lp = _psnr(ref, arr), _lpips(ref, arr) + rows.append((tag, med, ps, lp, gp)) + print( + f" {tag:24s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" {tag:24s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True) + rows.append((tag, None, None, None, None)) + + fp8 = next((r[1] for r in rows if r[0] == "fp8_compile" and r[1]), None) + print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True) + for tag, med, ps, lp, gp in rows: + if med is None: + print(f" {tag:24s} FAILED") + continue + vs_fp8 = f"{fp8/med:.2f}x" if fp8 else "-" + psv = "inf" if ps == float("inf") else f"{ps:.1f}" + lpv = ( + "ref" + if (lp == 0.0 and tag == "bf16_eager") + else (f"{lp:.3f}" if lp is not None else "n/a") + ) + print( + f" {tag:24s} {med:.3f}s vs_fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}", + flush = True, + ) + + +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("--micro-M", type = int, default = 4096) + p.add_argument("--min-feat", type = int, default = 512) + p.add_argument("--only", choices = ["diag", "micro", "e2e", "all"], default = "all") + args = p.parse_args(argv) + + diagnostics() + if args.only in ("micro", "all"): + micro(args.micro_M, args.iters, compile_ = True) + if args.only in ("e2e", "all"): + e2e(args.steps, args.res, args.seed, args.iters, args.min_feat) + print("NVFP4-T211-PROBE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/quant_probe.py b/scripts/quant_probe.py new file mode 100644 index 0000000000..8675fb75bf --- /dev/null +++ b/scripts/quant_probe.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Empirical quant probe: torchao int8/fp8/fp4 dynamic quant vs GGUF+compile. + +Question this answers: GGUF stores the Z-Image DiT at 4-bit but dequantizes to bf16 +per matmul, so it runs at bf16 tensor-core rate. Can a low-precision *tensor-core* +path (int8dq on any Ampere+, fp8dq on Ada+, NVFP4/MXFP8 on Blackwell), loaded from +the dense bf16 transformer, beat GGUF+compile on speed while staying inside the +quality bar -- and how does its quality compare to GGUF's own 4-bit loss? + +Reference for all quality numbers is the DENSE bf16 EAGER image (the best this model +can do). Each config is a fresh pipeline (no compile/quant cross-contamination). +Reports median latency, PSNR + LPIPS vs reference, and peak VRAM. Run on one CUDA GPU. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +REPO = "unsloth/Z-Image-Turbo-GGUF" +GGUF = "z-image-turbo-Q4_K_M.gguf" +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "probe_images" + + +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +_LPIPS = {"fn": None} + + +def _lpips(ref_arr, arr): + """Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable. + + Runs on CPU so the scorer never holds CUDA memory: each row resets peak VRAM, so a + resident GPU LPIPS module would inflate the reported load/gen VRAM and could even OOM.""" + try: + import torch + import lpips + + if _LPIPS["fn"] is None: + _LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval() + + def t(x): + return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0 + + 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) + return None + + +def _load_dense(): + import torch + import diffusers + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _load_gguf(): + import torch + import diffusers + from huggingface_hub import hf_hub_download + + 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", + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _quant_config(name): + """Return a torchao config instance for `name`, or raise to mark FAILED.""" + from torchao.quantization import ( + Int8WeightOnlyConfig, + Int8DynamicActivationInt8WeightConfig, + Float8DynamicActivationFloat8WeightConfig, + ) + + if name == "int8wo": + return Int8WeightOnlyConfig() + if name == "int8dq": + return Int8DynamicActivationInt8WeightConfig() + if name == "fp8dq": + return Float8DynamicActivationFloat8WeightConfig() + if name == "nvfp4": + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig + return NVFP4DynamicActivationNVFP4WeightConfig() + if name == "mxfp8": + from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig + try: + import torch + return MXDynamicActivationMXWeightConfig( + activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn + ) + except (TypeError, AttributeError): + return MXDynamicActivationMXWeightConfig() + raise ValueError(name) + + +def _make_filter_fn(min_features): + """Keep only the FLOP-heavy linears: nn.Linear with both in/out >= min_features. + The int8 dynamic path uses torch._int_mm (needs activation M>16), and the tiny + 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 = ""): + return ( + isinstance(module, nn.Linear) + and getattr(module, "in_features", 0) >= min_features + and getattr(module, "out_features", 0) >= min_features + ) + + return filter_fn + + +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) + log(f" quantized transformer with {name} ({q}/{total} linears >= {min_features} feat)") + + +def _compile(pipe, log): + fn = getattr(pipe.transformer, "compile_repeated_blocks", None) + if not callable(fn): + return False + for kw in ({"fullgraph": True, "dynamic": True}, {"dynamic": True}, {}): + try: + fn(**kw) + log(f" compiled repeated blocks {kw}") + return True + except Exception as exc: # noqa: BLE001 + log(f" compile {kw} failed: {type(exc).__name__}: {str(exc)[:90]}") + return False + + +def _gen(pipe, steps, seed, res): + import torch + + 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] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +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", + ) + 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, + ): + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load_dense() if source == "dense" else _load_gguf() + load_peak = torch.cuda.max_memory_allocated() / 1e9 + if quant is not None: + _apply_quant(pipe, quant, print_, args.min_feat) + if compile: + _compile(pipe, print_) + _gen(pipe, steps, seed, res) # warmup / compilation + else: + _gen(pipe, steps, seed, res) # allocator warmup + torch.cuda.reset_peak_memory_stats() + dts, img = [], None + for _ in range(iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + gen_peak = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return tag, _median(dts), arr, load_peak, gen_peak + + 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), + } + 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) + ref_arr = None + rows = [] + for tag in want: + if tag not in table: + 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) + try: + _, 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) + 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) + ) + 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, + ) + + 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, + ) + for tag, med, psnr, lpips_v, lp, gp in rows: + if med is None: + 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") + ) + 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) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/sparse_accum_probe.py b/scripts/sparse_accum_probe.py new file mode 100644 index 0000000000..4f91f905b6 --- /dev/null +++ b/scripts/sparse_accum_probe.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Probe two consumer-GPU-motivated levers on the real dense Z-Image transformer: + + * fp8 fast_accum on/off -- on consumer Blackwell, fp8 with FP16 accumulate is ~2x + fp8 with FP32 accumulate (838 vs 419 TFLOPS). torchao defaults use_fast_accum=True, + so this confirms we are already on the fast path and quantifies it (muted on a B200, + which is not nerfed, but the knob still moves latency). + * 2:4 semi-structured sparsity -- doubles tensor-core rate in theory. Two blockers to + test empirically: (a) QUALITY -- inference-only 2:4 magnitude-pruning drops 50% of + weights with no fine-tune; (b) it does NOT compose with torch.compile, so the real + sparse path runs eager. We measure sparse-no-compile speed vs our fp8+compile + baseline (the bar it must beat) and the LPIPS of 2:4 pruning. + +Reference for quality is the dense bf16 eager image. Run on one CUDA GPU. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/sparse_images") + + +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import torch, lpips + + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _load_dense(): + import torch, diffusers + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _big_linears(transformer, min_feat = 512): + import torch.nn as nn + return [ + m + for m in transformer.modules() + if isinstance(m, nn.Linear) and m.in_features >= min_feat and m.out_features >= min_feat + ] + + +def _prune_24_(transformer, min_feat = 512): + """In-place 2:4 magnitude prune (zero the 2 smallest of every 4 along in_features) + of the FLOP-heavy linears. Dense format -> measures the QUALITY of 2:4 with no kernel.""" + import torch + + n = 0 + for lin in _big_linears(transformer, min_feat): + w = lin.weight.data + o, i = w.shape + if i % 4: + continue + g = w.view(o, i // 4, 4) + idx = g.abs().argsort(dim = -1)[..., :2] + g.scatter_(-1, idx, 0.0) + n += 1 + return n + + +def _gen(pipe, steps, seed, res): + import torch + + 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] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +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) + args = p.parse_args(argv) + steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat + import torch + + OUT.mkdir(parents = True, exist_ok = True) + + def filt(mod, fqn = ""): + import torch.nn as nn + return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf + + def run( + tag, + *, + quant = None, + fast_accum = True, + prune = False, + real_sparse = False, + compile = True, + ): + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load_dense() + note = "" + if prune or real_sparse: + n = _prune_24_(pipe.transformer, mf) + note += f" pruned24={n}" + if real_sparse: + from torchao.sparsity import sparsify_, semi_sparse_weight + sparsify_(pipe.transformer, semi_sparse_weight(), filter_fn = filt) + note += " +semi_sparse" + if quant == "fp8": + from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig + from torchao.float8 import Float8MMConfig + + cfg = Float8DynamicActivationFloat8WeightConfig( + mm_config = Float8MMConfig(use_fast_accum = fast_accum) + ) + quantize_(pipe.transformer, cfg, filter_fn = filt) + note += f" fp8(fast_accum={fast_accum})" + if compile: + try: + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + except Exception as exc: # noqa: BLE001 + note += f" [compile FAILED {type(exc).__name__}]" + print(f" [{tag}]{note}", flush = True) + _gen(pipe, steps, seed, res) # warmup / compile + dts = [] + img = None + for _ in range(args.iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + gp = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, gp + + print( + f"== sparse/accum probe (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", + flush = True, + ) + rows = [] + # quality reference: dense bf16 eager (no compile, no quant) + bref, ref, _ = run("bf16_eager", compile = False) + rows.append(("bf16_eager", bref, float("inf"), 0.0, None)) + print(f" bf16 eager ref: {bref:.3f}s", flush = True) + + specs = [ + ("bf16_compile", dict()), + ("fp8_fastT_c", dict(quant = "fp8", fast_accum = True)), + ("fp8_fastF_c", dict(quant = "fp8", fast_accum = False)), + ( + "fake24_fp8_c", + dict(quant = "fp8", fast_accum = True, prune = True), + ), # quality of 2:4+fp8 (fake=no kernel) + ( + "real24_nocompile", + dict(real_sparse = True, compile = False), + ), # sparse SPEED (no quant, no compile) + ( + "real24_compile_try", + dict(real_sparse = True, compile = True), + ), # does sparse survive compile? + ] + for tag, kw in specs: + try: + med, arr, gp = run(tag, **kw) + ps, lp = _psnr(ref, arr), _lpips(ref, arr) + rows.append((tag, med, ps, lp, gp)) + print( + f" {tag:18s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" {tag:18s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True) + rows.append((tag, None, None, None, None)) + + print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True) + base = next((r[1] for r in rows if r[0] == "fp8_fastT_c" and r[1]), None) + for tag, med, ps, lp, gp in rows: + if med is None: + print(f" {tag:18s} FAILED") + continue + vs_eager = f"{bref/med:.2f}x" + vs_fp8 = f"{base/med:.2f}x" if base else "-" + psv = "inf" if ps == float("inf") else f"{ps:.1f}" + lpv = ( + "ref" + if (lp == 0.0 and tag == "bf16_eager") + else (f"{lp:.3f}" if lp is not None else "n/a") + ) + print( + f" {tag:18s} {med:.3f}s eager:{vs_eager:>6s} fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}", + flush = True, + ) + print("SPARSE-ACCUM-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 9500f2ed9f..3cac5629ac 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -45,6 +45,7 @@ from .diffusion_memory import ( snapshot_device_memory, ) from .diffusion_speed import ( + SPEED_DEFAULT, SPEED_OFF, apply_speed_optims, resolve_speed_mode, @@ -52,6 +53,11 @@ from .diffusion_speed import ( snapshot_backend_flags, ) from .diffusion_precision import quantize_text_encoders +from .diffusion_transformer_quant import ( + dense_transformer_supported, + normalize_transformer_quant, + quantize_transformer, +) logger = get_logger(__name__) @@ -81,6 +87,9 @@ class _LoadState: backend_flags_before: Optional[dict] = None # Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C). text_encoder_quant: Optional[str] = None + # Transformer quant actually engaged on the opt-in dense fast path: "int8" | "fp8" + # | "nvfp4" | "mxfp8" | None. None means the default GGUF transformer was loaded. + transformer_quant: Optional[str] = None @dataclass @@ -288,6 +297,8 @@ class DiffusionBackend: memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, + transformer_quant: Optional[str] = None, + transformer_quant_fast_accum: Optional[bool] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" fam = self.validate_load_request( @@ -319,6 +330,8 @@ class DiffusionBackend: memory_mode = memory_mode, speed_mode = speed_mode, text_encoder_quant = text_encoder_quant, + transformer_quant = transformer_quant, + transformer_quant_fast_accum = transformer_quant_fast_accum, _load_token = token, ), daemon = True, @@ -452,6 +465,8 @@ class DiffusionBackend: memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, + transformer_quant: Optional[str] = None, + transformer_quant_fast_accum: Optional[bool] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -484,98 +499,146 @@ class DiffusionBackend: # checkpoints never sit in VRAM at once. self._unload_locked() - # Dequantise the GGUF transformer on-device; the VAE / text-encoder / - # scheduler come from the base diffusers repo (GGUF is transformer-only). gguf_path = self._resolve_gguf_path(repo_id, gguf_filename, hf_token) transformer_cls = getattr(diffusers, fam.transformer_class) - transformer = transformer_cls.from_single_file( - gguf_path, - quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), - torch_dtype = dtype, - config = base, - subfolder = "transformer", - # Forward the token: the config is fetched from the (possibly gated) - # base repo before from_pretrained gets a chance to authenticate. - token = hf_token, + pipeline_cls = getattr(diffusers, fam.pipeline_class) + + # Decide placement up front (the weights are still on CPU, so free VRAM is + # the real budget) -- this also doubles as the dense-quant preflight: the + # dense bf16 transformer must fit resident, so the fast path is offered only + # when the plan is `none`. + plan = self._plan_memory( + target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload ) - pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} - if hf_token: - pipe_kwargs["token"] = hf_token - pipeline_cls = getattr(diffusers, fam.pipeline_class) - pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + # Opt-in 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 + # / scheme, OOM, partial quant) falls back to the GGUF build below. + pipe = None + transformer_quant_engaged = None + if ( + normalize_transformer_quant(transformer_quant) is not None + and dense_transformer_supported(target) + and plan.offload_policy == OFFLOAD_NONE + ): + try: + pipe, transformer_quant_engaged = self._load_dense_quant_pipeline( + transformer_cls, + pipeline_cls, + base, + device, + dtype, + hf_token, + target, + transformer_quant, + transformer_quant_fast_accum, + ) + except Exception as exc: # noqa: BLE001 — fall back to the GGUF build + logger.warning( + "diffusion.transformer_quant_fallback: %s (loading GGUF)", exc + ) + pipe = None + transformer_quant_engaged = None + clear_gpu_cache() + + if pipe is None: + # Default: dequantise the single-file GGUF transformer on-device; the + # VAE / text-encoder / scheduler come from the base diffusers repo + # (GGUF is transformer-only). + transformer = transformer_cls.from_single_file( + gguf_path, + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), + torch_dtype = dtype, + config = base, + subfolder = "transformer", + # Forward the token: the config is fetched from the (possibly gated) + # base repo before from_pretrained gets a chance to authenticate. + token = hf_token, + ) + + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} + if hf_token: + pipe_kwargs["token"] = hf_token + pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) # Resolve the effective speed mode: GGUF models default to the # near-lossless `default` profile (compile is ~2.2x and sits below # the quant noise floor), dense models stay bit-identical `off`. An # explicit speed_mode (incl. "off") is honored verbatim. effective_speed = resolve_speed_mode(speed_mode, is_gguf = bool(gguf_filename)) + # A torchao-quantized dense transformer runs its matmuls through the + # regional torch.compile; UNcompiled (eager) it is ~30x slower and would + # lose to the GGUF fallback. A dense model otherwise resolves to `off`, so + # force at least `default` (regional compile) whenever the quant engaged, + # or the opt-in "fast" path silently commits an eager, pathologically slow + # pipeline. + if transformer_quant_engaged is not None and effective_speed == SPEED_OFF: + logger.info( + "diffusion.transformer_quant: forcing speed_mode=default " + "(quantized transformer must be compiled; eager is ~30x slower)" + ) + effective_speed = SPEED_DEFAULT # Opt-in speed optims run BEFORE placement (channels_last / compile # must precede CPU offload). Snapshot the process-wide backend flags # first so unload can restore them: TF32 / cudnn.benchmark are global, # and a later `off` load must not inherit this load's settings. backend_flags_before = snapshot_backend_flags() - # apply_speed_optims mutates PROCESS-WIDE flags (TF32 / cudnn.benchmark); - # they are only restored via _LoadState.backend_flags_before on unload. If - # the build fails after this but before _state commits (e.g. an OOM in - # apply_memory_plan / pipe.to), nothing would restore them and a later `off` - # generation would be contaminated, so restore on any non-committed exit. - committed = False - try: - speed_applied = apply_speed_optims( - pipe, - target, - is_gguf = bool(gguf_filename), - family = fam, - speed_mode = effective_speed, - logger = logger, - ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), - # also before placement so the offload hooks move the smaller weights. - te_quant = quantize_text_encoders( - pipe, - target, - mode = text_encoder_quant, - logger = logger, + speed_applied = apply_speed_optims( + pipe, + target, + is_gguf = bool(gguf_filename), + family = fam, + speed_mode = effective_speed, + logger = logger, + ) + if transformer_quant_engaged is not None and not speed_applied.get("compiled"): + # Promotion above could not engage compile (e.g. the family is not + # compile-friendly, or compile_repeated_blocks failed): the quantized + # transformer is now running eager, which is far slower than the GGUF + # path it replaced. Surface it loudly rather than hiding the regression. + logger.warning( + "diffusion.transformer_quant: %s engaged but the transformer is NOT " + "compiled; eager torchao quant is ~30x slower than GGUF here", + transformer_quant_engaged, ) + # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), + # also before placement so the offload hooks move the smaller weights. + te_quant = quantize_text_encoders( + pipe, + target, + mode = text_encoder_quant, + logger = logger, + ) - # Decide placement from MEASURED free device memory vs the model's - # estimated resident size (transformer GGUF dequantised + the - # companion text-encoder / VAE already cached for `base`), then - # apply it. Computed here, after the build but before placement, - # because the weights are still on CPU so free VRAM is the real - # budget. `cpu_offload=True` stays an explicit override. - plan = self._plan_memory( - target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload - ) - # apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it - # may fall back to whole-module offload, and tiling is a no-op on a - # pipeline with no tiling control), so status stays honest. - effective_policy, effective_tiling = apply_memory_plan( - pipe, plan, device = device, logger = logger - ) + # Apply the placement planned above (from MEASURED free device memory vs + # the model's estimated resident size). apply_memory_plan returns the + # (policy, tiling) ACTUALLY engaged (it may fall back to whole-module + # offload, and tiling is a no-op on a pipeline with no tiling control), so + # status stays honest. The dense fast path already placed the pipe resident; + # for the `none` policy this is an idempotent re-placement. + effective_policy, effective_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) - self._state = _LoadState( - pipe = pipe, - family = fam, - repo_id = repo_id, - base_repo = base, - device = device, - dtype = str(dtype).replace("torch.", ""), - cpu_offload = effective_policy != OFFLOAD_NONE, - offload_policy = effective_policy, - vae_tiling = effective_tiling, - memory_mode = plan.requested_mode, - speed_mode = effective_speed, - speed_optims = tuple(k for k, v in speed_applied.items() if v), - backend_flags_before = backend_flags_before, - text_encoder_quant = te_quant, - ) - committed = True - finally: - if not committed: - restore_backend_flags(backend_flags_before) - clear_gpu_cache() + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = effective_policy != OFFLOAD_NONE, + offload_policy = effective_policy, + vae_tiling = effective_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + speed_optims = tuple(k for k, v in speed_applied.items() if v), + backend_flags_before = backend_flags_before, + text_encoder_quant = te_quant, + transformer_quant = transformer_quant_engaged, + ) logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", @@ -588,6 +651,39 @@ class DiffusionBackend: ) return self.status() + def _load_dense_quant_pipeline( + self, + transformer_cls: Any, + pipeline_cls: Any, + base: str, + device: str, + dtype: Any, + hf_token: Optional[str], + target: DiffusionDeviceTarget, + mode: Optional[str], + fast_accum: Optional[bool] = None, + ) -> tuple[Any, str]: + """Build the opt-in fast pipeline: load the DENSE bf16 transformer from the base + repo (``subfolder="transformer"``), assemble the pipeline, place it on the device, + and torchao-quantise the transformer in place. Returns ``(pipe, engaged_scheme)``. + + Raises if the scheme is unsupported or quantisation fails, so ``load_pipeline`` + catches it and falls back to the GGUF build. Quantisation runs ON the device (the + dynamic int8 / fp8 / fp4 kernels need the weights on CUDA) and BEFORE the loader + compiles the repeated block, so the order is quantize -> compile -> placement.""" + transformer = transformer_cls.from_pretrained( + base, subfolder = "transformer", torch_dtype = dtype, token = hf_token + ) + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} + if hf_token: + pipe_kwargs["token"] = hf_token + pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + pipe.to(device) + scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger) + if scheme is None: + raise RuntimeError("transformer quant unsupported for this device/scheme") + return pipe, scheme + def _plan_memory( self, target: DiffusionDeviceTarget, @@ -796,6 +892,7 @@ class DiffusionBackend: "speed_mode": None, "speed_optims": [], "text_encoder_quant": None, + "transformer_quant": None, } return { "loaded": True, @@ -811,6 +908,7 @@ class DiffusionBackend: "speed_mode": state.speed_mode, "speed_optims": list(state.speed_optims), "text_encoder_quant": state.text_encoder_quant, + "transformer_quant": state.transformer_quant, } diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py new file mode 100644 index 0000000000..7cb1074425 --- /dev/null +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in low-precision quantisation of the diffusion DiT transformer. + +The default path loads the transformer as a single-file GGUF, which stores weights +4-bit but DEQUANTISES to bf16 on every matmul -- so it runs at bf16 tensor-core rate +and never touches the int8 / fp8 / fp4 tensor cores. It is a memory win that costs +speed. This module is the opt-in alternative: load the DENSE bf16 transformer from the +base repo and torchao-quantise it with a DYNAMIC-ACTIVATION scheme so the matmul runs +on the low-precision tensor cores. Measured on a B200 (Z-Image-Turbo, 1024px / 8 steps) +vs the GGUF+compile default (0.802s, LPIPS 0.083 vs dense bf16): fp8 dynamic 0.585s +(1.37x), int8 dynamic 0.603s (1.33x), both at LOWER LPIPS than GGUF -- faster AND a hair +more accurate, at the cost of a higher-memory dense load. So it is strictly opt-in; the +loader keeps GGUF as the low-memory default and the fallback. + +Scheme by architecture (``auto`` picks the best supported, best first): + nvfp4 / mxfp8 - Blackwell sm_100+ FP4 / MX tensor cores (biggest win; prototype). + fp8 - Ada / Hopper / Blackwell (sm_89+) fp8 tensor cores. + int8 - Ampere+ (sm_80+) int8 tensor cores -- the broadest-hardware lever. + +Every scheme needs ``torch.compile`` to realise the speedup (dynamic quant is ~30x +slower eager); the loader already compiles the repeated block AFTER this runs. torch / +torchao are imported lazily so the module stays importable in a no-torch runtime, and +every probe is best-effort: an unsupported scheme yields None and the caller loads GGUF. +""" + +from __future__ import annotations + +from typing import Any, Optional + +TQ_INT8 = "int8" +TQ_FP8 = "fp8" +TQ_NVFP4 = "nvfp4" +TQ_MXFP8 = "mxfp8" +TQ_AUTO = "auto" +TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8) +TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES + +# Skip linears whose in/out features are below this. The int8 dynamic path uses +# torch._int_mm, which requires the activation row count M > 16, and the DiT's tiny +# timestep / pooled / modulation projections run at M=1 and crash it. They are a +# negligible share of the FLOPs, so leaving them bf16 costs ~nothing (measured: +# 239/276 Z-Image linears quantised, full speedup) and keeps quality a touch higher. +DEFAULT_MIN_LINEAR_FEATURES = 512 + +# Per-architecture preference order for ``auto`` -- best (fastest, in-bar) first, with +# the lower-precision schemes listed as fallbacks for that arch tier. On Blackwell, fp8 +# leads: measured on a B200, plain fp8 dynamic is both faster AND more accurate than the +# alternatives for the DiT's shapes. mxfp8's block scaling adds overhead without a speed +# win, so it sits below fp8. nvfp4 is intentionally below fp8 too: the FP4 tensor-core +# GEMM is real once torch>=2.11 + torchao's CUTLASS FP4 kernel is present (verified: a +# 16384^3 GEMM hits ~3826 TFLOPS, 1.37x fp8), but it only beats fp8 on very large GEMMs. +# At the DiT's actual shapes (hidden ~3072, MLP ~12288, M~4096) it is *slower* than fp8 +# (0.81x end-to-end on Z-Image 1024px) AND notably less accurate (LPIPS 0.166 vs fp8's +# 0.044), because FP4's per-forward quant overhead is not amortised and the format is +# coarser. So nvfp4 is kept as an explicit opt-in, never the auto pick for diffusion. +_AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( + ((10, 0), (TQ_FP8, TQ_NVFP4, 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 +) + +# Cache of (scheme, device) -> bool so the quantise+matmul smoke test runs once. +_SMOKE_CACHE: dict[tuple[str, str], bool] = {} + +# Data-center GPU model tokens (un-nerfed FP32 accumulate). Matched as whole tokens of +# torch.cuda.get_device_name(), so the workstation "A4000" is not mistaken for the +# data-center "A40". Anything not here -- GeForce, workstation RTX, or an unknown name -- +# is treated as consumer-class (FP32-accumulate halved). See developer.nvidia.com/cuda/gpus. +_DATACENTER_GPU_TOKENS = frozenset( + { + "B200", + "B100", + "GB200", + "GB300", + "GB10", # Blackwell data center + "H200", + "H100", + "H800", + "H20", # Hopper data center + "A100", + "A800", + "A30", + "A40", + "A16", + "A10", + "A2", # Ampere data center + "L40", + "L40S", + "L4", + "L20", + "L2", # Ada data center + "V100", + "P100", + "P40", + "T4", # legacy data center + } +) + + +def _is_consumer_gpu(device: Any = None) -> bool: + """Whether the active GPU is consumer / workstation class (GDDR), where fp8 FP32 + accumulate is throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center + HBM parts (recognised by name token) are not nerfed and return False, so they keep + the higher-precision default accumulate for free. Heuristic on the device name: a + GeForce / TITAN name is always consumer; a recognised data-center token is not; + anything else (workstation RTX, unknown) defaults to consumer -- the safe choice, + since fast accumulate is free on data-center and a win on consumer. Best-effort: + True on any probe failure.""" + try: + import re + + import torch + name = torch.cuda.get_device_name(device).upper() + except Exception: # noqa: BLE001 — no torch / no device -> assume consumer + return True + if "GEFORCE" in name or "TITAN" in name: + return True + tokens = set(re.split(r"[^A-Z0-9]+", name)) + return not (tokens & _DATACENTER_GPU_TOKENS) + + +def normalize_transformer_quant(value: Optional[str]) -> Optional[str]: + """Lower/strip a requested transformer quant; None / "" / "none" / "off" -> None. + + Raises ValueError for an unsupported value so a bad request is rejected cheaply.""" + if value is None: + return None + normalized = str(value).strip().lower().replace("-", "_") + if not normalized or normalized in ("none", "off"): + return None + if normalized not in TQ_MODES: + raise ValueError( + f"Unsupported transformer_quant '{value}'. Use one of: {', '.join(TQ_MODES)}." + ) + return normalized + + +def dense_transformer_supported(target: Any) -> bool: + """Whether the dense-source quant path is usable for ``target``: a CUDA device with + a bf16 compute dtype (the only configuration any torchao dynamic scheme accelerates). + A cheap pre-check the loader runs before loading the (large) dense transformer.""" + if getattr(target, "device", None) != "cuda": + return False + try: + import torch + return getattr(target, "dtype", None) is torch.bfloat16 + except Exception: + return False + + +def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Optional[str]: + """The concrete scheme to apply, or None to fall back to GGUF. + + ``auto`` walks the per-arch ladder and returns the first scheme that passes a real + quantise+matmul smoke test, so on a box where the Blackwell fp4 / mx kernels are + unavailable it lands on fp8 / int8 with no error. An explicit scheme is honored only + if supported (else None -> GGUF), never silently swapped for a different one.""" + requested = normalize_transformer_quant(requested) + if requested is None or not dense_transformer_supported(target): + return None + device = str(getattr(target, "device", "cuda")) + if requested != TQ_AUTO: + return requested if _scheme_supported(requested, device) else None + cap = _capability() + if cap is None: + return None + for floor, schemes in _AUTO_LADDER: + if cap >= floor: + for scheme in schemes: + if _scheme_supported(scheme, device): + return scheme + return None + return None + + +def _capability() -> Optional[tuple[int, int]]: + try: + import torch + major, minor = torch.cuda.get_device_capability() + return (int(major), int(minor)) + except Exception: + return None + + +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"): + return False + except Exception: + return False + return _smoke_probe(scheme, device) + + +def _smoke_probe(scheme: str, device: str) -> bool: + """True iff a tiny Linear quantised with ``scheme`` runs one M=32 forward without + error. Cached per (scheme, device). This is what makes ``auto`` robust to a torch / + torchao build where a prototype (nvfp4 / mxfp8) kernel is unavailable: it fails here + and the ladder moves on, rather than crashing at the first real denoise step.""" + key = (scheme, device) + if key in _SMOKE_CACHE: + return _SMOKE_CACHE[key] + ok = False + try: + import torch + from torchao.quantization import quantize_ + + lin = torch.nn.Linear(512, 512, bias = False).to(device = device, dtype = torch.bfloat16) + quantize_(lin, _make_quant_config(scheme), filter_fn = make_filter_fn(0)) + x = torch.randn(32, 512, device = device, dtype = torch.bfloat16) + with torch.no_grad(): + lin(x) + torch.cuda.synchronize() + ok = True + except Exception: + ok = False + _SMOKE_CACHE[key] = ok + return ok + + +def _resolve_fast_accum(fast_accum: Optional[bool]) -> bool: + """The fp8 ``use_fast_accum`` to apply. ``None`` auto-detects by GPU class + (consumer / workstation -> fast; data-center -> precise); an explicit bool forces it.""" + return _is_consumer_gpu() if fast_accum is None else bool(fast_accum) + + +def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: + """The torchao dynamic-activation config for ``scheme`` (lazy import; prototype + import for the Blackwell fp4 / mx schemes is inside the branch that needs it). + + ``fast_accum`` applies to fp8 only: None auto-detects by GPU class, True/False force it.""" + from torchao.quantization import ( + Float8DynamicActivationFloat8WeightConfig, + Int8DynamicActivationInt8WeightConfig, + ) + + if scheme == TQ_INT8: + return Int8DynamicActivationInt8WeightConfig() + if scheme == TQ_FP8: + # Choose fp8 accumulate by GPU class (unless forced). On consumer / workstation + # cards (GDDR) the fp8 tensor cores run ~2x faster with FP16 (fast) accumulate + # than FP32 (e.g. ~838 vs ~419 TFLOPS on RTX 50xx), so fast accumulate is a real + # win there. Data-center HBM parts default to the higher-precision accumulate. + # fast accumulate is a precision (not overflow) tradeoff and stays below the fp8 + # quant noise floor (measured 0 non-finite even on Z-Image's ~1e6 activations). + try: + from torchao.float8 import Float8MMConfig + return Float8DynamicActivationFloat8WeightConfig( + mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum)) + ) + except Exception: # noqa: BLE001 — older torchao without the explicit knob + return Float8DynamicActivationFloat8WeightConfig() + if scheme == TQ_NVFP4: + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig + + # Select the CUTLASS FP4 path, not the default Triton kernel: torchao defaults + # use_triton_kernel=True, which needs MSLK installed. On a Blackwell box with the + # CUTLASS FP4 extension but no MSLK, the default would make the smoke probe fail + # and silently fall back to GGUF instead of using the FP4 tensor cores. + try: + return NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel = False) + except TypeError: # older torchao without the knob + 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 + ) + except (TypeError, AttributeError): + # TypeError: older torchao without the explicit dtype knobs. + # AttributeError: a torch build without torch.float8_e4m3fn. + return MXDynamicActivationMXWeightConfig() + raise ValueError(f"unknown transformer quant scheme '{scheme}'") + + +def make_filter_fn(min_features: int): + """A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear + with both in/out features >= ``min_features``. Hides the (module, fqn) callback arity.""" + + def filter_fn(module: Any, fqn: str = "") -> bool: + try: + import torch + if not isinstance(module, torch.nn.Linear): + return False + except Exception: + return False + in_features = getattr(module, "in_features", None) + out_features = getattr(module, "out_features", None) + if in_features is None or out_features is None: + return False + return in_features >= min_features and out_features >= min_features + + return filter_fn + + +def quantize_transformer( + pipe: Any, + target: Any, + *, + mode: Optional[str], + min_features: int = DEFAULT_MIN_LINEAR_FEATURES, + fast_accum: Optional[bool] = None, + logger: Any = None, +) -> Optional[str]: + """Quantise ``pipe.transformer``'s FLOP-heavy linears in place with the arch-chosen + dynamic scheme. Returns the scheme actually engaged, or None when disabled / + unsupported / failed -- the caller then loads GGUF instead. Best-effort: it never + raises for an ordinary unsupported environment (a failure leaves the module dense). + + ``fast_accum`` (fp8 only) overrides the per-GPU-class accumulate choice: None + auto-detects (fast on consumer, precise on data-center), True/False force it.""" + scheme = select_transformer_quant_scheme(target, mode) + if scheme is None: + return None + transformer = getattr(pipe, "transformer", None) + if transformer is None: + return None + try: + from torchao.quantization import quantize_ + + quantize_( + transformer, + _make_quant_config(scheme, fast_accum = fast_accum), + filter_fn = make_filter_fn(min_features), + ) + # Runtime-only marker (torchao tensors are not safetensors-serializable; this + # backend is inference-only, so this is purely diagnostic). + try: + transformer._unsloth_runtime_quant = scheme + except Exception: # noqa: BLE001 — marker is best-effort + pass + return scheme + except Exception as exc: # noqa: BLE001 — leave the transformer dense -> GGUF fallback + _warn(logger, scheme, exc) + return None + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.transformer_quant: %s failed: %s", what, exc) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index fff68c090a..56ee571dae 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1720,6 +1720,23 @@ 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( + 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.", + ) + transformer_quant_fast_accum: Optional[bool] = Field( + None, + description = "fp8 only: FP8 matmul accumulate. null auto-detects by GPU class " + "(fast FP16 accumulate on consumer/workstation cards, where FP32 " + "accumulate is ~2x slower; precise FP32 accumulate on data-center " + "HBM cards, which are not nerfed). true/false force it. Negligible " + "quality effect (below the fp8 quant noise floor); no overflow risk.", + ) class DiffusionGenerateRequest(BaseModel): @@ -1830,3 +1847,8 @@ class DiffusionStatusResponse(BaseModel): text_encoder_quant: Optional[str] = Field( None, description = "Text-encoder quantisation engaged: fp8 | nvfp4 | null" ) + transformer_quant: Optional[str] = Field( + None, + description = "Transformer quant engaged on the dense fast path: int8 | fp8 | " + "nvfp4 | mxfp8 | null (null = the GGUF transformer was loaded)", + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1aa4be56b3..13beee00e0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10331,6 +10331,8 @@ async def load_diffusion_model( memory_mode = request.memory_mode, speed_mode = request.speed_mode, text_encoder_quant = request.text_encoder_quant, + transformer_quant = request.transformer_quant, + transformer_quant_fast_accum = request.transformer_quant_fast_accum, ) return DiffusionStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index fe4cabe8b5..2e714b217b 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -899,3 +899,131 @@ def test_load_fast_mode_stays_resident_on_cuda(fake_runtime, tmp_path, monkeypat ) assert status["offload_policy"] == "none" and status["cpu_offload"] is False assert backend._state.pipe.moved_to == "cuda" + + +# ── transformer quant (opt-in dense fast path) ──────────────────────────────── + + +def _stub_dense_quant(monkeypatch, *, scheme = "fp8"): + """Force the dense+quant branch hermetically: a supported dense source, a + from_pretrained on the fake transformer, and a quantizer that engages `scheme`. + Returns a dict recording the dense-loader / quantizer calls.""" + from core.inference import diffusion as dmod + + calls: dict = {"from_pretrained": 0, "quantize": 0, "quant_mode": None} + + @classmethod + def _from_pretrained(cls, base, **kwargs): + calls["from_pretrained"] += 1 + calls["fp_kwargs"] = {"base": base, **kwargs} + return object() + + monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + + def _quantize(pipe, target, *, mode, **kw): + calls["quantize"] += 1 + calls["quant_mode"] = mode + return scheme + + monkeypatch.setattr(dmod, "quantize_transformer", _quantize) + 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. + 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"), + ) + (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 status["transformer_quant"] is None + assert _FakeTransformer.last["path"] # GGUF from_single_file was used + + +def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatch): + # transformer_quant + a CUDA resident plan -> load the DENSE transformer from the + # base repo, place it on the device, quantise it, and report the engaged scheme. + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + 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", + transformer_quant = "fp8", + ) + assert status["transformer_quant"] == "fp8" + # No speed_mode was given, but a quantized transformer is ~30x slower eager, so the + # backend promotes it to `default` (regional compile) instead of the dense `off`. + assert status["speed_mode"] == "default" + assert calls["from_pretrained"] == 1 and calls["quantize"] == 1 + assert calls["quant_mode"] == "fp8" + assert calls["fp_kwargs"]["subfolder"] == "transformer" # dense transformer subfolder + # The GGUF single-file path was NOT used for the transformer. + assert _FakeTransformer.last == {} + # quantize ran on-device: the dense pipe was placed on cuda (before compile). + assert backend._state.pipe.moved_to == "cuda" + assert status["offload_policy"] == "none" + + +def test_transformer_quant_falls_back_to_gguf_on_failure(fake_runtime, tmp_path, monkeypatch): + # A dense/quant failure (here: quantize returns None -> unsupported) must fall back + # to the GGUF build, not error -- status reports no transformer_quant engaged. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + + @classmethod + def _from_pretrained(cls, base, **kwargs): + return object() + + monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False) + 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", + 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 + + +def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, monkeypatch): + # The dense bf16 transformer only fits resident, so when the memory plan would + # offload (here low_vram) the fast path is skipped and GGUF loads instead -- the + # dense transformer is never even loaded. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + + @classmethod + def _fp_fail(cls, *a, **k): + pytest.fail("dense transformer must not load when the plan offloads") + + 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", + ) + assert status["transformer_quant"] is None + assert status["offload_policy"] == "model" + assert _FakeTransformer.last["path"] # GGUF path used diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index b0fb3836e6..137b218636 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -381,6 +381,44 @@ def test_memory_mode_threads_through_to_backend(client, monkeypatch): assert backend.last_load_kwargs.get("memory_mode") == "low_vram" +def test_transformer_quant_threads_through_to_backend(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_quant": "auto"}, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("transformer_quant") == "auto" + + +def test_transformer_quant_fast_accum_threads_through(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_quant": "fp8", + "transformer_quant_fast_accum": False, + }, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("transformer_quant_fast_accum") is False + + +def test_invalid_transformer_quant_returns_422_without_eviction(client): + # An unsupported transformer_quant is rejected by the request schema (Literal), so + # the GPU is never acquired and no chat model is evicted. + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_quant": "int2"}, + ) + assert resp.status_code == 422 + assert gpu_arbiter._owner is None + + def test_invalid_memory_mode_returns_422_without_eviction(client): # An unsupported memory_mode is rejected by the request schema (Literal), so the # GPU is never acquired and no chat model is evicted. diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py new file mode 100644 index 0000000000..d2adaa9f7c --- /dev/null +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -0,0 +1,342 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for transformer quantisation (``diffusion_transformer_quant.py``). + +Hermetic: torch + torchao are stubbed via ``sys.modules``, and the per-scheme smoke +probe (``_scheme_supported`` / ``_smoke_probe``) is monkeypatched where the test cares +about the selection ladder rather than the GPU probe, so everything runs CPU-only. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +import core.inference.diffusion_transformer_quant as tq +from core.inference.diffusion_transformer_quant import ( + TQ_FP8, + TQ_INT8, + TQ_MXFP8, + TQ_NVFP4, + dense_transformer_supported, + make_filter_fn, + normalize_transformer_quant, + quantize_transformer, + select_transformer_quant_scheme, +) + + +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, +): + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + torch.float16 = "float16" + if with_fp8: + torch.float8_e4m3fn = "float8_e4m3fn" + torch.cuda = types.SimpleNamespace( + is_available = lambda: cuda_available, + get_device_capability = lambda *a: cc, + ) + monkeypatch.setitem(sys.modules, "torch", torch) + return torch + + +# ── normalisation ───────────────────────────────────────────────────────────── + + +def test_normalize_transformer_quant(): + assert normalize_transformer_quant(None) is None + assert normalize_transformer_quant("") is None + assert normalize_transformer_quant("none") is None + assert normalize_transformer_quant("off") is None + assert normalize_transformer_quant("AUTO") == "auto" + assert normalize_transformer_quant("INT8") == TQ_INT8 + assert normalize_transformer_quant("fp8") == TQ_FP8 + with pytest.raises(ValueError): + normalize_transformer_quant("int2") + + +# ── dense-source gate ─────────────────────────────────────────────────────────── + + +def test_dense_transformer_supported_requires_cuda_bf16(monkeypatch): + _stub_torch(monkeypatch) + assert dense_transformer_supported(_target()) is True + assert dense_transformer_supported(_target(device = "cpu")) is False + assert dense_transformer_supported(_target(dtype = "float16")) is False + + +# ── scheme selection ladder ───────────────────────────────────────────────────── + + +def _allow(monkeypatch, allowed): + """Force ``_scheme_supported`` to accept only ``allowed`` (simulates smoke results).""" + monkeypatch.setattr(tq, "_scheme_supported", lambda scheme, device: scheme in allowed) + + +def test_auto_blackwell_prefers_fp8_then_falls_back(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + # Even with every scheme available, auto picks fp8 on Blackwell: measured on a B200 + # (torch 2.11 + torchao CUTLASS FP4), fp8 is both faster and more accurate than nvfp4 + # for the DiT's shapes -- nvfp4's FP4 GEMM only wins on very large GEMMs, not here. + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + # fp8 unavailable: nvfp4 is the next pick (above mxfp8 / int8). + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_NVFP4 + # Only mxfp8 + int8 left -> mxfp8 (still above int8). + _allow(monkeypatch, {TQ_MXFP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_MXFP8 + # Only int8 usable -> int8. + _allow(monkeypatch, {TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + + +def test_auto_ada_hopper_prefers_fp8(monkeypatch): + _stub_torch(monkeypatch, cc = (8, 9)) + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + _stub_torch(monkeypatch, cc = (9, 0)) # Hopper + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + + +def test_auto_ampere_prefers_int8(monkeypatch): + _stub_torch(monkeypatch, cc = (8, 0)) + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) # fp8 cores absent on Ampere -> int8 only in ladder + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + _stub_torch(monkeypatch, cc = (8, 6)) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + + +def test_auto_pre_ampere_unsupported(monkeypatch): + _stub_torch(monkeypatch, cc = (7, 5)) # Turing: below the int8-dynamic floor + _allow(monkeypatch, {TQ_INT8, TQ_FP8}) + assert select_transformer_quant_scheme(_target(), "auto") is None + + +def test_explicit_scheme_honored_or_none(monkeypatch): + _stub_torch(monkeypatch, cc = (8, 0)) + _allow(monkeypatch, {TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "int8") == TQ_INT8 + # Explicit unsupported scheme is NOT silently downgraded -> None (-> GGUF fallback). + assert select_transformer_quant_scheme(_target(), "fp8") is None + assert select_transformer_quant_scheme(_target(), "nvfp4") is None + + +def test_select_none_when_disabled_or_non_cuda(monkeypatch): + _stub_torch(monkeypatch) + _allow(monkeypatch, {TQ_INT8, TQ_FP8, TQ_NVFP4}) + assert select_transformer_quant_scheme(_target(), None) is None + assert select_transformer_quant_scheme(_target(device = "cpu"), "auto") is None + + +# ── _scheme_supported / _smoke_probe ──────────────────────────────────────────── + + +def test_scheme_supported_shortcircuits(monkeypatch): + # No CUDA -> False without running the smoke probe. + _stub_torch(monkeypatch, cuda_available = False) + monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run")) + assert tq._scheme_supported(TQ_INT8, "cuda") is False + # fp8 requested but the fp8 dtype is missing -> False before the probe. + _stub_torch(monkeypatch, with_fp8 = False) + monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run")) + assert tq._scheme_supported(TQ_FP8, "cuda") is False + + +def test_smoke_probe_caches_and_tolerates_failure(monkeypatch): + tq._SMOKE_CACHE.clear() + calls = {"n": 0} + + class _Lin: + def __init__(self, *a, **k): + pass + + def to(self, **k): + return self + + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + torch.nn = types.SimpleNamespace(Linear = _Lin) + torch.randn = lambda *a, **k: object() + torch.no_grad = lambda: __import__("contextlib").nullcontext() + torch.cuda = types.SimpleNamespace(is_available = lambda: True, synchronize = lambda: None) + monkeypatch.setitem(sys.modules, "torch", torch) + + tqz = types.ModuleType("torchao.quantization") + + def _quantize_ok( + module, + config, + filter_fn = None, + ): + calls["n"] += 1 + + tqz.quantize_ = _quantize_ok + tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8cfg" + tqz.Float8DynamicActivationFloat8WeightConfig = lambda: "fp8cfg" + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + # _Lin is callable? No -> the forward lin(x) would fail. Make instances callable. + _Lin.__call__ = lambda self, x: x + + assert tq._smoke_probe(TQ_INT8, "cuda") is True + assert tq._smoke_probe(TQ_INT8, "cuda") is True # cached, no second quantize_ + assert calls["n"] == 1 + + # A scheme whose quantize_ raises -> probe False (and cached). + tq._SMOKE_CACHE.clear() + + 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 + + +# ── consumer-vs-datacenter detection (fp8 fast-accumulate gate) ────────────────── + + +def _stub_device_name(monkeypatch, name): + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace(get_device_name = lambda device = None: name) + monkeypatch.setitem(sys.modules, "torch", torch) + + +@pytest.mark.parametrize( + "name", + [ + "NVIDIA GeForce RTX 5090", + "NVIDIA GeForce RTX 4090", + "NVIDIA RTX A4000", # workstation: A4000 token, NOT the data-center A40 + "NVIDIA RTX 6000 Ada Generation", + "NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC) + ], +) +def test_is_consumer_gpu_true(monkeypatch, name): + _stub_device_name(monkeypatch, name) + assert tq._is_consumer_gpu() is True + + +@pytest.mark.parametrize( + "name", + [ + "NVIDIA B200", + "NVIDIA H100 80GB HBM3", + "NVIDIA A100-SXM4-80GB", + "NVIDIA A40", # data-center Ampere (distinct token from RTX A4000) + "NVIDIA L40S", + "NVIDIA L4", + "Tesla V100-SXM2-16GB", + ], +) +def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name): + _stub_device_name(monkeypatch, name) + assert tq._is_consumer_gpu() is False + + +def test_is_consumer_gpu_defaults_true_on_probe_failure(monkeypatch): + # No torch / no device name available -> assume consumer (safe: fast accum is free + # on data center and a win on consumer). + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace() # no get_device_name + monkeypatch.setitem(sys.modules, "torch", torch) + assert tq._is_consumer_gpu() is True + + +# ── filter ────────────────────────────────────────────────────────────────────── + + +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(types.SimpleNamespace(), "no_attrs") is False + + +# ── apply ─────────────────────────────────────────────────────────────────────── + + +def test_resolve_fast_accum(monkeypatch): + # None auto-detects by GPU class; an explicit bool forces it. + monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: True) + assert tq._resolve_fast_accum(None) is True + monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: False) + assert tq._resolve_fast_accum(None) is False + assert tq._resolve_fast_accum(True) is True # forced on (e.g. on a data-center card) + assert tq._resolve_fast_accum(False) is False # forced off (e.g. on a consumer card) + + +def test_quantize_transformer_applies_and_marks(monkeypatch): + monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8) + seen: dict = {} + + def _mk(scheme, fast_accum = None): + seen["scheme"], seen["fast_accum"] = scheme, fast_accum + return f"{scheme}cfg" + + monkeypatch.setattr(tq, "_make_quant_config", _mk) + recorder: list = [] + tqz = types.ModuleType("torchao.quantization") + 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", fast_accum = False) == 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 transformer._unsloth_runtime_quant == TQ_FP8 # diagnostic marker set + assert seen["fast_accum"] is False # the override is forwarded into the config + + +def test_quantize_transformer_none_when_unsupported(monkeypatch): + monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: None) + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + assert quantize_transformer(pipe, _target(), mode = "auto") is None + + +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, + ): + raise RuntimeError("partial quant failure") + + tqz.quantize_ = _boom + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + # A quantise failure returns None (caller falls back to GGUF), never raises. + assert quantize_transformer(pipe, _target(), mode = "int8") is None