From d127be20a50323c0bf7d343c3ba4c1913b0884cb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 11:03:37 +0000 Subject: [PATCH] Studio diffusion (Phase 14) review round 2: align helper name with the stack Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the identical helper already present higher in the diffusion stack (Phase 16). The helper definition, the runtime quantiser call, and the offline builder are now byte-identical to that version, so the two branches no longer introduce a divergent name for the same single-source-of-truth and the stack merges without a conflict on this fix. No behavior change. --- scripts/build_prequant_checkpoint.py | 12 ++++++------ .../core/inference/diffusion_transformer_quant.py | 15 ++++++++------- .../tests/test_diffusion_transformer_quant.py | 9 +++++---- 3 files changed, 19 insertions(+), 17 deletions(-) diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index 98fab2b49e..366822de94 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -57,7 +57,7 @@ def main(argv = None) -> int: from core.inference.diffusion_transformer_quant import ( TQ_SCHEMES, _make_quant_config, - int8_exclude_name_tokens, + exclude_tokens_for_scheme, make_filter_fn, ) from torchao.quantization import quantize_ @@ -79,15 +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) - # Use the SAME int8 M=1 exclusion as the runtime quantiser (single source of truth): - # otherwise an int8 prequant checkpoint quantises the AdaLN modulation / conditioning - # embedders and reintroduces the torch._int_mm M=1 crash when loaded via - # transformer_prequant_path. fp8/fp4/mx get an empty exclusion (artifacts unchanged). + # 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 = int8_exclude_name_tokens(scheme) + args.min_features, exclude_name_tokens = exclude_tokens_for_scheme(scheme) ), ) diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 8a9de50cdc..3431bbee27 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -63,12 +63,13 @@ _INT8_EXCLUDE_NAME_TOKENS = ( ) -def int8_exclude_name_tokens(scheme: str) -> tuple[str, ...]: - """Filter exclusions for ``scheme``: the M=1 AdaLN-modulation / conditioning-embedder - linears for int8 (they crash ``torch._int_mm``, which needs M > 16), empty for every - other scheme (fp8 / fp4 / mx use ``scaled_mm``, no M limit). The single source of truth - shared by the runtime quantiser and the offline prequant builder, so a prequant artifact's - quantised-layer set matches the runtime exactly (no reintroduced M=1 crash).""" +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 () @@ -369,7 +370,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(scheme) + 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 a019fe8744..03cea6c53a 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -357,17 +357,18 @@ def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch): assert keep(big(), "") is True -def test_int8_exclude_name_tokens_shared_by_runtime_and_builder(): +def test_exclude_tokens_for_scheme_shared_by_runtime_and_builder(): # The runtime quantiser and the offline prequant builder must apply the SAME int8 # exclusion, or an int8 prequant artifact quantises the M=1 modulation/embedder linears # and reintroduces the torch._int_mm crash. int8 gets the exclusion; others get none. from core.inference.diffusion_transformer_quant import ( _INT8_EXCLUDE_NAME_TOKENS, - int8_exclude_name_tokens, + exclude_tokens_for_scheme, ) - assert int8_exclude_name_tokens(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS + + assert exclude_tokens_for_scheme(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS for scheme in (TQ_FP8, TQ_NVFP4, TQ_MXFP8): - assert int8_exclude_name_tokens(scheme) == () + assert exclude_tokens_for_scheme(scheme) == () # ── apply ───────────────────────────────────────────────────────────────────────