diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index 18dd064650..98fab2b49e 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, + int8_exclude_name_tokens, make_filter_fn, ) from torchao.quantization import quantize_ @@ -78,7 +79,17 @@ 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)) + # 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). + quantize_( + transformer, + _make_quant_config(scheme), + filter_fn = make_filter_fn( + args.min_features, exclude_name_tokens = int8_exclude_name_tokens(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 fcaf66f6d3..570419aefa 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -62,6 +62,15 @@ _INT8_EXCLUDE_NAME_TOKENS = ( "pooled", ) + +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).""" + 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 +368,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 = int8_exclude_name_tokens(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 8f6ef6594d..6a6becd09f 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -357,6 +357,20 @@ 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(): + # 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, + ) + + assert int8_exclude_name_tokens(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS + for scheme in (TQ_FP8, TQ_NVFP4, TQ_MXFP8): + assert int8_exclude_name_tokens(scheme) == () + + # ── apply ───────────────────────────────────────────────────────────────────────