36df317293
Comment-only pass over the Python this PR touches: drop what the code already says, collapse multi-line explanations that still read on one line, and keep the reasoning that is not recoverable from the code. No code, docstring semantics or behaviour changes; verified with an AST comparison against the previous revision, and the backend suite is unchanged (same 37 environment failures as before: the API integration tests that need a live keyed server, the flash-attn install hooks, and the GPU memory fields).
88 lines
3.5 KiB
Python
88 lines
3.5 KiB
Python
# 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}")
|
|
# A modulation/embedder Linear sits outside the repeated blocks (no numeric index in its fqn),
|
|
# or has an AdaLN out==k*in (k>=3) 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())
|