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

Merge branch 'diffusion-phase8-quant' into diffusion-phase9-prequant

Browse source 
# Conflicts:
#	scripts/compare_engines.py
This commit is contained in:
Daniel Han 2026-07-01 11:29:36 +00:00
commit 4fec0bb79c
282 changed files with 22928 additions and 2814 deletions
Download patch file Download diff file Expand all files Collapse all files

14
scripts/compare_engines.py
View file

@ -40,9 +40,13 @@ def bench_pytorch(repo, gguf, resolutions, steps, seed, iters):
backend = DiffusionBackend()
for speed in ("off", "default"):
backend.begin_load(repo, gguf_filename = gguf, speed_mode = speed)
deadline = time.time() + 1800 # 30 min: a stuck download/load must not hang forever
while backend.load_progress().get("phase") != "ready":
if backend.load_progress().get("phase") == "error":
raise RuntimeError(backend.load_progress())
prog = backend.load_progress()
if prog.get("phase") == "error":
raise RuntimeError(prog)
if time.time() > deadline:
raise TimeoutError(f"load timed out (last progress: {prog})")
time.sleep(0.5)
for res in resolutions:
@ -128,11 +132,13 @@ def main(argv = None) -> int:
)
p.add_argument(
"--vae",
default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/flux_vae/ae.safetensors",
default = None,
help = "VAE safetensors for sd.cpp (required when benchmarking the sd.cpp engine)",
)
p.add_argument(
"--llm",
default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/qwen3_te/Qwen3-4B-Instruct-2507-Q4_K_M.gguf",
default = None,
help = "text-encoder GGUF for sd.cpp (required when benchmarking the sd.cpp engine)",
)
p.add_argument("--resolutions", default = "512,1024")
p.add_argument("--steps", type = int, default = 8)

19
scripts/diffusion_bench.py
View file

@ -142,8 +142,10 @@ def _psnr(ref_png: Path, cand_png: Path) -> float:
import numpy as np
from PIL import Image
a = np.asarray(Image.open(ref_png).convert("RGB"), dtype = np.float64)
b = np.asarray(Image.open(cand_png).convert("RGB"), dtype = np.float64)
with Image.open(ref_png) as im_a:
a = np.asarray(im_a.convert("RGB"), dtype = np.float64)
with Image.open(cand_png) as im_b:
b = np.asarray(im_b.convert("RGB"), dtype = np.float64)
if a.shape != b.shape:
# Different geometry means the comparison is meaningless; report worst case.
return 0.0
@ -368,9 +370,14 @@ def _compare(args: argparse.Namespace) -> int:
print(" refusing noisy comparison (pass --force-compare to override).", flush = True)
return 2
# PSNR vs the stored reference image.
# PSNR vs the stored reference image. The baseline stores an absolute reference_png,
# which breaks if the baseline directory was copied/moved, so fall back to reference.png
# next to the baseline JSON. A still-missing reference is a failure below, not a silent
# pass -- otherwise the benchmark would report PASS having done no image comparison.
ref_png = Path(baseline.get("accuracy", {}).get("reference_png", ""))
psnr = _psnr(ref_png, args._image_out) if ref_png.exists() else float("nan")
if not ref_png.is_file():
ref_png = baseline_path.parent / "reference.png"
psnr = _psnr(ref_png, args._image_out) if ref_png.is_file() else float("nan")
base_gen = baseline.get("generate", {})
cur_gen = metrics["generate"]
@ -402,7 +409,9 @@ def _compare(args: argparse.Namespace) -> int:
)
if base_peak and cur_peak and vram_reg > args.max_vram_regression:
failures.append(f"peak VRAM +{vram_reg * 100:.1f}% > {args.max_vram_regression * 100:.0f}%")
if not math.isnan(psnr) and psnr < args.min_psnr:
if math.isnan(psnr):
failures.append("PSNR reference image missing; cannot verify output quality")
elif psnr < args.min_psnr:
failures.append(f"PSNR {psnr:.2f}dB < {args.min_psnr:.1f}dB (output changed)")
if failures:

17
scripts/diffusion_quality.py
View file

