514f4c60fe
- Wan2.2-A14B step cache: pin the balanced FBCache threshold to 0.08 even when quant is active (per-family override in diffusion_cache.py). Auto-fp8 made the generic quant promotion (0.12) the family's effective default at pairwise LPIPS 0.128, over the 0.08 quality gate the balanced preset is held to. Measured operating point with fp8 actually engaged (1280x720/81f/50 steps, B200): fb@0.08 = 1.08x at 0.129 vs the old fb@0.12 = 2.58x at 0.181; documented in the preset table. Explicit thresholds and the fast preset are unaffected. - MagCache curves: validated the shipped 33-frame calibrations at the production 121-frame default for hunyuanvideo-1.5-720p, hunyuanvideo-1.5 (480p) and wan2.2-ti2v-5b. Fresh 121-frame calibrations differ by <= 0.024 max abs entry and produce byte-identical frames at the auto presets (hv720 quality 1.69x at LPIPS 0.042, hv480 quality 1.66x at 0.018, wan5b balanced 1.74x at 0.026, all pairwise vs the same-load uncached stack), so the curves ship unchanged with the frame-count transfer documented next to them. - Dual-GPU CFG parallelism: the secondary-device pick now prefers a device whose name and compute capability match the primary, and the gate declines a mismatched pair in auto mode (eager kernel selection is arch-dependent, so the advertised bit-identity cannot hold across different GPU models); an explicit cfg_parallel=on proceeds but is downgraded to lossless=False with a warning. - A14B expert step cache is now all-or-none, mirroring the transactional quant loop: a mixed outcome (cache engaged on one expert but not the other) is rolled back and reported uncached with the failure reason, on both the load path and the generation-time auto toggle. - Partial torchao quantization is no longer reported as dense: after an in-place quantize_/caster failure, the DiT / text encoder / VAE is scanned for leftover torchao tensor-subclass parameters and the load fails with a clear error when any are found (a half-quantized module cannot run as dense, and offload's Module.to() crashes on torchao tensors). Failures that swapped nothing keep the best-effort dense fallback. - Cleanup: apply_attention_backend / apply_speed_optims / the attention trim are called once on the pipe (they already fan out over every DiT internally), so the second A14B expert no longer passes through them twice; the stale dual-DiT helper comment is rewritten to match the two helper shapes. Tests: device-identity picker/gate/lossy-plan coverage, per-family threshold pin scoping, all-or-none rollback in both failure directions, and partial-quant detection for all three quant modules.
807 lines
37 KiB
Python
807 lines
37 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
|
|
|
|
"""Unit tests for transformer quantisation (``diffusion_transformer_quant.py``).
|
|
|
|
Hermetic: torch + torchao are stubbed via ``sys.modules``, and the per-scheme smoke
|
|
probe (``_scheme_supported`` / ``_smoke_probe``) is monkeypatched where the test cares
|
|
about the selection ladder rather than the GPU probe, so everything runs CPU-only.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
import core.inference.diffusion_transformer_quant as tq
|
|
from core.inference.diffusion_transformer_quant import (
|
|
TQ_FP8,
|
|
TQ_INT8,
|
|
TQ_MXFP8,
|
|
TQ_NVFP4,
|
|
dense_transformer_supported,
|
|
make_filter_fn,
|
|
normalize_transformer_quant,
|
|
quantize_transformer,
|
|
select_transformer_quant_scheme,
|
|
)
|
|
|
|
|
|
def _target(*, device = "cuda", dtype = "bfloat16"):
|
|
return types.SimpleNamespace(device = device, dtype = dtype)
|
|
|
|
|
|
def _stub_torch(
|
|
monkeypatch,
|
|
*,
|
|
cc = (10, 0),
|
|
with_fp8 = True,
|
|
cuda_available = True,
|
|
device_name = "NVIDIA B200",
|
|
):
|
|
torch = types.ModuleType("torch")
|
|
torch.bfloat16 = "bfloat16"
|
|
torch.float16 = "float16"
|
|
if with_fp8:
|
|
torch.float8_e4m3fn = "float8_e4m3fn"
|
|
torch.cuda = types.SimpleNamespace(
|
|
is_available = lambda: cuda_available,
|
|
get_device_capability = lambda *a: cc,
|
|
# data-center name by default so the ladder tests get the data-center order;
|
|
# consumer tests pass a GeForce name (or monkeypatch _is_consumer_gpu).
|
|
get_device_name = lambda *a: device_name,
|
|
)
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
return torch
|
|
|
|
|
|
# ── normalisation ─────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_normalize_transformer_quant():
|
|
assert normalize_transformer_quant(None) is None
|
|
assert normalize_transformer_quant("") is None
|
|
assert normalize_transformer_quant("none") is None
|
|
assert normalize_transformer_quant("off") is None
|
|
assert normalize_transformer_quant("AUTO") == "auto"
|
|
assert normalize_transformer_quant("INT8") == TQ_INT8
|
|
assert normalize_transformer_quant("fp8") == TQ_FP8
|
|
with pytest.raises(ValueError):
|
|
normalize_transformer_quant("int2")
|
|
|
|
|
|
# ── dense-source gate ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_dense_transformer_supported_requires_cuda_bf16(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
assert dense_transformer_supported(_target()) is True
|
|
assert dense_transformer_supported(_target(device = "cpu")) is False
|
|
assert dense_transformer_supported(_target(dtype = "float16")) is False
|
|
|
|
|
|
# ── scheme selection ladder ─────────────────────────────────────────────────────
|
|
|
|
|
|
def _allow(monkeypatch, allowed):
|
|
"""Force ``_scheme_supported`` to accept only ``allowed`` (simulates smoke results)."""
|
|
monkeypatch.setattr(tq, "_scheme_supported", lambda scheme, device: scheme in allowed)
|
|
|
|
|
|
def test_auto_blackwell_prefers_fp8_then_falls_back(monkeypatch):
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
# Even with every scheme available, auto picks fp8 on Blackwell: measured on a B200
|
|
# (torch 2.11 + torchao CUTLASS FP4), fp8 is both faster and more accurate than nvfp4
|
|
# for the DiT's shapes -- nvfp4's FP4 GEMM only wins on very large GEMMs, not here.
|
|
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
|
# fp8 unavailable: nvfp4 is DISABLED in the auto ladder (explicit opt-in only), so even
|
|
# though the hardware supports nvfp4 here, auto skips it and picks mxfp8 next.
|
|
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_MXFP8
|
|
# Only mxfp8 + int8 left -> mxfp8 (still above int8).
|
|
_allow(monkeypatch, {TQ_MXFP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_MXFP8
|
|
# Only int8 usable -> int8.
|
|
_allow(monkeypatch, {TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
|
|
|
|
|
def test_auto_consumer_blackwell_prefers_int8(monkeypatch):
|
|
# Consumer Blackwell (RTX 50xx): fp8 FP32-accumulate is throughput-halved while int8 is
|
|
# full-rate, so auto prefers int8 even though fp8 is available (the data-center default).
|
|
_stub_torch(monkeypatch, cc = (10, 0), device_name = "NVIDIA GeForce RTX 5090")
|
|
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
|
# int8 unavailable -> falls back to the rest of the tier (fp8 next).
|
|
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
|
|
|
|
|
def test_auto_consumer_ada_prefers_int8(monkeypatch):
|
|
# Consumer Ada (RTX 4090): int8 runs ~2x fp8's nerfed FP32-accumulate rate.
|
|
_stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA GeForce RTX 4090")
|
|
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
|
|
|
|
|
def test_auto_workstation_unknown_prefers_int8(monkeypatch):
|
|
# Unknown / workstation name -> treated as consumer (the safe default) -> int8 first.
|
|
_stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA RTX A5000")
|
|
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
|
|
|
|
|
def test_auto_professional_rtx_prefers_fp8(monkeypatch):
|
|
# Professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) are classified datacenter
|
|
# by the rest of the backend, so auto keeps fp8 first (not int8) -- matching llama_cpp.
|
|
for device_name, cc in (
|
|
("NVIDIA RTX PRO 6000 Blackwell Server Edition", (10, 0)),
|
|
("NVIDIA RTX 6000 Ada Generation", (8, 9)),
|
|
):
|
|
_stub_torch(monkeypatch, cc = cc, device_name = device_name)
|
|
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
|
|
|
|
|
def test_auto_ada_hopper_prefers_fp8(monkeypatch):
|
|
# Data-center Ada (L40S) / Hopper (H100): not nerfed -> fp8 first.
|
|
_stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA L40S")
|
|
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
|
_stub_torch(monkeypatch, cc = (9, 0), device_name = "NVIDIA H100 80GB HBM3") # Hopper
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
|
|
|
|
|
def test_auto_ampere_prefers_int8(monkeypatch):
|
|
_stub_torch(monkeypatch, cc = (8, 0))
|
|
_allow(monkeypatch, {TQ_FP8, TQ_INT8}) # fp8 cores absent on Ampere -> int8 only in ladder
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
|
_stub_torch(monkeypatch, cc = (8, 6))
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8
|
|
|
|
|
|
def test_auto_pre_ampere_unsupported(monkeypatch):
|
|
_stub_torch(monkeypatch, cc = (7, 5)) # Turing: below the int8-dynamic floor
|
|
_allow(monkeypatch, {TQ_INT8, TQ_FP8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") is None
|
|
|
|
|
|
def test_explicit_scheme_honored_or_none(monkeypatch):
|
|
_stub_torch(monkeypatch, cc = (8, 0))
|
|
_allow(monkeypatch, {TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "int8") == TQ_INT8
|
|
# Explicit unsupported scheme is NOT silently downgraded -> None (-> GGUF fallback).
|
|
assert select_transformer_quant_scheme(_target(), "fp8") is None
|
|
assert select_transformer_quant_scheme(_target(), "nvfp4") is None
|
|
|
|
|
|
def test_select_none_when_disabled_or_non_cuda(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
_allow(monkeypatch, {TQ_INT8, TQ_FP8, TQ_NVFP4})
|
|
assert select_transformer_quant_scheme(_target(), None) is None
|
|
assert select_transformer_quant_scheme(_target(device = "cpu"), "auto") is None
|
|
|
|
|
|
# ── _scheme_supported / _smoke_probe ────────────────────────────────────────────
|
|
|
|
|
|
def test_scheme_supported_shortcircuits(monkeypatch):
|
|
# No CUDA -> False without running the smoke probe.
|
|
_stub_torch(monkeypatch, cuda_available = False)
|
|
monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run"))
|
|
assert tq._scheme_supported(TQ_INT8, "cuda") is False
|
|
# fp8 requested but the fp8 dtype is missing -> False before the probe.
|
|
_stub_torch(monkeypatch, with_fp8 = False)
|
|
monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run"))
|
|
assert tq._scheme_supported(TQ_FP8, "cuda") is False
|
|
|
|
|
|
def test_smoke_probe_caches_and_tolerates_failure(monkeypatch):
|
|
tq._SMOKE_CACHE.clear()
|
|
calls = {"n": 0}
|
|
|
|
class _Lin:
|
|
def __init__(self, *a, **k):
|
|
pass
|
|
|
|
def to(self, **k):
|
|
return self
|
|
|
|
torch = types.ModuleType("torch")
|
|
torch.bfloat16 = "bfloat16"
|
|
torch.nn = types.SimpleNamespace(Linear = _Lin)
|
|
torch.randn = lambda *a, **k: object()
|
|
torch.no_grad = lambda: __import__("contextlib").nullcontext()
|
|
torch.cuda = types.SimpleNamespace(is_available = lambda: True, synchronize = lambda: None)
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
|
|
tqz = types.ModuleType("torchao.quantization")
|
|
|
|
def _quantize_ok(
|
|
module,
|
|
config,
|
|
filter_fn = None,
|
|
):
|
|
calls["n"] += 1
|
|
|
|
tqz.quantize_ = _quantize_ok
|
|
tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8cfg"
|
|
tqz.Float8DynamicActivationFloat8WeightConfig = lambda: "fp8cfg"
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
|
# _Lin is callable? No -> the forward lin(x) would fail. Make instances callable.
|
|
_Lin.__call__ = lambda self, x: x
|
|
|
|
assert tq._smoke_probe(TQ_INT8, "cuda") is True
|
|
assert tq._smoke_probe(TQ_INT8, "cuda") is True # cached, no second quantize_
|
|
assert calls["n"] == 1
|
|
|
|
# A scheme whose quantize_ raises -> probe False (and cached).
|
|
tq._SMOKE_CACHE.clear()
|
|
|
|
def _quantize_boom(
|
|
module,
|
|
config,
|
|
filter_fn = None,
|
|
):
|
|
raise RuntimeError("kernel unavailable")
|
|
|
|
tqz.quantize_ = _quantize_boom
|
|
assert tq._smoke_probe(TQ_FP8, "cuda") is False
|
|
|
|
|
|
# ── consumer-vs-datacenter detection (fp8 fast-accumulate gate) ──────────────────
|
|
|
|
|
|
def _stub_device_name(monkeypatch, name):
|
|
torch = types.ModuleType("torch")
|
|
torch.cuda = types.SimpleNamespace(get_device_name = lambda device = None: name)
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"name",
|
|
[
|
|
"NVIDIA GeForce RTX 5090",
|
|
"NVIDIA GeForce RTX 4090",
|
|
"NVIDIA RTX A4000", # workstation: A4000 token, NOT the data-center A40
|
|
"NVIDIA RTX A5000", # workstation: A5000 token, not professional/datacenter
|
|
"NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC)
|
|
],
|
|
)
|
|
def test_is_consumer_gpu_true(monkeypatch, name):
|
|
_stub_device_name(monkeypatch, name)
|
|
assert tq._is_consumer_gpu() is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"name",
|
|
[
|
|
"NVIDIA B200",
|
|
"NVIDIA B300", # Blackwell Ultra (matches llama_cpp datacenter regex)
|
|
"NVIDIA GH200 480GB", # Grace-Hopper superchip (was misread as consumer)
|
|
"NVIDIA H100 80GB HBM3",
|
|
"NVIDIA A100-SXM4-80GB",
|
|
"NVIDIA A40", # data-center Ampere (distinct token from RTX A4000)
|
|
"NVIDIA L40S",
|
|
"NVIDIA L4",
|
|
"Tesla V100-SXM2-16GB",
|
|
"NVIDIA RTX PRO 6000 Blackwell Server Edition", # professional -> datacenter-class
|
|
"NVIDIA RTX 6000 Ada Generation", # professional -> datacenter-class
|
|
],
|
|
)
|
|
def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name):
|
|
_stub_device_name(monkeypatch, name)
|
|
assert tq._is_consumer_gpu() is False
|
|
|
|
|
|
def test_is_consumer_gpu_defaults_true_on_probe_failure(monkeypatch):
|
|
# No torch / no device name available -> assume consumer (safe: fast accum is free
|
|
# on data center and a win on consumer).
|
|
torch = types.ModuleType("torch")
|
|
torch.cuda = types.SimpleNamespace() # no get_device_name
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
assert tq._is_consumer_gpu() is True
|
|
|
|
|
|
# ── filter ──────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_make_filter_fn(monkeypatch):
|
|
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)
|
|
assert keep(_Lin(1024, 4096), "blocks.0.attn.to_q") is True
|
|
assert keep(_Lin(256, 4096), "time_proj") is False # small in_features -> skip
|
|
assert keep(_Lin(4096, 256), "out_proj") is False # small out_features -> skip
|
|
assert keep(object(), "not_linear") is False # non-Linear -> skip
|
|
assert keep(types.SimpleNamespace(), "no_attrs") is False
|
|
|
|
|
|
def test_require_bf16_schemes_excludes_nvfp4():
|
|
# fp8 and mxfp8 assert a bf16 weight (torchao 0.17 / B200: "PerRow quantization only works for
|
|
# bfloat16 ..." and "Only supporting bf16 out dtype ..."), so they gate on it; nvfp4 quantises an
|
|
# fp32 weight fine, so it is NOT gated (leaving its large fp32 projections quantised, not dense).
|
|
from core.inference.diffusion_transformer_quant import (
|
|
_REQUIRE_BF16_SCHEMES,
|
|
TQ_FP8,
|
|
TQ_MXFP8,
|
|
TQ_NVFP4,
|
|
TQ_INT8,
|
|
)
|
|
|
|
assert TQ_FP8 in _REQUIRE_BF16_SCHEMES
|
|
assert TQ_MXFP8 in _REQUIRE_BF16_SCHEMES
|
|
assert TQ_NVFP4 not in _REQUIRE_BF16_SCHEMES
|
|
assert TQ_INT8 not in _REQUIRE_BF16_SCHEMES
|
|
|
|
|
|
def test_make_filter_fn_require_bf16_skips_non_bf16(monkeypatch):
|
|
# fp8 / mxfp8 assert a bf16 weight, so require_bf16 must skip a fp32 Linear (which Wan / Hunyuan
|
|
# video DiTs keep) while keeping the bf16 ones -- otherwise a single fp32 layer raises inside
|
|
# quantize_ and no-ops the whole pass. int8 and nvfp4 leave it off (they quantise fp32 fine).
|
|
torch = types.ModuleType("torch")
|
|
torch.bfloat16, torch.float32 = "bf16", "fp32"
|
|
|
|
class _Lin:
|
|
def __init__(self, i, o, dtype):
|
|
self.in_features, self.out_features = i, o
|
|
self.weight = types.SimpleNamespace(dtype = dtype)
|
|
|
|
torch.nn = types.SimpleNamespace(Linear = _Lin)
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
|
|
gated = make_filter_fn(512, require_bf16 = True)
|
|
assert gated(_Lin(1024, 4096, torch.bfloat16), "blocks.0.attn.to_q") is True
|
|
assert gated(_Lin(1024, 4096, torch.float32), "blocks.0.attn.to_q") is False # fp32 -> skip
|
|
assert gated(types.SimpleNamespace(in_features = 1024, out_features = 4096), "no_weight") is False
|
|
# int8 (require_bf16 off, the default) still quantises the fp32 linear.
|
|
assert make_filter_fn(512)(_Lin(1024, 4096, torch.float32), "blocks.0.attn.to_q") is True
|
|
|
|
|
|
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 exclusion, or a
|
|
# prequant artifact quantises a layer the runtime path skips (int8's M=1 modulation/embedder
|
|
# linears -> torch._int_mm crash; fp8's zero-padded-conditioning embedder -> black frames). int8
|
|
# gets the family-independent exclusion; the scaled_mm schemes exclude nothing WITHOUT a family.
|
|
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) == ()
|
|
|
|
|
|
def test_exclude_tokens_for_scheme():
|
|
# The shared scheme->exclusion decision used by BOTH the runtime quantise path and the offline
|
|
# prequant-checkpoint builder, so a checkpoint built ahead of time skips exactly the layers the
|
|
# runtime path skips (offline == runtime). int8 excludes the M=1 modulation / embedder tokens on
|
|
# every family; the scaled_mm schemes exclude nothing by default.
|
|
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
|
|
assert exclude_tokens_for_scheme(TQ_FP8) == ()
|
|
assert exclude_tokens_for_scheme(TQ_NVFP4) == ()
|
|
assert exclude_tokens_for_scheme(TQ_MXFP8) == ()
|
|
|
|
|
|
def test_exclude_tokens_for_scheme_family():
|
|
# Per-family fp8 exclusion: on Wan the per-row fp8 activation scale divides by a zero
|
|
# padding-token row in condition_embedder.text_embedder (-> inf -> black frames), so fp8 keeps
|
|
# condition_embedder bf16 while the 30-block stack stays fp8. int8 is unaffected by family (it
|
|
# tolerates zero rows), and an unknown / Hunyuan family gets no fp8 exclusion.
|
|
from core.inference.diffusion_transformer_quant import (
|
|
_INT8_EXCLUDE_NAME_TOKENS,
|
|
exclude_tokens_for_scheme,
|
|
)
|
|
|
|
assert exclude_tokens_for_scheme(TQ_FP8, "wan2.2-ti2v-5b") == ("condition_embedder",)
|
|
assert exclude_tokens_for_scheme(TQ_FP8, "wan2.2-t2v-a14b") == ("condition_embedder",)
|
|
assert exclude_tokens_for_scheme(TQ_MXFP8, "wan2.2-ti2v-5b") == ("condition_embedder",)
|
|
# Hunyuan is not localisable (fp8 stays denied), and an unknown family gets nothing.
|
|
assert exclude_tokens_for_scheme(TQ_FP8, "hunyuanvideo-1.5") == ()
|
|
assert exclude_tokens_for_scheme(TQ_FP8, "z-image") == ()
|
|
# int8 tolerates zero ROWS (per-token; no divide), so most families need no extra skip...
|
|
assert exclude_tokens_for_scheme(TQ_INT8, "wan2.2-ti2v-5b") == _INT8_EXCLUDE_NAME_TOKENS
|
|
# ...but the attention trim shrinks HunyuanVideo-1.5's text streams to their VALID token
|
|
# counts, where the int8 dynamic path fails two ways (both measured): a zero-token (M=0)
|
|
# input passes through UNPROJECTED (byt5/image embedders on t2v -> cond-type add crash),
|
|
# and torch._int_mm requires M > 16 (the ~6-token empty negative prompt crashes the
|
|
# TokenRefiner and the blocks' context-stream projections). All the text-stream linears
|
|
# must stay bf16; the M ~ 32k video-stream linears keep the int8 coverage.
|
|
from core.inference.diffusion_transformer_quant import _HUNYUAN15_INT8_EXCLUDES
|
|
|
|
for fam in ("hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"):
|
|
assert (
|
|
exclude_tokens_for_scheme(TQ_INT8, fam)
|
|
== _INT8_EXCLUDE_NAME_TOKENS + _HUNYUAN15_INT8_EXCLUDES
|
|
)
|
|
assert "context_embedder" in _HUNYUAN15_INT8_EXCLUDES # covers context_embedder_2 too
|
|
|
|
|
|
# ── apply ───────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_resolve_fast_accum(monkeypatch):
|
|
# None auto-detects by GPU class; an explicit bool forces it.
|
|
monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: True)
|
|
assert tq._resolve_fast_accum(None) is True
|
|
monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: False)
|
|
assert tq._resolve_fast_accum(None) is False
|
|
assert tq._resolve_fast_accum(True) is True # forced on (e.g. on a data-center card)
|
|
assert tq._resolve_fast_accum(False) is False # forced off (e.g. on a consumer card)
|
|
|
|
|
|
def test_fp8_config_uses_per_row_granularity():
|
|
"""FP8 must use PerRow (per-token activation + per-channel weight) scaling. torchao's
|
|
default is per-TENSOR: on a DiT with extreme activation outliers (z-image's ~6.6e4) one
|
|
outlier forces a tensor-wide scale that pushes normal values below fp8 resolution and the
|
|
denoise collapses to noise. This is the regression guard for that fix (validated on B200:
|
|
per-tensor fp8 = noise, per-row fp8 = matches bf16)."""
|
|
torchao_quant = pytest.importorskip("torchao.quantization")
|
|
per_row = getattr(torchao_quant, "PerRow", None)
|
|
if per_row is None:
|
|
pytest.skip("torchao build without PerRow granularity")
|
|
cfg = tq._make_quant_config(TQ_FP8)
|
|
gran = getattr(cfg, "granularity", None)
|
|
assert gran is not None, "fp8 config must set an explicit granularity, not torchao's default"
|
|
grans = gran if isinstance(gran, (list, tuple)) else [gran]
|
|
assert grans and all(isinstance(g, per_row) for g in grans), f"expected all PerRow, got {gran}"
|
|
|
|
|
|
def test_quantize_transformer_applies_and_marks(monkeypatch):
|
|
monkeypatch.setattr(
|
|
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8
|
|
)
|
|
seen: dict = {}
|
|
|
|
def _mk(scheme, fast_accum = None):
|
|
seen["scheme"], seen["fast_accum"] = scheme, fast_accum
|
|
return f"{scheme}cfg"
|
|
|
|
monkeypatch.setattr(tq, "_make_quant_config", _mk)
|
|
recorder: list = []
|
|
tqz = types.ModuleType("torchao.quantization")
|
|
tqz.quantize_ = lambda module, config, filter_fn = None: recorder.append(
|
|
(module, config, filter_fn)
|
|
)
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
|
|
|
transformer = types.SimpleNamespace()
|
|
pipe = types.SimpleNamespace(transformer = transformer)
|
|
assert quantize_transformer(pipe, _target(), mode = "fp8", fast_accum = False) == TQ_FP8
|
|
assert len(recorder) == 1 and recorder[0][0] is transformer and recorder[0][1] == "fp8cfg"
|
|
assert callable(recorder[0][2]) # a filter_fn was passed
|
|
assert transformer._unsloth_runtime_quant == TQ_FP8 # diagnostic marker set
|
|
assert seen["fast_accum"] is False # the override is forwarded into the config
|
|
|
|
|
|
def test_quantize_transformer_none_when_unsupported(monkeypatch):
|
|
monkeypatch.setattr(
|
|
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: None
|
|
)
|
|
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
|
assert quantize_transformer(pipe, _target(), mode = "auto") is None
|
|
|
|
|
|
def test_quantize_transformer_tolerates_failure(monkeypatch):
|
|
monkeypatch.setattr(
|
|
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_INT8
|
|
)
|
|
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: "cfg")
|
|
tqz = types.ModuleType("torchao.quantization")
|
|
|
|
def _boom(
|
|
module,
|
|
config,
|
|
filter_fn = None,
|
|
):
|
|
raise RuntimeError("partial quant failure")
|
|
|
|
tqz.quantize_ = _boom
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
|
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
|
# A quantise failure returns None (caller falls back to GGUF), never raises.
|
|
assert quantize_transformer(pipe, _target(), mode = "int8") is None
|
|
|
|
|
|
# ── family scheme deny (measured model-level breakage) ────────────────────────
|
|
|
|
|
|
def test_family_deny_auto_skips_fp8_for_qwen(monkeypatch):
|
|
# B200 with every scheme available: auto must NOT pick fp8 / nvfp4 / mxfp8 for the
|
|
# Qwen DiT (per-row fp8 renders black frames on it; see _FAMILY_SCHEME_DENY) and
|
|
# falls through the ladder to int8, which measures excellent on Qwen.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image") == TQ_INT8
|
|
assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image-edit") == TQ_INT8
|
|
|
|
|
|
def test_family_deny_refuses_explicit_fp8_for_qwen(monkeypatch):
|
|
# An explicit fp8 request on qwen-image returns None (same contract as an
|
|
# unsupported scheme: the caller builds the GGUF pipeline instead). int8 stays
|
|
# honored on qwen, and fp8 stays honored on families outside the deny table.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "fp8", family = "qwen-image") is None
|
|
assert select_transformer_quant_scheme(_target(), "int8", family = "qwen-image") == TQ_INT8
|
|
assert select_transformer_quant_scheme(_target(), "fp8", family = "z-image") == TQ_FP8
|
|
|
|
|
|
def test_family_allows_fp8_for_wan(monkeypatch):
|
|
# B200 with every scheme available: the Wan fp8 black frame was root-caused to a single input
|
|
# embedder dividing by a zero padding-token row (fixed by the condition_embedder exclude, not a
|
|
# deny), so auto now picks fp8 -- the ladder head -- for both the 5B TI2V and the A14B MoE.
|
|
# mxfp8 / nvfp4 stay denied (same per-row scaled_mm family, not separately validated), so with
|
|
# only those + int8 available auto still lands on int8.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto", family = "wan2.2-ti2v-5b") == TQ_FP8
|
|
assert select_transformer_quant_scheme(_target(), "auto", family = "wan2.2-t2v-a14b") == TQ_FP8
|
|
_allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_INT8}) # fp8 unavailable -> denied mx/nvfp4 skipped
|
|
assert select_transformer_quant_scheme(_target(), "auto", family = "wan2.2-ti2v-5b") == TQ_INT8
|
|
|
|
|
|
def test_family_deny_wan_fp8_allowed_mxfp8_nvfp4_refused(monkeypatch):
|
|
# An explicit fp8 request on a Wan family is now honored (made safe by the condition_embedder
|
|
# exclude); int8 stays honored; mxfp8 / nvfp4 stay refused (return None -> GGUF fallback) until
|
|
# separately validated on the Wan activation profile.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow(monkeypatch, {TQ_FP8, TQ_MXFP8, TQ_NVFP4, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "fp8", family = "wan2.2-ti2v-5b") == TQ_FP8
|
|
assert select_transformer_quant_scheme(_target(), "int8", family = "wan2.2-ti2v-5b") == TQ_INT8
|
|
assert select_transformer_quant_scheme(_target(), "mxfp8", family = "wan2.2-ti2v-5b") is None
|
|
assert select_transformer_quant_scheme(_target(), "nvfp4", family = "wan2.2-ti2v-5b") is None
|
|
|
|
|
|
def test_family_deny_auto_skips_fp8_for_hunyuan(monkeypatch):
|
|
# HunyuanVideo-1.5 DiT also renders black frames on fp8 (measured, LPIPS 0.82); both the
|
|
# 480p and 720p repacks deny fp8/mxfp8/nvfp4 and fall to int8. ltx-2 is NOT denied (its
|
|
# fp8 measures clean), so the deny is per family, not a blanket video rule.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto", family = "hunyuanvideo-1.5") == TQ_INT8
|
|
assert (
|
|
select_transformer_quant_scheme(_target(), "auto", family = "hunyuanvideo-1.5-720p")
|
|
== TQ_INT8
|
|
)
|
|
# ltx-2 keeps the ladder head (fp8) -- it is not a black-frame family.
|
|
assert select_transformer_quant_scheme(_target(), "auto", family = "ltx-2") == TQ_FP8
|
|
|
|
|
|
def test_family_deny_no_family_keeps_ladder(monkeypatch):
|
|
# Without a family (or an unknown one) the ladder is unchanged: fp8 first on B200.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
|
assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8
|
|
assert select_transformer_quant_scheme(_target(), "auto", family = "sdxl") == TQ_FP8
|
|
|
|
|
|
def test_is_int8_memory_fallback(monkeypatch):
|
|
# True only where AUTO lands on int8 as a black-frame/denied FALLBACK on a data-center,
|
|
# fp8-capable GPU (Hunyuan: fp8 denied -> int8), where dense+compile beats int8. The loader
|
|
# uses this (plus a 'dense fits resident' check) to prefer dense over int8. False everywhere
|
|
# int8 is a legitimate accelerator: fp8 families, consumer GPUs, and pre-Ada (no fp8) parts.
|
|
from core.inference.diffusion_transformer_quant import is_int8_memory_fallback
|
|
|
|
# data-center Blackwell, all schemes available: Hunyuan denies fp8/mx/nvfp4 -> auto int8 -> True.
|
|
_stub_torch(monkeypatch, cc = (10, 0), device_name = "NVIDIA B200")
|
|
_allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
|
|
assert is_int8_memory_fallback(_target(), "hunyuanvideo-1.5") is True
|
|
assert is_int8_memory_fallback(_target(), "hunyuanvideo-1.5-720p") is True
|
|
# Wan / LTX resolve to fp8 (a speed win), not int8 -> False (keep quantising).
|
|
assert is_int8_memory_fallback(_target(), "wan2.2-ti2v-5b") is False
|
|
assert is_int8_memory_fallback(_target(), "wan2.2-t2v-a14b") is False
|
|
assert is_int8_memory_fallback(_target(), "ltx-2") is False
|
|
assert is_int8_memory_fallback(_target(), None) is False
|
|
|
|
# Consumer GPU (fp8 accumulate halved -> int8 can be a speed win): never prefer dense.
|
|
_stub_torch(monkeypatch, cc = (10, 0), device_name = "NVIDIA GeForce RTX 5090")
|
|
_allow(monkeypatch, {TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8})
|
|
assert is_int8_memory_fallback(_target(), "hunyuanvideo-1.5") is False
|
|
|
|
# Ampere data-center (sm_80, no fp8 tensor cores): int8 is the genuine accelerator -> False.
|
|
_stub_torch(monkeypatch, cc = (8, 0), device_name = "NVIDIA A100")
|
|
_allow(monkeypatch, {TQ_INT8})
|
|
assert is_int8_memory_fallback(_target(), "hunyuanvideo-1.5") is False
|
|
|
|
|
|
def test_quantize_transformer_threads_family(monkeypatch):
|
|
# quantize_transformer passes the family down to the selector, so a denied
|
|
# (family, scheme) pair never reaches torchao.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow(monkeypatch, {TQ_FP8, TQ_INT8})
|
|
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
|
called = {}
|
|
tqz = types.ModuleType("torchao.quantization")
|
|
|
|
def _quantize(
|
|
module,
|
|
config,
|
|
filter_fn = None,
|
|
):
|
|
called["scheme"] = True
|
|
|
|
tqz.quantize_ = _quantize
|
|
tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8-cfg"
|
|
tqz.Float8DynamicActivationFloat8WeightConfig = lambda **kw: "fp8-cfg"
|
|
tqz.PerRow = lambda: "per-row"
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
|
assert quantize_transformer(pipe, _target(), mode = "fp8", family = "qwen-image") is None
|
|
assert called == {}
|
|
|
|
|
|
def test_quantize_transformer_fp8_wan_excludes_condition_embedder(monkeypatch):
|
|
# The Wan fp8 fix is a FILTER exclusion, not a deny: quantize_transformer must thread the family
|
|
# into the filter so condition_embedder (the embedder that reads the zero-padded text and would
|
|
# divide by a zero row scale) stays bf16 while the 30-block stack is quantised on fp8. Capture
|
|
# the filter that reaches torchao and check it on representative fully-qualified names.
|
|
monkeypatch.setattr(
|
|
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8
|
|
)
|
|
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg")
|
|
|
|
# torch stub with a real nn.Linear class so the captured filter's isinstance + bf16 gate runs.
|
|
torch = types.ModuleType("torch")
|
|
torch.bfloat16 = "bfloat16"
|
|
|
|
class _Linear:
|
|
def __init__(
|
|
self,
|
|
inf,
|
|
outf,
|
|
dtype = "bfloat16",
|
|
):
|
|
self.in_features, self.out_features = inf, outf
|
|
self.weight = types.SimpleNamespace(dtype = dtype)
|
|
|
|
torch.nn = types.SimpleNamespace(Linear = _Linear)
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
|
|
captured: dict = {}
|
|
tqz = types.ModuleType("torchao.quantization")
|
|
tqz.quantize_ = lambda module, config, filter_fn = None: captured.update(fn = filter_fn)
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
|
|
|
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
|
assert quantize_transformer(pipe, _target(), mode = "fp8", family = "wan2.2-ti2v-5b") == TQ_FP8
|
|
filt = captured["fn"]
|
|
big = _Linear(4096, 5120) # a FLOP-heavy bf16 linear (passes min_features + bf16 gate)
|
|
# condition_embedder.* is kept bf16 (the divide-by-zero origin); the block stack is fp8.
|
|
assert filt(big, "condition_embedder.text_embedder.linear_1") is False
|
|
assert filt(big, "condition_embedder.time_embedder.linear_1") is False
|
|
assert filt(big, "blocks.0.attn1.to_q") is True
|
|
assert filt(big, "blocks.0.ffn.net.0.proj") is True
|
|
assert (
|
|
filt(big, "blocks.0.attn2.to_k") is True
|
|
) # cross-attn K/V stay fp8 (embedder bias rescues rows)
|
|
|
|
|
|
# ── partial in-place quant detection (torchao_quantized_param_fqns) ────────────────
|
|
class _TorchaoLikeTensor:
|
|
"""Stands in for a torchao tensor subclass: detection keys on the class's module
|
|
path, so the fake just claims a torchao __module__."""
|
|
|
|
|
|
_TorchaoLikeTensor.__module__ = "torchao.dtypes.affine_quantized_tensor"
|
|
|
|
|
|
class _MutableDiT:
|
|
"""A fake transformer whose quantize_ pass 'swapped' one weight before failing."""
|
|
|
|
def __init__(self):
|
|
self._swapped = False
|
|
|
|
def named_parameters(self):
|
|
if self._swapped:
|
|
yield ("blocks.0.attn.to_q.weight", _TorchaoLikeTensor())
|
|
yield ("blocks.1.attn.to_q.weight", types.SimpleNamespace())
|
|
|
|
|
|
def test_torchao_param_scan_detects_swapped_weights():
|
|
dit = _MutableDiT()
|
|
assert tq.torchao_quantized_param_fqns(dit) == []
|
|
dit._swapped = True
|
|
assert tq.torchao_quantized_param_fqns(dit) == ["blocks.0.attn.to_q.weight"]
|
|
# Unscannable object -> no leftovers reported (best-effort, the pre-check path).
|
|
assert tq.torchao_quantized_param_fqns(object()) == []
|
|
|
|
|
|
def test_quantize_transformer_partial_failure_raises(monkeypatch):
|
|
# quantize_ swaps weights module-by-module, so a mid-pass exception (OOM on a 14B
|
|
# DiT, a layer a kernel rejects) can leave earlier layers quantized. That module
|
|
# cannot run as dense (offload's Module.to() crashes on torchao tensors), so the
|
|
# load must FAIL with a clear error instead of reporting a dense fallback.
|
|
monkeypatch.setattr(
|
|
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8
|
|
)
|
|
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg")
|
|
tqz = types.ModuleType("torchao.quantization")
|
|
|
|
def _convert_one_then_boom(module, config, filter_fn = None):
|
|
module._swapped = True # the in-place swap of the first submodule
|
|
raise RuntimeError("OOM mid-conversion")
|
|
|
|
tqz.quantize_ = _convert_one_then_boom
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
|
pipe = types.SimpleNamespace(transformer = _MutableDiT())
|
|
with pytest.raises(RuntimeError, match = "partially quantized"):
|
|
quantize_transformer(pipe, _target(), mode = "fp8")
|
|
|
|
|
|
def test_quantize_transformer_clean_failure_still_falls_back_dense(monkeypatch):
|
|
# A failure that swapped NOTHING keeps the best-effort contract: dense fallback.
|
|
monkeypatch.setattr(
|
|
tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8
|
|
)
|
|
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg")
|
|
tqz = types.ModuleType("torchao.quantization")
|
|
|
|
def _boom(module, config, filter_fn = None):
|
|
raise RuntimeError("failed before any swap")
|
|
|
|
tqz.quantize_ = _boom
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
|
pipe = types.SimpleNamespace(transformer = _MutableDiT())
|
|
assert quantize_transformer(pipe, _target(), mode = "fp8") is None
|