diff --git a/scripts/quant_probe.py b/scripts/quant_probe.py index 777f945cc4..6d69e7e8e7 100644 --- a/scripts/quant_probe.py +++ b/scripts/quant_probe.py @@ -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 diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 64957f1031..38e03e7ced 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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( diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 0db57917df..c192e5c8c0 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -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: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index e68a057a46..634f322639 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -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" diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 23f1012175..c89f562c62 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -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())