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

Studio diffusion (Phase 9): pre-quantized transformer loading

Browse source 
The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on
the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs
13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline
(scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton
on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True)
the quantized weights, so the dense bf16 never touches the GPU.

Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching
GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same
torchao config + min_features filter the runtime path uses, applied ahead of time.

New core/inference/diffusion_prequant.py (resolve_prequant_source +
load_prequantized_transformer, best-effort, lazy imports). diffusion.py
_load_dense_quant_pipeline tries the pre-quant source first and falls back to the
dense materialise+quantise path, then to GGUF, so the default is unchanged.
DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty
prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for
the resolver, the meta-init+assign loader, and the backend branch selection +
fallbacks; GPU verification via scripts/verify_prequant_backend.py.
This commit is contained in:
Daniel Han 2026-06-26 11:23:20 +00:00
commit b90f833469
11 changed files with 968 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

127
scripts/build_prequant_checkpoint.py Normal file
View file

@ -0,0 +1,127 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Build a pre-quantized transformer checkpoint for the Studio diffusion fast path.
Quantise a model's dense bf16 DiT transformer ONCE and save the quantized state dict, so
the backend can load the already-quantized weights at runtime (meta-init +
load_state_dict(assign=True)) instead of materialising the dense bf16 on the GPU. That
drops the transformer GPU load peak ~2x and the download ~2x for fp8 (measured on Z-Image:
12.9 -> 6.3 GB peak, 12 -> 6.28 GB on disk), with bit-identical output -- it is the exact
same torchao config + min_features filter the runtime path uses, applied ahead of time.
Run on one CUDA (Blackwell / Ada / Hopper) GPU. fp8 works on torch 2.9+; the FP4/MX schemes
need the newer kernels (see scripts/nvfp4_t211_probe.py).
python scripts/build_prequant_checkpoint.py \
--base Tongyi-MAI/Z-Image-Turbo --family z-image --scheme fp8 \
--out outputs/quant_research/prequant_fp8/transformer_fp8.pt [--upload-repo ORG/REPO]
"""
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"
def main(argv=None) -> int:
p = argparse.ArgumentParser()
p.add_argument("--base", required=True, help="diffusers base repo (carries the transformer subfolder)")
p.add_argument("--family", required=True, help="diffusion family name/alias (e.g. z-image)")
p.add_argument("--scheme", required=True, help="quant scheme: int8 | fp8 | nvfp4 | mxfp8")
p.add_argument("--out", required=True, help="output .pt path for the checkpoint")
p.add_argument("--min-features", type=int, default=512)
p.add_argument("--dtype", default="bfloat16", choices=["bfloat16"])
p.add_argument("--hf-token", default=None)
p.add_argument("--upload-repo", default=None, help="optional HF repo id to upload the checkpoint to")
p.add_argument("--upload-revision", default=None)
args = p.parse_args(argv)
sys.path.insert(0, str(BACKEND))
import torch
import torchao
import diffusers
from core.inference.diffusion_families import detect_family
from core.inference.diffusion_prequant import PREQUANT_FORMAT, prequant_filename
# Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant).
from core.inference.diffusion_transformer_quant import (
TQ_SCHEMES,
_make_quant_config,
make_filter_fn,
)
from torchao.quantization import quantize_
scheme = args.scheme.strip().lower()
if scheme not in TQ_SCHEMES:
print(f"error: --scheme must be one of {TQ_SCHEMES} (not 'auto')", flush=True)
return 2
fam = detect_family(args.base, override=args.family)
if fam is None:
print(f"error: unknown family '{args.family}'", flush=True)
return 2
transformer_cls = getattr(diffusers, fam.transformer_class)
print(f"== build prequant ({fam.name}/{scheme}, min_feat={args.min_features}) ==", flush=True)
print(f" loading dense transformer from {args.base} (subfolder=transformer) ...", flush=True)
t0 = time.time()
transformer = transformer_cls.from_pretrained(
args.base, subfolder="transformer", torch_dtype=torch.bfloat16, token=args.hf_token
).to("cuda")
print(f" quantising in place ({scheme}) ...", flush=True)
quantize_(transformer, _make_quant_config(scheme), filter_fn=make_filter_fn(args.min_features))
# Move the state dict to CPU for a portable, GPU-free artifact.
state_dict = {
k: (v.detach().to("cpu") if hasattr(v, "detach") else v)
for k, v in transformer.state_dict().items()
}
ckpt = {
"format": PREQUANT_FORMAT,
"metadata": {
"base_model_id": args.base,
"family": fam.name,
"scheme": scheme,
"min_features": args.min_features,
"torch_dtype": args.dtype,
"quant_backend": "torchao",
"transformer_class": fam.transformer_class,
"torch_version": torch.__version__,
"torchao_version": getattr(torchao, "__version__", "?"),
"diffusers_version": diffusers.__version__,
},
"state_dict": state_dict,
}
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
torch.save(ckpt, out)
size_gb = out.stat().st_size / 1e9
print(f" saved {out} ({size_gb:.2f} GB) in {time.time() - t0:.0f}s", flush=True)
print(f" metadata: {ckpt['metadata']}", flush=True)
if args.upload_repo:
from huggingface_hub import HfApi
dest = prequant_filename(scheme)
print(f" uploading -> {args.upload_repo}:{dest} ...", flush=True)
api = HfApi(token=args.hf_token)
api.create_repo(args.upload_repo, exist_ok=True)
api.upload_file(
path_or_fileobj=str(out),
path_in_repo=dest,
repo_id=args.upload_repo,
revision=args.upload_revision,
)
print(f" uploaded {dest} to {args.upload_repo}", flush=True)
print("BUILD-PREQUANT-DONE", flush=True)
return 0
if __name__ == "__main__":
sys.exit(main())

170
scripts/prequant_probe.py Normal file
View file

@ -0,0 +1,170 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Does a *pre-quantized* checkpoint fix the dense-quant load-VRAM spike?
The current fast-transformer path materialises the dense bf16 transformer on the GPU
and quantises it in place -> ~2x the GGUF load peak. This probe checks the fix: quantise
once, ``torch.save`` the quantized state dict, then load it onto an empty (meta) model
with ``load_state_dict(assign=True)`` so the bf16 never touches the GPU.
Modes (run each in its own process so peak VRAM is clean):
build -- load dense bf16, quantize_ fp8, torch.save the state dict + on-disk size.
baseline -- current path: from_pretrained bf16 -> quantize_ on GPU. Report load peak + gen.
prequant -- meta-init -> load_state_dict(saved, assign=True) -> cuda. Report load peak + gen.
Run on one CUDA (Blackwell) GPU. Reference image for LPIPS is the baseline path."""
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"
ROOT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research")
CKPT = ROOT / "prequant_fp8" / "transformer_fp8_state.pt"
OUT = ROOT / "prequant_images"
MIN_FEAT = 512
def _filt(mod, fqn=""):
import torch.nn as nn
return isinstance(mod, nn.Linear) and mod.in_features >= MIN_FEAT and mod.out_features >= MIN_FEAT
def _fp8_cfg():
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
return Float8DynamicActivationFloat8WeightConfig()
def _build():
import torch
import diffusers
from torchao.quantization import quantize_
torch.cuda.reset_peak_memory_stats()
t = diffusers.ZImageTransformer2DModel.from_pretrained(
BASE, subfolder="transformer", torch_dtype=torch.bfloat16).to("cuda")
quantize_(t, _fp8_cfg(), filter_fn=_filt)
CKPT.parent.mkdir(parents=True, exist_ok=True)
sd = t.state_dict()
# move to cpu for a portable, gpu-free checkpoint
sd = {k: (v.detach().to("cpu") if hasattr(v, "detach") else v) for k, v in sd.items()}
torch.save(sd, CKPT)
sz = CKPT.stat().st_size / 1e9
peak = torch.cuda.max_memory_allocated() / 1e9
print(f"[build] saved {CKPT.name} on-disk={sz:.2f} GB build_gpu_peak={peak:.1f} GB", flush=True)
return 0
def _make_pipe_from_transformer(t):
import diffusers
import torch
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 _baseline(steps, seed, res):
import torch
import diffusers
from torchao.quantization import quantize_
torch.cuda.reset_peak_memory_stats()
t = diffusers.ZImageTransformer2DModel.from_pretrained(
BASE, subfolder="transformer", torch_dtype=torch.bfloat16).to("cuda")
quantize_(t, _fp8_cfg(), filter_fn=_filt)
load_peak = torch.cuda.max_memory_allocated() / 1e9
pipe = _make_pipe_from_transformer(t)
img, dt = _gen(pipe, steps, seed, res) # warmup
img, dt = _gen(pipe, steps, seed, res)
OUT.mkdir(parents=True, exist_ok=True)
img.save(OUT / "baseline.png")
print(f"[baseline] transformer_load_gpu_peak={load_peak:.1f} GB gen={dt:.3f}s", flush=True)
return 0
def _prequant(steps, seed, res):
import torch
import diffusers
from accelerate import init_empty_weights
if not CKPT.exists():
print(f"[prequant] missing checkpoint {CKPT}; run --mode build first", flush=True)
return 1
torch.cuda.reset_peak_memory_stats()
cfg = diffusers.ZImageTransformer2DModel.load_config(BASE, subfolder="transformer")
with init_empty_weights():
t = diffusers.ZImageTransformer2DModel.from_config(cfg)
sd = torch.load(CKPT, weights_only=False, map_location="cpu")
missing, unexpected = t.load_state_dict(sd, strict=False, assign=True)
# any param/buffer still on meta (e.g. non-persistent buffers) -> materialise on cuda
leftover = [n for n, p in t.named_parameters() if p.is_meta] + [n for n, b in t.named_buffers() if b.is_meta]
if leftover:
print(f"[prequant] {len(leftover)} meta leftovers (non-persistent buffers): {leftover[:4]}", flush=True)
t = t.to_empty(device="cuda") # fallback path; re-loads sd below
t.load_state_dict(sd, strict=False, assign=True)
t = t.to(torch.bfloat16).to("cuda")
load_peak = torch.cuda.max_memory_allocated() / 1e9
print(f"[prequant] missing={len(missing)} unexpected={len(unexpected)} "
f"transformer_load_gpu_peak={load_peak:.1f} GB", flush=True)
pipe = _make_pipe_from_transformer(t)
img, dt = _gen(pipe, steps, seed, res) # warmup
img, dt = _gen(pipe, steps, seed, res)
OUT.mkdir(parents=True, exist_ok=True)
img.save(OUT / "prequant.png")
# LPIPS vs baseline if present
bpath = OUT / "baseline.png"
lp = None
if bpath.exists():
try:
import lpips
from PIL import Image
fn = lpips.LPIPS(net="alex", verbose=False).cuda().eval()
def tt(p):
a = np.array(Image.open(p).convert("RGB"))
return (torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda()
with torch.no_grad():
lp = float(fn(tt(bpath), tt(OUT / "prequant.png")).item())
except Exception as exc: # noqa: BLE001
print(f" (lpips: {type(exc).__name__})", flush=True)
print(f"[prequant] gen={dt:.3f}s LPIPS_vs_baseline={lp}", flush=True)
return 0
def main(argv=None) -> int:
p = argparse.ArgumentParser()
p.add_argument("--mode", choices=["build", "baseline", "prequant"], 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)
if args.mode == "build":
return _build()
if args.mode == "baseline":
return _baseline(args.steps, args.seed, args.res)
return _prequant(args.steps, args.seed, args.res)
if __name__ == "__main__":
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
rc = main()
print("PREQUANT-PROBE-DONE", flush=True)
sys.exit(rc)

126
scripts/verify_prequant_backend.py Normal file
View file

@ -0,0 +1,126 @@
# 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 sys
import time
from pathlib import Path
import numpy as np
BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"
BASE = "Tongyi-MAI/Z-Image-Turbo"
CKPT = "/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/prequant_fp8/transformer_fp8.pt"
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/prequant_verify_images")
logging.basicConfig(level=logging.INFO, format="%(message)s")
LOGGER = logging.getLogger("verify_prequant")
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))
import torch
import diffusers
from core.inference.diffusion_prequant import 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":
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)
img, dt = _gen(pipe, steps, seed, res) # warmup
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)
ref_path = OUT / "runtime.png"
if mode == "prequant" and ref_path.exists():
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)
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())