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

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

Browse source 
* Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict

Phase 1 of porting the richer diffusion stack onto the image-generation backend.

- Add a compartmentalized device/dtype policy module (diffusion_device.py)
  resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA
  capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or
  fp32, never a silent fp16 that renders a black image.
- Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved
  float16 to float32 for those families so they do not produce black images.
- Split the backend locks: a generation holds only _generate_lock, so status,
  unload, and a new load are never blocked by a long denoise. Add per-generation
  cancellation via callback_on_step_end so an eviction or a superseding load
  preempts a running generation; a replacement load waits for it to stop before
  allocating, so two pipelines never sit in VRAM at once.
- Validate a load request before the GPU handoff so an unloadable pick never
  evicts a working chat model, and reject missing local paths up front.
- Add CPU-only tests for the device policy, dtype guard, lock split and
  cancellation, and validate-before-evict, plus a GPU benchmark/regression
  script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR
  against a saved reference.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy

Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and
VAE tiling/slicing from measured free device memory vs the model's estimated
resident footprint, then applies it to the built pipeline. auto stays resident
when the model fits (byte-identical to the prior resident path), and falls to
whole-module offload when tight; fast/balanced/low_vram are explicit overrides.
Sequential submodule offload is unreliable for GGUF transformers on diffusers
0.38, so it falls back to whole-module offload and status reports the policy
actually engaged.

Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with
no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM
47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost.

73 prior + 35 new CPU tests pass.

* Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling

Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level,
use_stream) that keeps the transformer flowing through the GPU a few blocks at a
time while the text encoder / VAE stay resident, and fix VAE tiling to drive the
VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the
pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so
status never overstates either, and group falls back to whole-module offload when
the transformer can't be streamed.

Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group
cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 ->
2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names
now match that tradeoff: balanced = stream the transformer, low_vram = offload
every component. auto picks group when the companions fit resident, else model.

112 CPU tests pass.

* Studio diffusion (Phase 5): image quality-vs-quant accuracy harness

Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold
prompt + seed fixed, render a grid with a reference quant (default BF16), then render
each candidate quant and measure drift from the reference. Records mean PSNR + SSIM
(pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity
(transformers, --clip), plus file size, latency, and peak VRAM, then prints a
quality-vs-cost table and recommends the smallest quant within a quality budget.
--selftest validates the metrics on synthetic images with no GPU or model.

Verified on Z-Image (B200): the table degrades monotonically with quant size
(Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat
(~0.34) -- quantization erodes fine detail far more than prompt adherence.

* Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32)

Add a speed_mode knob (off by default, so the render path stays bit-identical):
default applies channels_last VAE + regional torch.compile of the denoiser's
repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional
compile is gated off for the GGUF transformer (dequantises per-op) and for families
flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image),
so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed
optims run before placement/offload, per the diffusers composition order. status now
reports speed_mode + the optims actually engaged.

Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last',
'tf32'], compile correctly skipped for GGUF; generation works in every mode.

121 CPU tests pass.

* Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting

Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3)
storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16
compute dtype while normalisations and embeddings stay full precision. Applied
before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder
dense). status reports which encoders were cast.

Verified on Z-Image (B200, balanced/group mode where the encoder stays resident):
generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload)
at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR
vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off
by default and documented as such, with the Phase 5 harness to size the cost.

127 CPU tests pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob)

Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant
(fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao
NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8
stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run
before placement; status reports the mode actually engaged. This is the lean
realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the
3045-line port.

Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the
bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE
option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and
both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default);
size it per model with the Phase 5 quality harness. diffusion_bench gains
--text-encoder-quant.

129 CPU tests pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac

Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the
chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm
/ XPU; this covers the hardware diffusers serves poorly, consuming the same
split GGUF assets Studio already curates.

- sd_cpp_args.py: pure sd-cli command builder. Maps the family to its
  text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1
  CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential)
  to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu /
  --vae-tiling / --diffusion-fa), so one user knob drives both engines.
- sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary()
  with the same precedence as the llama finder (env override, then the Studio
  install root, then in-tree, then PATH), an is_available/version probe, and a
  one-shot subprocess generate that streams progress and returns the PNG.
  runtime_env() prepends the binary's directory to the platform library path
  so a prebuilt's bundled libstable-diffusion.so resolves.
  select_diffusion_engine() is the pure routing decision (GPU backends to
  diffusers, CPU/MPS to native when present).
- install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt
  (macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the
  Studio install root. resolve_release_asset() is a pure, unit-tested
  host-to-asset matrix.
- scripts/sd_cpp_smoke.py: end-to-end native generation harness.

Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine,
routing, runtime env, and the installer resolver. Full diffusion suite 166
passing.

Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both
generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group
offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the
dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine

Builds on Phase 4's native stable-diffusion.cpp engine, extending it from
text-to-image to the wider feature surface, since sd.cpp supports all of these
through the binary already. Pure command-builder additions plus one engine
method, so the txt2img path is unchanged.

- sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img +
  strength make a run img2img, adding mask makes it inpaint, ref_images drives
  FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and
  lora_dir + the <lora:name:weight> prompt syntax select LoRAs. New
  SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run
  mode (input image + esrgan model, no prompt / text encoders).
- sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so
  generate() (now carrying the conditioning flags) and a new upscale() reuse
  the same streaming / error / output-check path.
- scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img /
  --strength / --upscale-model / --upscale-repeats.

Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the
upscale builder and its validation, and the engine's img2img + upscale paths.
Full diffusion suite 176 passing.

Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the
init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale
(512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing
coherent images. Video and the diffusers-path feature wiring are deferred.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 7): accuracy-preserving speed pass

Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy
bug and a dead-on-arrival speed path; this fixes both and adds the lossless /
near-lossless wins, all measured on a B200.

Correctness:
- TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32
  process-wide and never restored them, so a later `off` load silently inherited
  TF32 and was no longer bit-identical. Added snapshot_backend_flags /
  restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer
  runs and restored on unload. Verified: load max -> unload -> load off is now
  byte-identical (PSNR inf) to a fresh off.
- sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and
  only checked the timeout after EOF, so a child stuck in model load / GPU init
  with no output ignored the timeout. Drained stdout on a reader thread with a
  wall-clock deadline. Added a silent-hang regression test.

Speed (diffusers path), near-lossless, opt-in tiers:
- Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and
  Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks
  compiles and runs ~2.2x faster on the GGUF Z-Image transformer on
  torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the
  block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB
  vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move
  output quality. Gate relaxed; default tier delivers it.
- cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs).
- torch.inference_mode() around the pipeline call (lossless, strictly faster than
  the no_grad diffusers uses internally).

