From 5ef000f912294dbf32236bb469210b009ed2a389 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 27 Jun 2026 07:32:18 +0000 Subject: [PATCH] Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears) The prequant-checkpoint builder applied the dense quant filter without the int8-only M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8 checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now produces a working checkpoint on every supported model, giving int8 (the consumer-preferred scheme) the same ~2x load-VRAM and download reduction fp8 already had. --- scripts/build_prequant_checkpoint.py | 11 ++++++++++- .../inference/diffusion_transformer_quant.py | 12 +++++++++++- .../tests/test_diffusion_transformer_quant.py | 16 ++++++++++++++++ 3 files changed, 37 insertions(+), 2 deletions(-) diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index 18dd064650..8312e02130 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -57,6 +57,7 @@ def main(argv = None) -> int: from core.inference.diffusion_transformer_quant import ( TQ_SCHEMES, _make_quant_config, + exclude_tokens_for_scheme, make_filter_fn, ) from torchao.quantization import quantize_ @@ -78,7 +79,15 @@ def main(argv = None) -> int: args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token ).to("cuda") print(f" quantising in place ({scheme}) ...", flush = True) - quantize_(transformer, _make_quant_config(scheme), filter_fn = make_filter_fn(args.min_features)) + # Mirror the runtime path EXACTLY (the offline == runtime, LPIPS-0 invariant): for int8 also + # skip the M=1 AdaLN-modulation / conditioning-embedder projections, else the saved checkpoint + # bakes them as int8 and crashes (torch._int_mm needs M>16) at the first denoise step on + # Flux / Qwen. fp8 / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns (). + quantize_( + transformer, + _make_quant_config(scheme), + filter_fn = make_filter_fn(args.min_features, exclude_name_tokens = exclude_tokens_for_scheme(scheme)), + ) # Move the state dict to CPU for a portable, GPU-free artifact. state_dict = { diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 173d0aaa86..e9bd1c6ff1 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -62,6 +62,16 @@ _INT8_EXCLUDE_NAME_TOKENS = ( "pooled", ) + +def exclude_tokens_for_scheme(scheme: str) -> tuple[str, ...]: + """Name tokens to exclude from quantisation for ``scheme``. int8 (torch._int_mm, M>16) + skips the M=1 modulation / conditioning-embedder projections (see _INT8_EXCLUDE_NAME_TOKENS); + every other scheme uses scaled_mm (no M limit) and excludes nothing. Shared by the runtime + quantise path and the offline prequant-checkpoint builder so the two never drift -- an int8 + checkpoint built offline must skip exactly the layers the runtime path skips, or it bakes the + M=1 projections as int8 and crashes at the first denoise step on Flux / Qwen.""" + return _INT8_EXCLUDE_NAME_TOKENS if scheme == TQ_INT8 else () + # 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 @@ -359,7 +369,7 @@ def quantize_transformer( # int8 (torch._int_mm, M>16) additionally skips the M=1 modulation / conditioning-embedder # projections; fp8 / fp4 / mx (scaled_mm) have no such limit and quantise everything. - exclude = _INT8_EXCLUDE_NAME_TOKENS if scheme == TQ_INT8 else () + exclude = exclude_tokens_for_scheme(scheme) quantize_( transformer, _make_quant_config(scheme, fast_accum = fast_accum), diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index 9ee0af7649..87111324c9 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -353,6 +353,22 @@ def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch): assert make_filter_fn(512)(big(), "transformer_blocks.0.norm1.linear") is True +def test_exclude_tokens_for_scheme(): + # The shared scheme->exclusion decision used by BOTH the runtime quantise path and the offline + # prequant-checkpoint builder, so an int8 checkpoint built ahead of time skips exactly the + # layers the runtime path skips (offline == runtime). int8 excludes the M=1 modulation / + # embedder tokens; every scaled_mm scheme excludes nothing. + from core.inference.diffusion_transformer_quant import ( + _INT8_EXCLUDE_NAME_TOKENS, + exclude_tokens_for_scheme, + ) + + assert exclude_tokens_for_scheme(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS + assert exclude_tokens_for_scheme(TQ_FP8) == () + assert exclude_tokens_for_scheme(TQ_NVFP4) == () + assert exclude_tokens_for_scheme(TQ_MXFP8) == () + + # ── apply ───────────────────────────────────────────────────────────────────────