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

Trim the comments across the diffusion backend

Browse source 
Comment-only pass over the Python this PR touches: drop what the code already
says, collapse multi-line explanations that still read on one line, and keep
the reasoning that is not recoverable from the code. No code, docstring
semantics or behaviour changes; verified with an AST comparison against the
previous revision, and the backend suite is unchanged (same 37 environment
failures as before: the API integration tests that need a live keyed server,
the flash-attn install hooks, and the GPU memory fields).
This commit is contained in:
Daniel Han 2026-07-26 20:30:46 +00:00
commit 36df317293
113 changed files with 3389 additions and 4087 deletions
Download patch file Download diff file Expand all files Collapse all files

34
scripts/build_prequant_checkpoint.py
View file

@ -83,24 +83,15 @@ def main(argv = None) -> int:
args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token
).to("cuda")
print(f" quantising in place ({scheme}) ...", flush = True)
# Mirror the runtime path EXACTLY (offline == runtime, LPIPS-0 invariant): for int8 also skip
# the M=1 AdaLN-modulation / conditioning-embedder projections, else the 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 (). Pass the
# family: int8 also carries PER-FAMILY exclusions (Qwen-Image's unpadded text stream runs at
# M = prompt tokens, so a short prompt breaks _int_mm), and the loader validates the baked
# list against exclude_tokens_for_scheme(scheme, metadata["family"]) -- so building with
# family=None both bakes the crashing text-stream linears and yields an artifact the runtime
# then rejects (silently falling back to the dense quantise this script exists to avoid).
# Mirror the runtime exclusions exactly: int8 also skips the M=1 modulation projections
# (torch._int_mm needs M>16) plus per-family ones; scaled_mm schemes skip none. family=None
# bakes the crashing linears and yields an artifact the runtime then rejects.
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) keeps its _keep_in_fp32_modules in
# fp32 even under torch_dtype=bf16, so quantising one raises inside quantize_ and aborts the
# pass. nvfp4 quantises fp32 fine, so it isn't gated. Runtime quantize_transformer gates on
# scheme membership; mirror it here so the offline checkpoint quantises the same layer set.
# fp8 / mxfp8 assert bf16 weights, so skip any non-bf16 Linear (mixed-precision DiTs keep some
# in fp32); nvfp4 handles fp32. Mirrors the runtime quantize_transformer gate.
require_bf16 = scheme in _REQUIRE_BF16_SCHEMES
# fp8 bakes the accumulate mode into the saved kernels; record the resolved choice so the
# loader can refuse a checkpoint whose baked value contradicts an explicit runtime request.
# fp8 bakes the accumulate mode into the saved kernels; record it so the loader can reject a
# checkpoint contradicting an explicit runtime request.
fast_accum = _resolve_fast_accum(None) if scheme == TQ_FP8 else None
quantize_(
transformer,
@ -112,7 +103,7 @@ def main(argv = None) -> int:
),
)
# Move the state dict to CPU for a portable, GPU-free artifact.
# CPU state dict 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()
@ -122,10 +113,8 @@ def main(argv = None) -> int:
"family": fam.name,
"scheme": scheme,
"min_features": args.min_features,
# The layers skipped for this scheme (int8's M=1 modulation projections; () for the
# scaled_mm schemes), whether non-bf16 Linears were skipped (the scaled_mm bf16 gate), and,
# for fp8, the baked accumulate mode. All let the loader reject a checkpoint that wouldn't
# match the runtime path.
# Skipped layers, whether non-bf16 Linears were skipped, and the fp8 accumulate mode: all let
# the loader reject a checkpoint that would not match the runtime path.
"exclude_name_tokens": list(exclude_name_tokens),
"require_bf16": require_bf16,
"fast_accum": fast_accum,
@ -136,8 +125,7 @@ def main(argv = None) -> int:
"torchao_version": getattr(torchao, "__version__", "?"),
"diffusers_version": diffusers.__version__,
}
# Record the fp8 granularity so the loader can reject a stale per-tensor checkpoint
# (the runtime now requires per-row; see FP8_GRANULARITY).
# fp8 granularity: lets the loader reject a stale per-tensor checkpoint (runtime needs per-row).
if scheme == TQ_FP8:
metadata["fp8_granularity"] = FP8_GRANULARITY
ckpt = {

11
scripts/build_te_prequant_checkpoint.py
View file

@ -57,8 +57,7 @@ def main(argv = None) -> int:
from core.inference.diffusion_precision import _cast_fp8
from core.inference.diffusion_te_prequant import TE_PREQUANT_FORMAT
# The family is metadata for forensics; detection lives in different modules per
# branch (diffusion_families vs video_families), so resolve best-effort by name.
# Family is forensic metadata; detection differs per branch, so resolve best-effort by name.
family = args.family.strip().lower()
subfolder = args.component if args.config_subfolder is None else args.config_subfolder
@ -70,9 +69,8 @@ def main(argv = None) -> int:
print(f" loading dense encoder from {args.base} (subfolder={subfolder!r}) ...", flush = True)
t0 = time.time()
config = transformers.AutoConfig.from_pretrained(args.base, **from_pretrained_kwargs)
# Prefer the checkpoint's own architecture (what the diffusers pipeline instantiates,
# e.g. Gemma3ForConditionalGeneration); AutoModel.from_config would give the bare base
# class and record a te_class whose state dict the pipeline cannot use.
# Prefer the checkpoint's own architecture; AutoModel.from_config gives the bare base class,
# whose state dict the pipeline cannot use.
arch = (getattr(config, "architectures", None) or [None])[0]
if arch and hasattr(transformers, arch):
encoder_cls_name = arch
@ -104,8 +102,7 @@ def main(argv = None) -> int:
"te_class": encoder_cls_name,
"torch_dtype": args.dtype,
"cast_backend": "diffusers_layerwise",
# str(): torch.__version__ is a TorchVersion object; pickling it into the
# checkpoint makes torch.load(weights_only=True) reject the whole artifact.
# str(): a pickled TorchVersion makes torch.load(weights_only=True) reject the artifact.
"torch_version": str(torch.__version__),
"transformers_version": str(transformers.__version__),
}