Memory path:
- VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers;
  the balanced (group) tier keeps exact slicing only, so it is now bit-identical to
  the resident image (verified PSNR inf) and slightly faster.
- Group offload adds non_blocking + record_stream on the CUDA stream path to
  overlap each block's H2D copy with compute (lossless; gated on the installed
  diffusers signature so older versions still work).

Native (sd.cpp) path:
- native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a
  near-lossless CUDA win that was previously only added on offload tiers; max also
  -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so
  it is never auto-on. Engine generate() merges it, de-duped against offload flags.

Default profile: a GGUF model with no explicit speed_mode now resolves to the
`default` profile (resolve_speed_mode), since compile's perturbation sits below the
quantisation noise floor and so does not reduce quality versus the dense reference;
out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models
stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is
always honored, so the byte-identical path remains one flag away and is the
regression reference.

Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/
perf_verify.py (the B200 verification above), and diffusion_bench.py gains
--speed-mode so the speed tiers are benchmarkable.

Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore,
GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags +
the engine de-dup, and the sd-cli silent-hang timeout.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks

The opt-in `max` speed tier now compiles the repeated block with
mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode:
Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer
cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are
deliberately avoided -- both crash on the regionally-compiled block (its static
output buffer is overwritten across denoise steps), measured.

Adds two reproducible benchmarks used to validate the optimization research:
- scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head.
- scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes.

Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen;
coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune);
FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source)

Add an opt-in transformer_quant mode that loads the dense bf16 transformer and
torchao-quantises it onto the low-precision tensor cores, instead of the GGUF
transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a
B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's
0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor.

GGUF+compile stays the low-memory default and the fallback. The mode is gated on
CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB);
any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged
reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe
(Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips
the tiny projections that crash int8's torch._int_mm. New module mirrors
diffusion_precision.py; quant runs before compile before placement.

184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route
coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py
is the standalone torchao lever probe.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity

Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838
TFLOPS with FP16 accumulate; bf16 209), so:
- fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on;
  pinning it guards consumer cards against a default change. On B200 it is identical
  speed and slightly better quality (LPIPS 0.050 vs 0.091).
- the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate).

2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py):
2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the
cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile
(our main ~2x). Documented as a dead end, not shipped.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe

scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image
generation and reports max-abs + non-finite counts for use_fast_accum True vs False.
Confirms fast accumulation is an accumulation-precision knob, not an overflow one:
across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16),
0 non-finite elements and identical max-abs for both modes.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override

Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want
fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and
prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match
on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40;
GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it.

Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s),
no overflow, quality below the quant noise floor. So the default leans to accuracy on
data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force)
lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off).

187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override
threading).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9

scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer.
Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with
use_triton_kernel=False (the default triton path needs the missing MSLK library), but
only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than
using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's
CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto
on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder

Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4
tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS,
2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the
diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both
slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166
vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is
correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add
scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

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

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in

load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config

Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.

quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.

* Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load

- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
  missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
  whole snapshot. restore_backend_flags restores each flag independently so one failure can't
  leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
  apply_speed_optims mutated the process-wide flags but before _state captured them for unload
  to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
  contaminated later off generations.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output

Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.

* Studio diffusion (Phase 9) review fixes: prequant safety + validation

- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
  resolves inside an operator-configured ALLOWLIST of directories
  (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
  once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
  any path a load request named (arbitrary code execution). realpath() blocks symlink
  escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
  checkpoint that quantised a different layer set is rejected instead of silently
  loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
  path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
  author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts

- diffusion_memory: when group offload is unavailable and the plan falls back to
  whole-module offload, enable VAE tiling (the group plan left it off, but the fallback
  is the low-VRAM path where the decode spike can OOM). Covers both the group and
  sequential fallback branches.
- perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a
  balanced bit-identity regression actually fails the check instead of exiting 0.
- compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and
  the load-progress poll has a 30 min deadline instead of looping forever on a hang.
- test for the group->model fallback enabling VAE tiling.

* Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path

- diffusion: a torchao-quantized transformer is committed only compiled. A dense model
  resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF
  it replaced), so when transformer_quant engaged and speed resolved to off, promote to
  default (regional compile); warn loudly if compile still does not engage.
- diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the
  CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK);
  otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF.
- nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path).
- test asserts the eager-quant -> default-compile promotion.

* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening

- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
  timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
  extract through a per-member containment check (Zip-Slip guard); expanduser the
  --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
  the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
  installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
  sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
  crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs

Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.

Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.

* Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats

Codex review on the native engine arg builder:

- build_sd_cpp_command emitted --width/--height unconditionally, so an
  img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
  the input. width/height are now Optional (None = unset): an image-conditioned
  run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
  the size from the input image (set_width_and_height_if_unset); a plain txt2img
  run with unset dims keeps the prior 1024x1024 default; explicit dims are always
  honored. width/height are read only by the builder, so the type change is local.

- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
  that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
  explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
  and emits the flag for any explicit value != 1.

Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)

* Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc

Codex review: the transformer_prequant_path field description still told operators
to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the
prior security fix made that variable a directory allowlist -- _allowed_prequant_roots
deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator
following the documented =1 would have every transformer_prequant_path request
silently refused. The description now states it must name one or more allowlisted
directories and that a bare on/off value is not accepted.

Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does
not say =1, and describes an allowlist/directory (guards against doc drift).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-01 11:37:53 -07:00 • committed by GitHub
commit ddb6084ac4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1317 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

132
scripts/build_prequant_checkpoint.py Normal file
View file

@ -0,0 +1,132 @@
# 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())

196
scripts/prequant_probe.py Normal file
View file

@ -0,0 +1,196 @@
# 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(__file__).resolve().parent.parent / "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)

152
scripts/verify_prequant_backend.py Normal file
View file

@ -0,0 +1,152 @@
# 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")
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())

80
studio/backend/core/inference/diffusion.py
View file

@ -53,10 +53,16 @@ from .diffusion_speed import (
snapshot_backend_flags,
)
from .diffusion_precision import quantize_text_encoders
from .diffusion_prequant import (
load_prequantized_transformer,
resolve_prequant_source,
)
from .diffusion_transformer_quant import (
DEFAULT_MIN_LINEAR_FEATURES,
dense_transformer_supported,
normalize_transformer_quant,
quantize_transformer,
select_transformer_quant_scheme,
)
logger = get_logger(__name__)
@ -299,6 +305,7 @@ class DiffusionBackend:
text_encoder_quant: Optional[str] = None,
transformer_quant: Optional[str] = None,
transformer_quant_fast_accum: Optional[bool] = None,
transformer_prequant_path: Optional[str] = None,
) -> dict[str, Any]:
"""Validate, then run the (slow) load on a daemon thread. Returns at once."""
fam = self.validate_load_request(
@ -332,6 +339,7 @@ class DiffusionBackend:
text_encoder_quant = text_encoder_quant,
transformer_quant = transformer_quant,
transformer_quant_fast_accum = transformer_quant_fast_accum,
transformer_prequant_path = transformer_prequant_path,
_load_token = token,
),
daemon = True,
@ -467,6 +475,7 @@ class DiffusionBackend:
text_encoder_quant: Optional[str] = None,
transformer_quant: Optional[str] = None,
transformer_quant_fast_accum: Optional[bool] = None,
transformer_prequant_path: Optional[str] = None,
_load_token: Optional[int] = None,
) -> dict[str, Any]:
# Validate first (cheap, no torch/diffusers) so a direct call with a bad
@ -534,6 +543,8 @@ class DiffusionBackend:
target,
transformer_quant,
transformer_quant_fast_accum,
fam = fam,
prequant_path = transformer_prequant_path,
)
except Exception as exc: # noqa: BLE001 — fall back to the GGUF build
logger.warning(
@ -662,27 +673,76 @@ class DiffusionBackend:
target: DiffusionDeviceTarget,
mode: Optional[str],
fast_accum: Optional[bool] = None,
*,
fam: Optional[DiffusionFamily] = None,
prequant_path: Optional[str] = None,
) -> tuple[Any, str]:
"""Build the opt-in fast pipeline: load the DENSE bf16 transformer from the base
repo (``subfolder="transformer"``), assemble the pipeline, place it on the device,
and torchao-quantise the transformer in place. Returns ``(pipe, engaged_scheme)``.
"""Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
Two ways to get the quantized transformer, in order:
1. Pre-quantized: if a checkpoint is configured for the chosen scheme (an explicit
``prequant_path`` or the family's hosted repo), load the already-quantized
weights onto the meta device and assign them in -- the dense bf16 never lands on
the GPU, so the load peak is ~half and the download is smaller.
2. Dense + quantise (fallback): load the DENSE bf16 transformer from the base repo,
place it on the device, and torchao-quantise it in place.
Raises if the scheme is unsupported or quantisation fails, so ``load_pipeline``
catches it and falls back to the GGUF build. Quantisation runs ON the device (the
dynamic int8 / fp8 / fp4 kernels need the weights on CUDA) and BEFORE the loader
compiles the repeated block, so the order is quantize -> compile -> placement."""
catches it and falls back to the GGUF build. Quantisation runs ON the device and
BEFORE the loader compiles the repeated block, so the order stays quantize ->
compile -> placement."""
# 1. Pre-quantized checkpoint, when one is configured for the resolved scheme.
scheme = select_transformer_quant_scheme(target, mode)
if scheme is not None and fam is not None:
source = resolve_prequant_source(fam, scheme, path_override = prequant_path)
if source is not None:
transformer = load_prequantized_transformer(
transformer_cls,
base,
source,
device = device,
dtype = dtype,
hf_token = hf_token,
scheme = scheme,
# Reject a checkpoint built with a different Linear filter than the
# dense path uses, so the prequant and runtime-quant models match.
min_features = DEFAULT_MIN_LINEAR_FEATURES,
logger = logger,
)
if transformer is not None:
pipe = self._assemble_pipe(
pipeline_cls, base, transformer, dtype, hf_token, device
)
return pipe, scheme
# 2. Fallback: materialise the dense bf16 transformer and quantise it on-device.
transformer = transformer_cls.from_pretrained(
base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
)
pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device)
scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger)
if scheme is None:
raise RuntimeError("transformer quant unsupported for this device/scheme")
return pipe, scheme
@staticmethod
def _assemble_pipe(
pipeline_cls: Any,
base: str,
transformer: Any,
dtype: Any,
hf_token: Optional[str],
device: str,
) -> Any:
"""Assemble the diffusers pipeline around ``transformer`` and place it on ``device``
(a no-op for an already-placed pre-quantized transformer; it moves the companions)."""
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer}
if hf_token:
pipe_kwargs["token"] = hf_token
pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
pipe.to(device)
scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger)
if scheme is None:
raise RuntimeError("transformer quant unsupported for this device/scheme")
return pipe, scheme
return pipe
def _plan_memory(
self,

14
studio/backend/core/inference/diffusion_families.py
View file

@ -39,6 +39,12 @@ class DiffusionFamily:
# regional torch.compile. Now consulted on the GGUF path too (compile runs on the
# GGUF transformer); all current families compile, so this stays True.
supports_torch_compile: bool = True
# Optional pre-quantized transformer checkpoints, as (scheme, repo_id) pairs (a
# hashable mapping). When the fast transformer_quant path resolves a scheme with a
# hosted checkpoint, the loader fetches the already-quantized weights instead of
# materialising the dense bf16 transformer on the GPU (much lower load VRAM + a
# smaller download). Empty until checkpoints are hosted -> behaviour is unchanged.
prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple)
# Keyed by architecture, not per model variant: a checkpoint's specific base repo
@ -123,6 +129,14 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
return base or fam.base_repo
def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]:
"""The hosted pre-quantized transformer repo for ``scheme`` in this family, or None."""
for entry_scheme, repo_id in fam.prequant_repos:
if entry_scheme == scheme:
return repo_id
return None
def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path:
"""Resolve ``gguf_filename`` to a file under ``repo_root``, rejecting escapes.

287
studio/backend/core/inference/diffusion_prequant.py Normal file
View file

@ -0,0 +1,287 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Load a *pre-quantized* transformer instead of quantising a dense one on the GPU.
The opt-in fast transformer_quant path (see ``diffusion_transformer_quant.py``) loads
the dense bf16 transformer and torchao-``quantize_``s it in place. That materialises the
full bf16 weights on the GPU before quantising, so the load peak is ~2x the GGUF's and it
pulls the full bf16 download. When a transformer has already been quantised once and saved
(``scripts/build_prequant_checkpoint.py``), this module loads those weights directly:
1. build the transformer skeleton on the ``meta`` device (no storage) via
``accelerate.init_empty_weights`` + ``from_config``;
2. ``load_state_dict(assign=True)`` the quantized state dict (the torchao weight subclass
tensors are assigned in, not copied), so the dense bf16 never touches the GPU;
3. move to the device.
Measured (B200, Z-Image fp8): transformer GPU load peak 12.9 -> 6.3 GB, download 12 ->
6.28 GB, output bit-identical (LPIPS 0.0). The checkpoint carries the exact same scheme +
``min_features`` as the runtime path, so the result is identical to quantising on the fly.
Best-effort and lazily imported throughout: a missing / mismatched / unreadable checkpoint
returns None and the caller falls back to the dense-quantise path (and then to GGUF). All
behaviour is gated on a configured source -- with nothing configured this module is inert.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Optional
# torch.save dict layout this module reads (and the build script writes). Bumped if the
# on-disk structure changes so an old/foreign artifact is rejected rather than mis-loaded.
PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1"
# Loading a checkpoint ends in ``torch.load(weights_only=False)``, which executes arbitrary
# code embedded in the pickle. A hosted family *repo* checkpoint is first-party and trusted,
# but a ``source.kind == "path"`` can originate from the ``transformer_prequant_path`` field
# of a load request -- i.e. an authenticated API caller naming an arbitrary local file.
# Unpickling that is remote code execution, so a request-supplied path is unpickled ONLY when
# it resolves inside an operator-configured ALLOWLIST of directories. A bare on/off toggle is
# deliberately NOT accepted as a wildcard: enabling local checkpoints for one trusted
# directory must never also permit unpickling any other path a request happens to name. The
# trusted hosted-repo path is unaffected.
ALLOW_LOCAL_PREQUANT_PATH_ENV = "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH"
_PREQUANT_TOGGLE_TOKENS = {"1", "true", "yes", "on", "0", "false", "no", "off"}
def _allowed_prequant_roots() -> list:
"""Operator-allowlisted directories whose pre-quant checkpoints may be unpickled.
Set ``UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH`` to one or more directories (separated by
``os.pathsep``). A bare truthy/falsey toggle is ignored on purpose -- it must name a
directory, so there is no "allow everything" mode."""
import os
raw = (os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV) or "").strip()
if not raw:
return []
roots = []
for part in raw.split(os.pathsep):
part = part.strip()
if not part or part.lower() in _PREQUANT_TOGGLE_TOKENS:
continue # a bare on/off value is not a directory -> never a wildcard allow
try:
roots.append(os.path.realpath(os.path.expanduser(part)))
except Exception: # noqa: BLE001 — a bad entry is simply not allowlisted
continue
return roots
def _local_prequant_path_allowed(path: str) -> bool:
"""True only when ``path`` resolves inside an operator-allowlisted directory; an
arbitrary request-supplied path is never unpickled. ``realpath`` first so a symlink
cannot point an allowlisted name at a file outside the allowed roots."""
import os
roots = _allowed_prequant_roots()
if not roots:
return False
try:
real = os.path.realpath(os.path.expanduser(path))
except Exception: # noqa: BLE001
return False
return any(real == r or real.startswith(r + os.sep) for r in roots)
@dataclass(frozen = True)
class PrequantSource:
"""Where a pre-quantized transformer checkpoint lives. ``kind`` is "path" (a local
file) or "repo" (a Hub repo id in ``location`` + ``filename`` inside it)."""
kind: str
location: str
filename: Optional[str] = None
def prequant_filename(scheme: str) -> str:
"""The conventional checkpoint filename for ``scheme`` inside a Hub repo."""
return f"transformer_{scheme}.pt"
def resolve_prequant_source(
fam: Any,
scheme: str,
*,
path_override: Optional[str] = None,
) -> Optional[PrequantSource]:
"""Resolve where the pre-quantized checkpoint for ``(fam, scheme)`` should come from.
Priority: (1) an explicit local ``path_override`` (testing / power users); (2) the
family's hosted repo for ``scheme``; (3) None -> no pre-quant, caller quantises dense.
Pure: no IO, no torch -- it only decides the source, the loader fetches it.
"""
override = (path_override or "").strip()
if override:
return PrequantSource(kind = "path", location = override, filename = None)
try:
from .diffusion_families import family_prequant_repo
repo_id = family_prequant_repo(fam, scheme)
except Exception: # noqa: BLE001 — a bad family object must not break the load
repo_id = None
if repo_id:
return PrequantSource(kind = "repo", location = repo_id, filename = prequant_filename(scheme))
return None
def load_prequantized_transformer(
transformer_cls: Any,
base: str,
source: PrequantSource,
*,
device: str,
dtype: Any,
hf_token: Optional[str] = None,
scheme: str,
min_features: Optional[int] = None,
logger: Any = None,
) -> Optional[Any]:
"""Load the pre-quantized transformer described by ``source`` onto ``device``.
Returns the placed, already-quantized transformer, or None on any problem (missing /
mismatched / unreadable checkpoint, or a meta-init the class does not support) so the
caller falls back to the dense-quantise path. Best-effort: never raises for an
ordinary unavailable artifact.
"""
try:
# weights_only=False (required below) executes pickle code, so a caller-supplied
# local path is unpickled ONLY when it resolves inside an operator-allowlisted
# directory. The hosted family repo is first-party and always allowed.
if source.kind == "path" and not _local_prequant_path_allowed(source.location):
_warn(
logger,
f"{scheme}:path",
RuntimeError(
"request-supplied local pre-quant path refused (unpickling an arbitrary "
f"file is unsafe); set {ALLOW_LOCAL_PREQUANT_PATH_ENV} to an allowlisted "
"directory containing trusted checkpoints to permit it",
),
)
return None
path = _resolve_checkpoint_path(source, hf_token)
if path is None:
return None
import torch
# torchao weight subclasses are not safetensors-serializable, so the checkpoint is
# a torch.save pickle. weights_only=False is required to rebuild those subclasses.
# The local-path branch is gated above; the repo branch is a first-party artifact.
ckpt = torch.load(path, weights_only = False, map_location = "cpu")
if not _validate_checkpoint(ckpt, scheme, base, logger, min_features = min_features):
return None
state_dict = ckpt["state_dict"]
config = transformer_cls.load_config(base, subfolder = "transformer", token = hf_token)
from accelerate import init_empty_weights
with init_empty_weights():
transformer = transformer_cls.from_config(config)
# assign=True swaps in the loaded (quantized) tensors rather than copying into the
# meta tensors (a copy into meta is a no-op); strict=True since the saved state
# dict is the full state dict of the same class (non-persistent buffers excluded).
transformer.load_state_dict(state_dict, strict = True, assign = True)
if _has_meta_tensors(transformer):
# A class with non-persistent buffers (computed in __init__, absent from the
# state dict) leaves those on meta. Rebuild on CPU so the buffers hold their
# real values, then re-assign the quantized weights. The dense bf16 lives in
# CPU RAM only -- the GPU still receives just the quantized footprint.
transformer = transformer_cls.from_config(config)
transformer.load_state_dict(state_dict, strict = True, assign = True)
transformer = transformer.to(device)
try: # diagnostic marker, mirrors the runtime-quant path
transformer._unsloth_runtime_quant = scheme
except Exception: # noqa: BLE001 — marker is best-effort
pass
if logger is not None:
logger.info(
"diffusion.prequant: loaded %s checkpoint (%s) onto %s",
scheme,
source.kind,
device,
)
return transformer
except Exception as exc: # noqa: BLE001 — fall back to the dense-quantise path
_warn(logger, f"{scheme}:{source.kind}", exc)
return None
def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) -> Optional[str]:
"""The local file path for ``source``, downloading from the Hub if needed; None if absent."""
if source.kind == "path":
import os
return source.location if os.path.isfile(source.location) else None
if source.kind == "repo":
from huggingface_hub import hf_hub_download
return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token)
return None
def _validate_checkpoint(
ckpt: Any,
scheme: str,
base: str,
logger: Any,
min_features: Optional[int] = None,
) -> bool:
"""Reject a checkpoint that is the wrong format / scheme / base model / filter.
``min_features`` (when given) is the runtime Linear-feature threshold: a checkpoint
built with a different ``--min-features`` quantises a different set of Linear layers,
so ``load_state_dict(assign=True)`` would silently install a model that does not match
what the dense path produces while status still reports the requested scheme. Reject it."""
if not isinstance(ckpt, dict) or ckpt.get("format") != PREQUANT_FORMAT:
_warn(logger, scheme, ValueError("unrecognised pre-quant checkpoint format"))
return False
if "state_dict" not in ckpt:
_warn(logger, scheme, ValueError("pre-quant checkpoint has no state_dict"))
return False
meta = ckpt.get("metadata") or {}
if meta.get("scheme") != scheme:
_warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}"))
return False
ckpt_base = meta.get("base_model_id")
if ckpt_base and base and not _same_base_model(ckpt_base, base):
_warn(logger, scheme, ValueError(f"checkpoint base {ckpt_base!r} != {base!r}"))
return False
if min_features is not None:
ckpt_min = meta.get("min_features")
if ckpt_min is not None and int(ckpt_min) != int(min_features):
_warn(
logger,
scheme,
ValueError(f"checkpoint min_features {ckpt_min!r} != runtime {min_features!r}"),
)
return False
return True
def _same_base_model(a: str, b: str) -> bool:
"""Tolerant compare of two base-model ids: an exact match, or the same final
path/repo segment (so a local path or a fork id matches the canonical repo, e.g.
``/models/Z-Image-Turbo`` vs ``Tongyi-MAI/Z-Image-Turbo``)."""
def _tail(x: str) -> str:
return x.replace("\\", "/").rstrip("/").split("/")[-1].lower()
return a == b or _tail(a) == _tail(b)
def _has_meta_tensors(module: Any) -> bool:
"""True if any parameter or buffer is still on the meta device after loading."""
from itertools import chain
try:
return any(
getattr(t, "is_meta", False) for t in chain(module.parameters(), module.buffers())
)
except Exception: # noqa: BLE001
return False
def _warn(logger: Any, what: str, exc: Exception) -> None:
if logger is not None:
logger.warning("diffusion.prequant: %s failed: %s", what, exc)

13
studio/backend/models/inference.py
View file

@ -1737,6 +1737,19 @@ class DiffusionLoadRequest(BaseModel):
"HBM cards, which are not nerfed). true/false force it. Negligible "
"quality effect (below the fp8 quant noise floor); no overflow risk.",
)
transformer_prequant_path: Optional[str] = Field(
None,
description = "Local path to a pre-quantized transformer checkpoint (built by "
"scripts/build_prequant_checkpoint.py) for the requested transformer_quant "
"scheme. Loads the already-quantized weights with the dense bf16 never on the "
"GPU (~half the load VRAM and a smaller download). null uses the family's hosted "
"checkpoint if configured, else quantises the dense transformer at load time. "
"Loading a local path unpickles the file (arbitrary code execution), so it is "
"ignored unless the path resolves inside a directory the operator allowlisted "
"via UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH (one or more directories, separated by "
"the OS path separator). A bare on/off value such as '1' is deliberately not "
"accepted -- it must name an allowed directory.",
)
class DiffusionGenerateRequest(BaseModel):

1
studio/backend/routes/inference.py
View file

@ -10333,6 +10333,7 @@ async def load_diffusion_model(
text_encoder_quant = request.text_encoder_quant,
transformer_quant = request.transformer_quant,
transformer_quant_fast_accum = request.transformer_quant_fast_accum,
transformer_prequant_path = request.transformer_prequant_path,
)
return DiffusionStatusResponse(**status_dict)
except (ValueError, FileNotFoundError) as exc:

75
studio/backend/tests/test_diffusion_backend.py
View file

@ -920,6 +920,10 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"):
monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False)
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
# Resolve the scheme without the real GPU smoke probe, and configure no pre-quant
# checkpoint so the dense materialise+quantise branch is the one exercised.
monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: scheme)
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None)
def _quantize(pipe, target, *, mode, **kw):
calls["quantize"] += 1
@ -974,6 +978,77 @@ def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatc
assert status["offload_policy"] == "none"
def test_transformer_quant_prequant_path_engaged(fake_runtime, tmp_path, monkeypatch):
# A configured pre-quant checkpoint -> load the already-quantized transformer directly;
# the dense from_pretrained and the on-device quantize_transformer are NOT used.
from core.inference import diffusion as dmod
backend = DiffusionBackend()
_force_cuda_target(backend, monkeypatch)
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8")
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object())
prequant_obj = object()
loaded: dict = {"n": 0}
def _load_prequant(transformer_cls, base, source, **kw):
loaded["n"] += 1
loaded["scheme"] = kw.get("scheme")
return prequant_obj
monkeypatch.setattr(dmod, "load_prequantized_transformer", _load_prequant)
@classmethod
def _fp_fail(cls, *a, **k):
pytest.fail("dense from_pretrained must not run when a prequant checkpoint loads")
monkeypatch.setattr(_FakeTransformer, "from_pretrained", _fp_fail, raising = False)
monkeypatch.setattr(
dmod,
"quantize_transformer",
lambda *a, **k: pytest.fail("quantize_transformer must not run on the prequant path"),
)
(tmp_path / "m.gguf").write_bytes(b"x")
status = backend.load_pipeline(
str(tmp_path),
gguf_filename = "m.gguf",
family_override = "z-image",
transformer_quant = "fp8",
transformer_prequant_path = str(tmp_path / "zimage_fp8.pt"),
)
assert status["transformer_quant"] == "fp8"
assert loaded["n"] == 1 and loaded["scheme"] == "fp8"
# The pre-quantized transformer object was assembled into the pipeline...
assert _FakePipeline.last.get("transformer") is prequant_obj
# ...and the GGUF single-file path was not used.
assert _FakeTransformer.last == {}
def test_transformer_quant_prequant_load_fails_falls_back_to_dense(
fake_runtime, tmp_path, monkeypatch
):
# A configured prequant source whose load returns None must fall back to the dense
# materialise+quantise path (not straight to GGUF), preserving the fast mode.
from core.inference import diffusion as dmod
backend = DiffusionBackend()
_force_cuda_target(backend, monkeypatch)
calls = _stub_dense_quant(monkeypatch, scheme = "fp8")
# Override the no-prequant default: a source resolves, but its load fails.
monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object())
monkeypatch.setattr(dmod, "load_prequantized_transformer", lambda *a, **k: None)
(tmp_path / "m.gguf").write_bytes(b"x")
status = backend.load_pipeline(
str(tmp_path),
gguf_filename = "m.gguf",
family_override = "z-image",
transformer_quant = "fp8",
)
assert status["transformer_quant"] == "fp8"
assert calls["from_pretrained"] == 1 and calls["quantize"] == 1 # dense path ran
assert _FakeTransformer.last == {} # GGUF not used
def test_transformer_quant_falls_back_to_gguf_on_failure(fake_runtime, tmp_path, monkeypatch):
# A dense/quant failure (here: quantize returns None -> unsupported) must fall back
# to the GGUF build, not error -- status reports no transformer_quant engaged.

348
studio/backend/tests/test_diffusion_prequant.py Normal file
View file

@ -0,0 +1,348 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Hermetic CPU tests for the pre-quantized transformer load path.
torch / accelerate are stubbed via ``sys.modules`` (the module under test imports them
lazily), and ``transformer_cls`` is a fake that records calls -- so the resolver, the
meta-init + ``load_state_dict(assign=True)`` flow, and the validation/fallback behaviour
are all exercised without CUDA, torchao, or a real diffusers model.
"""
from __future__ import annotations
import contextlib
import sys
import types
import core.inference.diffusion_prequant as pq
from core.inference.diffusion_families import DiffusionFamily
from core.inference.diffusion_prequant import (
PREQUANT_FORMAT,
PrequantSource,
load_prequantized_transformer,
resolve_prequant_source,
)
# ── resolve_prequant_source ──────────────────────────────────────────────────────
def _fam(prequant_repos = ()):
return DiffusionFamily(
name = "z-image",
pipeline_class = "ZImagePipeline",
transformer_class = "ZImageTransformer2DModel",
base_repo = "Tongyi-MAI/Z-Image-Turbo",
prequant_repos = prequant_repos,
)
def test_resolve_path_override_wins():
fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),))
src = resolve_prequant_source(fam, "fp8", path_override = "/tmp/local.pt")
assert src == PrequantSource(kind = "path", location = "/tmp/local.pt", filename = None)
def test_resolve_family_repo_by_scheme():
fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"), ("int8", "org/hosted-int8")))
src = resolve_prequant_source(fam, "int8")
assert src.kind == "repo" and src.location == "org/hosted-int8"
assert src.filename == "transformer_int8.pt"
def test_resolve_wrong_scheme_is_none():
fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),))
assert resolve_prequant_source(fam, "int8") is None
def test_resolve_nothing_configured_is_none():
assert resolve_prequant_source(_fam(), "fp8") is None
assert resolve_prequant_source(_fam(), "fp8", path_override = "") is None
# ── load_prequantized_transformer ────────────────────────────────────────────────
class _FakeTransformer:
calls: dict = {}
def __init__(self):
self.assigned = None
self.moved = None
@classmethod
def load_config(cls, base, **kw):
cls.calls["load_config"] = {"base": base, **kw}
return {"cfg": True}
@classmethod
def from_config(cls, config):
cls.calls["from_config"] = config
return cls()
@classmethod
def from_pretrained(cls, *a, **k): # the dense path -- must never run here
cls.calls["from_pretrained"] = True
raise AssertionError("from_pretrained must not be called on the prequant path")
def load_state_dict(
self,
sd,
strict = True,
assign = False,
):
_FakeTransformer.calls["load_state_dict"] = {"strict": strict, "assign": assign}
self.assigned = sd
def parameters(self):
return []
def buffers(self):
return []
def to(self, device):
self.moved = device
return self
def _stub_torch_accelerate(
monkeypatch,
ckpt,
*,
load_raises = False,
):
torch = types.ModuleType("torch")
def _load(
path,
weights_only = False,
map_location = None,
):
if load_raises:
raise RuntimeError("corrupt checkpoint")
return ckpt
torch.load = _load
monkeypatch.setitem(sys.modules, "torch", torch)
accelerate = types.ModuleType("accelerate")
accelerate.init_empty_weights = lambda: contextlib.nullcontext()
monkeypatch.setitem(sys.modules, "accelerate", accelerate)
def _good_ckpt(scheme = "fp8", base = "Tongyi-MAI/Z-Image-Turbo"):
return {
"format": PREQUANT_FORMAT,
"metadata": {"scheme": scheme, "base_model_id": base},
"state_dict": {"weight": object()},
}
def _load(
monkeypatch,
tmp_path,
ckpt,
*,
scheme = "fp8",
load_raises = False,
exists = True,
allow_local = True,
):
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises)
# The local-path branch is opt-in via a directory ALLOWLIST (it unpickles an arbitrary
# file); these tests exercise the load mechanics, so allowlist tmp_path (where ckpt.pt
# lives) unless a test is checking the gate.
if allow_local:
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path))
else:
monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
path = tmp_path / "ckpt.pt"
if exists:
path.write_bytes(b"x")
source = PrequantSource(kind = "path", location = str(path), filename = None)
return load_prequantized_transformer(
_FakeTransformer,
"Tongyi-MAI/Z-Image-Turbo",
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = scheme,
logger = None,
)
def test_load_meta_init_and_assign(monkeypatch, tmp_path):
t = _load(monkeypatch, tmp_path, _good_ckpt())
assert t is not None
# meta-init path was used, not the dense from_pretrained.
assert "from_config" in _FakeTransformer.calls
assert "from_pretrained" not in _FakeTransformer.calls
# assign=True is the whole point (copy into meta is a no-op).
assert _FakeTransformer.calls["load_state_dict"] == {"strict": True, "assign": True}
assert t.moved == "cuda"
assert t._unsloth_runtime_quant == "fp8"
def test_load_missing_file_is_none(monkeypatch, tmp_path):
assert _load(monkeypatch, tmp_path, _good_ckpt(), exists = False) is None
def test_load_torch_load_raises_is_none(monkeypatch, tmp_path):
assert _load(monkeypatch, tmp_path, _good_ckpt(), load_raises = True) is None
def test_load_format_mismatch_is_none(monkeypatch, tmp_path):
bad = _good_ckpt()
bad["format"] = "something_else"
assert _load(monkeypatch, tmp_path, bad) is None
def test_load_scheme_mismatch_is_none(monkeypatch, tmp_path):
# checkpoint built for int8, but fp8 was requested.
assert _load(monkeypatch, tmp_path, _good_ckpt(scheme = "int8"), scheme = "fp8") is None
def test_load_base_mismatch_is_none(monkeypatch, tmp_path):
assert _load(monkeypatch, tmp_path, _good_ckpt(base = "other/model")) is None
# ── local-path opt-in gate (RCE guard) ───────────────────────────────────────────
def test_load_local_path_refused_by_default(monkeypatch, tmp_path):
# A valid checkpoint at a real file is still refused: torch.load must never run on a
# request-supplied path without the operator opt-in.
called = {"load": False}
def _explode(*a, **k):
called["load"] = True
raise AssertionError("torch.load must not run on a refused local path")
torch = types.ModuleType("torch")
torch.load = _explode
monkeypatch.setitem(sys.modules, "torch", torch)
monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
path = tmp_path / "ckpt.pt"
path.write_bytes(b"x")
source = PrequantSource(kind = "path", location = str(path), filename = None)
result = load_prequantized_transformer(
_FakeTransformer,
"Tongyi-MAI/Z-Image-Turbo",
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
logger = None,
)
assert result is None
assert called["load"] is False
def test_load_local_path_allowed_with_optin(monkeypatch, tmp_path):
assert _load(monkeypatch, tmp_path, _good_ckpt(), allow_local = True) is not None
def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path):
# The hosted-repo branch is first-party and trusted: it loads with no opt-in env set.
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, _good_ckpt())
monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
downloaded = tmp_path / "transformer_fp8.pt"
downloaded.write_bytes(b"x")
hub = types.ModuleType("huggingface_hub")
hub.hf_hub_download = lambda repo_id, filename, token = None: str(downloaded)
monkeypatch.setitem(sys.modules, "huggingface_hub", hub)
source = PrequantSource(kind = "repo", location = "org/hosted-fp8", filename = "transformer_fp8.pt")
result = load_prequantized_transformer(
_FakeTransformer,
"Tongyi-MAI/Z-Image-Turbo",
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
logger = None,
)
assert result is not None
def test_load_local_path_outside_allowlist_refused(monkeypatch, tmp_path):
# Even with the opt-in set, a path OUTSIDE every allowlisted directory must not be
# unpickled: enabling one trusted dir is not a wildcard for arbitrary request paths.
called = {"load": False}
def _explode(*a, **k):
called["load"] = True
raise AssertionError("torch.load must not run on a path outside the allowlist")
torch = types.ModuleType("torch")
torch.load = _explode
monkeypatch.setitem(sys.modules, "torch", torch)
allowed = tmp_path / "allowed"
allowed.mkdir()
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(allowed))
outside = tmp_path / "evil.pt" # a real file, but outside the allowlisted dir
outside.write_bytes(b"x")
source = PrequantSource(kind = "path", location = str(outside), filename = None)
result = load_prequantized_transformer(
_FakeTransformer,
"Tongyi-MAI/Z-Image-Turbo",
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
logger = None,
)
assert result is None
assert called["load"] is False
def test_load_min_features_mismatch_is_none(monkeypatch, tmp_path):
# A checkpoint built with a different --min-features quantises a different Linear set,
# so it must be rejected when the runtime threshold is supplied.
ckpt = _good_ckpt()
ckpt["metadata"]["min_features"] = 256 # built with 256, runtime asks for 512
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, ckpt)
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path))
path = tmp_path / "ckpt.pt"
path.write_bytes(b"x")
source = PrequantSource(kind = "path", location = str(path), filename = None)
result = load_prequantized_transformer(
_FakeTransformer,
"Tongyi-MAI/Z-Image-Turbo",
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
min_features = 512,
logger = None,
)
assert result is None
def test_load_base_fork_tail_matches(monkeypatch, tmp_path):
# A local path / fork id with the same final segment as the canonical base is accepted.
ckpt = _good_ckpt(base = "Tongyi-MAI/Z-Image-Turbo")
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, ckpt)
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path))
path = tmp_path / "ckpt.pt"
path.write_bytes(b"x")
source = PrequantSource(kind = "path", location = str(path), filename = None)
result = load_prequantized_transformer(
_FakeTransformer,
"/local/models/Z-Image-Turbo", # different prefix, same tail
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
logger = None,
)
assert result is not None

