diff --git a/scripts/int8_linear_probe.py b/scripts/int8_linear_probe.py new file mode 100644 index 0000000000..6a832553a0 --- /dev/null +++ b/scripts/int8_linear_probe.py @@ -0,0 +1,79 @@ +# 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 bf031d13d0..173d0aaa86 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -37,13 +37,31 @@ 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", +) + # 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 @@ -286,9 +304,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: @@ -301,7 +321,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 any(tok in name for tok in exclude_name_tokens): + return False + return True return filter_fn @@ -331,10 +357,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 = _INT8_EXCLUDE_NAME_TOKENS if scheme == TQ_INT8 else () 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 dd9a12c60d..9ee0af7649 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -309,6 +309,50 @@ 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 + + # ── apply ───────────────────────────────────────────────────────────────────────