ca8f415a7e
The transformer and text encoder auto-quantize; the VAE stayed dense. VAEs are convolutional, so torchao int8 (Linear/2D-only) does not apply, but Float8DynamicActivationFloat8WeightConfig quantizes Conv2d/Conv3d weights with PerTensor granularity (auto-skipping convs whose channels are not a multiple of 16, so the 3-channel RGB head stays dense). New diffusion_vae_quant.py offers two schemes: fp8_dynamic (torchao conv compute fp8, cc>=8.9, resident) and fp8 (diffusers layerwise storage cast, any conv, survives offload); no int8 (no Conv3d int8 kernel). select_vae_quant_scheme walks (fp8_dynamic, fp8) with a live conv smoke probe, an offload gate, a per-family deny list, and a force_fp32 gate; the image + video loaders map unset vae_quant to auto, skip the vae_force_fp32 Wan families, and record the engaged scheme. Guards _align_vae_dtype to skip the img2img/inpaint re-cast when the VAE is quantized (its fp8 tensor subclasses reject .to(dtype=)). Verified on a B200: %16 Conv2d/Conv3d/Linear -> Float8Tensor, conv_out dense, forward runs.
379 lines
17 KiB
Python
379 lines
17 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 VAE quantisation (``diffusion_vae_quant.py``).
|
|
|
|
Hermetic: torch + the diffusers / torchao casters are stubbed via ``sys.modules`` so
|
|
gating, the conv filter, and the apply path run without a GPU, real diffusers, or real
|
|
torchao. Mirrors tests/test_diffusion_precision.py's stubbing style.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
import core.inference.diffusion_vae_quant as vq
|
|
from core.inference.diffusion_vae_quant import (
|
|
VAE_QUANT_AUTO,
|
|
VAE_QUANT_FP8,
|
|
VAE_QUANT_FP8_DYNAMIC,
|
|
_cast_vae_fp8,
|
|
_cast_vae_fp8_dynamic,
|
|
normalize_vae_quant,
|
|
quantize_vae,
|
|
select_vae_quant_scheme,
|
|
vae_quant_supported,
|
|
)
|
|
|
|
|
|
def _target(*, device = "cuda", dtype = "bfloat16", cc = (10, 0)):
|
|
return types.SimpleNamespace(device = device, dtype = dtype, _cc = cc)
|
|
|
|
|
|
class _Weight:
|
|
"""A stand-in for a conv / linear weight tensor: exposes ``.shape`` and ``.dim()``."""
|
|
|
|
def __init__(self, shape):
|
|
self.shape = shape
|
|
|
|
def dim(self):
|
|
return len(self.shape)
|
|
|
|
|
|
def _stub_torch(monkeypatch, *, with_fp8 = True, cc = (10, 0)):
|
|
torch = types.ModuleType("torch")
|
|
torch.bfloat16 = "bfloat16"
|
|
torch.float16 = "float16"
|
|
if with_fp8:
|
|
torch.float8_e4m3fn = "float8_e4m3fn"
|
|
# The conv filter isinstance-checks nn.Linear / nn.Conv2d / nn.Conv3d, and the layerwise
|
|
# caster reads torch.float8_e4m3fn, so the stub torch must expose both.
|
|
torch.nn = types.SimpleNamespace(
|
|
Linear = type("Linear", (), {}),
|
|
Conv2d = type("Conv2d", (), {}),
|
|
Conv3d = type("Conv3d", (), {}),
|
|
)
|
|
torch.cuda = types.SimpleNamespace(
|
|
get_device_capability = lambda *a: cc,
|
|
synchronize = lambda *a, **k: None,
|
|
is_available = lambda: True,
|
|
)
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
return torch
|
|
|
|
|
|
def _stub_torchao(monkeypatch, captured):
|
|
# torchao's fp8_dynamic conv config + quantize_. Records the (config, filter_fn) so the
|
|
# PerTensor granularity and the conv filter closure can be asserted.
|
|
tq = types.ModuleType("torchao.quantization")
|
|
tq.quantize_ = lambda module, config, filter_fn = None: captured.update(
|
|
module = module, config = config, filter_fn = filter_fn
|
|
)
|
|
tq.Float8DynamicActivationFloat8WeightConfig = lambda granularity = None: ("fp8dyn", granularity)
|
|
tq.PerTensor = lambda: "pertensor"
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tq)
|
|
return tq
|
|
|
|
|
|
def _stub_diffusers(monkeypatch, recorder):
|
|
hooks = types.ModuleType("diffusers.hooks")
|
|
casting = types.ModuleType("diffusers.hooks.layerwise_casting")
|
|
casting.DEFAULT_SKIP_MODULES_PATTERN = ("norm",)
|
|
hooks.apply_layerwise_casting = lambda module, **kw: recorder.append(("fp8", module, kw))
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting)
|
|
|
|
|
|
def _stub_capability(monkeypatch, cc):
|
|
"""Stub the transformer module's ``_capability`` that select_vae_quant_scheme imports."""
|
|
dtq = types.ModuleType("core.inference.diffusion_transformer_quant")
|
|
dtq._capability = lambda: cc
|
|
monkeypatch.setitem(sys.modules, "core.inference.diffusion_transformer_quant", dtq)
|
|
return dtq
|
|
|
|
|
|
def _allow_vae(monkeypatch, allowed):
|
|
"""Force vae_quant_supported to accept only ``allowed`` (simulates the hardware gate)."""
|
|
monkeypatch.setattr(vq, "vae_quant_supported", lambda target, mode: mode in allowed)
|
|
|
|
|
|
# ── normalisation ─────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_normalize_vae_quant():
|
|
assert normalize_vae_quant(None) is None
|
|
assert normalize_vae_quant("") is None
|
|
assert normalize_vae_quant("none") is None
|
|
# "off" disables (like the TE / transformer normalisers) -> dense.
|
|
assert normalize_vae_quant("off") is None
|
|
# "auto" passes through for select_vae_quant_scheme to resolve.
|
|
assert normalize_vae_quant("AUTO") == VAE_QUANT_AUTO
|
|
assert normalize_vae_quant("FP8") == VAE_QUANT_FP8
|
|
# Hyphens fold to underscores so "fp8-dynamic" is accepted.
|
|
assert normalize_vae_quant("FP8-Dynamic") == VAE_QUANT_FP8_DYNAMIC
|
|
# int8 / nvfp4 have no VAE scheme -> rejected.
|
|
with pytest.raises(ValueError):
|
|
normalize_vae_quant("int8")
|
|
with pytest.raises(ValueError):
|
|
normalize_vae_quant("nvfp4")
|
|
|
|
|
|
# ── gating ────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_vae_quant_supported_fp8_requires_cuda_bf16_and_fp8(monkeypatch):
|
|
_stub_torch(monkeypatch, with_fp8 = True, cc = (8, 9))
|
|
assert vae_quant_supported(_target(), VAE_QUANT_FP8) is True
|
|
assert vae_quant_supported(_target(device = "cpu"), VAE_QUANT_FP8) is False
|
|
assert vae_quant_supported(_target(dtype = "float16"), VAE_QUANT_FP8) is False
|
|
# No fp8 dtype at all -> unsupported.
|
|
_stub_torch(monkeypatch, with_fp8 = False, cc = (8, 9))
|
|
assert vae_quant_supported(_target(), VAE_QUANT_FP8) is False
|
|
|
|
|
|
def test_vae_quant_supported_fp8_dynamic_requires_sm89(monkeypatch):
|
|
# Compute fp8 conv (torch._scaled_mm) needs fp8-GEMM silicon: Ada sm_89+ / Hopper / Blackwell.
|
|
_stub_torch(monkeypatch, cc = (8, 9))
|
|
assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is True
|
|
_stub_torch(monkeypatch, cc = (9, 0))
|
|
assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is True
|
|
# Ampere (8.0) has no fp8 GEMM.
|
|
_stub_torch(monkeypatch, cc = (8, 0))
|
|
assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is False
|
|
|
|
|
|
# ── auto ladder (select_vae_quant_scheme) ───────────────────────────────────────
|
|
|
|
|
|
def test_select_datacenter_prefers_fp8_dynamic(monkeypatch):
|
|
# Data-center fp8-GEMM silicon with a passing conv probe: fp8_dynamic leads the ladder.
|
|
_stub_capability(monkeypatch, (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
|
|
monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True)
|
|
assert select_vae_quant_scheme(_target(), "auto", family = "flux.1") == VAE_QUANT_FP8_DYNAMIC
|
|
|
|
|
|
def test_select_offload_uses_layerwise_fp8(monkeypatch):
|
|
# Under offload the torchao fp8_dynamic mode (rejects Module.to()) is skipped BEFORE the
|
|
# probe -> layerwise fp8. The probe must not even run.
|
|
_stub_capability(monkeypatch, (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
|
|
monkeypatch.setattr(
|
|
vq, "_vae_fp8_dynamic_probe", lambda device: pytest.fail("probe must not run under offload")
|
|
)
|
|
assert (
|
|
select_vae_quant_scheme(_target(), "auto", offload_active = True) == VAE_QUANT_FP8
|
|
)
|
|
|
|
|
|
def test_select_force_fp32_stays_dense(monkeypatch):
|
|
# A force-fp32 (Wan) family never quantises, for auto or an explicit request.
|
|
_stub_capability(monkeypatch, (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
|
|
assert select_vae_quant_scheme(_target(), "auto", force_fp32 = True) is None
|
|
assert select_vae_quant_scheme(_target(), "fp8", force_fp32 = True) is None
|
|
|
|
|
|
def test_select_family_deny_skips_scheme(monkeypatch):
|
|
_stub_capability(monkeypatch, (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
|
|
monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True)
|
|
monkeypatch.setattr(
|
|
vq, "_VAE_FAMILY_SCHEME_DENY", {"badfam": frozenset({VAE_QUANT_FP8_DYNAMIC})}
|
|
)
|
|
# fp8_dynamic denied for this family -> falls to layerwise fp8.
|
|
assert select_vae_quant_scheme(_target(), "auto", family = "badfam") == VAE_QUANT_FP8
|
|
|
|
|
|
def test_select_no_capability_is_none(monkeypatch):
|
|
_stub_capability(monkeypatch, None)
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
|
|
assert select_vae_quant_scheme(_target(), "auto") is None
|
|
|
|
|
|
def test_select_support_gate_falls_to_fp8(monkeypatch):
|
|
# fp8_dynamic hardware-unsupported (only fp8 allowed) -> layerwise fp8.
|
|
_stub_capability(monkeypatch, (8, 9))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8})
|
|
monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True)
|
|
assert select_vae_quant_scheme(_target(), "auto") == VAE_QUANT_FP8
|
|
|
|
|
|
def test_select_smoke_failure_falls_to_fp8(monkeypatch):
|
|
# fp8_dynamic is hardware-supported but its conv probe fails -> layerwise fp8.
|
|
_stub_capability(monkeypatch, (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
|
|
monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: False)
|
|
assert select_vae_quant_scheme(_target(), "auto") == VAE_QUANT_FP8
|
|
|
|
|
|
def test_select_explicit_passthrough(monkeypatch):
|
|
# An explicit request is returned as-is (quantize_vae re-gates it); no ladder walk,
|
|
# so no transformer-module stub is needed.
|
|
assert select_vae_quant_scheme(_target(), "fp8") == VAE_QUANT_FP8
|
|
assert select_vae_quant_scheme(_target(), "fp8_dynamic") == VAE_QUANT_FP8_DYNAMIC
|
|
assert select_vae_quant_scheme(_target(), None) is None
|
|
assert select_vae_quant_scheme(_target(), "none") is None
|
|
|
|
|
|
def test_select_explicit_family_deny_returns_none(monkeypatch):
|
|
monkeypatch.setattr(vq, "_VAE_FAMILY_SCHEME_DENY", {"badfam": frozenset({VAE_QUANT_FP8})})
|
|
assert select_vae_quant_scheme(_target(), "fp8", family = "badfam") is None
|
|
|
|
|
|
# ── conv-aware filter (_cast_vae_fp8_dynamic) ───────────────────────────────────
|
|
|
|
|
|
def test_fp8_dynamic_conv_filter(monkeypatch):
|
|
# The real filter closure the PerTensor fp8 conv config receives: %16-channel Conv2d/Conv3d
|
|
# and Linear quantise; off-%16 channels and the conv_out / norm_out head stay dense.
|
|
torch = _stub_torch(monkeypatch)
|
|
captured: dict = {}
|
|
_stub_torchao(monkeypatch, captured)
|
|
_cast_vae_fp8_dynamic(object(), _target())
|
|
# PerTensor granularity (NOT the DiT's per-row) is what the config was built with.
|
|
assert captured["config"] == ("fp8dyn", "pertensor")
|
|
ff = captured["filter_fn"]
|
|
nn = torch.nn
|
|
|
|
def _mod(cls, shape):
|
|
m = cls()
|
|
m.weight = _Weight(shape)
|
|
return m
|
|
|
|
# Conv2d (4D) / Conv3d (5D) with both channel dims a multiple of 16 quantise.
|
|
assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.up.0.resnets.0.conv1") is True
|
|
assert ff(_mod(nn.Conv3d, (64, 64, 3, 3, 3)), "decoder.mid_block.conv3d") is True
|
|
# nn.Linear (mid-block attention projections) also quantise.
|
|
assert ff(_mod(nn.Linear, (512, 512)), "decoder.mid_block.attentions.0.to_q") is True
|
|
# Channels not a multiple of 16 excluded (torchao would skip them regardless): the RGB
|
|
# in/out head (C=3) and any off-16 dim.
|
|
assert ff(_mod(nn.Conv2d, (128, 3, 3, 3)), "encoder.conv_in") is False
|
|
assert ff(_mod(nn.Conv2d, (24, 128, 3, 3)), "decoder.up.1.upsamplers.0.conv") is False
|
|
# conv_out / proj_out / norm_out excluded by NAME even with %16 channels.
|
|
assert ff(_mod(nn.Conv2d, (16, 128, 3, 3)), "decoder.conv_out") is False
|
|
assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.conv_norm_out") is False
|
|
assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.norm_out.conv") is False
|
|
assert ff(_mod(nn.Linear, (512, 512)), "decoder.proj_out") is False
|
|
# A non-conv/linear module (e.g. a GroupNorm) is excluded outright.
|
|
class _GroupNorm:
|
|
pass
|
|
|
|
gn = _GroupNorm()
|
|
gn.weight = _Weight((128,))
|
|
assert ff(gn, "decoder.mid_block.resnets.0.norm1") is False
|
|
# A weight with dim() < 2 (a 1D param) is excluded.
|
|
assert ff(_mod(nn.Conv2d, (128,)), "decoder.some.bias_only") is False
|
|
|
|
|
|
def test_cast_vae_fp8_layerwise_skips_head_and_norms(monkeypatch):
|
|
# The layerwise storage cast passes the decoder head + norm tokens through to diffusers'
|
|
# skip_modules_pattern (on top of the diffusers default), so they stay dense.
|
|
_stub_torch(monkeypatch)
|
|
recorder: list = []
|
|
_stub_diffusers(monkeypatch, recorder)
|
|
vae = object()
|
|
_cast_vae_fp8(vae, _target())
|
|
assert len(recorder) == 1
|
|
_, mod, kw = recorder[0]
|
|
assert mod is vae
|
|
assert kw["storage_dtype"] == "float8_e4m3fn"
|
|
assert kw["compute_dtype"] == "bfloat16"
|
|
skip = kw["skip_modules_pattern"]
|
|
# The diffusers default is preserved and the keep-dense tokens are appended.
|
|
assert "norm" in skip
|
|
for tok in ("conv_out", "proj_out", "conv_norm_out", "norm_out"):
|
|
assert tok in skip
|
|
|
|
|
|
# ── apply (quantize_vae) ────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_quantize_vae_disabled_returns_none(monkeypatch):
|
|
pipe = types.SimpleNamespace(vae = object())
|
|
assert quantize_vae(pipe, _target(), mode = None) is None
|
|
assert quantize_vae(pipe, _target(), mode = "none") is None
|
|
|
|
|
|
def test_quantize_vae_force_fp32_stays_dense(monkeypatch):
|
|
# A force-fp32 (Wan) family never casts, for an explicit scheme or auto.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("must not cast"))
|
|
monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("must not cast"))
|
|
pipe = types.SimpleNamespace(vae = object())
|
|
assert quantize_vae(pipe, _target(), mode = "fp8", force_fp32 = True) is None
|
|
assert quantize_vae(pipe, _target(), mode = "auto", force_fp32 = True) is None
|
|
|
|
|
|
def test_quantize_vae_offload_skips_fp8_dynamic(monkeypatch):
|
|
# Explicit fp8_dynamic under offload is skipped (torchao tensors reject Module.to());
|
|
# layerwise fp8 still engages.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
|
|
monkeypatch.setattr(
|
|
vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("torchao must not run under offload")
|
|
)
|
|
pipe = types.SimpleNamespace(vae = object())
|
|
assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", offload_active = True) is None
|
|
fp8_calls: list = []
|
|
monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: fp8_calls.append(v))
|
|
assert quantize_vae(pipe, _target(), mode = "fp8", offload_active = True) == VAE_QUANT_FP8
|
|
assert len(fp8_calls) == 1
|
|
|
|
|
|
def test_quantize_vae_unsupported_hw_is_noop(monkeypatch):
|
|
# An explicit scheme on hardware that does not support it applies nothing.
|
|
_stub_torch(monkeypatch, cc = (8, 0))
|
|
_allow_vae(monkeypatch, set())
|
|
monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("must not cast"))
|
|
pipe = types.SimpleNamespace(vae = object())
|
|
assert quantize_vae(pipe, _target(), mode = "fp8") is None
|
|
|
|
|
|
def test_quantize_vae_none_vae_is_noop(monkeypatch):
|
|
# A pipeline with no VAE attribute is a best-effort no-op even when the mode is supported.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8})
|
|
pipe = types.SimpleNamespace() # no .vae
|
|
assert quantize_vae(pipe, _target(), mode = "fp8") is None
|
|
|
|
|
|
def test_quantize_vae_explicit_fp8_applies(monkeypatch):
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8})
|
|
calls: list = []
|
|
monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: calls.append(v))
|
|
vae = object()
|
|
pipe = types.SimpleNamespace(vae = vae)
|
|
assert quantize_vae(pipe, _target(), mode = "fp8") == VAE_QUANT_FP8
|
|
assert calls == [vae]
|
|
|
|
|
|
def test_quantize_vae_tolerates_caster_failure(monkeypatch):
|
|
# The caster raising leaves the VAE dense (best-effort) -> None.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8})
|
|
|
|
def _boom(v, t):
|
|
raise RuntimeError("fp8 unsupported for this layer")
|
|
|
|
monkeypatch.setattr(vq, "_cast_vae_fp8", _boom)
|
|
pipe = types.SimpleNamespace(vae = object())
|
|
assert quantize_vae(pipe, _target(), mode = "fp8") is None
|
|
|
|
|
|
def test_quantize_vae_auto_resolves_and_applies(monkeypatch):
|
|
# End-to-end: mode="auto" resolves via the ladder then applies the resolved caster.
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
_stub_capability(monkeypatch, (10, 0))
|
|
_allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
|
|
monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device: True)
|
|
calls: list = []
|
|
monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: calls.append(v))
|
|
vae = object()
|
|
pipe = types.SimpleNamespace(vae = vae)
|
|
assert quantize_vae(pipe, _target(), mode = "auto", family = "flux.1") == VAE_QUANT_FP8_DYNAMIC
|
|
assert calls == [vae]
|