29
studio/backend/tests/test_diffusion_routes.py
View file

@ -408,6 +408,35 @@ def test_transformer_quant_fast_accum_threads_through(client, monkeypatch):
assert backend.last_load_kwargs.get("transformer_quant_fast_accum") is False
def test_transformer_prequant_path_threads_through(client, monkeypatch):
backend = _FakeBackend()
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
resp = client.post(
"/api/inference/images/load",
json = {
"model_path": "x/z-image",
"gguf_filename": "q.gguf",
"transformer_quant": "fp8",
"transformer_prequant_path": "/data/zimage_fp8.pt",
},
)
assert resp.status_code == 200
assert backend.last_load_kwargs.get("transformer_prequant_path") == "/data/zimage_fp8.pt"
def test_prequant_path_doc_describes_allowlist_not_toggle():
# The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a
# directory allowlist, not a =1 toggle (diffusion_prequant._allowed_prequant_roots
# drops bare on/off tokens), so operators following the doc don't get every
# request silently refused.
from models.inference import DiffusionLoadRequest
desc = DiffusionLoadRequest.model_fields["transformer_prequant_path"].description
assert "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH" in desc
assert "=1" not in desc
assert "allowlist" in desc.lower() or "director" in desc.lower()
def test_invalid_transformer_quant_returns_422_without_eviction(client):
# An unsupported transformer_quant is rejected by the request schema (Literal), so
# the GPU is never acquired and no chat model is evicted.