12
scripts/diffusion_bench.py
View file

@ -41,10 +41,8 @@ from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional
# The backend package lives at unsloth/studio/backend; this file is at
# unsloth/scripts/diffusion_bench.py. Put the backend root on sys.path so
# ``core.inference.diffusion`` imports as the server does. (The backend import is deferred
# into main() so --help never triggers torch.)
# Put the backend root on sys.path so ``core.inference.diffusion`` imports as the server does.
# (The backend import is deferred into main() so --help never triggers torch.)
_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend"
if str(_BACKEND_ROOT) not in sys.path:
sys.path.insert(0, str(_BACKEND_ROOT))
@ -370,10 +368,8 @@ 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. The baseline stores an absolute reference_png, which
# breaks if the baseline dir 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 -- else the
# benchmark reports PASS having done no image comparison.
# PSNR vs the stored reference. reference_png is absolute, so fall back to reference.png beside
# the baseline JSON; a still-missing reference fails below rather than passing silently.
ref_png = Path(baseline.get("accuracy", {}).get("reference_png", ""))
if not ref_png.is_file():
ref_png = baseline_path.parent / "reference.png"

14
scripts/diffusion_quality.py
View file

@ -67,9 +67,8 @@ def _to_rgb(path_or_img: Any) -> Any:
return np.asarray(img.convert("RGB"), dtype = np.float64)
# Finite PSNR (dB) a perfect (inf) sample is capped to when averaged with imperfect ones, so a
# lossless render counts as excellent without hiding diverged samples. Well above the ~37 dB
# compile and ~21 dB quant noise floors this harness reports.
# Finite PSNR (dB) cap for perfect (inf) samples, so a lossless render averages as excellent
# without hiding diverged ones. Well above the ~37 dB compile and ~21 dB quant noise floors.
_PERFECT_MATCH_PSNR = 100.0
@ -192,8 +191,7 @@ 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).
# Local paths: stat directly, since the Hub lookup returns None and _recommend would drop them.
try:
local = Path(repo).expanduser()
if local.is_dir():
@ -286,10 +284,8 @@ def _compare(
clip_sim.append(clip.image_similarity(img, ref))
def _mean(xs: list[float]) -> Optional[float]:
# +inf marks an identical render (reference vs itself, or a lossless quant/offload)
# scoring PSNR=inf -- the case this harness verifies. Report inf ONLY when every sample is
# inf; a mix means some renders diverged, so a bare inf would mask them. Cap the perfect
# ones to a high finite PSNR and average so the drift still shows. (Only PSNR is ever inf.)
# +inf marks an identical render. Report inf only when every sample is inf; otherwise cap the
# perfect ones to a high finite PSNR so partial drift still shows. (Only PSNR is ever inf.)
if not xs:
return None
if all(x == math.inf for x in xs):

5
scripts/fbcache_flux_probe.py
View file

@ -101,9 +101,8 @@ def run(
from diffusers.hooks import apply_first_block_cache
apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = threshold))
if compile_:
# FBCache's per-step decision is a graph break, so a cached run must compile with
# fullgraph=False (mirroring production); fullgraph=True would fail the warmup compile and
# the row would fall back to an eager cached run, producing misleading speedups.
# FBCache's per-step decision is a graph break, so cached runs compile with fullgraph=False as
# production does; fullgraph=True would fail warmup and silently fall back to eager.
fullgraph = threshold is None
try:
pipe.transformer.compile_repeated_blocks(fullgraph = fullgraph, dynamic = True)

