diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index 18dd064650..366822de94 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,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)) + # 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/scripts/int8_linear_probe.py b/scripts/int8_linear_probe.py new file mode 100644 index 0000000000..5cef1b68e7 --- /dev/null +++ b/scripts/int8_linear_probe.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Examine which Linear layers the int8 dense-quant filter would select, to find the large M=1 +modulation/embedder projections that crash torch._int_mm (M>16). Loads each transformer on the +META device (no weights, no GPU) from its base-repo config, lists nn.Linear fqn/in/out, and marks +those that pass min_features=512. CPU-only, fast.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + +# (label, transformer_class, base_repo) +MODELS = [ + ("flux.1-dev", "FluxTransformer2DModel", "black-forest-labs/FLUX.1-dev"), + ("qwen-image", "QwenImageTransformer2DModel", "Qwen/Qwen-Image"), + ("z-image", "ZImageTransformer2DModel", "Tongyi-MAI/Z-Image-Turbo"), + ("flux.2-klein-4b", "Flux2Transformer2DModel", "black-forest-labs/FLUX.2-klein-4B"), +] +MIN = 512 + + +def main() -> int: + import diffusers + import torch + from accelerate import init_empty_weights + + tok = os.environ.get("HF_TOKEN") + for label, cls_name, base in MODELS: + cls = getattr(diffusers, cls_name, None) + if cls is None: + print(f"\n### {label}: {cls_name} NOT in diffusers") + continue + try: + cfg = cls.load_config(base, subfolder = "transformer", token = tok) + with init_empty_weights(): + model = cls.from_config(cfg) + except Exception as e: # noqa: BLE001 + print(f"\n### {label}: load failed {type(e).__name__}: {e}") + continue + lins = [(n, m) for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)] + selected = [(n, m) for n, m in lins if m.in_features >= MIN and m.out_features >= MIN] + print(f"\n### {label}: {len(lins)} Linear, {len(selected)} pass min_features={MIN}") + # Heuristic: a modulation/embedder Linear is one OUTSIDE the repeated transformer blocks, + # i.e. its fqn does not contain a numeric block index, OR out==k*in (k>=3) AdaLN shape. + sus = [] + for n, m in selected: + depth_idx = any(p.isdigit() for p in n.split(".")) + ratio = m.out_features / m.in_features if m.in_features else 0 + tag = [] + if not depth_idx: + tag.append("NO-BLOCK-IDX") + if ratio >= 3: + tag.append(f"out={ratio:.0f}xin") + if any( + t in n.lower() + for t in ("norm", "embed", "time", "guidance", "modulation", "adaln", "cond") + ): + tag.append("NAME") + if tag: + sus.append((n, m.in_features, m.out_features, ",".join(tag))) + # Print the distinct fqn shapes (collapse block indices to {i}) + import re + + seen = {} + for n, i, o, tag in sus: + key = re.sub(r"\.\d+\.", ".{i}.", n) + seen.setdefault((key, i, o, tag), 0) + seen[(key, i, o, tag)] += 1 + print(f" SUSPECT (M=1 risk) distinct patterns:") + for (key, i, o, tag), cnt in sorted(seen.items()): + print(f" [{cnt:>3}x] {key:55s} {i:>6}->{o:<6} [{tag}]") + # Also show a few non-suspect selected names for contrast (the real FLOP linears) + good = [ + n + for n, m in selected + if (n, m.in_features, m.out_features) not in {(s[0], s[1], s[2]) for s in sus} + ][:6] + print(f" kept-for-int8 examples: {good}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index cb6873a661..9e2e56a2a5 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -37,13 +37,42 @@ 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. +# Skip linears whose in/out features are below this. A small share of the FLOPs, so leaving +# them bf16 costs ~nothing and keeps quality a touch higher. DEFAULT_MIN_LINEAR_FEATURES = 512 +# int8-ONLY name exclusions. The int8 dynamic path uses torch._int_mm, which requires the +# activation row count M > 16. A DiT's AdaLN *modulation* projections and its timestep / +# guidance / pooled-text *conditioning embedders* are computed once from the [batch, dim] +# conditioning vector (M = batch = 1), not per token -- so they crash _int_mm even though +# their feature dims are large (e.g. Flux's norm1.linear 3072->18432, Qwen's img_mod.1, Flux.2's +# *_modulation.linear). min_features does NOT catch them (they are big), so int8 also skips any +# Linear whose fqn matches one of these tokens. They are a negligible share of the FLOPs (run at +# M=1, once per block), so int8 keeps the full speedup on the attention/FFN layers (M = seq). +# fp8 / nvfp4 / mxfp8 use scaled_mm (no M>16 limit) and quantise these layers fine, so the +# exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq) +# are deliberately NOT excluded. +_INT8_EXCLUDE_NAME_TOKENS = ( + "norm", # AdaLN modulation .linear (norm1 / norm1_context / norm / norm_out) + "_mod", # Qwen img_mod / txt_mod + "modulation", # Flux.2 double/single_stream_modulation + "timestep_embed", + "guidance_embed", + "time_text_embed", # Flux/Qwen time_text_embed.* (pooled-text + timestep); NOT context_embedder + "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 @@ -307,9 +336,11 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: 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 make_filter_fn(min_features: int, exclude_name_tokens: tuple[str, ...] = ()): + """A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear with both + in/out features >= ``min_features`` AND whose fully-qualified name contains none of + ``exclude_name_tokens`` (used by int8 to skip the M=1 modulation / conditioning-embedder + projections that crash ``torch._int_mm``). Hides the (module, fqn) callback arity.""" def filter_fn(module: Any, fqn: str = "") -> bool: try: @@ -322,7 +353,13 @@ def make_filter_fn(min_features: int): 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 + if in_features < min_features or out_features < min_features: + return False + if exclude_name_tokens: + name = fqn.lower() if fqn else "" + if any(tok in name for tok in exclude_name_tokens): + return False + return True return filter_fn @@ -352,10 +389,13 @@ def quantize_transformer( try: from torchao.quantization import quantize_ + # 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 = exclude_tokens_for_scheme(scheme) quantize_( transformer, _make_quant_config(scheme, fast_accum = fast_accum), - filter_fn = make_filter_fn(min_features), + filter_fn = make_filter_fn(min_features, exclude_name_tokens = exclude), ) # Runtime-only marker (torchao tensors are not safetensors-serializable; this # backend is inference-only, so this is purely diagnostic). diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index d527723b2e..a418367e98 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -325,6 +325,67 @@ def test_make_filter_fn(monkeypatch): assert keep(types.SimpleNamespace(), "no_attrs") is False +def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch): + # The int8 path skips the large M=1 AdaLN modulation / conditioning-embedder projections + # (they crash torch._int_mm's M>16), while keeping the attention / FFN compute layers and + # the sequence embedders. fp8 (no exclusion) keeps everything. + from core.inference.diffusion_transformer_quant import _INT8_EXCLUDE_NAME_TOKENS + + 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, exclude_name_tokens = _INT8_EXCLUDE_NAME_TOKENS) + big = lambda: _Lin(3072, 18432) # noqa: E731 — large enough to pass min_features + # Excluded (M=1 modulation / conditioning embedders), despite large features: + for fqn in ( + "transformer_blocks.0.norm1.linear", + "transformer_blocks.0.norm1_context.linear", + "single_transformer_blocks.0.norm.linear", + "norm_out.linear", + "transformer_blocks.0.img_mod.1", + "transformer_blocks.0.txt_mod.1", + "double_stream_modulation_img.linear", + "time_text_embed.timestep_embedder.linear_2", + "time_text_embed.guidance_embedder.linear_2", + "time_guidance_embed.timestep_embedder.linear_2", + ): + assert keep(big(), fqn) is False, fqn + # Kept (M=seq compute layers + sequence embedders), NOT matched by the modulation tokens: + for fqn in ( + "transformer_blocks.0.attn.to_q", + "transformer_blocks.0.ff.net.0.proj", + "single_transformer_blocks.0.proj_mlp", + "single_transformer_blocks.0.attn.to_qkv_mlp_proj", + "context_embedder", # "context" contains "text" -> must NOT be excluded + "txt_in", + ): + assert keep(big(), fqn) is True, fqn + # Without the exclusion (fp8 path), the modulation layer is kept. + assert make_filter_fn(512)(big(), "transformer_blocks.0.norm1.linear") is True + # A None / empty fqn must not crash the exclusion check (defensive against the callback + # passing no name); with no name nothing matches the exclusion tokens -> kept. + assert keep(big(), None) is True + assert keep(big(), "") is True + + +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, + exclude_tokens_for_scheme, + ) + assert exclude_tokens_for_scheme(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS + for scheme in (TQ_FP8, TQ_NVFP4, TQ_MXFP8): + assert exclude_tokens_for_scheme(scheme) == () + + # ── apply ───────────────────────────────────────────────────────────────────────