diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index 9766467351..66f5f58a89 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -86,8 +86,10 @@ def main(argv = None) -> int: # Mirror the runtime path EXACTLY (the offline == runtime, LPIPS-0 invariant): for int8 also # skip the M=1 AdaLN-modulation / conditioning-embedder projections, else the saved checkpoint # bakes them as int8 and crashes (torch._int_mm needs M>16) at the first denoise step on - # Flux / Qwen. fp8 / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns (). - exclude_name_tokens = exclude_tokens_for_scheme(scheme) + # Flux / Qwen. fp8 / fp4 / mx use scaled_mm (no M limit) and exclude nothing EXCEPT on families + # whose padded conditioning would divide by a zero row scale (e.g. Wan's condition_embedder); + # pass the family so the offline set matches the runtime one exactly. + exclude_name_tokens = exclude_tokens_for_scheme(scheme, fam.name) # fp8 and mxfp8 assert a bf16 weight, so their filter must skip any non-bf16 Linear the # transformer keeps: a mixed-precision DiT (Wan / Hunyuan) retains its _keep_in_fp32_modules in # fp32 even under torch_dtype=bf16, so quantising one would raise inside quantize_ and abort the diff --git a/scripts/fp8_layer_ablation.py b/scripts/fp8_layer_ablation.py new file mode 100644 index 0000000000..3862e251db --- /dev/null +++ b/scripts/fp8_layer_ablation.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Per-layer fp8 ablation probe for video DiTs (Wan / Hunyuan): find which layers -- if any -- +break under production per-row fp8 (torch._scaled_mm), or prove the failure is systemic. + +Motivation: the dense video default auto-quantises the DiT; on Blackwell the auto ladder leads +with fp8. On Wan2.2 and HunyuanVideo-1.5 that renders every frame BLACK (int8 is clean); the +shipped fix denies fp8 for those families -> int8. The question this probe answers: is the black +frame caused by a small, nameable set of outlier layers (so we could keep fp8 on the rest, mixed), +or is it systemic to the _scaled_mm path (activation dynamic-quant overflow / accumulate) -- in +which case per-layer exclusion cannot help and int8 stays the right call? + +Method (single-forward proxy; cheap): run one bf16 generation, capture a REAL mid-schedule DiT +input tuple via a forward_pre_hook, then compare the dense bf16 DiT output on that tuple to the +output after quantising layer SUBSETS with the production fp8 config. Because it reuses +``_make_quant_config`` / ``make_filter_fn`` from diffusion_transformer_quant, the GEMM path is +byte-identical to production (no MSLK). + +Phase 0 (this file's default) is the gate: + 0a proxy-validity : fp8-ALL must reproduce the black signal in one forward (cosine collapse or + non-finite output). If fp8-ALL is CLEAN in one forward, the black frame is + multi-step / cache compounding -> single-forward bisection cannot localise it + -> STOP, keep int8. + 0b mechanism : (i) per-Linear finiteness hooks locate the first inf/NaN (a localisable + overflow layer); (ii) flip use_fast_accum True/False -- if it flips + black->clean the fix is a one-line accumulate flag, no per-layer work. +It also records per-Linear input outlier stats (within-token spread + amax) on the bf16 forward, +ready for the Phase 2 outlier ranking if a bucket localises. + +Example: + CUDA_VISIBLE_DEVICES=1 python scripts/fp8_layer_ablation.py --family wan2.2-ti2v-5b --phase 0 +""" + +from __future__ import annotations + +import argparse +import copy +import gc +import json +import os +import sys +import time +from pathlib import Path +from typing import Any, Optional + +os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1") +os.environ.setdefault("HF_HOME", "/mnt/disks/unslothai/ubuntu/workspace_81/BACKUP_05/temp/hf_cache") + +_REPO_ROOT = Path(__file__).resolve().parent.parent +_BACKEND_ROOT = _REPO_ROOT / "studio" / "backend" +for _p in (str(_BACKEND_ROOT), str(_REPO_ROOT / "scripts")): + if _p not in sys.path: + sys.path.insert(0, _p) + +# Reuse the benchmark's real pipeline-build + family plumbing so the load is identical. +from video_speedmem_bench import _FAMILIES, PROMPT, _build_pipe, _import_diffusers # noqa: E402 +from core.inference.diffusion_transformer_quant import ( # noqa: E402 + DEFAULT_MIN_LINEAR_FEATURES, + TQ_FP8, + _REQUIRE_BF16_SCHEMES, + _make_quant_config, + make_filter_fn, +) + + +# ── cuda helpers ──────────────────────────────────────────────────────────────── +def _empty() -> None: + import torch + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _detach_to_cpu(obj: Any) -> Any: + """Recursively detach + move tensors in a nested args/kwargs structure to CPU.""" + import torch + + if torch.is_tensor(obj): + return obj.detach().to("cpu") + if isinstance(obj, tuple): + return tuple(_detach_to_cpu(o) for o in obj) + if isinstance(obj, list): + return [_detach_to_cpu(o) for o in obj] + if isinstance(obj, dict): + return {k: _detach_to_cpu(v) for k, v in obj.items()} + return obj + + +def _to_device(obj: Any, device: str) -> Any: + """Recursively move tensors to ``device`` (dtypes preserved -- timestep stays int).""" + import torch + + if torch.is_tensor(obj): + return obj.to(device) + if isinstance(obj, tuple): + return tuple(_to_device(o, device) for o in obj) + if isinstance(obj, list): + return [_to_device(o, device) for o in obj] + if isinstance(obj, dict): + return {k: _to_device(v, device) for k, v in obj.items()} + return obj + + +def _extract_tensor(out: Any): + import torch + + if torch.is_tensor(out): + return out + s = getattr(out, "sample", None) + if torch.is_tensor(s): + return s + if isinstance(out, (list, tuple)) and out and torch.is_tensor(out[0]): + return out[0] + return out + + +# ── forward-tuple capture ──────────────────────────────────────────────────────── +class _Stop(Exception): + pass + + +def _capture_forward_tuple(pipe, *, steps, width, height, num_frames, guidance, seed, + capture_call, gvg): + """Run a bf16 generation and grab the ``capture_call``-th transformer forward's + (args, kwargs) to CPU, then abort. This is a REAL mid-schedule DiT input.""" + import torch + + holder: dict[str, Any] = {} + calls = [0] + + def _pre(mod, args, kwargs): + calls[0] += 1 + if calls[0] == capture_call: + holder["args"] = _detach_to_cpu(args) + holder["kwargs"] = _detach_to_cpu(kwargs) + raise _Stop() + return None + + handle = pipe.transformer.register_forward_pre_hook(_pre, with_kwargs=True) + g = torch.Generator(device="cuda").manual_seed(seed) + kwargs = dict( + prompt=PROMPT, width=width, height=height, num_frames=num_frames, + num_inference_steps=steps, generator=g, + ) + if gvg: + guider = getattr(pipe, "guider", None) + if guider is not None and hasattr(guider, "guidance_scale"): + try: + guider.guidance_scale = guidance + except Exception: + pass + else: + kwargs["guidance_scale"] = guidance + try: + pipe(**kwargs) + except _Stop: + pass + finally: + handle.remove() + if "kwargs" not in holder: + raise RuntimeError( + f"transformer forward never reached call #{capture_call} " + f"(saw {calls[0]}); lower --capture-call" + ) + return holder + + +def _forward_out(module, tup: dict, device: str): + import torch + + args = _to_device(tup.get("args", ()), device) + kwargs = _to_device(tup.get("kwargs", {}), device) + with torch.no_grad(): + out = module(*args, **kwargs) + return _extract_tensor(out) + + +# ── scoring ────────────────────────────────────────────────────────────────────── +def _score(ref, cand) -> dict: + """Compare a candidate DiT output to the bf16 reference. Non-finite entries are the + strongest black-frame signal; cosine / relL2 / norm-ratio are computed on finite entries.""" + import torch + + reff = ref.detach().flatten().float() + canf = cand.detach().flatten().float() + fin_c = torch.isfinite(canf) + frac_nonfinite = float(1.0 - fin_c.float().mean().item()) + mask = fin_c & torch.isfinite(reff) + if mask.any(): + a = canf[mask] + b = reff[mask] + cos = float(torch.nn.functional.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item()) + rel = float((torch.linalg.vector_norm(a - b) / (torch.linalg.vector_norm(b) + 1e-12)).item()) + norm_ratio = float((torch.linalg.vector_norm(a) / (torch.linalg.vector_norm(b) + 1e-12)).item()) + else: + cos, rel, norm_ratio = 0.0, float("inf"), 0.0 + return { + "cosine": round(cos, 5), + "relL2": round(rel, 5), + "frac_nonfinite": round(frac_nonfinite, 6), + "norm_ratio": round(norm_ratio, 5), + "cand_amax": round(float(canf[fin_c].abs().max().item()) if fin_c.any() else float("inf"), 3), + } + + +def _verdict(sc: dict) -> str: + """Black-frame proxy verdict from a score dict.""" + if sc["frac_nonfinite"] > 0: + return "BROKEN(non-finite)" + if sc["cosine"] < 0.5 or sc["norm_ratio"] < 0.3 or sc["norm_ratio"] > 3.0: + return "BROKEN(collapse)" + if sc["cosine"] < 0.98: + return "DEGRADED" + return "CLEAN" + + +# ── quant helpers ────────────────────────────────────────────────────────────── +def _fp8_filter(min_features: int, exclude_tokens=(), only_tokens=()): + """A production fp8 filter (require_bf16) optionally narrowed: + - exclude_tokens: skip linears whose fqn contains any token (keep them bf16) + - only_tokens : quantise ONLY linears whose fqn contains any token.""" + base = make_filter_fn(min_features, exclude_name_tokens=exclude_tokens, + require_bf16=(TQ_FP8 in _REQUIRE_BF16_SCHEMES)) + + def filt(module, fqn: str = "") -> bool: + if not base(module, fqn): + return False + name = (fqn or "").lower() + if exclude_tokens and any(t in name for t in exclude_tokens): + return False + if only_tokens and not any(t in name for t in only_tokens): + return False + return True + + return filt + + +def _ablate(dense_cpu, tup, *, min_features, fast_accum=None, exclude_tokens=(), only_tokens=(), + instrument=False): + """Deepcopy the dense DiT -> GPU -> fp8-quantise the selected subset -> one forward. + Returns (cpu_output_tensor, finiteness_records|None).""" + import torch + from torchao.quantization import quantize_ + + m = copy.deepcopy(dense_cpu).to("cuda") + filt = _fp8_filter(min_features, exclude_tokens=exclude_tokens, only_tokens=only_tokens) + quantize_(m, _make_quant_config(TQ_FP8, fast_accum=fast_accum), filter_fn=filt) + + records = None + hooks = [] + if instrument: + records = [] + + def mk(nm): + def hook(mod, inp, out): + o = out if torch.is_tensor(out) else (out[0] if isinstance(out, (list, tuple)) and out else None) + if o is None: + return + of = o.detach().float() + in_amax = None + in_frac_zero_rows = None + if inp and torch.is_tensor(inp[0]): + xin = inp[0].detach().float() + in_amax = float(xin.abs().max().item()) + if xin.dim() >= 2: + ra = xin.reshape(-1, xin.shape[-1]).abs().amax(dim=-1) + in_frac_zero_rows = round(float((ra < 1e-6).float().mean().item()), 4) + records.append({ + "name": nm, + "out_finite": bool(torch.isfinite(of).all().item()), + "out_amax": float(of.abs().max().item()) if torch.isfinite(of).any() else float("inf"), + "in_amax": in_amax, + "in_frac_zero_rows": in_frac_zero_rows, + }) + return hook + + for nm, mod in m.named_modules(): + if isinstance(mod, torch.nn.Linear): + hooks.append(mod.register_forward_hook(mk(nm))) + + out = _forward_out(m, tup, "cuda") + out = out.detach().to("cpu") + for h in hooks: + h.remove() + del m + _empty() + return out, records + + +# ── bf16 reference + input outlier stats (for Phase 2 ranking) ─────────────────── +# ── Phase 1 buckets (fqn substring tokens), keyed by DiT kind ──────────────────── +# Grounded in Phase 0: on Wan the inf originates in the layers whose ROW dimension is the +# padded text sequence (condition_embedder.text_embedder + cross-attn K/V that project +# encoder_hidden_states); the video self-attention (attn1) + attn2.to_q (video query) + FFN +# stay finite. "textpath" is the hypothesised minimal exclude set. +_BUCKETS: dict[str, dict[str, tuple[str, ...]]] = { + "wan": { + "textpath": ("text_embedder", "attn2.to_k", "attn2.to_v", "attn2.add_k", "attn2.add_v"), + "condition_embedder": ("condition_embedder",), + "attn2_all": ("attn2.",), + "attn2_kv": ("attn2.to_k", "attn2.to_v", "attn2.add_k", "attn2.add_v"), + "attn1": ("attn1.",), + "ffn": ("ffn.",), + }, + "hunyuan": { + # candidate fix: the three input embedders consuming zero-padded / zero conditioning + # (image_embeds all-zero for T2V, encoder_hidden_states_2 ByT5 all-zero, pooled + # time_text_embed). "context_embedder" substring also matches "context_embedder_2". + "embedders": ("context_embedder", "image_embedder"), + "context_embedder": ("context_embedder",), # text refiner + context_embedder_2 + "image_embedder": ("image_embedder",), + "context_embedder_2": ("context_embedder_2",), + "proj_out": ("proj_out",), + "main_blocks": ("transformer_blocks.",), # ONLY main blocks -> should be CLEAN + }, +} + + +def _dit_kind(family: str) -> str: + return "hunyuan" if "hunyuan" in family.lower() else "wan" + + +def _zero_row_diag(tup: dict) -> dict: + """Directly test the divide-by-zero hypothesis: fraction of near-zero (padding) rows in + the captured tensors. A per-row fp8 activation scale = row_amax / 448, so a row_amax of 0 + yields scale 0 -> x/0 = inf. Padded text sequences (encoder_hidden_states) are the suspect.""" + import torch + + diag = {} + for key, val in tup.get("kwargs", {}).items(): + if torch.is_tensor(val) and val.dim() >= 2 and val.is_floating_point(): + x = val.float().reshape(-1, val.shape[-1]) + row_amax = x.abs().amax(dim=-1) + n = int(row_amax.numel()) + n_zero = int((row_amax < 1e-6).sum().item()) + diag[key] = { + "rows": n, + "zero_amax_rows": n_zero, + "frac_zero_rows": round(n_zero / n, 4) if n else 0.0, + "min_row_amax": round(float(row_amax.min().item()), 8), + } + return diag + + +def _bf16_reference_and_stats(dense_gpu, tup, *, min_features): + """Reference output + per-Linear input outlier stats on the bf16 forward. ``spread`` = + mean over tokens of (per-token channel amax / per-token channel median-abs): PR#150's + outlier proxy. Only linears the fp8 filter would quantise are recorded.""" + import torch + + keep = _fp8_filter(min_features) # which linears fp8 would touch + stats: list[dict] = [] + hooks = [] + + def mk(nm): + def hook(mod, inp, out): + if not inp or not torch.is_tensor(inp[0]): + return + x = inp[0].detach().float() + x2 = x.reshape(-1, x.shape[-1]).abs() # [tokens, channels] + amax_tok = x2.amax(dim=-1) + med_tok = x2.median(dim=-1).values + spread = float((amax_tok / (med_tok + 1e-9)).mean().item()) + stats.append({ + "name": nm, + "in_amax": float(x2.max().item()), + "spread": round(spread, 2), + "in_features": int(mod.in_features), + "out_features": int(mod.out_features), + }) + return hook + + for nm, mod in dense_gpu.named_modules(): + if isinstance(mod, torch.nn.Linear) and keep(mod, nm): + hooks.append(mod.register_forward_hook(mk(nm))) + + ref = _forward_out(dense_gpu, tup, "cuda") + ref = ref.detach().to("cpu") + for h in hooks: + h.remove() + stats.sort(key=lambda s: s["spread"], reverse=True) + return ref, stats + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--family", default="wan2.2-ti2v-5b", choices=sorted(_FAMILIES)) + ap.add_argument("--phase", type=int, default=0, choices=[0, 1, 2]) + ap.add_argument("--only", default="", help="phase 2: comma tokens -- fp8 ONLY these linears") + ap.add_argument("--exclude", default="", help="phase 2: comma tokens -- fp8 all EXCEPT these") + ap.add_argument("--steps", type=int, default=20) + ap.add_argument("--num-frames", type=int, default=25) + ap.add_argument("--width", type=int, default=512) + ap.add_argument("--height", type=int, default=320) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--capture-call", type=int, default=8, + help="which transformer forward call to capture (~mid schedule)") + ap.add_argument("--min-features", type=int, default=DEFAULT_MIN_LINEAR_FEATURES) + ap.add_argument("--out", default="outputs/fp8_ablation") + args = ap.parse_args(argv) + + import torch + + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + + spec = _FAMILIES[args.family] + repo = spec["repo"] + force_fp32 = spec.get("vae_force_fp32", False) + guidance = spec.get("guidance", 5.0) + + from core.inference.video_families import detect_video_family + + fam_obj = detect_video_family(repo) + gvg = bool(getattr(fam_obj, "guidance_via_guider", False)) + + print(f"== fp8 layer ablation: family={args.family} repo={repo} phase={args.phase} ==", flush=True) + t0 = time.perf_counter() + pipe = _build_pipe(repo, force_fp32) + print(f"[load] pipe built in {time.perf_counter()-t0:.1f}s", flush=True) + + tup = _capture_forward_tuple( + pipe, steps=args.steps, width=args.width, height=args.height, + num_frames=args.num_frames, guidance=guidance, seed=args.seed, + capture_call=args.capture_call, gvg=gvg, + ) + ks = {k: (tuple(v.shape) if torch.is_tensor(v) else type(v).__name__) + for k, v in tup.get("kwargs", {}).items()} + print(f"[capture] call #{args.capture_call} kwargs={ks}", flush=True) + + # Keep the dense bf16 DiT (still on GPU) as reference source; copy to CPU as the ablation seed. + dense_gpu = pipe.transformer + ref, stats = _bf16_reference_and_stats(dense_gpu, tup, min_features=args.min_features) + print(f"[bf16 ref] out shape={tuple(ref.shape)} amax={float(ref.float().abs().max()):.3f} " + f"n_quantized_linears={len(stats)}", flush=True) + + dense_cpu = copy.deepcopy(dense_gpu).to("cpu") + # Free the pipeline (VAE/text encoders) to leave the GPU for one transformer at a time. + del pipe, dense_gpu + _empty() + + report: dict[str, Any] = {"family": args.family, "repo": repo, "capture_call": args.capture_call, + "input_shapes": ks, "ref_amax": round(float(ref.float().abs().max()), 3), + "zero_row_diag": _zero_row_diag(tup)} + print(f"[zero-row diag] {report['zero_row_diag']}", flush=True) + + if args.phase == 0: + # ── Phase 0a: proxy validity -- fp8-ALL (production accum = None auto-detect) ── + out_all, _ = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=None) + sc_all = _score(ref, out_all) + report["fp8_all_autoaccum"] = {**sc_all, "verdict": _verdict(sc_all)} + print(f"[0a fp8-ALL auto-accum] {sc_all} -> {_verdict(sc_all)}", flush=True) + + # ── Phase 0b-ii: accumulate flip ── + for fa in (False, True): + o, _ = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=fa) + sc = _score(ref, o) + report[f"fp8_all_fast_accum_{fa}"] = {**sc, "verdict": _verdict(sc)} + print(f"[0b accum={fa}] {sc} -> {_verdict(sc)}", flush=True) + + # ── Phase 0b-i: per-Linear finiteness (instrumented fp8-ALL forward) ── + _, records = _ablate(dense_cpu, tup, min_features=args.min_features, fast_accum=None, + instrument=True) + if records: + nonfinite = [r for r in records if not r["out_finite"]] + report["first_nonfinite_linears"] = nonfinite[:10] + report["n_nonfinite_linears"] = len(nonfinite) + finite_sorted = sorted([r for r in records if r["out_finite"]], + key=lambda r: r["out_amax"], reverse=True) + report["top_out_amax_linears"] = finite_sorted[:10] + print(f"[0b finiteness] {len(nonfinite)}/{len(records)} linears non-finite; " + f"first={[r['name'] for r in nonfinite[:5]]}", flush=True) + print(f"[0b top out_amax] {[(r['name'], round(r['out_amax'],1)) for r in finite_sorted[:5]]}", + flush=True) + + report["top_spread_linears"] = stats[:20] + print(f"[stats] top-spread linears: " + f"{[(s['name'], s['spread']) for s in stats[:8]]}", flush=True) + + va = report["fp8_all_autoaccum"]["verdict"] + if va.startswith("BROKEN"): + print("\n[GATE] fp8-ALL reproduces the black signal in one forward -> proxy VALID; " + "proceed to Phase 1 bucket ablation.", flush=True) + elif va == "CLEAN": + print("\n[GATE] fp8-ALL is CLEAN in one forward but full gen is black -> failure is " + "multi-step/cache; single-forward bisection cannot localise -> STOP, keep int8.", + flush=True) + else: + print(f"\n[GATE] fp8-ALL verdict={va}: borderline; inspect scores.", flush=True) + + elif args.phase == 1: # ── Phase 1: bucket ablation (necessity + sufficiency) ── + kind = _dit_kind(args.family) + buckets = _BUCKETS[kind] + results: dict[str, Any] = {} + + # baseline: fp8-ALL (should be BROKEN, matching Phase 0) + o, _ = _ablate(dense_cpu, tup, min_features=args.min_features) + results["fp8_all"] = {**_score(ref, o)} + results["fp8_all"]["verdict"] = _verdict(results["fp8_all"]) + print(f"[1 baseline fp8-ALL] {results['fp8_all']['verdict']} {results['fp8_all']}", flush=True) + + for bname, toks in buckets.items(): + # NECESSITY: fp8 everything EXCEPT this bucket -- if CLEAN, this bucket is the culprit. + o, _ = _ablate(dense_cpu, tup, min_features=args.min_features, exclude_tokens=toks) + sc_ex = _score(ref, o) + # SUFFICIENCY: fp8 ONLY this bucket -- if BROKEN, this bucket alone reproduces damage. + o2, _ = _ablate(dense_cpu, tup, min_features=args.min_features, only_tokens=toks) + sc_only = _score(ref, o2) + results[bname] = { + "exclude": {**sc_ex, "verdict": _verdict(sc_ex)}, + "only": {**sc_only, "verdict": _verdict(sc_only)}, + } + print(f"[1 {bname}] EXCLUDE->{_verdict(sc_ex)} (cos {sc_ex['cosine']}, " + f"nf {sc_ex['frac_nonfinite']}) | ONLY->{_verdict(sc_only)} " + f"(cos {sc_only['cosine']}, nf {sc_only['frac_nonfinite']})", flush=True) + + report["phase1"] = results + + else: # ── Phase 2: instrument a specific only/exclude set, find first overflow ── + only = tuple(t.strip() for t in args.only.split(",") if t.strip()) + exclude = tuple(t.strip() for t in args.exclude.split(",") if t.strip()) + o, records = _ablate(dense_cpu, tup, min_features=args.min_features, + only_tokens=only, exclude_tokens=exclude, instrument=True) + sc = _score(ref, o) + report["phase2"] = {"only": only, "exclude": exclude, **sc, "verdict": _verdict(sc)} + print(f"[2 only={only} exclude={exclude}] {sc} -> {_verdict(sc)}", flush=True) + if records: + # first non-finite linears in execution order, with their input amax: an in_amax of + # ~0 at the first-broken layer confirms a zero-amax (padding) input row -> scale 0 -> inf. + nonfinite = [r for r in records if not r["out_finite"]] + report["phase2_first_nonfinite"] = nonfinite[:15] + report["phase2_n_nonfinite"] = len(nonfinite) + print(f"[2 finiteness] {len(nonfinite)}/{len(records)} quantised-path linears non-finite", + flush=True) + for r in nonfinite[:8]: + print(f" first-inf {r['name']} in_amax={r['in_amax']} " + f"in_frac_zero_rows={r['in_frac_zero_rows']} out_amax={r['out_amax']}", + flush=True) + + dest = out / f"phase{args.phase}_{args.family}.json" + with open(dest, "w", encoding="utf-8") as fh: + json.dump(report, fh, indent=2) + print(f"wrote {dest}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/video_speedmem_bench.py b/scripts/video_speedmem_bench.py index 882d61a15d..4bd2825835 100644 --- a/scripts/video_speedmem_bench.py +++ b/scripts/video_speedmem_bench.py @@ -154,6 +154,23 @@ def _frames_to_arrays(output) -> list: return arrs +def _mean_luma(arrs: list) -> Optional[float]: + """Mean Rec.601 luma over all frames (0-255). ~0 == black frames (the fp8 failure signal).""" + import numpy as np + + if not arrs: + return None + vals = [] + for a in arrs: + a = np.asarray(a).astype(np.float32) + if a.ndim == 3 and a.shape[-1] >= 3: + luma = 0.299 * a[..., 0] + 0.587 * a[..., 1] + 0.114 * a[..., 2] + else: + luma = a + vals.append(float(luma.mean())) + return round(sum(vals) / len(vals), 3) if vals else None + + def _mean_lpips(ref_arrs: list, arrs: list) -> Optional[float]: """Mean per-frame LPIPS over the min common frame count.""" if not ref_arrs or not arrs: @@ -217,6 +234,12 @@ _CONFIGS: dict[str, dict[str, Any]] = { "te_fbcache": dict(te="auto", vae="none", dit="none", speed="default", attn="auto", cache="auto"), "ditfp8_fbcache": dict(te="none", vae="none", dit="auto", speed="default", attn="auto", cache="auto"), "ditint8_fbcache_prod": dict(te="none", vae="none", dit="int8", speed="default", attn="auto", cache="auto"), + # mixed-fp8 vs int8 head-to-head (Phase 3): the DiT-quant accuracy comparison on Wan/Hunyuan. + # fp8 here goes through the production quantize_transformer family exclude (input embedders + # kept bf16), so it is only non-black if the mixed-fp8 wiring is live. cache on AND off. + "ditfp8mixed_nocache": dict(te="none", vae="none", dit="fp8", speed="default", attn="native", cache="off"), + "ditfp8mixed_fbcache": dict(te="none", vae="none", dit="fp8", speed="default", attn="auto", cache="auto"), + "ditint8_nocache": dict(te="none", vae="none", dit="int8", speed="default", attn="native", cache="off"), } @@ -441,6 +464,7 @@ def _run_config(name: str, cfg: dict, *, family: str, steps: int, width: int, he "gen_latency_s": round(_median(dts), 3), "per_step_ms": round(_median(steps_ms), 1), "n_frames": len(arrs), + "mean_luma": _mean_luma(arrs), } del pipe _empty() diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index 903e323d0f..ecce54efab 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -319,7 +319,11 @@ def _validate_checkpoint( ckpt_excludes = meta.get("exclude_name_tokens") if ckpt_excludes is not None: from .diffusion_transformer_quant import exclude_tokens_for_scheme - expected = tuple(exclude_tokens_for_scheme(scheme)) + # The exclude set is derived from scheme AND family (fp8 keeps a family's padded-conditioning + # embedder bf16); use the checkpoint's recorded family so an artifact baked under an older + # token list (e.g. a Wan fp8 checkpoint built before the condition_embedder exclude) is + # rejected and re-quantised rather than loaded with a stale, black-framing layer set. + expected = tuple(exclude_tokens_for_scheme(scheme, meta.get("family"))) if tuple(ckpt_excludes) != expected: _warn( logger, diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index bf3a80cdd5..cea2b4ad99 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -83,14 +83,45 @@ _INT8_EXCLUDE_NAME_TOKENS = ( ) -def exclude_tokens_for_scheme(scheme: str) -> tuple[str, ...]: - """Name tokens to exclude from quantisation for ``scheme``. int8 (torch._int_mm, M>16) - skips the M=1 modulation / conditioning-embedder projections (see _INT8_EXCLUDE_NAME_TOKENS); - every other scheme uses scaled_mm (no M limit) and excludes nothing. Shared by the runtime - quantise path and the offline prequant-checkpoint builder so the two never drift -- an int8 - checkpoint built offline must skip exactly the layers the runtime path skips, or it bakes the - M=1 projections as int8 and crashes at the first denoise step on Flux / Qwen.""" - return _INT8_EXCLUDE_NAME_TOKENS if scheme == TQ_INT8 else () +# fp8 (and the other per-row scaled_mm schemes) PER-FAMILY name exclusions. Per-row fp8 scales +# each activation ROW by row_amax / 448; a row whose amax is 0 -- a PADDING token in a +# zero-padded conditioning sequence -- yields scale 0 and x / 0 = inf, collapsing the DiT to +# black frames (measured on B200 with the production torch._scaled_mm path via +# scripts/fp8_layer_ablation.py). The fix is to keep the specific Linear that first consumes the +# raw zero-padded sequence in bf16; its bias makes every downstream row non-zero, so the rest of +# the DiT is fp8-clean. This is family-specific because the offending layer's name is: +# wan2.2 (WanTransformer3DModel): condition_embedder.text_embedder reads the raw 512-token text +# (~all padding for a short prompt). Excluding the whole condition_embedder (a handful of +# one-off M=1/M=seq embedders, negligible FLOPs + memory) leaves the entire 30-block +# attn1/attn2/ffn stack on fp8 -- measured fp8-except-condition_embedder cosine 0.99975 vs bf16, +# 0 non-finite, vs fp8-everywhere = 100% non-finite. int8 (per-token; a zero row rounds to 0, +# no divide) needs no exclusion, which is why it was always clean and is the fallback. +# HunyuanVideo-1.5 is deliberately absent: its MMDiT masks the padding text tokens to zero INSIDE +# every block, so the per-block context stream (add_*_proj / to_add_out / ff_context) regenerates +# zero rows in every layer -- no small exclude set exists, so it stays fp8-denied -> int8 (see +# _FAMILY_SCHEME_DENY). Applies to every per-row scaled_mm scheme (fp8 today; mxfp8 / nvfp4 inherit +# it if ever un-denied for these families), never to int8. +_FP8_FAMILY_EXCLUDE_NAME_TOKENS: dict[str, tuple[str, ...]] = { + "wan2.2-ti2v-5b": ("condition_embedder",), + "wan2.2-t2v-a14b": ("condition_embedder",), # same DiT class + padded-text conditioning +} + + +def exclude_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tuple[str, ...]: + """Name tokens to exclude from quantisation for ``scheme`` (optionally family-specific). + + int8 (torch._int_mm, M>16) skips the M=1 modulation / conditioning-embedder projections (see + _INT8_EXCLUDE_NAME_TOKENS) on every family. The per-row scaled_mm schemes (fp8 / mxfp8 / nvfp4) + exclude nothing by default, but on families whose zero-padded conditioning sequence would divide + by a zero row scale they skip the offending input embedder (see _FP8_FAMILY_EXCLUDE_NAME_TOKENS). + ``family=None`` preserves the historical behaviour (int8 tokens, or () otherwise). Shared by the + runtime quantise path and the offline prequant-checkpoint builder so the two never drift -- a + checkpoint built offline must skip exactly the layers the runtime path skips, or it bakes a layer + that then crashes (int8 M=1 -> _int_mm) or infs (fp8 padding row -> scaled_mm) at the first + denoise step.""" + if scheme == TQ_INT8: + return _INT8_EXCLUDE_NAME_TOKENS + return _FP8_FAMILY_EXCLUDE_NAME_TOKENS.get(str(family or "").strip().lower(), ()) # Per-architecture preference order for ``auto`` -- best (fastest, in-bar) first, with @@ -153,13 +184,21 @@ _AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( _FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = { "qwen-image": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), "qwen-image-edit": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), # same DiT + activations - # Wan2.2 video DiTs (WanTransformer3DModel): fp8 renders black frames (measured); both the - # 5B TI2V and the A14B MoE share the DiT class + activation profile, so both deny -> int8. - "wan2.2-ti2v-5b": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), - "wan2.2-t2v-a14b": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), + # Wan2.2 video DiTs (WanTransformer3DModel): plain fp8 rendered black frames, but the failure + # is a SINGLE input embedder (condition_embedder) dividing by a zero padding-token row -- the + # 30-block compute stack is fp8-clean. fp8 is therefore allowed here and made safe by the + # per-family exclude (_FP8_FAMILY_EXCLUDE_NAME_TOKENS keeps condition_embedder bf16). mxfp8 / + # nvfp4 remain denied (same per-row scaled_mm family, not separately validated) so auto still + # skips them; both the 5B TI2V and the A14B MoE share the DiT class + activation profile. + "wan2.2-ti2v-5b": frozenset({TQ_MXFP8, TQ_NVFP4}), + "wan2.2-t2v-a14b": frozenset({TQ_MXFP8, TQ_NVFP4}), # HunyuanVideo-1.5 DiT (HunyuanVideo15Transformer3DModel): fp8 renders black frames (measured, - # LPIPS 0.82); int8 is clean. The 480p and 720p repacks share the DiT + activations, so both - # deny -> int8. (ltx-2 fp8 measures clean on the same stack, so it is intentionally absent.) + # LPIPS 0.82) and -- unlike Wan -- the failure is NOT confinable to an input embedder: its MMDiT + # masks padding text tokens to zero inside every block, so the per-block context stream + # (add_*_proj / to_add_out / ff_context) regenerates zero rows layer after layer (measured: fp8 + # on only the main blocks is still 100% non-finite). No small exclude set exists, so fp8 stays + # denied and auto lands on int8 (clean, per-token). The 480p and 720p repacks share the DiT + + # activations, so both deny. (ltx-2 fp8 measures clean on the same stack, so it is absent.) "hunyuanvideo-1.5": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), "hunyuanvideo-1.5-720p": frozenset({TQ_FP8, TQ_MXFP8, TQ_NVFP4}), } @@ -502,12 +541,14 @@ def quantize_transformer( try: from torchao.quantization import quantize_ - # int8 (torch._int_mm, M>16) additionally skips the M=1 modulation / conditioning-embedder - # projections; fp8 / fp4 / mx (scaled_mm) have no such limit and quantise everything -- but - # fp8 and mxfp8 assert a bf16 weight, so on a mixed-precision DiT (Wan / Hunyuan keep some - # fp32 linears) they must skip the non-bf16 ones or the whole pass raises and no-ops. nvfp4 - # quantises fp32 weights fine, so it is not gated (see _REQUIRE_BF16_SCHEMES). - exclude = exclude_tokens_for_scheme(scheme) + # int8 (torch._int_mm, M>16) skips the M=1 modulation / conditioning-embedder projections; + # fp8 / fp4 / mx (scaled_mm) have no M limit but, on a family whose zero-padded conditioning + # sequence would divide a per-row activation scale by a zero row, skip that one input + # embedder (Wan's condition_embedder -> _FP8_FAMILY_EXCLUDE_NAME_TOKENS). fp8 and mxfp8 also + # assert a bf16 weight, so on a mixed-precision DiT (Wan / Hunyuan keep some fp32 linears) + # they must skip the non-bf16 ones or the whole pass raises and no-ops. nvfp4 quantises fp32 + # weights fine, so it is not gated (see _REQUIRE_BF16_SCHEMES). + exclude = exclude_tokens_for_scheme(scheme, family) quantize_( transformer, _make_quant_config(scheme, fast_accum = fast_accum), diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index 8b11cbd98f..96b1c9eca4 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -286,6 +286,26 @@ def test_load_exclude_tokens_match_ok(monkeypatch, tmp_path): assert _load(monkeypatch, tmp_path, ckpt, scheme = "int8") is not None +def test_load_fp8_family_exclude_match_ok(monkeypatch, tmp_path): + # A Wan fp8 checkpoint keeps condition_embedder bf16 (the zero-padded-text divide-by-zero + # origin); the validator derives the expected set from scheme AND the recorded family, so a + # checkpoint baked with that exclude validates and loads. + ckpt = _good_ckpt(scheme = "fp8") + ckpt["metadata"]["family"] = "wan2.2-ti2v-5b" + ckpt["metadata"]["exclude_name_tokens"] = ["condition_embedder"] + assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is not None + + +def test_load_fp8_family_exclude_stale_is_none(monkeypatch, tmp_path): + # An OLD Wan fp8 checkpoint baked before the condition_embedder exclude (empty token list) would + # quantise condition_embedder and render black frames; the family-aware validator rejects it so + # the loader re-quantises dense with the correct exclude instead of loading the stale artifact. + ckpt = _good_ckpt(scheme = "fp8") + ckpt["metadata"]["family"] = "wan2.2-ti2v-5b" + ckpt["metadata"]["exclude_name_tokens"] = [] + assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is None + + def test_load_require_bf16_mismatch_is_none(monkeypatch, tmp_path): # An fp8 (scaled_mm) checkpoint built WITHOUT the bf16 gate quantised a different layer set # than the runtime filter now produces, so it must be rejected rather than loaded. diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index b35bbc7a7d..ff4da40337 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -416,9 +416,10 @@ def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch): def test_exclude_tokens_for_scheme_shared_by_runtime_and_builder(): - # The runtime quantiser and the offline prequant builder must apply the SAME int8 - # exclusion, or an int8 prequant artifact quantises the M=1 modulation/embedder linears - # and reintroduces the torch._int_mm crash. int8 gets the exclusion; others get none. + # The runtime quantiser and the offline prequant builder must apply the SAME exclusion, or a + # prequant artifact quantises a layer the runtime path skips (int8's M=1 modulation/embedder + # linears -> torch._int_mm crash; fp8's zero-padded-conditioning embedder -> black frames). int8 + # gets the family-independent exclusion; the scaled_mm schemes exclude nothing WITHOUT a family. from core.inference.diffusion_transformer_quant import ( _INT8_EXCLUDE_NAME_TOKENS, exclude_tokens_for_scheme, @@ -430,9 +431,9 @@ def test_exclude_tokens_for_scheme_shared_by_runtime_and_builder(): def test_exclude_tokens_for_scheme(): # The shared scheme->exclusion decision used by BOTH the runtime quantise path and the offline - # prequant-checkpoint builder, so an int8 checkpoint built ahead of time skips exactly the - # layers the runtime path skips (offline == runtime). int8 excludes the M=1 modulation / - # embedder tokens; every scaled_mm scheme excludes nothing. + # prequant-checkpoint builder, so a checkpoint built ahead of time skips exactly the layers the + # runtime path skips (offline == runtime). int8 excludes the M=1 modulation / embedder tokens on + # every family; the scaled_mm schemes exclude nothing by default. from core.inference.diffusion_transformer_quant import ( _INT8_EXCLUDE_NAME_TOKENS, exclude_tokens_for_scheme, @@ -444,6 +445,26 @@ def test_exclude_tokens_for_scheme(): assert exclude_tokens_for_scheme(TQ_MXFP8) == () +def test_exclude_tokens_for_scheme_family(): + # Per-family fp8 exclusion: on Wan the per-row fp8 activation scale divides by a zero + # padding-token row in condition_embedder.text_embedder (-> inf -> black frames), so fp8 keeps + # condition_embedder bf16 while the 30-block stack stays fp8. int8 is unaffected by family (it + # tolerates zero rows), and an unknown / Hunyuan family gets no fp8 exclusion. + from core.inference.diffusion_transformer_quant import ( + _INT8_EXCLUDE_NAME_TOKENS, + exclude_tokens_for_scheme, + ) + + assert exclude_tokens_for_scheme(TQ_FP8, "wan2.2-ti2v-5b") == ("condition_embedder",) + assert exclude_tokens_for_scheme(TQ_FP8, "wan2.2-t2v-a14b") == ("condition_embedder",) + assert exclude_tokens_for_scheme(TQ_MXFP8, "wan2.2-ti2v-5b") == ("condition_embedder",) + # Hunyuan is not localisable (fp8 stays denied), and an unknown family gets nothing. + assert exclude_tokens_for_scheme(TQ_FP8, "hunyuanvideo-1.5") == () + assert exclude_tokens_for_scheme(TQ_FP8, "z-image") == () + # int8 is family-independent (its zero-row handling needs no per-family skip). + assert exclude_tokens_for_scheme(TQ_INT8, "wan2.2-ti2v-5b") == _INT8_EXCLUDE_NAME_TOKENS + + # ── apply ─────────────────────────────────────────────────────────────────────── @@ -554,25 +575,30 @@ def test_family_deny_refuses_explicit_fp8_for_qwen(monkeypatch): assert select_transformer_quant_scheme(_target(), "fp8", family = "z-image") == TQ_FP8 -def test_family_deny_auto_skips_fp8_for_wan(monkeypatch): - # B200 with every scheme available: auto must NOT pick fp8 / nvfp4 / mxfp8 for the Wan - # video DiT (per-row fp8 renders black frames on it, measured; see _FAMILY_SCHEME_DENY) - # and falls through the ladder to int8, which is clean on Wan. Both the 5B TI2V and the - # A14B MoE share the WanTransformer3DModel activation profile, so both deny to int8. +def test_family_allows_fp8_for_wan(monkeypatch): + # B200 with every scheme available: the Wan fp8 black frame was root-caused to a single input + # embedder dividing by a zero padding-token row (fixed by the condition_embedder exclude, not a + # deny), so auto now picks fp8 -- the ladder head -- for both the 5B TI2V and the A14B MoE. + # mxfp8 / nvfp4 stay denied (same per-row scaled_mm family, not separately validated), so with + # only those + int8 available auto still lands on int8. _stub_torch(monkeypatch, cc = (10, 0)) _allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto", family = "wan2.2-ti2v-5b") == TQ_FP8 + assert select_transformer_quant_scheme(_target(), "auto", family = "wan2.2-t2v-a14b") == TQ_FP8 + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_INT8}) # fp8 unavailable -> denied mx/nvfp4 skipped assert select_transformer_quant_scheme(_target(), "auto", family = "wan2.2-ti2v-5b") == TQ_INT8 - assert select_transformer_quant_scheme(_target(), "auto", family = "wan2.2-t2v-a14b") == TQ_INT8 -def test_family_deny_refuses_explicit_fp8_for_wan(monkeypatch): - # An explicit fp8 request on a Wan family returns None (same GGUF-fallback contract as - # qwen); int8 stays honored on Wan, and fp8 stays honored on video families outside the - # deny table (e.g. an untested family keeps the default ladder until validated). +def test_family_deny_wan_fp8_allowed_mxfp8_nvfp4_refused(monkeypatch): + # An explicit fp8 request on a Wan family is now honored (made safe by the condition_embedder + # exclude); int8 stays honored; mxfp8 / nvfp4 stay refused (return None -> GGUF fallback) until + # separately validated on the Wan activation profile. _stub_torch(monkeypatch, cc = (10, 0)) - _allow(monkeypatch, {TQ_FP8, TQ_INT8}) - assert select_transformer_quant_scheme(_target(), "fp8", family = "wan2.2-ti2v-5b") is None + _allow(monkeypatch, {TQ_FP8, TQ_MXFP8, TQ_NVFP4, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "fp8", family = "wan2.2-ti2v-5b") == TQ_FP8 assert select_transformer_quant_scheme(_target(), "int8", family = "wan2.2-ti2v-5b") == TQ_INT8 + assert select_transformer_quant_scheme(_target(), "mxfp8", family = "wan2.2-ti2v-5b") is None + assert select_transformer_quant_scheme(_target(), "nvfp4", family = "wan2.2-ti2v-5b") is None def test_family_deny_auto_skips_fp8_for_hunyuan(monkeypatch): @@ -618,3 +644,42 @@ def test_quantize_transformer_threads_family(monkeypatch): monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) assert quantize_transformer(pipe, _target(), mode = "fp8", family = "qwen-image") is None assert called == {} + + +def test_quantize_transformer_fp8_wan_excludes_condition_embedder(monkeypatch): + # The Wan fp8 fix is a FILTER exclusion, not a deny: quantize_transformer must thread the family + # into the filter so condition_embedder (the embedder that reads the zero-padded text and would + # divide by a zero row scale) stays bf16 while the 30-block stack is quantised on fp8. Capture + # the filter that reaches torchao and check it on representative fully-qualified names. + monkeypatch.setattr( + tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8 + ) + monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg") + + # torch stub with a real nn.Linear class so the captured filter's isinstance + bf16 gate runs. + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + + class _Linear: + def __init__(self, inf, outf, dtype = "bfloat16"): + self.in_features, self.out_features = inf, outf + self.weight = types.SimpleNamespace(dtype = dtype) + + torch.nn = types.SimpleNamespace(Linear = _Linear) + monkeypatch.setitem(sys.modules, "torch", torch) + + captured: dict = {} + tqz = types.ModuleType("torchao.quantization") + tqz.quantize_ = lambda module, config, filter_fn = None: captured.update(fn = filter_fn) + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + assert quantize_transformer(pipe, _target(), mode = "fp8", family = "wan2.2-ti2v-5b") == TQ_FP8 + filt = captured["fn"] + big = _Linear(4096, 5120) # a FLOP-heavy bf16 linear (passes min_features + bf16 gate) + # condition_embedder.* is kept bf16 (the divide-by-zero origin); the block stack is fp8. + assert filt(big, "condition_embedder.text_embedder.linear_1") is False + assert filt(big, "condition_embedder.time_embedder.linear_1") is False + assert filt(big, "blocks.0.attn1.to_q") is True + assert filt(big, "blocks.0.ffn.net.0.proj") is True + assert filt(big, "blocks.0.attn2.to_k") is True # cross-attn K/V stay fp8 (embedder bias rescues rows)