unsloth/scripts/verify_prequant_backend.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""GPU verification of the Phase 9 pre-quantized load path through the real backend code.

Exercises the actual product functions (``load_prequantized_transformer`` and the runtime
``quantize_transformer``), not a reimplementation:

  prequant  -- load the checkpoint built by build_prequant_checkpoint.py via the real
               ``load_prequantized_transformer`` (meta-init + assign), measure GPU load peak,
               generate.
  runtime   -- the existing path: from_pretrained dense bf16 -> ``quantize_transformer`` on
               device, measure GPU load peak, generate (the LPIPS reference).

Asserts the prequant load peak is far below the dense one and the images match (LPIPS ~0).
Run each mode in its own process for a clean peak. One CUDA GPU."""

from __future__ import annotations

import argparse
import logging
import os
import sys
import time
from pathlib import Path

import numpy as np

_REPO = Path(__file__).resolve().parent.parent
_RESEARCH = _REPO / "outputs" / "quant_research"
BACKEND = _REPO / "studio" / "backend"
BASE = "Tongyi-MAI/Z-Image-Turbo"
CKPT = os.environ.get("PREQUANT_CKPT", str(_RESEARCH / "prequant_fp8" / "transformer_fp8.pt"))
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path(os.environ.get("PREQUANT_OUT_DIR", str(_RESEARCH / "prequant_verify_images")))

logging.basicConfig(level = logging.INFO, format = "%(message)s")
LOGGER = logging.getLogger("verify_prequant")

# Prequant must match runtime numerics (LPIPS) and beat the dense load peak by this fraction.
LPIPS_MAX = 0.02
PREQUANT_PEAK_MAX_FRACTION = 0.75
_RUNTIME_PEAK_FILE = OUT / "runtime_peak.txt"


def _target(dtype):
    import types
    return types.SimpleNamespace(device = "cuda", dtype = dtype)


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 _lpips(ref, arr):
    try:
        import lpips, torch

        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(fn(t(ref), t(arr)).item())
    except Exception as exc:  # noqa: BLE001
        print(f"  (lpips: {type(exc).__name__})", flush = True)
        return None


def run(mode, steps, seed, res):
    sys.path.insert(0, str(BACKEND))
    os.environ["UNSLOTH_NVFP4_DIFFUSION"] = "1"
    import torch
    import diffusers
    from core.inference.diffusion_prequant import (
        ALLOW_LOCAL_PREQUANT_PATH_ENV,
        PrequantSource,
        load_prequantized_transformer,
    )
    from core.inference.diffusion_transformer_quant import quantize_transformer

    OUT.mkdir(parents = True, exist_ok = True)
    transformer_cls = diffusers.ZImageTransformer2DModel
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.empty_cache()

    if mode == "prequant":
        # Local checkpoints are refused unless allowlisted;
        # CKPT is operator-supplied and trusted.
        ckpt_dir = os.path.dirname(os.path.realpath(CKPT))
        existing = os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV, "")
        os.environ[ALLOW_LOCAL_PREQUANT_PATH_ENV] = (
            ckpt_dir if not existing else existing + os.pathsep + ckpt_dir
        )
        source = PrequantSource(kind = "path", location = CKPT, filename = None)
        transformer = load_prequantized_transformer(
            transformer_cls,
            BASE,
            source,
            device = "cuda",
            dtype = torch.bfloat16,
            hf_token = None,
            scheme = "fp8",
            logger = LOGGER,
        )
        if transformer is None:
            print("prequant load FAILED (returned None)", flush = True)
            return 1
        pipe = diffusers.ZImagePipeline.from_pretrained(
            BASE, torch_dtype = torch.bfloat16, transformer = transformer
        )
        pipe.to("cuda")
        load_peak = torch.cuda.max_memory_allocated() / 1e9
        marker = getattr(transformer, "_unsloth_runtime_quant", None)
        print(f"[prequant] load_gpu_peak={load_peak:.1f} GB  marker={marker}", flush = True)
    else:  # runtime
        transformer = transformer_cls.from_pretrained(
            BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
        ).to("cuda")
        pipe = diffusers.ZImagePipeline.from_pretrained(
            BASE, torch_dtype = torch.bfloat16, transformer = transformer
        )
        pipe.to("cuda")
        scheme = quantize_transformer(pipe, _target(torch.bfloat16), mode = "fp8", logger = LOGGER)
        load_peak = torch.cuda.max_memory_allocated() / 1e9
        print(f"[runtime] engaged={scheme}  load_gpu_peak={load_peak:.1f} GB", flush = True)
        # Persist the dense reference peak so a later prequant run can enforce its VRAM win.
        _RUNTIME_PEAK_FILE.write_text(f"{load_peak:.6f}")

    img, dt = _gen(pipe, steps, seed, res)
    img, dt = _gen(pipe, steps, seed, res)
    img.save(OUT / f"{mode}.png")
    print(f"[{mode}] gen={dt:.3f}s saved {mode}.png", flush = True)

    if mode != "prequant":
        return 0

    # Enforce both invariants so a broken checkpoint fails loudly instead of merely completing.
    ref_path = OUT / "runtime.png"
    if not ref_path.exists():
        print("FAIL: runtime reference image missing; run --mode runtime first", flush = True)
        return 1
    from PIL import Image

    lp = _lpips(np.array(Image.open(ref_path).convert("RGB")), np.array(img))
    print(f"[prequant] LPIPS_vs_runtime={lp}", flush = True)
    if lp is None:
        print("FAIL: LPIPS could not be computed (install lpips)", flush = True)
        return 1
    if lp > LPIPS_MAX:
        print(f"FAIL: LPIPS {lp:.4f} > {LPIPS_MAX} (prequant diverged from runtime)", flush = True)
        return 1
    if _RUNTIME_PEAK_FILE.exists():
        try:
            runtime_peak = float(_RUNTIME_PEAK_FILE.read_text().strip())
        except ValueError:
            runtime_peak = 0.0
        if runtime_peak > 0.0 and load_peak > runtime_peak * PREQUANT_PEAK_MAX_FRACTION:
            print(
                f"FAIL: prequant load peak {load_peak:.1f} GB not below "
                f"{PREQUANT_PEAK_MAX_FRACTION:.0%} of dense {runtime_peak:.1f} GB",
                flush = True,
            )
            return 1
    return 0


def main(argv = None) -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--mode", choices = ["prequant", "runtime"], required = True)
    p.add_argument("--steps", type = int, default = 8)
    p.add_argument("--res", type = int, default = 1024)
    p.add_argument("--seed", type = int, default = 42)
    args = p.parse_args(argv)
    rc = run(args.mode, args.steps, args.seed, args.res)
    print("VERIFY-PREQUANT-DONE", flush = True)
    return rc


if __name__ == "__main__":
    sys.exit(main())