Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) (#6694)
* 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 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 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.) * [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
GitHub
parent
6b9b1c72d3
commit
f43910a4a4
13 changed files with 2243 additions and 74 deletions
|
|
@ -218,6 +218,10 @@ def _run(args: argparse.Namespace) -> dict[str, Any]:
|
|||
memory_mode = args.memory_mode,
|
||||
speed_mode = args.speed_mode,
|
||||
text_encoder_quant = args.text_encoder_quant,
|
||||
transformer_quant = args.transformer_quant,
|
||||
transformer_quant_fast_accum = {"auto": None, "on": True, "off": False}[
|
||||
args.fp8_fast_accum
|
||||
],
|
||||
)
|
||||
_wait_for_load(backend)
|
||||
_cuda_sync()
|
||||
|
|
@ -299,6 +303,8 @@ def _run(args: argparse.Namespace) -> dict[str, Any]:
|
|||
"speed_mode": args.speed_mode,
|
||||
"cpu_offload": args.cpu_offload,
|
||||
"text_encoder_quant": args.text_encoder_quant,
|
||||
"transformer_quant": args.transformer_quant,
|
||||
"fp8_fast_accum": args.fp8_fast_accum,
|
||||
},
|
||||
}
|
||||
|
||||
|
|
@ -465,6 +471,21 @@ def _build_parser() -> argparse.ArgumentParser:
|
|||
choices = ["fp8", "nvfp4"],
|
||||
help = "quantise the companion text encoder (fp8 or nvfp4)",
|
||||
)
|
||||
p.add_argument(
|
||||
"--transformer-quant",
|
||||
default = None,
|
||||
choices = ["auto", "int8", "fp8", "nvfp4", "mxfp8"],
|
||||
help = "opt-in fast transformer: load the DENSE bf16 transformer and torchao-"
|
||||
"quantise it onto the low-precision tensor cores (faster than GGUF, higher "
|
||||
"VRAM). auto picks per GPU; falls back to GGUF if unsupported / no VRAM",
|
||||
)
|
||||
p.add_argument(
|
||||
"--fp8-fast-accum",
|
||||
default = "auto",
|
||||
choices = ["auto", "on", "off"],
|
||||
help = "fp8 accumulate: auto picks by GPU class (fast on consumer, precise on "
|
||||
"data-center); on/off force it",
|
||||
)
|
||||
p.add_argument(
|
||||
"--cpu-offload", action = "store_true", help = "legacy: force whole-module CPU offload"
|
||||
)
|
||||
|
|
|
|||
115
scripts/fp8_overflow_check.py
Normal file
115
scripts/fp8_overflow_check.py
Normal file
|
|
@ -0,0 +1,115 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Empirically check that fp8 dynamic quant with fast accumulation does not overflow.
|
||||
|
||||
Hooks every quantised Linear's output during a real Z-Image generation and reports the
|
||||
global max |output| and any non-finite (Inf/NaN) count, for use_fast_accum True vs False.
|
||||
The concern fast_accum raises is accumulation *precision*, not overflow (the accumulator
|
||||
stays FP32-range and torchao's dynamic per-row scale keeps FP8 inputs <= 448); this proves
|
||||
it on the real model, including Z-Image's large (~9e5) activation peaks. Run on one CUDA GPU.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
BASE = "Tongyi-MAI/Z-Image-Turbo"
|
||||
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
|
||||
|
||||
|
||||
def _load_dense():
|
||||
import torch, diffusers
|
||||
|
||||
t = diffusers.ZImageTransformer2DModel.from_pretrained(
|
||||
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
|
||||
)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _run(fast_accum, steps, res, seed, mf):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
|
||||
from torchao.float8 import Float8MMConfig
|
||||
|
||||
pipe = _load_dense()
|
||||
|
||||
def filt(mod, fqn = ""):
|
||||
return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf
|
||||
|
||||
quantize_(
|
||||
pipe.transformer,
|
||||
Float8DynamicActivationFloat8WeightConfig(
|
||||
mm_config = Float8MMConfig(use_fast_accum = fast_accum)
|
||||
),
|
||||
filter_fn = filt,
|
||||
)
|
||||
|
||||
stats = {"max_abs": 0.0, "nonfinite": 0, "hooked": 0}
|
||||
|
||||
def hook(mod, inp, out):
|
||||
t = out[0] if isinstance(out, tuple) else out
|
||||
if not torch.is_tensor(t):
|
||||
return
|
||||
finite = torch.isfinite(t)
|
||||
nf = int((~finite).sum().item())
|
||||
stats["nonfinite"] += nf
|
||||
m = float(t[finite].abs().max().item()) if finite.any() else float("inf")
|
||||
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.
|
||||
for m in pipe.transformer.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
m.register_forward_hook(hook)
|
||||
stats["hooked"] += 1
|
||||
|
||||
g = torch.Generator(device = "cuda").manual_seed(seed)
|
||||
img = pipe(
|
||||
prompt = PROMPT,
|
||||
width = res,
|
||||
height = res,
|
||||
num_inference_steps = steps,
|
||||
guidance_scale = 0.0,
|
||||
generator = g,
|
||||
).images[0]
|
||||
import numpy as np
|
||||
|
||||
arr = np.array(img)
|
||||
img_finite = bool(np.isfinite(arr).all())
|
||||
del pipe
|
||||
torch.cuda.empty_cache()
|
||||
return stats, img_finite
|
||||
|
||||
|
||||
def main(argv = None) -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--steps", type = int, default = 4)
|
||||
p.add_argument("--res", type = int, default = 512)
|
||||
p.add_argument("--seed", type = int, default = 42)
|
||||
p.add_argument("--min-feat", type = int, default = 512)
|
||||
args = p.parse_args(argv)
|
||||
|
||||
print(f"== fp8 overflow check (Z-Image dense, {args.res}px, {args.steps} steps) ==", flush = True)
|
||||
for fast in (True, False):
|
||||
stats, img_finite = _run(fast, args.steps, args.res, args.seed, args.min_feat)
|
||||
print(
|
||||
f" fast_accum={str(fast):5s} hooked_linears={stats['hooked']:3d} "
|
||||
f"max|linear_out|={stats['max_abs']:.1f} nonfinite_elems={stats['nonfinite']} "
|
||||
f"image_all_finite={img_finite}",
|
||||
flush = True,
|
||||
)
|
||||
print("FP8-OVERFLOW-CHECK-DONE", flush = True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
|
||||
sys.exit(main())
|
||||
186
scripts/nvfp4_probe.py
Normal file
186
scripts/nvfp4_probe.py
Normal file
|
|
@ -0,0 +1,186 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Probe NVFP4 via torchao with use_triton_kernel=False (no MSLK) on the real dense
|
||||
Z-Image transformer: is it a genuine FP4-tensor-core speedup over fp8, and is quality
|
||||
in-bar? Reference for LPIPS is dense bf16 eager. Run on one CUDA (Blackwell) GPU."""
|
||||
|
||||
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"
|
||||
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_images"
|
||||
|
||||
|
||||
def _psnr(a, b):
|
||||
mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
|
||||
return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse))
|
||||
|
||||
|
||||
_LP = {"fn": None}
|
||||
|
||||
|
||||
def _lpips(ref, arr):
|
||||
try:
|
||||
import torch, lpips
|
||||
|
||||
if _LP["fn"] is None:
|
||||
_LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval()
|
||||
|
||||
def t(x):
|
||||
return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda()
|
||||
|
||||
with torch.no_grad():
|
||||
return float(_LP["fn"](t(ref), t(arr)).item())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" (lpips: {type(exc).__name__})", flush = True)
|
||||
return None
|
||||
|
||||
|
||||
def _load_dense():
|
||||
import torch, diffusers
|
||||
|
||||
t = diffusers.ZImageTransformer2DModel.from_pretrained(
|
||||
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
|
||||
)
|
||||
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 _median(xs):
|
||||
return sorted(xs)[len(xs) // 2]
|
||||
|
||||
|
||||
def main(argv = None) -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--steps", type = int, default = 8)
|
||||
p.add_argument("--res", type = int, default = 1024)
|
||||
p.add_argument("--seed", type = int, default = 42)
|
||||
p.add_argument("--iters", type = int, default = 3)
|
||||
p.add_argument("--min-feat", type = int, default = 512)
|
||||
p.add_argument("--out-dir", default = None, help = "image output dir (default: repo outputs/)")
|
||||
args = p.parse_args(argv)
|
||||
steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
global OUT
|
||||
if args.out_dir:
|
||||
OUT = Path(args.out_dir).expanduser()
|
||||
OUT.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
def filt(mod, fqn = ""):
|
||||
return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf
|
||||
|
||||
def run(
|
||||
tag,
|
||||
*,
|
||||
cfg = None,
|
||||
compile = True,
|
||||
):
|
||||
torch.compiler.reset()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
pipe = _load_dense()
|
||||
if cfg is not None:
|
||||
from torchao.quantization import quantize_
|
||||
quantize_(pipe.transformer, cfg, filter_fn = filt)
|
||||
if compile:
|
||||
try:
|
||||
pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(
|
||||
f" [{tag}] compile failed: {type(exc).__name__}: {str(exc)[:90]}", flush = True
|
||||
)
|
||||
_gen(pipe, steps, seed, res) # warmup / compile
|
||||
dts, img = [], None
|
||||
for _ in range(args.iters):
|
||||
img, dt = _gen(pipe, steps, seed, res)
|
||||
dts.append(dt)
|
||||
gp = torch.cuda.max_memory_allocated() / 1e9
|
||||
arr = np.array(img)
|
||||
img.save(OUT / f"{tag}.png")
|
||||
del pipe
|
||||
torch.cuda.empty_cache()
|
||||
return _median(dts), arr, gp
|
||||
|
||||
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8
|
||||
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV
|
||||
|
||||
print(f"== nvfp4 probe (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", flush = True)
|
||||
bref, ref, _ = run("bf16_eager", cfg = None, compile = False)
|
||||
print(f" bf16 eager ref: {bref:.3f}s", flush = True)
|
||||
rows = [("bf16_eager", bref, float("inf"), 0.0, None)]
|
||||
|
||||
specs = [
|
||||
("bf16_compile", None, True),
|
||||
("fp8_compile", FP8(), True),
|
||||
("nvfp4_notriton_compile", NV(use_triton_kernel = False), True),
|
||||
("nvfp4_notriton_eager", NV(use_triton_kernel = False), False),
|
||||
]
|
||||
for tag, cfg, comp in specs:
|
||||
try:
|
||||
med, arr, gp = run(tag, cfg = cfg, compile = comp)
|
||||
ps, lp = _psnr(ref, arr), _lpips(ref, arr)
|
||||
rows.append((tag, med, ps, lp, gp))
|
||||
print(
|
||||
f" {tag:24s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G",
|
||||
flush = True,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
print(f" {tag:24s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True)
|
||||
rows.append((tag, None, None, None, None))
|
||||
|
||||
fp8 = next((r[1] for r in rows if r[0] == "fp8_compile" and r[1]), None)
|
||||
print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True)
|
||||
for tag, med, ps, lp, gp in rows:
|
||||
if med is None:
|
||||
print(f" {tag:24s} FAILED")
|
||||
continue
|
||||
vs_fp8 = f"{fp8/med:.2f}x" if fp8 else "-"
|
||||
psv = "inf" if ps == float("inf") else f"{ps:.1f}"
|
||||
lpv = (
|
||||
"ref"
|
||||
if (lp == 0.0 and tag == "bf16_eager")
|
||||
else (f"{lp:.3f}" if lp is not None else "n/a")
|
||||
)
|
||||
print(
|
||||
f" {tag:24s} {med:.3f}s vs_fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}",
|
||||
flush = True,
|
||||
)
|
||||
print("NVFP4-PROBE-DONE", flush = True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
|
||||
sys.exit(main())
|
||||
301
scripts/nvfp4_t211_probe.py
Normal file
301
scripts/nvfp4_t211_probe.py
Normal file
|
|
@ -0,0 +1,301 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""torch>=2.11 NVFP4 probe. Three parts:
|
||||
|
||||
A. diagnostics -- torch/torchao versions, cpp-extension load state, device.
|
||||
B. GEMM micro -- isolated per-linear forward latency (bf16 / fp8 / nvfp4-cutlass /
|
||||
nvfp4-triton) at Z-Image-like shapes, to measure raw FP4
|
||||
tensor-core throughput free of pipeline overhead.
|
||||
C. end-to-end -- real dense Z-Image transformer, latency + LPIPS + PSNR + VRAM,
|
||||
reference = dense bf16 eager.
|
||||
|
||||
Run on one CUDA (Blackwell) GPU. This is the experiment that decides whether NVFP4
|
||||
becomes a genuine speedup once torch>=2.11 + torchao's CUTLASS FP4 GEMM is present."""
|
||||
|
||||
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"
|
||||
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_t211_images")
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------- diag
|
||||
def diagnostics() -> None:
|
||||
import torch
|
||||
import torchao
|
||||
|
||||
print("== A. diagnostics ==", flush = True)
|
||||
print(f" torch {torch.__version__}", flush = True)
|
||||
print(f" torchao {torchao.__version__}", flush = True)
|
||||
print(f" cuda {torch.version.cuda}", flush = True)
|
||||
if torch.cuda.is_available():
|
||||
print(
|
||||
f" device {torch.cuda.get_device_name(0)} sm{torch.cuda.get_device_capability(0)}",
|
||||
flush = True,
|
||||
)
|
||||
print(f" torch.ops.torchao present: {hasattr(torch.ops, 'torchao')}", flush = True)
|
||||
print(
|
||||
f" fp4 primitives: e2m1={hasattr(torch, 'float4_e2m1fn_x2')} "
|
||||
f"e8m0={hasattr(torch, 'float8_e8m0fnu')} _scaled_mm={hasattr(torch, '_scaled_mm')}",
|
||||
flush = True,
|
||||
)
|
||||
# torchao prints "Skipping import of cpp extensions ..." to stderr at import on torch<2.11.
|
||||
# On 2.11 that line is absent -> the CUTLASS FP4 GEMM extension is live.
|
||||
print(
|
||||
" (no 'Skipping import of cpp extensions' line above => cpp/CUTLASS ext loaded)",
|
||||
flush = True,
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------- micro
|
||||
def _configs():
|
||||
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8
|
||||
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV
|
||||
return {
|
||||
"bf16": None,
|
||||
"fp8": FP8(),
|
||||
"nvfp4_cutlass": NV(use_triton_kernel = False),
|
||||
"nvfp4_triton": NV(use_triton_kernel = True),
|
||||
}
|
||||
|
||||
|
||||
def _bench_linear(K, N, M, cfg, iters, compile_):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torchao.quantization import quantize_
|
||||
|
||||
torch.compiler.reset()
|
||||
torch.cuda.empty_cache()
|
||||
m = nn.Sequential(nn.Linear(K, N, bias = False)).cuda().to(torch.bfloat16)
|
||||
if cfg is not None:
|
||||
quantize_(m, cfg)
|
||||
fn = torch.compile(m, fullgraph = True, dynamic = False) if compile_ else m
|
||||
x = torch.randn(M, K, device = "cuda", dtype = torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
for _ in range(3): # warmup / compile
|
||||
fn(x)
|
||||
torch.cuda.synchronize()
|
||||
dts = []
|
||||
for _ in range(iters):
|
||||
t0 = time.perf_counter()
|
||||
fn(x)
|
||||
torch.cuda.synchronize()
|
||||
dts.append(time.perf_counter() - t0)
|
||||
del m, fn, x
|
||||
torch.cuda.empty_cache()
|
||||
med = sorted(dts)[len(dts) // 2]
|
||||
tflops = 2.0 * M * K * N / med / 1e12
|
||||
return med, tflops
|
||||
|
||||
|
||||
def micro(M, iters, compile_):
|
||||
print(f"\n== B. GEMM micro (M={M}, compile={compile_}, iters={iters}) ==", flush = True)
|
||||
# (K, N): qkv-ish, mlp-up, mlp-down for a ~3072-dim DiT
|
||||
shapes = [(3072, 3072), (3072, 12288), (12288, 3072)]
|
||||
cfgs = _configs()
|
||||
for K, N in shapes:
|
||||
print(f" shape K={K} N={N}:", flush = True)
|
||||
base_ms = None
|
||||
fp8_ms = None
|
||||
for name, cfg in cfgs.items():
|
||||
try:
|
||||
med, tfl = _bench_linear(K, N, M, cfg, iters, compile_)
|
||||
ms = med * 1e3
|
||||
if name == "bf16":
|
||||
base_ms = ms
|
||||
if name == "fp8":
|
||||
fp8_ms = ms
|
||||
vs_bf16 = f"{base_ms/ms:.2f}x" if base_ms else "-"
|
||||
vs_fp8 = f"{fp8_ms/ms:.2f}x" if fp8_ms else "-"
|
||||
print(
|
||||
f" {name:16s} {ms:7.3f} ms {tfl:7.1f} TFLOPS vs_bf16={vs_bf16:>6s} vs_fp8={vs_fp8:>6s}",
|
||||
flush = True,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" {name:16s} FAILED: {type(exc).__name__}: {str(exc)[:120]}", flush = True)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------- e2e
|
||||
def _psnr(a, b):
|
||||
mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
|
||||
return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse))
|
||||
|
||||
|
||||
_LP = {"fn": None}
|
||||
|
||||
|
||||
def _lpips(ref, arr):
|
||||
try:
|
||||
import lpips
|
||||
import torch
|
||||
|
||||
if _LP["fn"] is None:
|
||||
_LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval()
|
||||
|
||||
def t(x):
|
||||
return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda()
|
||||
|
||||
with torch.no_grad():
|
||||
return float(_LP["fn"](t(ref), t(arr)).item())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" (lpips: {type(exc).__name__})", flush = True)
|
||||
return None
|
||||
|
||||
|
||||
def _load_dense():
|
||||
import diffusers
|
||||
import torch
|
||||
|
||||
t = diffusers.ZImageTransformer2DModel.from_pretrained(
|
||||
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
|
||||
)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
|
||||
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 _median(xs):
|
||||
return sorted(xs)[len(xs) // 2]
|
||||
|
||||
|
||||
def e2e(steps, res, seed, iters, mf):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
OUT.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
def filt(mod, fqn = ""):
|
||||
return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf
|
||||
|
||||
def run(
|
||||
tag,
|
||||
*,
|
||||
cfg = None,
|
||||
compile = True,
|
||||
):
|
||||
torch.compiler.reset()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
pipe = _load_dense()
|
||||
if cfg is not None:
|
||||
from torchao.quantization import quantize_
|
||||
quantize_(pipe.transformer, cfg, filter_fn = filt)
|
||||
if compile:
|
||||
try:
|
||||
pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(
|
||||
f" [{tag}] compile failed: {type(exc).__name__}: {str(exc)[:90]}", flush = True
|
||||
)
|
||||
_gen(pipe, steps, seed, res) # warmup / compile
|
||||
dts, img = [], None
|
||||
for _ in range(iters):
|
||||
img, dt = _gen(pipe, steps, seed, res)
|
||||
dts.append(dt)
|
||||
gp = torch.cuda.max_memory_allocated() / 1e9
|
||||
arr = np.array(img)
|
||||
img.save(OUT / f"{tag}.png")
|
||||
del pipe
|
||||
torch.cuda.empty_cache()
|
||||
return _median(dts), arr, gp
|
||||
|
||||
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV
|
||||
from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8
|
||||
|
||||
print(
|
||||
f"\n== C. end-to-end (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", flush = True
|
||||
)
|
||||
bref, ref, _ = run("bf16_eager", cfg = None, compile = False)
|
||||
print(f" bf16 eager ref: {bref:.3f}s", flush = True)
|
||||
rows = [("bf16_eager", bref, float("inf"), 0.0, None)]
|
||||
|
||||
specs = [
|
||||
("bf16_compile", None, True),
|
||||
("fp8_compile", FP8(), True),
|
||||
("nvfp4_cutlass_compile", NV(use_triton_kernel = False), True),
|
||||
("nvfp4_triton_compile", NV(use_triton_kernel = True), True),
|
||||
]
|
||||
for tag, cfg, comp in specs:
|
||||
try:
|
||||
med, arr, gp = run(tag, cfg = cfg, compile = comp)
|
||||
ps, lp = _psnr(ref, arr), _lpips(ref, arr)
|
||||
rows.append((tag, med, ps, lp, gp))
|
||||
print(
|
||||
f" {tag:24s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G",
|
||||
flush = True,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
print(f" {tag:24s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True)
|
||||
rows.append((tag, None, None, None, None))
|
||||
|
||||
fp8 = next((r[1] for r in rows if r[0] == "fp8_compile" and r[1]), None)
|
||||
print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True)
|
||||
for tag, med, ps, lp, gp in rows:
|
||||
if med is None:
|
||||
print(f" {tag:24s} FAILED")
|
||||
continue
|
||||
vs_fp8 = f"{fp8/med:.2f}x" if fp8 else "-"
|
||||
psv = "inf" if ps == float("inf") else f"{ps:.1f}"
|
||||
lpv = (
|
||||
"ref"
|
||||
if (lp == 0.0 and tag == "bf16_eager")
|
||||
else (f"{lp:.3f}" if lp is not None else "n/a")
|
||||
)
|
||||
print(
|
||||
f" {tag:24s} {med:.3f}s vs_fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}",
|
||||
flush = True,
|
||||
)
|
||||
|
||||
|
||||
def main(argv = None) -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--steps", type = int, default = 8)
|
||||
p.add_argument("--res", type = int, default = 1024)
|
||||
p.add_argument("--seed", type = int, default = 42)
|
||||
p.add_argument("--iters", type = int, default = 3)
|
||||
p.add_argument("--micro-M", type = int, default = 4096)
|
||||
p.add_argument("--min-feat", type = int, default = 512)
|
||||
p.add_argument("--only", choices = ["diag", "micro", "e2e", "all"], default = "all")
|
||||
args = p.parse_args(argv)
|
||||
|
||||
diagnostics()
|
||||
if args.only in ("micro", "all"):
|
||||
micro(args.micro_M, args.iters, compile_ = True)
|
||||
if args.only in ("e2e", "all"):
|
||||
e2e(args.steps, args.res, args.seed, args.iters, args.min_feat)
|
||||
print("NVFP4-T211-PROBE-DONE", flush = True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
|
||||
sys.exit(main())
|
||||
318
scripts/quant_probe.py
Normal file
318
scripts/quant_probe.py
Normal file
|
|
@ -0,0 +1,318 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Empirical quant probe: torchao int8/fp8/fp4 dynamic quant vs GGUF+compile.
|
||||
|
||||
Question this answers: GGUF stores the Z-Image DiT at 4-bit but dequantizes to bf16
|
||||
per matmul, so it runs at bf16 tensor-core rate. Can a low-precision *tensor-core*
|
||||
path (int8dq on any Ampere+, fp8dq on Ada+, NVFP4/MXFP8 on Blackwell), loaded from
|
||||
the dense bf16 transformer, beat GGUF+compile on speed while staying inside the
|
||||
quality bar -- and how does its quality compare to GGUF's own 4-bit loss?
|
||||
|
||||
Reference for all quality numbers is the DENSE bf16 EAGER image (the best this model
|
||||
can do). Each config is a fresh pipeline (no compile/quant cross-contamination).
|
||||
Reports median latency, PSNR + LPIPS vs reference, and peak VRAM. Run on one CUDA GPU.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
REPO = "unsloth/Z-Image-Turbo-GGUF"
|
||||
GGUF = "z-image-turbo-Q4_K_M.gguf"
|
||||
BASE = "Tongyi-MAI/Z-Image-Turbo"
|
||||
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
|
||||
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "probe_images"
|
||||
|
||||
|
||||
def _psnr(a, b):
|
||||
mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
|
||||
return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse))
|
||||
|
||||
|
||||
_LPIPS = {"fn": None}
|
||||
|
||||
|
||||
def _lpips(ref_arr, arr):
|
||||
"""Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable.
|
||||
|
||||
Runs on CPU so the scorer never holds CUDA memory: each row resets peak VRAM, so a
|
||||
resident GPU LPIPS module would inflate the reported load/gen VRAM and could even OOM."""
|
||||
try:
|
||||
import torch
|
||||
import lpips
|
||||
|
||||
if _LPIPS["fn"] is None:
|
||||
_LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval()
|
||||
|
||||
def t(x):
|
||||
return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
|
||||
|
||||
with torch.no_grad():
|
||||
return float(_LPIPS["fn"](t(ref_arr), t(arr)).item())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" (lpips unavailable: {type(exc).__name__}: {str(exc)[:80]})", flush = True)
|
||||
return None
|
||||
|
||||
|
||||
def _load_dense():
|
||||
import torch
|
||||
import diffusers
|
||||
|
||||
t = diffusers.ZImageTransformer2DModel.from_pretrained(
|
||||
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
|
||||
)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _load_gguf():
|
||||
import torch
|
||||
import diffusers
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
t = diffusers.ZImageTransformer2DModel.from_single_file(
|
||||
hf_hub_download(REPO, GGUF),
|
||||
quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = torch.bfloat16),
|
||||
torch_dtype = torch.bfloat16,
|
||||
config = BASE,
|
||||
subfolder = "transformer",
|
||||
)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _quant_config(name):
|
||||
"""Return a torchao config instance for `name`, or raise to mark FAILED."""
|
||||
from torchao.quantization import (
|
||||
Int8WeightOnlyConfig,
|
||||
Int8DynamicActivationInt8WeightConfig,
|
||||
Float8DynamicActivationFloat8WeightConfig,
|
||||
)
|
||||
|
||||
if name == "int8wo":
|
||||
return Int8WeightOnlyConfig()
|
||||
if name == "int8dq":
|
||||
return Int8DynamicActivationInt8WeightConfig()
|
||||
if name == "fp8dq":
|
||||
return Float8DynamicActivationFloat8WeightConfig()
|
||||
if name == "nvfp4":
|
||||
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig
|
||||
return NVFP4DynamicActivationNVFP4WeightConfig()
|
||||
if name == "mxfp8":
|
||||
from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
|
||||
try:
|
||||
import torch
|
||||
return MXDynamicActivationMXWeightConfig(
|
||||
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
|
||||
)
|
||||
except (TypeError, AttributeError):
|
||||
return MXDynamicActivationMXWeightConfig()
|
||||
raise ValueError(name)
|
||||
|
||||
|
||||
def _make_filter_fn(min_features):
|
||||
"""Keep only the FLOP-heavy linears: nn.Linear with both in/out >= min_features.
|
||||
The int8 dynamic path uses torch._int_mm (needs activation M>16), and the tiny
|
||||
timestep/pooled projections (in_features=256) run at M=1 and crash it -- skip them."""
|
||||
import torch.nn as nn
|
||||
|
||||
def filter_fn(module, fqn = ""):
|
||||
return (
|
||||
isinstance(module, nn.Linear)
|
||||
and getattr(module, "in_features", 0) >= min_features
|
||||
and getattr(module, "out_features", 0) >= min_features
|
||||
)
|
||||
|
||||
return filter_fn
|
||||
|
||||
|
||||
def _apply_quant(pipe, name, log, min_features):
|
||||
import torch.nn as nn
|
||||
from torchao.quantization import quantize_
|
||||
|
||||
cfg = _quant_config(name)
|
||||
total = sum(1 for m in pipe.transformer.modules() if isinstance(m, nn.Linear))
|
||||
filt = _make_filter_fn(min_features)
|
||||
q = sum(1 for n, m in pipe.transformer.named_modules() if filt(m, n))
|
||||
quantize_(pipe.transformer, cfg, filter_fn = filt)
|
||||
log(f" quantized transformer with {name} ({q}/{total} linears >= {min_features} feat)")
|
||||
|
||||
|
||||
def _compile(pipe, log):
|
||||
fn = getattr(pipe.transformer, "compile_repeated_blocks", None)
|
||||
if not callable(fn):
|
||||
return False
|
||||
for kw in ({"fullgraph": True, "dynamic": True}, {"dynamic": True}, {}):
|
||||
try:
|
||||
fn(**kw)
|
||||
log(f" compiled repeated blocks {kw}")
|
||||
return True
|
||||
except Exception as exc: # noqa: BLE001
|
||||
log(f" compile {kw} failed: {type(exc).__name__}: {str(exc)[:90]}")
|
||||
return False
|
||||
|
||||
|
||||
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 _median(xs):
|
||||
return sorted(xs)[len(xs) // 2]
|
||||
|
||||
|
||||
def main(argv = None) -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--steps", type = int, default = 8)
|
||||
p.add_argument("--res", type = int, default = 1024)
|
||||
p.add_argument("--seed", type = int, default = 42)
|
||||
p.add_argument("--iters", type = int, default = 3)
|
||||
p.add_argument(
|
||||
"--min-feat",
|
||||
type = int,
|
||||
default = 512,
|
||||
help = "only quantize Linear with in&out features >= this (int8 _int_mm needs M>16)",
|
||||
)
|
||||
p.add_argument(
|
||||
"--configs",
|
||||
default = "bf16,bf16_c,gguf_c,int8dq_c,fp8dq_c,nvfp4_c,mxfp8_c,int8wo_c",
|
||||
help = "comma list; suffix _c = +compile",
|
||||
)
|
||||
args = p.parse_args(argv)
|
||||
steps, res, seed, iters = args.steps, args.res, args.seed, args.iters
|
||||
|
||||
import torch
|
||||
|
||||
OUT.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
def run(
|
||||
tag,
|
||||
*,
|
||||
source,
|
||||
quant = None,
|
||||
compile = False,
|
||||
):
|
||||
torch.compiler.reset()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
pipe = _load_dense() if source == "dense" else _load_gguf()
|
||||
load_peak = torch.cuda.max_memory_allocated() / 1e9
|
||||
if quant is not None:
|
||||
_apply_quant(pipe, quant, print_, args.min_feat)
|
||||
if compile:
|
||||
_compile(pipe, print_)
|
||||
_gen(pipe, steps, seed, res) # warmup / compilation
|
||||
else:
|
||||
_gen(pipe, steps, seed, res) # allocator warmup
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
dts, img = [], None
|
||||
for _ in range(iters):
|
||||
img, dt = _gen(pipe, steps, seed, res)
|
||||
dts.append(dt)
|
||||
gen_peak = torch.cuda.max_memory_allocated() / 1e9
|
||||
arr = np.array(img)
|
||||
img.save(OUT / f"{tag}.png")
|
||||
del pipe
|
||||
torch.cuda.empty_cache()
|
||||
return tag, _median(dts), arr, load_peak, gen_peak
|
||||
|
||||
print_ = lambda s: print(s, flush = True) # noqa: E731
|
||||
|
||||
# config table: tag -> (source, quant, compile)
|
||||
table = {
|
||||
"bf16": ("dense", None, False),
|
||||
"bf16_c": ("dense", None, True),
|
||||
"gguf_c": ("gguf", None, True),
|
||||
"int8wo_c": ("dense", "int8wo", True),
|
||||
"int8dq_c": ("dense", "int8dq", True),
|
||||
"fp8dq_c": ("dense", "fp8dq", True),
|
||||
"nvfp4_c": ("dense", "nvfp4", True),
|
||||
"mxfp8_c": ("dense", "mxfp8", True),
|
||||
}
|
||||
want = [c.strip() for c in args.configs.split(",") if c.strip()]
|
||||
|
||||
print(f"== quant probe (Z-Image-Turbo, {res}px, {steps} steps, seed {seed}) ==", flush = True)
|
||||
ref_arr = None
|
||||
rows = []
|
||||
for tag in want:
|
||||
if tag not in table:
|
||||
print(f" {tag}: unknown config, skipping", flush = True)
|
||||
continue
|
||||
source, quant, compile = table[tag]
|
||||
print(f"-- {tag} (source={source} quant={quant} compile={compile}) --", flush = True)
|
||||
try:
|
||||
_, med, arr, lp, gp = run(tag, source = source, quant = quant, compile = compile)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
import traceback
|
||||
|
||||
print(f" {tag:10s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True)
|
||||
traceback.print_exc()
|
||||
rows.append((tag, None, None, None, None, None))
|
||||
continue
|
||||
if ref_arr is None and tag == "bf16":
|
||||
ref_arr = arr
|
||||
psnr = _psnr(ref_arr, arr) if ref_arr is not None else None
|
||||
lpips_v = (
|
||||
_lpips(ref_arr, arr)
|
||||
if (ref_arr is not None and tag != "bf16")
|
||||
else (0.0 if tag == "bf16" else None)
|
||||
)
|
||||
rows.append((tag, med, psnr, lpips_v, lp, gp))
|
||||
ps = f"{psnr:.1f}dB" if psnr is not None else "n/a"
|
||||
lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a"
|
||||
print(
|
||||
f" {tag:10s} {med:.3f}s PSNR={ps:>7s} LPIPS={lps:>6s} loadVRAM={lp:.1f}G genVRAM={gp:.1f}G",
|
||||
flush = True,
|
||||
)
|
||||
|
||||
base = next((r[1] for r in rows if r[0] == "bf16" and r[1]), None)
|
||||
gguf = next((r[1] for r in rows if r[0] == "gguf_c" and r[1]), None)
|
||||
print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True)
|
||||
print(
|
||||
f"{'config':10s} {'sec':>7s} {'vs_bf16':>8s} {'vs_gguf':>8s} {'PSNR':>8s} {'LPIPS':>7s} {'loadG':>6s} {'genG':>6s}",
|
||||
flush = True,
|
||||
)
|
||||
for tag, med, psnr, lpips_v, lp, gp in rows:
|
||||
if med is None:
|
||||
print(f"{tag:10s} {'FAILED':>7s}", flush = True)
|
||||
continue
|
||||
vb = f"{base/med:.2f}x" if base else "-"
|
||||
vg = f"{gguf/med:.2f}x" if gguf else "-"
|
||||
ps = (
|
||||
f"{psnr:.1f}"
|
||||
if psnr is not None and psnr != float("inf")
|
||||
else ("inf" if psnr == float("inf") else "n/a")
|
||||
)
|
||||
lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a"
|
||||
print(
|
||||
f"{tag:10s} {med:>7.3f} {vb:>8s} {vg:>8s} {ps:>8s} {lps:>7s} {lp:>6.1f} {gp:>6.1f}",
|
||||
flush = True,
|
||||
)
|
||||
print("QUANT-PROBE-DONE", flush = True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
|
||||
sys.exit(main())
|
||||
251
scripts/sparse_accum_probe.py
Normal file
251
scripts/sparse_accum_probe.py
Normal file
|
|
@ -0,0 +1,251 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Probe two consumer-GPU-motivated levers on the real dense Z-Image transformer:
|
||||
|
||||
* fp8 fast_accum on/off -- on consumer Blackwell, fp8 with FP16 accumulate is ~2x
|
||||
fp8 with FP32 accumulate (838 vs 419 TFLOPS). torchao defaults use_fast_accum=True,
|
||||
so this confirms we are already on the fast path and quantifies it (muted on a B200,
|
||||
which is not nerfed, but the knob still moves latency).
|
||||
* 2:4 semi-structured sparsity -- doubles tensor-core rate in theory. Two blockers to
|
||||
test empirically: (a) QUALITY -- inference-only 2:4 magnitude-pruning drops 50% of
|
||||
weights with no fine-tune; (b) it does NOT compose with torch.compile, so the real
|
||||
sparse path runs eager. We measure sparse-no-compile speed vs our fp8+compile
|
||||
baseline (the bar it must beat) and the LPIPS of 2:4 pruning.
|
||||
|
||||
Reference for quality is the dense bf16 eager image. Run on one CUDA GPU.
|
||||
"""
|
||||
|
||||
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"
|
||||
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/sparse_images")
|
||||
|
||||
|
||||
def _psnr(a, b):
|
||||
mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
|
||||
return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse))
|
||||
|
||||
|
||||
_LP = {"fn": None}
|
||||
|
||||
|
||||
def _lpips(ref, arr):
|
||||
try:
|
||||
import torch, lpips
|
||||
|
||||
if _LP["fn"] is None:
|
||||
_LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval()
|
||||
|
||||
def t(x):
|
||||
return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda()
|
||||
|
||||
with torch.no_grad():
|
||||
return float(_LP["fn"](t(ref), t(arr)).item())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" (lpips: {type(exc).__name__})", flush = True)
|
||||
return None
|
||||
|
||||
|
||||
def _load_dense():
|
||||
import torch, diffusers
|
||||
|
||||
t = diffusers.ZImageTransformer2DModel.from_pretrained(
|
||||
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
|
||||
)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _big_linears(transformer, min_feat = 512):
|
||||
import torch.nn as nn
|
||||
return [
|
||||
m
|
||||
for m in transformer.modules()
|
||||
if isinstance(m, nn.Linear) and m.in_features >= min_feat and m.out_features >= min_feat
|
||||
]
|
||||
|
||||
|
||||
def _prune_24_(transformer, min_feat = 512):
|
||||
"""In-place 2:4 magnitude prune (zero the 2 smallest of every 4 along in_features)
|
||||
of the FLOP-heavy linears. Dense format -> measures the QUALITY of 2:4 with no kernel."""
|
||||
import torch
|
||||
|
||||
n = 0
|
||||
for lin in _big_linears(transformer, min_feat):
|
||||
w = lin.weight.data
|
||||
o, i = w.shape
|
||||
if i % 4:
|
||||
continue
|
||||
g = w.view(o, i // 4, 4)
|
||||
idx = g.abs().argsort(dim = -1)[..., :2]
|
||||
g.scatter_(-1, idx, 0.0)
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
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 _median(xs):
|
||||
return sorted(xs)[len(xs) // 2]
|
||||
|
||||
|
||||
def main(argv = None) -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--steps", type = int, default = 8)
|
||||
p.add_argument("--res", type = int, default = 1024)
|
||||
p.add_argument("--seed", type = int, default = 42)
|
||||
p.add_argument("--iters", type = int, default = 3)
|
||||
p.add_argument("--min-feat", type = int, default = 512)
|
||||
args = p.parse_args(argv)
|
||||
steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat
|
||||
import torch
|
||||
|
||||
OUT.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
def filt(mod, fqn = ""):
|
||||
import torch.nn as nn
|
||||
return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf
|
||||
|
||||
def run(
|
||||
tag,
|
||||
*,
|
||||
quant = None,
|
||||
fast_accum = True,
|
||||
prune = False,
|
||||
real_sparse = False,
|
||||
compile = True,
|
||||
):
|
||||
torch.compiler.reset()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
pipe = _load_dense()
|
||||
note = ""
|
||||
if prune or real_sparse:
|
||||
n = _prune_24_(pipe.transformer, mf)
|
||||
note += f" pruned24={n}"
|
||||
if real_sparse:
|
||||
from torchao.sparsity import sparsify_, semi_sparse_weight
|
||||
sparsify_(pipe.transformer, semi_sparse_weight(), filter_fn = filt)
|
||||
note += " +semi_sparse"
|
||||
if quant == "fp8":
|
||||
from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig
|
||||
from torchao.float8 import Float8MMConfig
|
||||
|
||||
cfg = Float8DynamicActivationFloat8WeightConfig(
|
||||
mm_config = Float8MMConfig(use_fast_accum = fast_accum)
|
||||
)
|
||||
quantize_(pipe.transformer, cfg, filter_fn = filt)
|
||||
note += f" fp8(fast_accum={fast_accum})"
|
||||
if compile:
|
||||
try:
|
||||
pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
note += f" [compile FAILED {type(exc).__name__}]"
|
||||
print(f" [{tag}]{note}", flush = True)
|
||||
_gen(pipe, steps, seed, res) # warmup / compile
|
||||
dts = []
|
||||
img = None
|
||||
for _ in range(args.iters):
|
||||
img, dt = _gen(pipe, steps, seed, res)
|
||||
dts.append(dt)
|
||||
gp = torch.cuda.max_memory_allocated() / 1e9
|
||||
arr = np.array(img)
|
||||
img.save(OUT / f"{tag}.png")
|
||||
del pipe
|
||||
torch.cuda.empty_cache()
|
||||
return _median(dts), arr, gp
|
||||
|
||||
print(
|
||||
f"== sparse/accum probe (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==",
|
||||
flush = True,
|
||||
)
|
||||
rows = []
|
||||
# quality reference: dense bf16 eager (no compile, no quant)
|
||||
bref, ref, _ = run("bf16_eager", compile = False)
|
||||
rows.append(("bf16_eager", bref, float("inf"), 0.0, None))
|
||||
print(f" bf16 eager ref: {bref:.3f}s", flush = True)
|
||||
|
||||
specs = [
|
||||
("bf16_compile", dict()),
|
||||
("fp8_fastT_c", dict(quant = "fp8", fast_accum = True)),
|
||||
("fp8_fastF_c", dict(quant = "fp8", fast_accum = False)),
|
||||
(
|
||||
"fake24_fp8_c",
|
||||
dict(quant = "fp8", fast_accum = True, prune = True),
|
||||
), # quality of 2:4+fp8 (fake=no kernel)
|
||||
(
|
||||
"real24_nocompile",
|
||||
dict(real_sparse = True, compile = False),
|
||||
), # sparse SPEED (no quant, no compile)
|
||||
(
|
||||
"real24_compile_try",
|
||||
dict(real_sparse = True, compile = True),
|
||||
), # does sparse survive compile?
|
||||
]
|
||||
for tag, kw in specs:
|
||||
try:
|
||||
med, arr, gp = run(tag, **kw)
|
||||
ps, lp = _psnr(ref, arr), _lpips(ref, arr)
|
||||
rows.append((tag, med, ps, lp, gp))
|
||||
print(
|
||||
f" {tag:18s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G",
|
||||
flush = True,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
print(f" {tag:18s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True)
|
||||
rows.append((tag, None, None, None, None))
|
||||
|
||||
print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True)
|
||||
base = next((r[1] for r in rows if r[0] == "fp8_fastT_c" and r[1]), None)
|
||||
for tag, med, ps, lp, gp in rows:
|
||||
if med is None:
|
||||
print(f" {tag:18s} FAILED")
|
||||
continue
|
||||
vs_eager = f"{bref/med:.2f}x"
|
||||
vs_fp8 = f"{base/med:.2f}x" if base else "-"
|
||||
psv = "inf" if ps == float("inf") else f"{ps:.1f}"
|
||||
lpv = (
|
||||
"ref"
|
||||
if (lp == 0.0 and tag == "bf16_eager")
|
||||
else (f"{lp:.3f}" if lp is not None else "n/a")
|
||||
)
|
||||
print(
|
||||
f" {tag:18s} {med:.3f}s eager:{vs_eager:>6s} fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}",
|
||||
flush = True,
|
||||
)
|
||||
print("SPARSE-ACCUM-DONE", flush = True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
|
||||
sys.exit(main())
|
||||
|
|
@ -45,6 +45,7 @@ from .diffusion_memory import (
|
|||
snapshot_device_memory,
|
||||
)
|
||||
from .diffusion_speed import (
|
||||
SPEED_DEFAULT,
|
||||
SPEED_OFF,
|
||||
apply_speed_optims,
|
||||
resolve_speed_mode,
|
||||
|
|
@ -52,6 +53,11 @@ from .diffusion_speed import (
|
|||
snapshot_backend_flags,
|
||||
)
|
||||
from .diffusion_precision import quantize_text_encoders
|
||||
from .diffusion_transformer_quant import (
|
||||
dense_transformer_supported,
|
||||
normalize_transformer_quant,
|
||||
quantize_transformer,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -81,6 +87,9 @@ class _LoadState:
|
|||
backend_flags_before: Optional[dict] = None
|
||||
# Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C).
|
||||
text_encoder_quant: Optional[str] = None
|
||||
# Transformer quant actually engaged on the opt-in dense fast path: "int8" | "fp8"
|
||||
# | "nvfp4" | "mxfp8" | None. None means the default GGUF transformer was loaded.
|
||||
transformer_quant: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -288,6 +297,8 @@ class DiffusionBackend:
|
|||
memory_mode: Optional[str] = None,
|
||||
speed_mode: Optional[str] = None,
|
||||
text_encoder_quant: Optional[str] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
transformer_quant_fast_accum: Optional[bool] = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Validate, then run the (slow) load on a daemon thread. Returns at once."""
|
||||
fam = self.validate_load_request(
|
||||
|
|
@ -319,6 +330,8 @@ class DiffusionBackend:
|
|||
memory_mode = memory_mode,
|
||||
speed_mode = speed_mode,
|
||||
text_encoder_quant = text_encoder_quant,
|
||||
transformer_quant = transformer_quant,
|
||||
transformer_quant_fast_accum = transformer_quant_fast_accum,
|
||||
_load_token = token,
|
||||
),
|
||||
daemon = True,
|
||||
|
|
@ -452,6 +465,8 @@ class DiffusionBackend:
|
|||
memory_mode: Optional[str] = None,
|
||||
speed_mode: Optional[str] = None,
|
||||
text_encoder_quant: Optional[str] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
transformer_quant_fast_accum: Optional[bool] = None,
|
||||
_load_token: Optional[int] = None,
|
||||
) -> dict[str, Any]:
|
||||
# Validate first (cheap, no torch/diffusers) so a direct call with a bad
|
||||
|
|
@ -484,98 +499,146 @@ class DiffusionBackend:
|
|||
# checkpoints never sit in VRAM at once.
|
||||
self._unload_locked()
|
||||
|
||||
# Dequantise the GGUF transformer on-device; the VAE / text-encoder /
|
||||
# scheduler come from the base diffusers repo (GGUF is transformer-only).
|
||||
gguf_path = self._resolve_gguf_path(repo_id, gguf_filename, hf_token)
|
||||
transformer_cls = getattr(diffusers, fam.transformer_class)
|
||||
transformer = transformer_cls.from_single_file(
|
||||
gguf_path,
|
||||
quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype),
|
||||
torch_dtype = dtype,
|
||||
config = base,
|
||||
subfolder = "transformer",
|
||||
# Forward the token: the config is fetched from the (possibly gated)
|
||||
# base repo before from_pretrained gets a chance to authenticate.
|
||||
token = hf_token,
|
||||
pipeline_cls = getattr(diffusers, fam.pipeline_class)
|
||||
|
||||
# Decide placement up front (the weights are still on CPU, so free VRAM is
|
||||
# the real budget) -- this also doubles as the dense-quant preflight: the
|
||||
# dense bf16 transformer must fit resident, so the fast path is offered only
|
||||
# when the plan is `none`.
|
||||
plan = self._plan_memory(
|
||||
target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload
|
||||
)
|
||||
|
||||
pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer}
|
||||
if hf_token:
|
||||
pipe_kwargs["token"] = hf_token
|
||||
pipeline_cls = getattr(diffusers, fam.pipeline_class)
|
||||
pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
|
||||
# Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it
|
||||
# (int8 / fp8 / fp4 tensor cores), which beats GGUF's bf16-rate per-matmul
|
||||
# dequant on both speed and quality, at the cost of a higher-memory dense
|
||||
# load. Gated on CUDA + bf16 + a resident fit; ANY failure (unsupported arch
|
||||
# / scheme, OOM, partial quant) falls back to the GGUF build below.
|
||||
pipe = None
|
||||
transformer_quant_engaged = None
|
||||
if (
|
||||
normalize_transformer_quant(transformer_quant) is not None
|
||||
and dense_transformer_supported(target)
|
||||
and plan.offload_policy == OFFLOAD_NONE
|
||||
):
|
||||
try:
|
||||
pipe, transformer_quant_engaged = self._load_dense_quant_pipeline(
|
||||
transformer_cls,
|
||||
pipeline_cls,
|
||||
base,
|
||||
device,
|
||||
dtype,
|
||||
hf_token,
|
||||
target,
|
||||
transformer_quant,
|
||||
transformer_quant_fast_accum,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — fall back to the GGUF build
|
||||
logger.warning(
|
||||
"diffusion.transformer_quant_fallback: %s (loading GGUF)", exc
|
||||
)
|
||||
pipe = None
|
||||
transformer_quant_engaged = None
|
||||
clear_gpu_cache()
|
||||
|
||||
if pipe is None:
|
||||
# Default: dequantise the single-file GGUF transformer on-device; the
|
||||
# VAE / text-encoder / scheduler come from the base diffusers repo
|
||||
# (GGUF is transformer-only).
|
||||
transformer = transformer_cls.from_single_file(
|
||||
gguf_path,
|
||||
quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype),
|
||||
torch_dtype = dtype,
|
||||
config = base,
|
||||
subfolder = "transformer",
|
||||
# Forward the token: the config is fetched from the (possibly gated)
|
||||
# base repo before from_pretrained gets a chance to authenticate.
|
||||
token = hf_token,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
# Resolve the effective speed mode: GGUF models default to the
|
||||
# near-lossless `default` profile (compile is ~2.2x and sits below
|
||||
# the quant noise floor), dense models stay bit-identical `off`. An
|
||||
# explicit speed_mode (incl. "off") is honored verbatim.
|
||||
effective_speed = resolve_speed_mode(speed_mode, is_gguf = bool(gguf_filename))
|
||||
# A torchao-quantized dense transformer runs its matmuls through the
|
||||
# regional torch.compile; UNcompiled (eager) it is ~30x slower and would
|
||||
# lose to the GGUF fallback. A dense model otherwise resolves to `off`, so
|
||||
# force at least `default` (regional compile) whenever the quant engaged,
|
||||
# or the opt-in "fast" path silently commits an eager, pathologically slow
|
||||
# pipeline.
|
||||
if transformer_quant_engaged is not None and effective_speed == SPEED_OFF:
|
||||
logger.info(
|
||||
"diffusion.transformer_quant: forcing speed_mode=default "
|
||||
"(quantized transformer must be compiled; eager is ~30x slower)"
|
||||
)
|
||||
effective_speed = SPEED_DEFAULT
|
||||
# Opt-in speed optims run BEFORE placement (channels_last / compile
|
||||
# must precede CPU offload). Snapshot the process-wide backend flags
|
||||
# first so unload can restore them: TF32 / cudnn.benchmark are global,
|
||||
# and a later `off` load must not inherit this load's settings.
|
||||
backend_flags_before = snapshot_backend_flags()
|
||||
# apply_speed_optims mutates PROCESS-WIDE flags (TF32 / cudnn.benchmark);
|
||||
# they are only restored via _LoadState.backend_flags_before on unload. If
|
||||
# the build fails after this but before _state commits (e.g. an OOM in
|
||||
# apply_memory_plan / pipe.to), nothing would restore them and a later `off`
|
||||
# generation would be contaminated, so restore on any non-committed exit.
|
||||
committed = False
|
||||
try:
|
||||
speed_applied = apply_speed_optims(
|
||||
pipe,
|
||||
target,
|
||||
is_gguf = bool(gguf_filename),
|
||||
family = fam,
|
||||
speed_mode = effective_speed,
|
||||
logger = logger,
|
||||
)
|
||||
# Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
|
||||
# also before placement so the offload hooks move the smaller weights.
|
||||
te_quant = quantize_text_encoders(
|
||||
pipe,
|
||||
target,
|
||||
mode = text_encoder_quant,
|
||||
logger = logger,
|
||||
speed_applied = apply_speed_optims(
|
||||
pipe,
|
||||
target,
|
||||
is_gguf = bool(gguf_filename),
|
||||
family = fam,
|
||||
speed_mode = effective_speed,
|
||||
logger = logger,
|
||||
)
|
||||
if transformer_quant_engaged is not None and not speed_applied.get("compiled"):
|
||||
# Promotion above could not engage compile (e.g. the family is not
|
||||
# compile-friendly, or compile_repeated_blocks failed): the quantized
|
||||
# transformer is now running eager, which is far slower than the GGUF
|
||||
# path it replaced. Surface it loudly rather than hiding the regression.
|
||||
logger.warning(
|
||||
"diffusion.transformer_quant: %s engaged but the transformer is NOT "
|
||||
"compiled; eager torchao quant is ~30x slower than GGUF here",
|
||||
transformer_quant_engaged,
|
||||
)
|
||||
# Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4),
|
||||
# also before placement so the offload hooks move the smaller weights.
|
||||
te_quant = quantize_text_encoders(
|
||||
pipe,
|
||||
target,
|
||||
mode = text_encoder_quant,
|
||||
logger = logger,
|
||||
)
|
||||
|
||||
# Decide placement from MEASURED free device memory vs the model's
|
||||
# estimated resident size (transformer GGUF dequantised + the
|
||||
# companion text-encoder / VAE already cached for `base`), then
|
||||
# apply it. Computed here, after the build but before placement,
|
||||
# because the weights are still on CPU so free VRAM is the real
|
||||
# budget. `cpu_offload=True` stays an explicit override.
|
||||
plan = self._plan_memory(
|
||||
target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload
|
||||
)
|
||||
# apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it
|
||||
# may fall back to whole-module offload, and tiling is a no-op on a
|
||||
# pipeline with no tiling control), so status stays honest.
|
||||
effective_policy, effective_tiling = apply_memory_plan(
|
||||
pipe, plan, device = device, logger = logger
|
||||
)
|
||||
# Apply the placement planned above (from MEASURED free device memory vs
|
||||
# the model's estimated resident size). apply_memory_plan returns the
|
||||
# (policy, tiling) ACTUALLY engaged (it may fall back to whole-module
|
||||
# offload, and tiling is a no-op on a pipeline with no tiling control), so
|
||||
# status stays honest. The dense fast path already placed the pipe resident;
|
||||
# for the `none` policy this is an idempotent re-placement.
|
||||
effective_policy, effective_tiling = apply_memory_plan(
|
||||
pipe, plan, device = device, logger = logger
|
||||
)
|
||||
|
||||
self._state = _LoadState(
|
||||
pipe = pipe,
|
||||
family = fam,
|
||||
repo_id = repo_id,
|
||||
base_repo = base,
|
||||
device = device,
|
||||
dtype = str(dtype).replace("torch.", ""),
|
||||
cpu_offload = effective_policy != OFFLOAD_NONE,
|
||||
offload_policy = effective_policy,
|
||||
vae_tiling = effective_tiling,
|
||||
memory_mode = plan.requested_mode,
|
||||
speed_mode = effective_speed,
|
||||
speed_optims = tuple(k for k, v in speed_applied.items() if v),
|
||||
backend_flags_before = backend_flags_before,
|
||||
text_encoder_quant = te_quant,
|
||||
)
|
||||
committed = True
|
||||
finally:
|
||||
if not committed:
|
||||
restore_backend_flags(backend_flags_before)
|
||||
clear_gpu_cache()
|
||||
self._state = _LoadState(
|
||||
pipe = pipe,
|
||||
family = fam,
|
||||
repo_id = repo_id,
|
||||
base_repo = base,
|
||||
device = device,
|
||||
dtype = str(dtype).replace("torch.", ""),
|
||||
cpu_offload = effective_policy != OFFLOAD_NONE,
|
||||
offload_policy = effective_policy,
|
||||
vae_tiling = effective_tiling,
|
||||
memory_mode = plan.requested_mode,
|
||||
speed_mode = effective_speed,
|
||||
speed_optims = tuple(k for k, v in speed_applied.items() if v),
|
||||
backend_flags_before = backend_flags_before,
|
||||
text_encoder_quant = te_quant,
|
||||
transformer_quant = transformer_quant_engaged,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s",
|
||||
|
|
@ -588,6 +651,39 @@ class DiffusionBackend:
|
|||
)
|
||||
return self.status()
|
||||
|
||||
def _load_dense_quant_pipeline(
|
||||
self,
|
||||
transformer_cls: Any,
|
||||
pipeline_cls: Any,
|
||||
base: str,
|
||||
device: str,
|
||||
dtype: Any,
|
||||
hf_token: Optional[str],
|
||||
target: DiffusionDeviceTarget,
|
||||
mode: Optional[str],
|
||||
fast_accum: Optional[bool] = 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)``.
|
||||
|
||||
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."""
|
||||
transformer = transformer_cls.from_pretrained(
|
||||
base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
|
||||
)
|
||||
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
|
||||
|
||||
def _plan_memory(
|
||||
self,
|
||||
target: DiffusionDeviceTarget,
|
||||
|
|
@ -796,6 +892,7 @@ class DiffusionBackend:
|
|||
"speed_mode": None,
|
||||
"speed_optims": [],
|
||||
"text_encoder_quant": None,
|
||||
"transformer_quant": None,
|
||||
}
|
||||
return {
|
||||
"loaded": True,
|
||||
|
|
@ -811,6 +908,7 @@ class DiffusionBackend:
|
|||
"speed_mode": state.speed_mode,
|
||||
"speed_optims": list(state.speed_optims),
|
||||
"text_encoder_quant": state.text_encoder_quant,
|
||||
"transformer_quant": state.transformer_quant,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
347
studio/backend/core/inference/diffusion_transformer_quant.py
Normal file
347
studio/backend/core/inference/diffusion_transformer_quant.py
Normal file
|
|
@ -0,0 +1,347 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Opt-in low-precision quantisation of the diffusion DiT transformer.
|
||||
|
||||
The default path loads the transformer as a single-file GGUF, which stores weights
|
||||
4-bit but DEQUANTISES to bf16 on every matmul -- so it runs at bf16 tensor-core rate
|
||||
and never touches the int8 / fp8 / fp4 tensor cores. It is a memory win that costs
|
||||
speed. This module is the opt-in alternative: load the DENSE bf16 transformer from the
|
||||
base repo and torchao-quantise it with a DYNAMIC-ACTIVATION scheme so the matmul runs
|
||||
on the low-precision tensor cores. Measured on a B200 (Z-Image-Turbo, 1024px / 8 steps)
|
||||
vs the GGUF+compile default (0.802s, LPIPS 0.083 vs dense bf16): fp8 dynamic 0.585s
|
||||
(1.37x), int8 dynamic 0.603s (1.33x), both at LOWER LPIPS than GGUF -- faster AND a hair
|
||||
more accurate, at the cost of a higher-memory dense load. So it is strictly opt-in; the
|
||||
loader keeps GGUF as the low-memory default and the fallback.
|
||||
|
||||
Scheme by architecture (``auto`` picks the best supported, best first):
|
||||
nvfp4 / mxfp8 - Blackwell sm_100+ FP4 / MX tensor cores (biggest win; prototype).
|
||||
fp8 - Ada / Hopper / Blackwell (sm_89+) fp8 tensor cores.
|
||||
int8 - Ampere+ (sm_80+) int8 tensor cores -- the broadest-hardware lever.
|
||||
|
||||
Every scheme needs ``torch.compile`` to realise the speedup (dynamic quant is ~30x
|
||||
slower eager); the loader already compiles the repeated block AFTER this runs. torch /
|
||||
torchao are imported lazily so the module stays importable in a no-torch runtime, and
|
||||
every probe is best-effort: an unsupported scheme yields None and the caller loads GGUF.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
TQ_INT8 = "int8"
|
||||
TQ_FP8 = "fp8"
|
||||
TQ_NVFP4 = "nvfp4"
|
||||
TQ_MXFP8 = "mxfp8"
|
||||
TQ_AUTO = "auto"
|
||||
TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8)
|
||||
TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES
|
||||
|
||||
# Skip linears whose in/out features are below this. The int8 dynamic path uses
|
||||
# torch._int_mm, which requires the activation row count M > 16, and the DiT's tiny
|
||||
# timestep / pooled / modulation projections run at M=1 and crash it. They are a
|
||||
# negligible share of the FLOPs, so leaving them bf16 costs ~nothing (measured:
|
||||
# 239/276 Z-Image linears quantised, full speedup) and keeps quality a touch higher.
|
||||
DEFAULT_MIN_LINEAR_FEATURES = 512
|
||||
|
||||
# Per-architecture preference order for ``auto`` -- best (fastest, in-bar) first, with
|
||||
# the lower-precision schemes listed as fallbacks for that arch tier. On Blackwell, fp8
|
||||
# leads: measured on a B200, plain fp8 dynamic is both faster AND more accurate than the
|
||||
# alternatives for the DiT's shapes. mxfp8's block scaling adds overhead without a speed
|
||||
# win, so it sits below fp8. nvfp4 is intentionally below fp8 too: the FP4 tensor-core
|
||||
# GEMM is real once torch>=2.11 + torchao's CUTLASS FP4 kernel is present (verified: a
|
||||
# 16384^3 GEMM hits ~3826 TFLOPS, 1.37x fp8), but it only beats fp8 on very large GEMMs.
|
||||
# At the DiT's actual shapes (hidden ~3072, MLP ~12288, M~4096) it is *slower* than fp8
|
||||
# (0.81x end-to-end on Z-Image 1024px) AND notably less accurate (LPIPS 0.166 vs fp8's
|
||||
# 0.044), because FP4's per-forward quant overhead is not amortised and the format is
|
||||
# coarser. So nvfp4 is kept as an explicit opt-in, never the auto pick for diffusion.
|
||||
_AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = (
|
||||
((10, 0), (TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8)), # Blackwell sm_100+
|
||||
((8, 9), (TQ_FP8, TQ_INT8)), # Ada sm_89 / Hopper sm_90
|
||||
((8, 0), (TQ_INT8,)), # Ampere sm_80 / sm_86
|
||||
)
|
||||
|
||||
# Cache of (scheme, device) -> bool so the quantise+matmul smoke test runs once.
|
||||
_SMOKE_CACHE: dict[tuple[str, str], bool] = {}
|
||||
|
||||
# Data-center GPU model tokens (un-nerfed FP32 accumulate). Matched as whole tokens of
|
||||
# torch.cuda.get_device_name(), so the workstation "A4000" is not mistaken for the
|
||||
# data-center "A40". Anything not here -- GeForce, workstation RTX, or an unknown name --
|
||||
# is treated as consumer-class (FP32-accumulate halved). See developer.nvidia.com/cuda/gpus.
|
||||
_DATACENTER_GPU_TOKENS = frozenset(
|
||||
{
|
||||
"B200",
|
||||
"B100",
|
||||
"GB200",
|
||||
"GB300",
|
||||
"GB10", # Blackwell data center
|
||||
"H200",
|
||||
"H100",
|
||||
"H800",
|
||||
"H20", # Hopper data center
|
||||
"A100",
|
||||
"A800",
|
||||
"A30",
|
||||
"A40",
|
||||
"A16",
|
||||
"A10",
|
||||
"A2", # Ampere data center
|
||||
"L40",
|
||||
"L40S",
|
||||
"L4",
|
||||
"L20",
|
||||
"L2", # Ada data center
|
||||
"V100",
|
||||
"P100",
|
||||
"P40",
|
||||
"T4", # legacy data center
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_consumer_gpu(device: Any = None) -> bool:
|
||||
"""Whether the active GPU is consumer / workstation class (GDDR), where fp8 FP32
|
||||
accumulate is throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center
|
||||
HBM parts (recognised by name token) are not nerfed and return False, so they keep
|
||||
the higher-precision default accumulate for free. Heuristic on the device name: a
|
||||
GeForce / TITAN name is always consumer; a recognised data-center token is not;
|
||||
anything else (workstation RTX, unknown) defaults to consumer -- the safe choice,
|
||||
since fast accumulate is free on data-center and a win on consumer. Best-effort:
|
||||
True on any probe failure."""
|
||||
try:
|
||||
import re
|
||||
|
||||
import torch
|
||||
name = torch.cuda.get_device_name(device).upper()
|
||||
except Exception: # noqa: BLE001 — no torch / no device -> assume consumer
|
||||
return True
|
||||
if "GEFORCE" in name or "TITAN" in name:
|
||||
return True
|
||||
tokens = set(re.split(r"[^A-Z0-9]+", name))
|
||||
return not (tokens & _DATACENTER_GPU_TOKENS)
|
||||
|
||||
|
||||
def normalize_transformer_quant(value: Optional[str]) -> Optional[str]:
|
||||
"""Lower/strip a requested transformer quant; None / "" / "none" / "off" -> None.
|
||||
|
||||
Raises ValueError for an unsupported value so a bad request is rejected cheaply."""
|
||||
if value is None:
|
||||
return None
|
||||
normalized = str(value).strip().lower().replace("-", "_")
|
||||
if not normalized or normalized in ("none", "off"):
|
||||
return None
|
||||
if normalized not in TQ_MODES:
|
||||
raise ValueError(
|
||||
f"Unsupported transformer_quant '{value}'. Use one of: {', '.join(TQ_MODES)}."
|
||||
)
|
||||
return normalized
|
||||
|
||||
|
||||
def dense_transformer_supported(target: Any) -> bool:
|
||||
"""Whether the dense-source quant path is usable for ``target``: a CUDA device with
|
||||
a bf16 compute dtype (the only configuration any torchao dynamic scheme accelerates).
|
||||
A cheap pre-check the loader runs before loading the (large) dense transformer."""
|
||||
if getattr(target, "device", None) != "cuda":
|
||||
return False
|
||||
try:
|
||||
import torch
|
||||
return getattr(target, "dtype", None) is torch.bfloat16
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Optional[str]:
|
||||
"""The concrete scheme to apply, or None to fall back to GGUF.
|
||||
|
||||
``auto`` walks the per-arch ladder and returns the first scheme that passes a real
|
||||
quantise+matmul smoke test, so on a box where the Blackwell fp4 / mx kernels are
|
||||
unavailable it lands on fp8 / int8 with no error. An explicit scheme is honored only
|
||||
if supported (else None -> GGUF), never silently swapped for a different one."""
|
||||
requested = normalize_transformer_quant(requested)
|
||||
if requested is None or not dense_transformer_supported(target):
|
||||
return None
|
||||
device = str(getattr(target, "device", "cuda"))
|
||||
if requested != TQ_AUTO:
|
||||
return requested if _scheme_supported(requested, device) else None
|
||||
cap = _capability()
|
||||
if cap is None:
|
||||
return None
|
||||
for floor, schemes in _AUTO_LADDER:
|
||||
if cap >= floor:
|
||||
for scheme in schemes:
|
||||
if _scheme_supported(scheme, device):
|
||||
return scheme
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _capability() -> Optional[tuple[int, int]]:
|
||||
try:
|
||||
import torch
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
return (int(major), int(minor))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _scheme_supported(scheme: str, device: str) -> bool:
|
||||
"""CUDA + (for fp8) the fp8 dtype + a cached quantise+matmul smoke test for ``scheme``."""
|
||||
try:
|
||||
import torch
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
if scheme == TQ_FP8 and not hasattr(torch, "float8_e4m3fn"):
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
return _smoke_probe(scheme, device)
|
||||
|
||||
|
||||
def _smoke_probe(scheme: str, device: str) -> bool:
|
||||
"""True iff a tiny Linear quantised with ``scheme`` runs one M=32 forward without
|
||||
error. Cached per (scheme, device). This is what makes ``auto`` robust to a torch /
|
||||
torchao build where a prototype (nvfp4 / mxfp8) kernel is unavailable: it fails here
|
||||
and the ladder moves on, rather than crashing at the first real denoise step."""
|
||||
key = (scheme, device)
|
||||
if key in _SMOKE_CACHE:
|
||||
return _SMOKE_CACHE[key]
|
||||
ok = False
|
||||
try:
|
||||
import torch
|
||||
from torchao.quantization import quantize_
|
||||
|
||||
lin = torch.nn.Linear(512, 512, bias = False).to(device = device, dtype = torch.bfloat16)
|
||||
quantize_(lin, _make_quant_config(scheme), filter_fn = make_filter_fn(0))
|
||||
x = torch.randn(32, 512, device = device, dtype = torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
lin(x)
|
||||
torch.cuda.synchronize()
|
||||
ok = True
|
||||
except Exception:
|
||||
ok = False
|
||||
_SMOKE_CACHE[key] = ok
|
||||
return ok
|
||||
|
||||
|
||||
def _resolve_fast_accum(fast_accum: Optional[bool]) -> bool:
|
||||
"""The fp8 ``use_fast_accum`` to apply. ``None`` auto-detects by GPU class
|
||||
(consumer / workstation -> fast; data-center -> precise); an explicit bool forces it."""
|
||||
return _is_consumer_gpu() if fast_accum is None else bool(fast_accum)
|
||||
|
||||
|
||||
def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
|
||||
"""The torchao dynamic-activation config for ``scheme`` (lazy import; prototype
|
||||
import for the Blackwell fp4 / mx schemes is inside the branch that needs it).
|
||||
|
||||
``fast_accum`` applies to fp8 only: None auto-detects by GPU class, True/False force it."""
|
||||
from torchao.quantization import (
|
||||
Float8DynamicActivationFloat8WeightConfig,
|
||||
Int8DynamicActivationInt8WeightConfig,
|
||||
)
|
||||
|
||||
if scheme == TQ_INT8:
|
||||
return Int8DynamicActivationInt8WeightConfig()
|
||||
if scheme == TQ_FP8:
|
||||
# Choose fp8 accumulate by GPU class (unless forced). On consumer / workstation
|
||||
# cards (GDDR) the fp8 tensor cores run ~2x faster with FP16 (fast) accumulate
|
||||
# than FP32 (e.g. ~838 vs ~419 TFLOPS on RTX 50xx), so fast accumulate is a real
|
||||
# win there. Data-center HBM parts default to the higher-precision accumulate.
|
||||
# fast accumulate is a precision (not overflow) tradeoff and stays below the fp8
|
||||
# quant noise floor (measured 0 non-finite even on Z-Image's ~1e6 activations).
|
||||
try:
|
||||
from torchao.float8 import Float8MMConfig
|
||||
return Float8DynamicActivationFloat8WeightConfig(
|
||||
mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum))
|
||||
)
|
||||
except Exception: # noqa: BLE001 — older torchao without the explicit knob
|
||||
return Float8DynamicActivationFloat8WeightConfig()
|
||||
if scheme == TQ_NVFP4:
|
||||
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig
|
||||
|
||||
# Select the CUTLASS FP4 path, not the default Triton kernel: torchao defaults
|
||||
# use_triton_kernel=True, which needs MSLK installed. On a Blackwell box with the
|
||||
# CUTLASS FP4 extension but no MSLK, the default would make the smoke probe fail
|
||||
# and silently fall back to GGUF instead of using the FP4 tensor cores.
|
||||
try:
|
||||
return NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel = False)
|
||||
except TypeError: # older torchao without the knob
|
||||
return NVFP4DynamicActivationNVFP4WeightConfig()
|
||||
if scheme == TQ_MXFP8:
|
||||
import torch
|
||||
from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
|
||||
try:
|
||||
return MXDynamicActivationMXWeightConfig(
|
||||
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
|
||||
)
|
||||
except (TypeError, AttributeError):
|
||||
# TypeError: older torchao without the explicit dtype knobs.
|
||||
# AttributeError: a torch build without torch.float8_e4m3fn.
|
||||
return MXDynamicActivationMXWeightConfig()
|
||||
raise ValueError(f"unknown transformer quant scheme '{scheme}'")
|
||||
|
||||
|
||||
def make_filter_fn(min_features: int):
|
||||
"""A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear
|
||||
with both in/out features >= ``min_features``. Hides the (module, fqn) callback arity."""
|
||||
|
||||
def filter_fn(module: Any, fqn: str = "") -> bool:
|
||||
try:
|
||||
import torch
|
||||
if not isinstance(module, torch.nn.Linear):
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
in_features = getattr(module, "in_features", None)
|
||||
out_features = getattr(module, "out_features", None)
|
||||
if in_features is None or out_features is None:
|
||||
return False
|
||||
return in_features >= min_features and out_features >= min_features
|
||||
|
||||
return filter_fn
|
||||
|
||||
|
||||
def quantize_transformer(
|
||||
pipe: Any,
|
||||
target: Any,
|
||||
*,
|
||||
mode: Optional[str],
|
||||
min_features: int = DEFAULT_MIN_LINEAR_FEATURES,
|
||||
fast_accum: Optional[bool] = None,
|
||||
logger: Any = None,
|
||||
) -> Optional[str]:
|
||||
"""Quantise ``pipe.transformer``'s FLOP-heavy linears in place with the arch-chosen
|
||||
dynamic scheme. Returns the scheme actually engaged, or None when disabled /
|
||||
unsupported / failed -- the caller then loads GGUF instead. Best-effort: it never
|
||||
raises for an ordinary unsupported environment (a failure leaves the module dense).
|
||||
|
||||
``fast_accum`` (fp8 only) overrides the per-GPU-class accumulate choice: None
|
||||
auto-detects (fast on consumer, precise on data-center), True/False force it."""
|
||||
scheme = select_transformer_quant_scheme(target, mode)
|
||||
if scheme is None:
|
||||
return None
|
||||
transformer = getattr(pipe, "transformer", None)
|
||||
if transformer is None:
|
||||
return None
|
||||
try:
|
||||
from torchao.quantization import quantize_
|
||||
|
||||
quantize_(
|
||||
transformer,
|
||||
_make_quant_config(scheme, fast_accum = fast_accum),
|
||||
filter_fn = make_filter_fn(min_features),
|
||||
)
|
||||
# Runtime-only marker (torchao tensors are not safetensors-serializable; this
|
||||
# backend is inference-only, so this is purely diagnostic).
|
||||
try:
|
||||
transformer._unsloth_runtime_quant = scheme
|
||||
except Exception: # noqa: BLE001 — marker is best-effort
|
||||
pass
|
||||
return scheme
|
||||
except Exception as exc: # noqa: BLE001 — leave the transformer dense -> GGUF fallback
|
||||
_warn(logger, scheme, exc)
|
||||
return None
|
||||
|
||||
|
||||
def _warn(logger: Any, what: str, exc: Exception) -> None:
|
||||
if logger is not None:
|
||||
logger.warning("diffusion.transformer_quant: %s failed: %s", what, exc)
|
||||
|
|
@ -1720,6 +1720,23 @@ class DiffusionLoadRequest(BaseModel):
|
|||
"memory-vs-quality tradeoff (shifts fine detail), not free; "
|
||||
"pairs well with balanced mode.",
|
||||
)
|
||||
transformer_quant: Optional[Literal["auto", "int8", "fp8", "nvfp4", "mxfp8"]] = Field(
|
||||
None,
|
||||
description = "Opt-in fast transformer: load the DENSE bf16 transformer instead "
|
||||
"of the GGUF and torchao-quantise it onto the low-precision tensor "
|
||||
"cores (faster than GGUF's bf16-rate dequant, at higher VRAM). auto "
|
||||
"picks the best for the GPU (Blackwell nvfp4/mxfp8, Ada/Hopper fp8, "
|
||||
"Ampere int8); an explicit scheme forces it. Needs CUDA + bf16 + room "
|
||||
"for the dense load; falls back to GGUF otherwise.",
|
||||
)
|
||||
transformer_quant_fast_accum: Optional[bool] = Field(
|
||||
None,
|
||||
description = "fp8 only: FP8 matmul accumulate. null auto-detects by GPU class "
|
||||
"(fast FP16 accumulate on consumer/workstation cards, where FP32 "
|
||||
"accumulate is ~2x slower; precise FP32 accumulate on data-center "
|
||||
"HBM cards, which are not nerfed). true/false force it. Negligible "
|
||||
"quality effect (below the fp8 quant noise floor); no overflow risk.",
|
||||
)
|
||||
|
||||
|
||||
class DiffusionGenerateRequest(BaseModel):
|
||||
|
|
@ -1830,3 +1847,8 @@ class DiffusionStatusResponse(BaseModel):
|
|||
text_encoder_quant: Optional[str] = Field(
|
||||
None, description = "Text-encoder quantisation engaged: fp8 | nvfp4 | null"
|
||||
)
|
||||
transformer_quant: Optional[str] = Field(
|
||||
None,
|
||||
description = "Transformer quant engaged on the dense fast path: int8 | fp8 | "
|
||||
"nvfp4 | mxfp8 | null (null = the GGUF transformer was loaded)",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10331,6 +10331,8 @@ async def load_diffusion_model(
|
|||
memory_mode = request.memory_mode,
|
||||
speed_mode = request.speed_mode,
|
||||
text_encoder_quant = request.text_encoder_quant,
|
||||
transformer_quant = request.transformer_quant,
|
||||
transformer_quant_fast_accum = request.transformer_quant_fast_accum,
|
||||
)
|
||||
return DiffusionStatusResponse(**status_dict)
|
||||
except (ValueError, FileNotFoundError) as exc:
|
||||
|
|
|
|||
|
|
@ -899,3 +899,131 @@ def test_load_fast_mode_stays_resident_on_cuda(fake_runtime, tmp_path, monkeypat
|
|||
)
|
||||
assert status["offload_policy"] == "none" and status["cpu_offload"] is False
|
||||
assert backend._state.pipe.moved_to == "cuda"
|
||||
|
||||
|
||||
# ── transformer quant (opt-in dense fast path) ────────────────────────────────
|
||||
|
||||
|
||||
def _stub_dense_quant(monkeypatch, *, scheme = "fp8"):
|
||||
"""Force the dense+quant branch hermetically: a supported dense source, a
|
||||
from_pretrained on the fake transformer, and a quantizer that engages `scheme`.
|
||||
Returns a dict recording the dense-loader / quantizer calls."""
|
||||
from core.inference import diffusion as dmod
|
||||
|
||||
calls: dict = {"from_pretrained": 0, "quantize": 0, "quant_mode": None}
|
||||
|
||||
@classmethod
|
||||
def _from_pretrained(cls, base, **kwargs):
|
||||
calls["from_pretrained"] += 1
|
||||
calls["fp_kwargs"] = {"base": base, **kwargs}
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False)
|
||||
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
|
||||
|
||||
def _quantize(pipe, target, *, mode, **kw):
|
||||
calls["quantize"] += 1
|
||||
calls["quant_mode"] = mode
|
||||
return scheme
|
||||
|
||||
monkeypatch.setattr(dmod, "quantize_transformer", _quantize)
|
||||
return calls
|
||||
|
||||
|
||||
def test_default_load_skips_dense_quant_path(fake_runtime, tmp_path, monkeypatch):
|
||||
# With no transformer_quant flag the GGUF path is taken and the dense gate is
|
||||
# never even consulted (short-circuit), so the default cannot regress.
|
||||
from core.inference import diffusion as dmod
|
||||
|
||||
monkeypatch.setattr(
|
||||
dmod,
|
||||
"dense_transformer_supported",
|
||||
lambda *a, **k: pytest.fail("dense path must not run without the flag"),
|
||||
)
|
||||
(tmp_path / "m.gguf").write_bytes(b"x")
|
||||
backend = DiffusionBackend()
|
||||
status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image")
|
||||
assert status["transformer_quant"] is None
|
||||
assert _FakeTransformer.last["path"] # GGUF from_single_file was used
|
||||
|
||||
|
||||
def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatch):
|
||||
# transformer_quant + a CUDA resident plan -> load the DENSE transformer from the
|
||||
# base repo, place it on the device, quantise it, and report the engaged scheme.
|
||||
backend = DiffusionBackend()
|
||||
_force_cuda_target(backend, monkeypatch)
|
||||
calls = _stub_dense_quant(monkeypatch, scheme = "fp8")
|
||||
(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"
|
||||
# No speed_mode was given, but a quantized transformer is ~30x slower eager, so the
|
||||
# backend promotes it to `default` (regional compile) instead of the dense `off`.
|
||||
assert status["speed_mode"] == "default"
|
||||
assert calls["from_pretrained"] == 1 and calls["quantize"] == 1
|
||||
assert calls["quant_mode"] == "fp8"
|
||||
assert calls["fp_kwargs"]["subfolder"] == "transformer" # dense transformer subfolder
|
||||
# The GGUF single-file path was NOT used for the transformer.
|
||||
assert _FakeTransformer.last == {}
|
||||
# quantize ran on-device: the dense pipe was placed on cuda (before compile).
|
||||
assert backend._state.pipe.moved_to == "cuda"
|
||||
assert status["offload_policy"] == "none"
|
||||
|
||||
|
||||
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.
|
||||
from core.inference import diffusion as dmod
|
||||
|
||||
backend = DiffusionBackend()
|
||||
_force_cuda_target(backend, monkeypatch)
|
||||
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
|
||||
|
||||
@classmethod
|
||||
def _from_pretrained(cls, base, **kwargs):
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False)
|
||||
monkeypatch.setattr(dmod, "quantize_transformer", lambda pipe, target, **kw: 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["loaded"] is True
|
||||
assert status["transformer_quant"] is None # fell back
|
||||
assert _FakeTransformer.last["path"] # GGUF from_single_file used
|
||||
|
||||
|
||||
def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, monkeypatch):
|
||||
# The dense bf16 transformer only fits resident, so when the memory plan would
|
||||
# offload (here low_vram) the fast path is skipped and GGUF loads instead -- the
|
||||
# dense transformer is never even loaded.
|
||||
from core.inference import diffusion as dmod
|
||||
|
||||
backend = DiffusionBackend()
|
||||
_force_cuda_target(backend, monkeypatch)
|
||||
monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True)
|
||||
|
||||
@classmethod
|
||||
def _fp_fail(cls, *a, **k):
|
||||
pytest.fail("dense transformer must not load when the plan offloads")
|
||||
|
||||
monkeypatch.setattr(_FakeTransformer, "from_pretrained", _fp_fail, raising = False)
|
||||
(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",
|
||||
memory_mode = "low_vram",
|
||||
)
|
||||
assert status["transformer_quant"] is None
|
||||
assert status["offload_policy"] == "model"
|
||||
assert _FakeTransformer.last["path"] # GGUF path used
|
||||
|
|
|
|||
|
|
@ -381,6 +381,44 @@ def test_memory_mode_threads_through_to_backend(client, monkeypatch):
|
|||
assert backend.last_load_kwargs.get("memory_mode") == "low_vram"
|
||||
|
||||
|
||||
def test_transformer_quant_threads_through_to_backend(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": "auto"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert backend.last_load_kwargs.get("transformer_quant") == "auto"
|
||||
|
||||
|
||||
def test_transformer_quant_fast_accum_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_quant_fast_accum": False,
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert backend.last_load_kwargs.get("transformer_quant_fast_accum") is False
|
||||
|
||||
|
||||
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.
|
||||
resp = client.post(
|
||||
"/api/inference/images/load",
|
||||
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_quant": "int2"},
|
||||
)
|
||||
assert resp.status_code == 422
|
||||
assert gpu_arbiter._owner is None
|
||||
|
||||
|
||||
def test_invalid_memory_mode_returns_422_without_eviction(client):
|
||||
# An unsupported memory_mode is rejected by the request schema (Literal), so the
|
||||
# GPU is never acquired and no chat model is evicted.
|
||||
|
|
|
|||
342
studio/backend/tests/test_diffusion_transformer_quant.py
Normal file
342
studio/backend/tests/test_diffusion_transformer_quant.py
Normal file
|
|
@ -0,0 +1,342 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit tests for transformer quantisation (``diffusion_transformer_quant.py``).
|
||||
|
||||
Hermetic: torch + torchao are stubbed via ``sys.modules``, and the per-scheme smoke
|
||||
probe (``_scheme_supported`` / ``_smoke_probe``) is monkeypatched where the test cares
|
||||
about the selection ladder rather than the GPU probe, so everything runs CPU-only.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
import core.inference.diffusion_transformer_quant as tq
|
||||
from core.inference.diffusion_transformer_quant import (
|
||||
TQ_FP8,
|
||||
TQ_INT8,
|
||||
TQ_MXFP8,
|
||||
TQ_NVFP4,
|
||||
dense_transformer_supported,
|
||||
make_filter_fn,
|
||||
normalize_transformer_quant,
|
||||
quantize_transformer,
|
||||
select_transformer_quant_scheme,
|
||||
)
|
||||
|
||||
|
||||
def _target(*, device = "cuda", dtype = "bfloat16"):
|
||||
return types.SimpleNamespace(device = device, dtype = dtype)
|
||||
|
||||
|
||||
def _stub_torch(
|
||||
monkeypatch,
|
||||
*,
|
||||
cc = (10, 0),
|
||||
with_fp8 = True,
|
||||
cuda_available = True,
|
||||
):
|
||||
torch = types.ModuleType("torch")
|
||||
torch.bfloat16 = "bfloat16"
|
||||
torch.float16 = "float16"
|
||||
if with_fp8:
|
||||
torch.float8_e4m3fn = "float8_e4m3fn"
|
||||
torch.cuda = types.SimpleNamespace(
|
||||
is_available = lambda: cuda_available,
|
||||
get_device_capability = lambda *a: cc,
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
return torch
|
||||
|
||||
|
||||
# ── normalisation ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_normalize_transformer_quant():
|
||||
assert normalize_transformer_quant(None) is None
|
||||
assert normalize_transformer_quant("") is None
|
||||
assert normalize_transformer_quant("none") is None
|
||||
assert normalize_transformer_quant("off") is None
|
||||
assert normalize_transformer_quant("AUTO") == "auto"
|
||||
assert normalize_transformer_quant("INT8") == TQ_INT8
|
||||
assert normalize_transformer_quant("fp8") == TQ_FP8
|
||||
with pytest.raises(ValueError):
|
||||
normalize_transformer_quant("int2")
|
||||
|
||||
|
||||
# ── dense-source gate ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_dense_transformer_supported_requires_cuda_bf16(monkeypatch):
|
||||
_stub_torch(monkeypatch)
|
||||
assert dense_transformer_supported(_target()) is True
|
||||
assert dense_transformer_supported(_target(device = "cpu")) is False
|
||||
assert dense_transformer_supported(_target(dtype = "float16")) is False
|
||||
|
||||
|
||||
# ── scheme selection ladder ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _allow(monkeypatch, allowed):
|
||||
"""Force ``_scheme_supported`` to accept only ``allowed`` (simulates smoke results)."""
|
||||
monkeypatch.setattr(tq, "_scheme_supported", lambda scheme, device: scheme in allowed)
|
||||
|
||||
|
||||
def test_auto_blackwell_prefers_fp8_then_falls_back(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
# Even with every scheme available, auto picks fp8 on Blackwell: measured on a B200
|
||||
# (torch 2.11 + torchao CUTLASS FP4), fp8 is both faster and more accurate than nvfp4
|
||||
# for the DiT's shapes -- nvfp4's FP4 GEMM only wins on very large GEMMs, not here.
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
||||
# fp8 unavailable: nvfp4 is the next pick (above mxfp8 / int8).
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_NVFP4
|
||||
# Only mxfp8 + int8 left -> mxfp8 (still above int8).
|
||||
_allow(monkeypatch, {TQ_MXFP8, TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_MXFP8
|
||||
# Only int8 usable -> int8.
|
||||
_allow(monkeypatch, {TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
||||
|
||||
|
||||
def test_auto_ada_hopper_prefers_fp8(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (8, 9))
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
||||
_stub_torch(monkeypatch, cc = (9, 0)) # Hopper
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
||||
|
||||
|
||||
def test_auto_ampere_prefers_int8(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (8, 0))
|
||||
_allow(monkeypatch, {TQ_FP8, TQ_INT8}) # fp8 cores absent on Ampere -> int8 only in ladder
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
||||
_stub_torch(monkeypatch, cc = (8, 6))
|
||||
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
||||
|
||||
|
||||
def test_auto_pre_ampere_unsupported(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (7, 5)) # Turing: below the int8-dynamic floor
|
||||
_allow(monkeypatch, {TQ_INT8, TQ_FP8})
|
||||
assert select_transformer_quant_scheme(_target(), "auto") is None
|
||||
|
||||
|
||||
def test_explicit_scheme_honored_or_none(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (8, 0))
|
||||
_allow(monkeypatch, {TQ_INT8})
|
||||
assert select_transformer_quant_scheme(_target(), "int8") == TQ_INT8
|
||||
# Explicit unsupported scheme is NOT silently downgraded -> None (-> GGUF fallback).
|
||||
assert select_transformer_quant_scheme(_target(), "fp8") is None
|
||||
assert select_transformer_quant_scheme(_target(), "nvfp4") is None
|
||||
|
||||
|
||||
def test_select_none_when_disabled_or_non_cuda(monkeypatch):
|
||||
_stub_torch(monkeypatch)
|
||||
_allow(monkeypatch, {TQ_INT8, TQ_FP8, TQ_NVFP4})
|
||||
assert select_transformer_quant_scheme(_target(), None) is None
|
||||
assert select_transformer_quant_scheme(_target(device = "cpu"), "auto") is None
|
||||
|
||||
|
||||
# ── _scheme_supported / _smoke_probe ────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_scheme_supported_shortcircuits(monkeypatch):
|
||||
# No CUDA -> False without running the smoke probe.
|
||||
_stub_torch(monkeypatch, cuda_available = False)
|
||||
monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run"))
|
||||
assert tq._scheme_supported(TQ_INT8, "cuda") is False
|
||||
# fp8 requested but the fp8 dtype is missing -> False before the probe.
|
||||
_stub_torch(monkeypatch, with_fp8 = False)
|
||||
monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run"))
|
||||
assert tq._scheme_supported(TQ_FP8, "cuda") is False
|
||||
|
||||
|
||||
def test_smoke_probe_caches_and_tolerates_failure(monkeypatch):
|
||||
tq._SMOKE_CACHE.clear()
|
||||
calls = {"n": 0}
|
||||
|
||||
class _Lin:
|
||||
def __init__(self, *a, **k):
|
||||
pass
|
||||
|
||||
def to(self, **k):
|
||||
return self
|
||||
|
||||
torch = types.ModuleType("torch")
|
||||
torch.bfloat16 = "bfloat16"
|
||||
torch.nn = types.SimpleNamespace(Linear = _Lin)
|
||||
torch.randn = lambda *a, **k: object()
|
||||
torch.no_grad = lambda: __import__("contextlib").nullcontext()
|
||||
torch.cuda = types.SimpleNamespace(is_available = lambda: True, synchronize = lambda: None)
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
|
||||
tqz = types.ModuleType("torchao.quantization")
|
||||
|
||||
def _quantize_ok(
|
||||
module,
|
||||
config,
|
||||
filter_fn = None,
|
||||
):
|
||||
calls["n"] += 1
|
||||
|
||||
tqz.quantize_ = _quantize_ok
|
||||
tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8cfg"
|
||||
tqz.Float8DynamicActivationFloat8WeightConfig = lambda: "fp8cfg"
|
||||
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
||||
# _Lin is callable? No -> the forward lin(x) would fail. Make instances callable.
|
||||
_Lin.__call__ = lambda self, x: x
|
||||
|
||||
assert tq._smoke_probe(TQ_INT8, "cuda") is True
|
||||
assert tq._smoke_probe(TQ_INT8, "cuda") is True # cached, no second quantize_
|
||||
assert calls["n"] == 1
|
||||
|
||||
# A scheme whose quantize_ raises -> probe False (and cached).
|
||||
tq._SMOKE_CACHE.clear()
|
||||
|
||||
def _quantize_boom(
|
||||
module,
|
||||
config,
|
||||
filter_fn = None,
|
||||
):
|
||||
raise RuntimeError("kernel unavailable")
|
||||
|
||||
tqz.quantize_ = _quantize_boom
|
||||
assert tq._smoke_probe(TQ_FP8, "cuda") is False
|
||||
|
||||
|
||||
# ── consumer-vs-datacenter detection (fp8 fast-accumulate gate) ──────────────────
|
||||
|
||||
|
||||
def _stub_device_name(monkeypatch, name):
|
||||
torch = types.ModuleType("torch")
|
||||
torch.cuda = types.SimpleNamespace(get_device_name = lambda device = None: name)
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name",
|
||||
[
|
||||
"NVIDIA GeForce RTX 5090",
|
||||
"NVIDIA GeForce RTX 4090",
|
||||
"NVIDIA RTX A4000", # workstation: A4000 token, NOT the data-center A40
|
||||
"NVIDIA RTX 6000 Ada Generation",
|
||||
"NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC)
|
||||
],
|
||||
)
|
||||
def test_is_consumer_gpu_true(monkeypatch, name):
|
||||
_stub_device_name(monkeypatch, name)
|
||||
assert tq._is_consumer_gpu() is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name",
|
||||
[
|
||||
"NVIDIA B200",
|
||||
"NVIDIA H100 80GB HBM3",
|
||||
"NVIDIA A100-SXM4-80GB",
|
||||
"NVIDIA A40", # data-center Ampere (distinct token from RTX A4000)
|
||||
"NVIDIA L40S",
|
||||
"NVIDIA L4",
|
||||
"Tesla V100-SXM2-16GB",
|
||||
],
|
||||
)
|
||||
def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name):
|
||||
_stub_device_name(monkeypatch, name)
|
||||
assert tq._is_consumer_gpu() is False
|
||||
|
||||
|
||||
def test_is_consumer_gpu_defaults_true_on_probe_failure(monkeypatch):
|
||||
# No torch / no device name available -> assume consumer (safe: fast accum is free
|
||||
# on data center and a win on consumer).
|
||||
torch = types.ModuleType("torch")
|
||||
torch.cuda = types.SimpleNamespace() # no get_device_name
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
assert tq._is_consumer_gpu() is True
|
||||
|
||||
|
||||
# ── filter ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_make_filter_fn(monkeypatch):
|
||||
class _Lin:
|
||||
def __init__(self, i, o):
|
||||
self.in_features, self.out_features = i, o
|
||||
|
||||
torch = types.ModuleType("torch")
|
||||
torch.nn = types.SimpleNamespace(Linear = _Lin)
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
|
||||
keep = make_filter_fn(512)
|
||||
assert keep(_Lin(1024, 4096), "blocks.0.attn.to_q") is True
|
||||
assert keep(_Lin(256, 4096), "time_proj") is False # small in_features -> skip
|
||||
assert keep(_Lin(4096, 256), "out_proj") is False # small out_features -> skip
|
||||
assert keep(object(), "not_linear") is False # non-Linear -> skip
|
||||
assert keep(types.SimpleNamespace(), "no_attrs") is False
|
||||
|
||||
|
||||
# ── apply ───────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_resolve_fast_accum(monkeypatch):
|
||||
# None auto-detects by GPU class; an explicit bool forces it.
|
||||
monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: True)
|
||||
assert tq._resolve_fast_accum(None) is True
|
||||
monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: False)
|
||||
assert tq._resolve_fast_accum(None) is False
|
||||
assert tq._resolve_fast_accum(True) is True # forced on (e.g. on a data-center card)
|
||||
assert tq._resolve_fast_accum(False) is False # forced off (e.g. on a consumer card)
|
||||
|
||||
|
||||
def test_quantize_transformer_applies_and_marks(monkeypatch):
|
||||
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8)
|
||||
seen: dict = {}
|
||||
|
||||
def _mk(scheme, fast_accum = None):
|
||||
seen["scheme"], seen["fast_accum"] = scheme, fast_accum
|
||||
return f"{scheme}cfg"
|
||||
|
||||
monkeypatch.setattr(tq, "_make_quant_config", _mk)
|
||||
recorder: list = []
|
||||
tqz = types.ModuleType("torchao.quantization")
|
||||
tqz.quantize_ = lambda module, config, filter_fn = None: recorder.append(
|
||||
(module, config, filter_fn)
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
||||
|
||||
transformer = types.SimpleNamespace()
|
||||
pipe = types.SimpleNamespace(transformer = transformer)
|
||||
assert quantize_transformer(pipe, _target(), mode = "fp8", fast_accum = False) == TQ_FP8
|
||||
assert len(recorder) == 1 and recorder[0][0] is transformer and recorder[0][1] == "fp8cfg"
|
||||
assert callable(recorder[0][2]) # a filter_fn was passed
|
||||
assert transformer._unsloth_runtime_quant == TQ_FP8 # diagnostic marker set
|
||||
assert seen["fast_accum"] is False # the override is forwarded into the config
|
||||
|
||||
|
||||
def test_quantize_transformer_none_when_unsupported(monkeypatch):
|
||||
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: None)
|
||||
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
||||
assert quantize_transformer(pipe, _target(), mode = "auto") is None
|
||||
|
||||
|
||||
def test_quantize_transformer_tolerates_failure(monkeypatch):
|
||||
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_INT8)
|
||||
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: "cfg")
|
||||
tqz = types.ModuleType("torchao.quantization")
|
||||
|
||||
def _boom(
|
||||
module,
|
||||
config,
|
||||
filter_fn = None,
|
||||
):
|
||||
raise RuntimeError("partial quant failure")
|
||||
|
||||
tqz.quantize_ = _boom
|
||||
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
||||
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
||||
# A quantise failure returns None (caller falls back to GGUF), never raises.
|
||||
assert quantize_transformer(pipe, _target(), mode = "int8") is None
|
||||
Loading…
Add table
Add a link
Reference in a new issue