@ -186,6 +186,18 @@ def _wait_for_load(backend: Any, timeout_s: int = 3600) -> None:
def _hf_file_size_mib(repo: str, filename: str) -> Optional[int]:
# A local model dir / file: stat it directly. The Hub lookup below returns None for
# a local path, which would drop every candidate from _recommend (file_size_mib None).
try:
local = Path(repo).expanduser()
if local.is_dir():
f = local / filename
if f.is_file():
return int(f.stat().st_size // (1024 * 1024))
elif local.is_file():
return int(local.stat().st_size // (1024 * 1024))
except Exception:
pass
try:
from huggingface_hub import HfApi
info = HfApi().model_info(repo, files_metadata = True, token = os.environ.get("HF_TOKEN"))
@ -268,6 +280,11 @@ def _compare(
clip_sim.append(clip.image_similarity(img, ref))
def _mean(xs: list[float]) -> Optional[float]:
# Preserve +inf: an identical render (reference vs itself, or a lossless
# quant/offload) scores PSNR=inf, which is exactly the case this harness
# verifies; dropping it as non-finite would print "-" instead of "inf".
if xs and any(x == math.inf for x in xs):
return math.inf
finite = [x for x in xs if math.isfinite(x)]
return round(sum(finite) / len(finite), 4) if finite else None

6
scripts/nvfp4_probe.py
View file

@ -16,7 +16,7 @@ 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_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_images"
def _psnr(a, b):
@ -84,11 +84,15 @@ def main(argv = None) -> int:
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 = ""):

145
scripts/nvfp4_t211_probe.py
View file

@ -32,30 +32,38 @@ 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)
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)
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)
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),
"nvfp4_cutlass": NV(use_triton_kernel = False),
"nvfp4_triton": NV(use_triton_kernel = True),
}
@ -66,11 +74,11 @@ def _bench_linear(K, N, M, cfg, iters, compile_):
torch.compiler.reset()
torch.cuda.empty_cache()
m = nn.Sequential(nn.Linear(K, N, bias=False)).cuda().to(torch.bfloat16)
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)
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)
@ -89,12 +97,12 @@ def _bench_linear(K, N, M, cfg, iters, compile_):
def micro(M, iters, compile_):
print(f"\n== B. GEMM micro (M={M}, compile={compile_}, iters={iters}) ==", flush=True)
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)
print(f" shape K={K} N={N}:", flush = True)
base_ms = None
fp8_ms = None
for name, cfg in cfgs.items():
@ -107,9 +115,12 @@ def micro(M, iters, compile_):
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)
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)
print(f" {name:16s} FAILED: {type(exc).__name__}: {str(exc)[:120]}", flush = True)
# ----------------------------------------------------------------------------- e2e
@ -125,8 +136,9 @@ def _lpips(ref, arr):
try:
import lpips
import torch
if _LP["fn"] is None:
_LP["fn"] = lpips.LPIPS(net="alex", verbose=False).cuda().eval()
_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()
@ -134,7 +146,7 @@ def _lpips(ref, arr):
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)
print(f" (lpips: {type(exc).__name__})", flush = True)
return None
@ -143,8 +155,9 @@ def _load_dense():
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)
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
)
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
pipe.to("cuda")
return pipe
@ -152,11 +165,17 @@ def _load_dense():
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,24 +188,31 @@ def e2e(steps, res, seed, iters, mf):
import torch
import torch.nn as nn
OUT.mkdir(parents=True, exist_ok=True)
OUT.mkdir(parents = True, exist_ok = True)
def filt(mod, fqn=""):
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):
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)
quantize_(pipe.transformer, cfg, filter_fn = filt)
if compile:
try:
pipe.transformer.compile_repeated_blocks(fullgraph=True, dynamic=True)
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)
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):
@ -202,58 +228,71 @@ def e2e(steps, res, seed, iters, mf):
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)
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),
("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)
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)
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)
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)
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)
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:
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")
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)
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)
print("NVFP4-T211-PROBE-DONE", flush = True)
return 0

7
scripts/perf_verify.py
View file

@ -151,7 +151,12 @@ def main(argv = None) -> int:
flush = True,
)
ok = (leak_psnr == float("inf")) and (def_t < off_t) and (_psnr(off_img, def_img) >= 30)
ok = (
(leak_psnr == float("inf"))
and (bal_psnr == float("inf")) # check 3: balanced must be bit-identical to off
and (def_t < off_t)
and (_psnr(off_img, def_img) >= 30)
)
print(f"\nPERF-VERIFY {'OK' if ok else 'CHECK'}", flush = True)
return 0 if ok else 1

14
scripts/quant_probe.py
View file

@ -27,7 +27,7 @@ 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("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/probe_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "probe_images"
def _psnr(a, b):
@ -39,17 +39,19 @@ _LPIPS = {"fn": None}
def _lpips(ref_arr, arr):
"""Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable."""
"""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).cuda().eval()
_LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval()
def t(x):
t = torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
return t.cuda()
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())
@ -111,7 +113,7 @@ def _quant_config(name):
return MXDynamicActivationMXWeightConfig(
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
)
except TypeError:
except (TypeError, AttributeError):
return MXDynamicActivationMXWeightConfig()
raise ValueError(name)

7
scripts/scan_packages.py
View file

@ -1208,9 +1208,10 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]:
HIGH,
package,
filename,
f"Python wheel ships large ({len(content) // 1024} KB) JS bundle "
"(uncommon; manually review)",
"",
# Size stays in evidence, not the check label, so the baseline key
# does not drift when a wheel's bundle grows by a few KB.
"Python wheel ships large JS bundle (uncommon; manually review)",
f"{len(content) // 1024} KB JS bundle",
)
)
return findings

2
scripts/scan_packages_baseline.json
View file

@ -1181,7 +1181,7 @@
{
"package": "tensorboard",
"file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js",
"check": "Python wheel ships large (1918 KB) JS bundle (uncommon; manually review)",
"check": "Python wheel ships large JS bundle (uncommon; manually review)",
"severity": "HIGH",
"evidence": ""
},