5
scripts/fp8_overflow_check.py
View file

@ -63,9 +63,8 @@ def _run(fast_accum, steps, res, seed, mf):
if m > stats["max_abs"]:
stats["max_abs"] = m
# Hook the quantised linears (where an fp8-accumulation overflow would surface). Run EAGER:
# forward hooks don't trace through torch.compile, and the fp8 fast-accum accumulation is
# identical compiled or eager -- compile only changes scheduling.
# Hook the quantised linears where an fp8-accumulation overflow surfaces. Run eager: forward
# hooks don't trace through torch.compile, and accumulation is identical either way.
for m in pipe.transformer.modules():
if isinstance(m, nn.Linear):
m.register_forward_hook(hook)

6
scripts/image_speedmem_bench.py
View file

@ -50,8 +50,7 @@ for _p in (str(_BACKEND_ROOT), str(_REPO_ROOT / "scripts")):
if _p not in sys.path:
sys.path.insert(0, _p)
# Fixed prompt set (the diffusion_quality.py defaults + one photographic subject) so the
# LPIPS mean is not hostage to a single composition.
# Fixed prompt set so the LPIPS mean is not hostage to a single composition.
PROMPTS = [
"A cozy reading nook by a rain-streaked window, warm lamplight, a cat asleep on a stack of books",
"A lone lighthouse on a rocky cliff at sunset, dramatic clouds, crashing waves, highly detailed",
@ -313,8 +312,7 @@ def _generate(
if "callback_on_step_end" in call_params:
kwargs["callback_on_step_end"] = _cb
last.clear()
# Production resets the step cache before every generation; without this the
# first step compares against the PREVIOUS prompt's final residual.
# Reset the step cache like production, else step 1 compares against the previous prompt.
_reset_step_cache(pipe)
_sync()
t0 = time.perf_counter()

4
scripts/int8_linear_probe.py
View file

@ -45,8 +45,8 @@ def main() -> int:
lins = [(n, m) for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)]
selected = [(n, m) for n, m in lins if m.in_features >= MIN and m.out_features >= MIN]
print(f"\n### {label}: {len(lins)} Linear, {len(selected)} pass min_features={MIN}")
# Heuristic: a modulation/embedder Linear is one OUTSIDE the repeated transformer blocks,
# i.e. its fqn does not contain a numeric block index, OR out==k*in (k>=3) AdaLN shape.
# A modulation/embedder Linear sits outside the repeated blocks (no numeric index in its fqn),
# or has an AdaLN out==k*in (k>=3) shape.
sus = []
for n, m in selected:
depth_idx = any(p.isdigit() for p in n.split("."))

4
scripts/nvfp4_t211_probe.py
View file

@ -47,8 +47,8 @@ def diagnostics() -> None:
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.
# torchao prints "Skipping import of cpp extensions ..." on torch<2.11; absence of that line
# on 2.11 means the CUTLASS FP4 GEMM extension is live.
print(
" (no 'Skipping import of cpp extensions' line above => cpp/CUTLASS ext loaded)",
flush = True,

15
scripts/perf_levers_probe.py
View file

@ -36,9 +36,8 @@ def _lpips(ref, arr):
import lpips
import torch
# Keep the metric model on CPU: caching it on CUDA leaves it resident across variants, and
# each run resets peak-memory stats, so its VRAM would be charged to (and reduce headroom
# for) every later variant's measurement.
# Keep the metric model on CPU: cached on CUDA it stays resident across variants and its VRAM
# is charged to every later measurement (each run resets peak-memory stats).
if _LP["fn"] is None:
_LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval()
@ -72,8 +71,8 @@ def _reset_inductor_flags():
ic.coordinate_descent_tuning = False
ic.coordinate_descent_check_all_directions = False
ic.epilogue_fusion = True
# Reset the int-mm fusion flag too, or it leaks from the inductor_flags variant into
# every later compiled row and the attention/fbcache measurements stop being isolated.
# Reset the int-mm fusion flag too, else it leaks from the inductor_flags variant into every
# later compiled row.
try:
ic.force_fuse_int_mm_with_mul = False
except Exception: # noqa: BLE001
@ -149,10 +148,8 @@ def run(
torch.cuda.empty_cache()
return None
else:
# set_attention_backend pins diffusers' PROCESS-WIDE active backend, and a fresh
# transformer's processors (backend None) inherit it. Force native for the no-attn variants
# so they aren't measured under a prior variant's kernel (e.g. fbcache with a leftover sage
# backend).
# set_attention_backend pins diffusers' process-wide backend and fresh processors inherit it,
# so force native for the no-attn variants (else they run under a prior variant's kernel).
try:
pipe.transformer.set_attention_backend("native")
except Exception as exc: # noqa: BLE001 — best-effort isolation

15
scripts/verify_prequant_backend.py
View file

@ -37,10 +37,8 @@ OUT = Path(os.environ.get("PREQUANT_OUT_DIR", str(_RESEARCH / "prequant_verify_i
logging.basicConfig(level = logging.INFO, format = "%(message)s")
LOGGER = logging.getLogger("verify_prequant")
# Prequant and runtime produce the SAME quantized weights, so their images must be
# near-identical; anything above this LPIPS means the prequant path diverged and the run fails.
# The prequant load peak must also sit clearly below the dense runtime peak (the point of the
# path); require at least this fractional headroom.
# Prequant and runtime produce the same quantized weights, so above this LPIPS the prequant
# path diverged. The prequant load peak must also sit this fraction below the dense peak.
LPIPS_MAX = 0.02
PREQUANT_PEAK_MAX_FRACTION = 0.75
_RUNTIME_PEAK_FILE = OUT / "runtime_peak.txt"
@ -102,9 +100,8 @@ def run(mode, steps, seed, res):
torch.cuda.empty_cache()
if mode == "prequant":
# A local checkpoint is refused unless its directory is allowlisted (unpickling an
# arbitrary file is unsafe). This verifier's CKPT is operator-supplied and trusted, so
# allowlist its directory here or the load returns None and measures nothing.
# A local checkpoint is refused unless its directory is allowlisted (unpickling is unsafe).
# CKPT is operator-supplied and trusted, so allowlist it or the load returns None.
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] = (
@ -153,8 +150,8 @@ def run(mode, steps, seed, res):
if mode != "prequant":
return 0
# Enforce the two invariants this verifier exists to check, so a broken prequant
# checkpoint fails loudly instead of passing just because generation completed.
# Enforce both invariants, so a broken prequant checkpoint fails loudly rather than passing
# just because generation completed.
ref_path = OUT / "runtime.png"
if not ref_path.exists():
print("FAIL: runtime reference image missing; run --mode runtime first", flush = True)

17
scripts/video_quality.py
View file

@ -69,8 +69,7 @@ DEFAULT_PROMPT = (
"splashing water, cinematic, camera tracking sideways"
)
# Finite PSNR (dB) an identical clip is capped to when averaged, matching
# scripts/diffusion_quality.py.
# Finite PSNR (dB) cap for an identical clip, matching scripts/diffusion_quality.py.
_PERFECT_MATCH_PSNR = 100.0
@ -159,11 +158,9 @@ def clip_metrics(
"""All frame metrics for one candidate clip vs the reference clip."""
import numpy as np
# The gate holds the requested shape (num_frames) fixed for reference and candidate, so both
# clips must decode to the same frame count. A shorter candidate is a truncated/corrupt render;
# comparing only the shared prefix would let good early frames mask the missing tail, so the
# mismatch is recorded and gated as FAIL (see verdict()) rather than dropped. No off-by-one is
# tolerated.
# num_frames is fixed for reference and candidate, so both clips must decode to the same frame
# count. A shorter candidate is truncated: comparing the shared prefix would let good early
# frames mask the missing tail, so the mismatch is gated as FAIL (see verdict()).
ref_count, cand_count = len(ref_frames), len(cand_frames)
frame_count_mismatch = ref_count != cand_count
n = min(ref_count, cand_count)
@ -212,8 +209,7 @@ def audio_metrics(ref_audio: Optional[Any], cand_audio: Optional[Any]) -> dict[s
return float(np.sqrt((arr**2).mean())) if arr.size else 0.0
ref_rms, cand_rms = _rms(ref_audio), _rms(cand_audio)
# NaN candidate audio compares False against any threshold, so call it out
# explicitly: a NaN track is a collapse, not a pass.
# NaN compares False against any threshold, so call it out: a NaN track is a collapse.
silent_collapse = (
ref_rms is not None
and ref_rms >= 1e-3
@ -459,8 +455,7 @@ def selftest() -> int:
shifted = clip_metrics(ref, make_clip(offset = 0.5))
check(shifted["ssim_mean"] < same["ssim_mean"], "content shift lowers ssim")
# A truncated render whose surviving prefix is pixel-identical must still FAIL
# on the frame-count mismatch alone, not PASS on the good early frames.
# A truncated render with a pixel-identical prefix must still FAIL on the frame-count mismatch.
truncated = clip_metrics(ref, make_clip()[: n // 2])
check(
truncated["frame_count_mismatch"] is True