Auto-quantize the VAE (image/video decoder) to fp8 (default on, gated)
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.
This commit is contained in:
Daniel Han
parent
e7168ef34c
commit
ca8f415a7e
7 changed files with 805 additions and 3 deletions
|
|
@ -91,6 +91,7 @@ from .diffusion_cache import (
|
|||
normalize_transformer_cache,
|
||||
)
|
||||
from .diffusion_precision import TE_QUANT_AUTO, normalize_te_quant, quantize_text_encoders
|
||||
from .diffusion_vae_quant import VAE_QUANT_AUTO, normalize_vae_quant, quantize_vae
|
||||
from .diffusion_prequant import (
|
||||
load_prequantized_transformer,
|
||||
resolve_prequant_source,
|
||||
|
|
@ -388,6 +389,10 @@ class _LoadState:
|
|||
backend_flags_before: Optional[dict] = None
|
||||
# Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C).
|
||||
text_encoder_quant: Optional[str] = None
|
||||
# VAE quantisation actually engaged: "fp8" (layerwise storage) | "fp8_dynamic" (torchao
|
||||
# conv compute) | None. When set, the resident VAE holds fp8 tensor subclasses, so the
|
||||
# img2img/inpaint _align_vae_dtype re-cast must be skipped (it would corrupt them).
|
||||
vae_quant: Optional[str] = None
|
||||
# Transformer quant actually engaged on the opt-in dense fast path: "int8" | "fp8"
|
||||
# | "nvfp4" | "mxfp8" | None. None means the default GGUF transformer was loaded.
|
||||
transformer_quant: Optional[str] = None
|
||||
|
|
@ -801,6 +806,7 @@ class DiffusionBackend:
|
|||
memory_mode: Optional[str] = None,
|
||||
speed_mode: Optional[str] = None,
|
||||
text_encoder_quant: Optional[str] = None,
|
||||
vae_quant: Optional[str] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
transformer_quant_fast_accum: Optional[bool] = None,
|
||||
transformer_prequant_path: Optional[str] = None,
|
||||
|
|
@ -849,6 +855,7 @@ class DiffusionBackend:
|
|||
memory_mode = memory_mode,
|
||||
speed_mode = speed_mode,
|
||||
text_encoder_quant = text_encoder_quant,
|
||||
vae_quant = vae_quant,
|
||||
transformer_quant = transformer_quant,
|
||||
transformer_quant_fast_accum = transformer_quant_fast_accum,
|
||||
transformer_prequant_path = transformer_prequant_path,
|
||||
|
|
@ -1167,6 +1174,7 @@ class DiffusionBackend:
|
|||
memory_mode: Optional[str] = None,
|
||||
speed_mode: Optional[str] = None,
|
||||
text_encoder_quant: Optional[str] = None,
|
||||
vae_quant: Optional[str] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
transformer_quant_fast_accum: Optional[bool] = None,
|
||||
transformer_prequant_path: Optional[str] = None,
|
||||
|
|
@ -1208,12 +1216,18 @@ class DiffusionBackend:
|
|||
normalize_attention_backend(attention_backend)
|
||||
normalize_transformer_cache(transformer_cache)
|
||||
normalize_te_quant(text_encoder_quant)
|
||||
normalize_vae_quant(vae_quant)
|
||||
# text_encoder_quant tri-state, mirroring transformer_quant: UNSET (None / "") -> auto,
|
||||
# which picks the best accurate TE scheme for this GPU + family (fp8_dynamic / int8 /
|
||||
# layerwise fp8) or stays dense when none qualifies. An explicit "none"/"off" pins the
|
||||
# encoder dense; an explicit scheme forces it. So the shipped default is auto.
|
||||
if text_encoder_quant is None or str(text_encoder_quant).strip() == "":
|
||||
text_encoder_quant = TE_QUANT_AUTO
|
||||
# vae_quant tri-state, same contract: UNSET -> auto (fp8_dynamic conv compute on resident
|
||||
# fp8-GEMM silicon that passes the conv probe, else layerwise fp8, else dense); none/off ->
|
||||
# dense; an explicit scheme forces it.
|
||||
if vae_quant is None or str(vae_quant).strip() == "":
|
||||
vae_quant = VAE_QUANT_AUTO
|
||||
# For a full pipeline the repo itself supplies every component, so it is its
|
||||
# own base; the single-file kinds resolve the companion base diffusers repo.
|
||||
base = (
|
||||
|
|
@ -1731,6 +1745,18 @@ class DiffusionBackend:
|
|||
offload_active = plan.offload_policy != OFFLOAD_NONE,
|
||||
logger = logger,
|
||||
)
|
||||
# Quantise the dense VAE (opt-in fp8 layerwise / fp8_dynamic torchao conv),
|
||||
# also before placement so the offload hooks move the smaller weights. The
|
||||
# image families do not force-fp32 their VAE; auto skips torchao under offload.
|
||||
vae_quant_engaged = quantize_vae(
|
||||
pipe,
|
||||
target,
|
||||
mode = vae_quant,
|
||||
family = fam.name,
|
||||
offload_active = plan.offload_policy != OFFLOAD_NONE,
|
||||
force_fp32 = False,
|
||||
logger = logger,
|
||||
)
|
||||
|
||||
# Apply the placement planned above (from MEASURED free device memory vs
|
||||
# the model's estimated resident size). apply_memory_plan returns the
|
||||
|
|
@ -1778,6 +1804,15 @@ class DiffusionBackend:
|
|||
if text_encoder_quant == TE_QUANT_AUTO
|
||||
else "requested",
|
||||
),
|
||||
"vae_quant": (
|
||||
vae_quant,
|
||||
vae_quant_engaged or "off",
|
||||
"dense (no accurate scheme for this GPU / disabled)"
|
||||
if vae_quant_engaged is None
|
||||
else "auto-selected for this GPU + family"
|
||||
if vae_quant == VAE_QUANT_AUTO
|
||||
else "requested",
|
||||
),
|
||||
"attention_backend": (
|
||||
attention_backend,
|
||||
attention_engaged or "native",
|
||||
|
|
@ -1823,6 +1858,7 @@ class DiffusionBackend:
|
|||
speed_optims = tuple(k for k, v in speed_applied.items() if v),
|
||||
backend_flags_before = backend_flags_before,
|
||||
text_encoder_quant = te_quant,
|
||||
vae_quant = vae_quant_engaged,
|
||||
transformer_quant = transformer_quant_engaged,
|
||||
attention_backend = attention_engaged,
|
||||
attention_request = attention_backend,
|
||||
|
|
@ -2231,7 +2267,9 @@ class DiffusionBackend:
|
|||
return pipe
|
||||
|
||||
@staticmethod
|
||||
def _align_vae_dtype(pipe: Any, denoiser_attr: str = "transformer") -> None:
|
||||
def _align_vae_dtype(
|
||||
pipe: Any, denoiser_attr: str = "transformer", vae_quant: Optional[str] = None
|
||||
) -> None:
|
||||
"""Cast the VAE to the denoiser's compute dtype before an image-conditioned
|
||||
call. The img2img/inpaint pipelines VAE-encode the input image at the text-
|
||||
encoder dtype (bf16), but a prior txt2img DECODE may have left the shared VAE
|
||||
|
|
@ -2239,7 +2277,14 @@ class DiffusionBackend:
|
|||
(bf16 image vs fp32 VAE). Re-aligning here is safe: our families run bf16 or
|
||||
fp32 only (the fp16 guard promotes fp16), and a later txt2img decode re-upcasts
|
||||
as needed. ``denoiser_attr`` is ``pipe.transformer`` for DiT families and
|
||||
``pipe.unet`` for SDXL. Best-effort; a no-op when already aligned."""
|
||||
``pipe.unet`` for SDXL. Best-effort; a no-op when already aligned.
|
||||
|
||||
Skipped when ``vae_quant`` engaged: a quantised VAE holds fp8 tensor subclasses
|
||||
that mishandle ``.to(dtype=...)`` (torchao rejects it), so the re-cast would
|
||||
corrupt the weights. The VAE already runs at the compute dtype under fp8, so the
|
||||
alignment is unnecessary there anyway."""
|
||||
if vae_quant is not None:
|
||||
return
|
||||
denoiser = getattr(pipe, denoiser_attr, None)
|
||||
vae = getattr(pipe, "vae", None)
|
||||
if denoiser is None or vae is None:
|
||||
|
|
@ -2724,7 +2769,9 @@ class DiffusionBackend:
|
|||
if init_pil is not None:
|
||||
# Keep the VAE encode dtype consistent with the input image.
|
||||
# state.family is always a DiffusionFamily, which defines denoiser_attr.
|
||||
self._align_vae_dtype(pipe, state.family.denoiser_attr)
|
||||
# A quantised VAE (fp8 tensor subclasses) must NOT be re-cast, so pass the
|
||||
# engaged scheme through to skip the re-align in that case.
|
||||
self._align_vae_dtype(pipe, state.family.denoiser_attr, state.vae_quant)
|
||||
|
||||
# Pipelines vary in which kwargs they accept (img2img derives size from the
|
||||
# input image and may reject width/height; a distilled pipe may take no
|
||||
|
|
@ -3007,6 +3054,7 @@ class DiffusionBackend:
|
|||
"speed_mode": None,
|
||||
"speed_optims": [],
|
||||
"text_encoder_quant": None,
|
||||
"vae_quant": None,
|
||||
"transformer_quant": None,
|
||||
"attention_backend": None,
|
||||
"transformer_cache": None,
|
||||
|
|
@ -3032,6 +3080,7 @@ class DiffusionBackend:
|
|||
"speed_mode": state.speed_mode,
|
||||
"speed_optims": list(state.speed_optims),
|
||||
"text_encoder_quant": state.text_encoder_quant,
|
||||
"vae_quant": state.vae_quant,
|
||||
"transformer_quant": state.transformer_quant,
|
||||
"attention_backend": state.attention_backend,
|
||||
"transformer_cache": state.transformer_cache,
|
||||
|
|
|
|||
303
studio/backend/core/inference/diffusion_vae_quant.py
Normal file
303
studio/backend/core/inference/diffusion_vae_quant.py
Normal file
|
|
@ -0,0 +1,303 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Auto low-precision casting of the diffusion pipeline's VAE (image / video decoder).
|
||||
|
||||
The transformer and the companion text encoder both quantise, but the VAE loads dense
|
||||
and is a real resident chunk (the FLUX.2 / Qwen conv VAEs run to a few GB, the video
|
||||
Conv3d VAEs more). Unlike the encoders it is CONVOLUTIONAL, so the Linear-only torchao
|
||||
int8 path does not apply -- there is no int8 Conv3d kernel. Exactly two schemes fit:
|
||||
|
||||
fp8_dynamic - torchao dynamic fp8 COMPUTE with PER-TENSOR granularity. Its
|
||||
Float8DynamicActivationFloat8WeightConfig quantises Conv2d (4D) and
|
||||
Conv3d (5D) weights, not just Linear, and torchao auto-skips any conv
|
||||
whose C_out / C_in is not a multiple of 16 (so the 3-channel RGB
|
||||
``conv_out`` head stays dense). Needs fp8-GEMM silicon (cc >= 8.9) and a
|
||||
resident (non-offloaded) VAE -- the fp8 tensor subclasses reject the
|
||||
Module.to() an offload hook uses.
|
||||
fp8 - diffusers layerwise casting: 8-bit (e4m3) STORAGE, upcast per layer to the
|
||||
compute dtype. Storage-only, so it runs on ANY conv (2D / 3D) on any
|
||||
fp8-capable card (cc >= 8.9) and survives group offload.
|
||||
|
||||
There is no int8 (no Conv3d int8 kernel) and no nvfp4 for the VAE. ``auto`` (the loader
|
||||
default) walks (fp8_dynamic, fp8): fp8_dynamic on resident data-center / Ada+ silicon that
|
||||
passes a live conv smoke probe, else layerwise fp8, else dense. ``none``/``off`` keeps the
|
||||
VAE dense bf16; an explicit scheme forces it (re-gated). A quantised VAE MUST NOT be
|
||||
``.to(dtype=...)``'d afterwards (the fp8 tensor subclasses mishandle it), so the loader
|
||||
skips the img2img/inpaint VAE re-align when a scheme engaged. torch / diffusers / torchao
|
||||
are imported lazily so the module stays importable in a no-torch runtime.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
VAE_QUANT_FP8 = "fp8"
|
||||
VAE_QUANT_FP8_DYNAMIC = "fp8_dynamic"
|
||||
VAE_QUANT_AUTO = "auto"
|
||||
# Concrete schemes (excludes "auto"): "auto" resolves to one of these via
|
||||
# select_vae_quant_scheme, and these are the values the casters dispatch on.
|
||||
VAE_QUANT_MODES = (VAE_QUANT_FP8, VAE_QUANT_FP8_DYNAMIC)
|
||||
|
||||
# The RGB decoder head + output normalisations stay dense. torchao's fp8_dynamic path
|
||||
# already skips the 3-channel ``conv_out`` via its C_out/C_in %16 rule, but the layerwise
|
||||
# fp8 path has no such rule, so name them explicitly for both. Substring-matched against the
|
||||
# module fqn: "conv_out"/"proj_out" (the pixel head), "conv_norm_out"/"norm_out" (the head's
|
||||
# group-norm), whose low-magnitude outputs the coarse fp8 grid would band.
|
||||
_VAE_KEEP_DENSE_TOKENS = ("conv_out", "proj_out", "conv_norm_out", "norm_out")
|
||||
|
||||
# Best-first ``auto`` order: fp8_dynamic (compute fp8 on the conv tensor cores) leads; layerwise
|
||||
# ``fp8`` (storage-only) is the universal fallback and the sole scheme that survives group offload.
|
||||
_VAE_AUTO_LADDER = (VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8)
|
||||
|
||||
# VAEs whose activation ranges break a scheme at the MODEL level (measured decoded-image
|
||||
# LPIPS / SSIM vs the dense bf16 VAE). Populated from the accuracy sweep; a denied scheme is
|
||||
# skipped by ``auto`` and refused when requested explicitly. Empty by default -- the
|
||||
# vae_force_fp32 video families are gated separately at the loader (they never quantise).
|
||||
_VAE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {}
|
||||
|
||||
# Cache of device -> bool for the fp8_dynamic conv smoke probe (run once per device).
|
||||
_VAE_DYNAMIC_PROBE_CACHE: dict[str, bool] = {}
|
||||
|
||||
|
||||
def normalize_vae_quant(value: Optional[str]) -> Optional[str]:
|
||||
"""Lower/strip a requested VAE quant; None / "" / "none" / "off" -> None,
|
||||
"auto" -> "auto" (resolved later by select_vae_quant_scheme).
|
||||
|
||||
Raises ValueError for an unsupported value so a bad request is rejected cheaply."""
|
||||
if value is None:
|
||||
return None
|
||||
normalized = str(value).strip().lower().replace("-", "_")
|
||||
if not normalized or normalized in ("none", "off"):
|
||||
return None
|
||||
if normalized == VAE_QUANT_AUTO:
|
||||
return VAE_QUANT_AUTO
|
||||
if normalized not in VAE_QUANT_MODES:
|
||||
raise ValueError(
|
||||
f"Unsupported vae_quant '{value}'. Use one of: "
|
||||
f"{', '.join((VAE_QUANT_AUTO,) + VAE_QUANT_MODES)}, none/off."
|
||||
)
|
||||
return normalized
|
||||
|
||||
|
||||
def vae_quant_supported(target: Any, mode: str) -> bool:
|
||||
"""Whether ``mode`` is usable for ``target``: a CUDA device with a bf16 compute dtype, plus
|
||||
the fp8 dtype (fp8 layerwise storage) and, for fp8_dynamic, fp8-GEMM silicon (Ada sm_89+ /
|
||||
Hopper / Blackwell). There is no int8 / nvfp4 VAE scheme, so both modes need fp8."""
|
||||
if getattr(target, "device", None) != "cuda":
|
||||
return False
|
||||
try:
|
||||
import torch
|
||||
|
||||
if getattr(target, "dtype", None) is not torch.bfloat16:
|
||||
return False
|
||||
if mode == VAE_QUANT_FP8:
|
||||
return hasattr(torch, "float8_e4m3fn")
|
||||
if mode == VAE_QUANT_FP8_DYNAMIC:
|
||||
# Compute fp8 (torch._scaled_mm on the conv weights) needs fp8-GEMM silicon.
|
||||
return hasattr(torch, "float8_e4m3fn") and torch.cuda.get_device_capability() >= (8, 9)
|
||||
except Exception:
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
def _vae_family_denied(family: Optional[str], scheme: str) -> bool:
|
||||
return scheme in _VAE_FAMILY_SCHEME_DENY.get((family or "").strip().lower(), frozenset())
|
||||
|
||||
|
||||
def _vae_fp8_dynamic_probe(device: str) -> bool:
|
||||
"""True iff torchao's fp8_dynamic CONV path runs on this build: quantise a tiny
|
||||
Conv2d(16, 16, 1) (channels a multiple of 16 so torchao does not skip it) with the
|
||||
PerTensor fp8 config and run one forward. Cached per device. This makes ``auto`` robust
|
||||
to a torchao build whose Float8 config lacks the conv path (the Linear-only fp8 the
|
||||
transformer probes does not prove the conv path works) -- it fails here and the ladder
|
||||
falls to layerwise fp8 rather than crashing at the first decode."""
|
||||
if device in _VAE_DYNAMIC_PROBE_CACHE:
|
||||
return _VAE_DYNAMIC_PROBE_CACHE[device]
|
||||
ok = False
|
||||
try:
|
||||
import torch
|
||||
from torch import nn
|
||||
from torchao.quantization import (
|
||||
Float8DynamicActivationFloat8WeightConfig,
|
||||
PerTensor,
|
||||
quantize_,
|
||||
)
|
||||
|
||||
conv = nn.Conv2d(16, 16, 1).to(device = device, dtype = torch.bfloat16)
|
||||
quantize_(
|
||||
conv,
|
||||
Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()),
|
||||
filter_fn = lambda m, fqn = "": isinstance(m, nn.Conv2d),
|
||||
)
|
||||
x = torch.randn(1, 16, 4, 4, device = device, dtype = torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
conv(x)
|
||||
torch.cuda.synchronize()
|
||||
ok = True
|
||||
except Exception:
|
||||
ok = False
|
||||
_VAE_DYNAMIC_PROBE_CACHE[device] = ok
|
||||
return ok
|
||||
|
||||
|
||||
def select_vae_quant_scheme(
|
||||
target: Any,
|
||||
requested: Optional[str],
|
||||
*,
|
||||
family: Optional[str] = None,
|
||||
offload_active: bool = False,
|
||||
force_fp32: bool = False,
|
||||
) -> Optional[str]:
|
||||
"""Resolve the concrete VAE scheme to apply, or None to stay dense bf16.
|
||||
|
||||
An explicit scheme is returned as-is unless family-denied (``quantize_vae`` re-gates it).
|
||||
``force_fp32`` (a family that runs its VAE in fp32) always stays dense. ``auto`` walks
|
||||
(fp8_dynamic, fp8) and returns the first that: survives the active offload policy (torchao
|
||||
fp8_dynamic tensors reject the Module.to() an offload hook uses -> only layerwise fp8 under
|
||||
offload), is not family-denied, is hardware-supported, and (fp8_dynamic only) passes a real
|
||||
conv smoke probe. Returns None when nothing qualifies (e.g. no CUDA / pre-Ada)."""
|
||||
requested = normalize_vae_quant(requested)
|
||||
if requested is None or force_fp32:
|
||||
return None
|
||||
if requested != VAE_QUANT_AUTO:
|
||||
return None if _vae_family_denied(family, requested) else requested
|
||||
from .diffusion_transformer_quant import _capability
|
||||
|
||||
if _capability() is None:
|
||||
return None
|
||||
device = str(getattr(target, "device", "cuda"))
|
||||
for scheme in _VAE_AUTO_LADDER:
|
||||
# torchao fp8_dynamic produces tensor subclasses that reject Module.to(); an offload hook
|
||||
# moves the VAE that way and hard-crashes, so under offload only layerwise fp8 engages.
|
||||
if offload_active and scheme == VAE_QUANT_FP8_DYNAMIC:
|
||||
continue
|
||||
if _vae_family_denied(family, scheme):
|
||||
continue
|
||||
if not vae_quant_supported(target, scheme):
|
||||
continue
|
||||
# fp8_dynamic additionally needs the torchao CONV fp8 kernel to actually run on this build.
|
||||
if scheme == VAE_QUANT_FP8_DYNAMIC and not _vae_fp8_dynamic_probe(device):
|
||||
continue
|
||||
return scheme
|
||||
return None
|
||||
|
||||
|
||||
def quantize_vae(
|
||||
pipe: Any,
|
||||
target: Any,
|
||||
*,
|
||||
mode: Optional[str],
|
||||
family: Optional[str] = None,
|
||||
offload_active: bool = False,
|
||||
force_fp32: bool = False,
|
||||
logger: Any = None,
|
||||
) -> Optional[str]:
|
||||
"""Quantise ``pipe.vae`` in place with ``mode`` (auto / fp8 / fp8_dynamic). Returns the scheme
|
||||
actually engaged, or None when disabled, unsupported, force-fp32, family-denied, or no VAE was
|
||||
cast. ``auto`` resolves to the best scheme for the GPU + family via ``select_vae_quant_scheme``;
|
||||
an explicit scheme is re-gated the same way. Under offload the torchao fp8_dynamic mode is
|
||||
skipped (its tensor subclasses reject the ``Module.to()`` an offload hook uses); layerwise fp8
|
||||
still engages. Best-effort: any failure leaves the VAE dense."""
|
||||
mode = normalize_vae_quant(mode)
|
||||
if mode is None:
|
||||
return None
|
||||
if mode == VAE_QUANT_AUTO:
|
||||
mode = select_vae_quant_scheme(
|
||||
target,
|
||||
VAE_QUANT_AUTO,
|
||||
family = family,
|
||||
offload_active = offload_active,
|
||||
force_fp32 = force_fp32,
|
||||
)
|
||||
if mode is None:
|
||||
return None
|
||||
else:
|
||||
# An explicit scheme re-runs the same gates ``auto`` applies in select_vae_quant_scheme.
|
||||
if force_fp32:
|
||||
_note(logger, f"vae '{mode}' skipped: family runs its VAE in fp32")
|
||||
return None
|
||||
if _vae_family_denied(family, mode):
|
||||
_note(logger, f"vae '{mode}' denied for family '{family}' (out-of-bar; staying dense)")
|
||||
return None
|
||||
# Layerwise fp8 is storage-only and streams fine under offload; the torchao fp8_dynamic
|
||||
# tensors reject Module.to(), so skip only that one when an offload policy is active.
|
||||
if offload_active and mode == VAE_QUANT_FP8_DYNAMIC:
|
||||
_note(
|
||||
logger,
|
||||
f"vae '{mode}' skipped under offload (torchao tensors reject Module.to()); "
|
||||
"pin a resident memory mode or use fp8",
|
||||
)
|
||||
return None
|
||||
if not vae_quant_supported(target, mode):
|
||||
return None
|
||||
vae = getattr(pipe, "vae", None)
|
||||
if vae is None:
|
||||
return None
|
||||
try:
|
||||
if mode == VAE_QUANT_FP8_DYNAMIC:
|
||||
_cast_vae_fp8_dynamic(vae, target)
|
||||
else:
|
||||
_cast_vae_fp8(vae, target)
|
||||
return mode
|
||||
except Exception as exc: # noqa: BLE001 — leave the VAE dense
|
||||
_warn(logger, mode, exc)
|
||||
return None
|
||||
|
||||
|
||||
def _cast_vae_fp8_dynamic(vae: Any, target: Any) -> None:
|
||||
# torchao dynamic fp8 COMPUTE with PER-TENSOR granularity (NOT the DiT's per-row config):
|
||||
# Float8DynamicActivationFloat8WeightConfig quantises Conv2d (4D) / Conv3d (5D) weights, not
|
||||
# just Linear. Both channel dims must be a multiple of 16 (torchao skips the conv otherwise,
|
||||
# which already leaves the 3-channel RGB head dense), and the decoder head / output norms are
|
||||
# excluded by name for good measure. A weight.dim() < 2 (e.g. a bias-only leaf) never matches.
|
||||
from torch import nn
|
||||
from torchao.quantization import (
|
||||
Float8DynamicActivationFloat8WeightConfig,
|
||||
PerTensor,
|
||||
quantize_,
|
||||
)
|
||||
|
||||
def filter_fn(module: Any, fqn: str = "") -> bool:
|
||||
if not isinstance(module, (nn.Linear, nn.Conv2d, nn.Conv3d)):
|
||||
return False
|
||||
weight = getattr(module, "weight", None)
|
||||
if weight is None or weight.dim() < 2:
|
||||
return False
|
||||
if weight.shape[0] % 16 != 0 or weight.shape[1] % 16 != 0:
|
||||
return False
|
||||
name = fqn.lower() if fqn else ""
|
||||
return not any(tok in name for tok in _VAE_KEEP_DENSE_TOKENS)
|
||||
|
||||
quantize_(vae, Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()), filter_fn = filter_fn)
|
||||
|
||||
|
||||
def _cast_vae_fp8(vae: Any, target: Any) -> None:
|
||||
import re
|
||||
|
||||
import torch
|
||||
from diffusers.hooks import apply_layerwise_casting
|
||||
from diffusers.hooks.layerwise_casting import DEFAULT_SKIP_MODULES_PATTERN
|
||||
|
||||
# diffusers' layerwise casting stores each supported leaf module's weights in fp8 and upcasts
|
||||
# them per forward. Storage-only, so it works on any conv (2D / 3D) and survives offload. Keep
|
||||
# the RGB decoder head / output norms dense (their low-magnitude pixels band on the coarse fp8
|
||||
# grid); the diffusers default pattern only covers pos/patch embeds. Names are literal
|
||||
# substrings (re.escape) matched against the module fqn.
|
||||
skip = tuple(DEFAULT_SKIP_MODULES_PATTERN) + tuple(re.escape(t) for t in _VAE_KEEP_DENSE_TOKENS)
|
||||
apply_layerwise_casting(
|
||||
vae,
|
||||
storage_dtype = torch.float8_e4m3fn,
|
||||
compute_dtype = target.dtype,
|
||||
skip_modules_pattern = skip,
|
||||
)
|
||||
|
||||
|
||||
def _warn(logger: Any, what: str, exc: Exception) -> None:
|
||||
if logger is not None:
|
||||
logger.warning("diffusion.vae_quant: (%s) failed: %s", what, exc)
|
||||
|
||||
|
||||
def _note(logger: Any, msg: str) -> None:
|
||||
if logger is not None:
|
||||
logger.info("diffusion.vae_quant: %s", msg)
|
||||
|
|
@ -79,6 +79,7 @@ from .diffusion_transformer_quant import (
|
|||
select_transformer_quant_scheme,
|
||||
)
|
||||
from .diffusion_precision import TE_QUANT_AUTO, normalize_te_quant, quantize_text_encoders
|
||||
from .diffusion_vae_quant import VAE_QUANT_AUTO, normalize_vae_quant, quantize_vae
|
||||
from .video_families import (
|
||||
VIDEO_CANCELLED_MSG,
|
||||
VIDEO_NOT_LOADED_MSG,
|
||||
|
|
@ -289,6 +290,10 @@ class _VideoLoadState:
|
|||
# The companion text encoder (UMT5 / Gemma3 / Qwen2.5-VL) loads dense bf16 and is often the
|
||||
# largest resident component; this shrinks it in place, mirroring the image backend.
|
||||
text_encoder_quant: Optional[str] = None
|
||||
# VAE quant actually engaged ("fp8" layerwise | "fp8_dynamic" torchao conv) or None. The
|
||||
# convolutional decoder (Conv2d/Conv3d) shrinks in place; the vae_force_fp32 families
|
||||
# (Wan) never quantise (force_fp32 -> dense). Mirrors the image backend's _LoadState.
|
||||
vae_quant: Optional[str] = None
|
||||
resolved: Optional[dict] = None
|
||||
|
||||
|
||||
|
|
@ -394,6 +399,7 @@ class VideoBackend:
|
|||
model_kind: Optional[str] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
text_encoder_quant: Optional[str] = None,
|
||||
vae_quant: Optional[str] = None,
|
||||
) -> VideoFamily:
|
||||
"""Cheap, network-free validation shared by the route and the load path."""
|
||||
kind = resolve_video_model_kind(gguf_filename, model_kind)
|
||||
|
|
@ -517,6 +523,8 @@ class VideoBackend:
|
|||
# Reject a malformed text_encoder_quant the same way (applies to any load kind: the dense
|
||||
# text encoder is resident for pipeline / gguf / single_file alike).
|
||||
normalize_te_quant(text_encoder_quant)
|
||||
# Same for vae_quant (the dense VAE is resident for every load kind).
|
||||
normalize_vae_quant(vae_quant)
|
||||
_ensure_mp4_encoder_available()
|
||||
return fam
|
||||
|
||||
|
|
@ -537,6 +545,7 @@ class VideoBackend:
|
|||
transformer_cache_threshold: Optional[float] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
text_encoder_quant: Optional[str] = None,
|
||||
vae_quant: Optional[str] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Validate, then run the (slow) load on a daemon thread. Returns at once."""
|
||||
|
|
@ -549,6 +558,7 @@ class VideoBackend:
|
|||
model_kind = model_kind,
|
||||
transformer_quant = transformer_quant,
|
||||
text_encoder_quant = text_encoder_quant,
|
||||
vae_quant = vae_quant,
|
||||
)
|
||||
with self._lock:
|
||||
if self._loading is not None and self._loading.error is None:
|
||||
|
|
@ -573,6 +583,7 @@ class VideoBackend:
|
|||
transformer_cache_threshold = transformer_cache_threshold,
|
||||
transformer_quant = transformer_quant,
|
||||
text_encoder_quant = text_encoder_quant,
|
||||
vae_quant = vae_quant,
|
||||
model_kind = model_kind,
|
||||
_load_token = token,
|
||||
),
|
||||
|
|
@ -886,6 +897,7 @@ class VideoBackend:
|
|||
transformer_cache_threshold: Optional[float] = None,
|
||||
transformer_quant: Optional[str] = None,
|
||||
text_encoder_quant: Optional[str] = None,
|
||||
vae_quant: Optional[str] = None,
|
||||
model_kind: Optional[str] = None,
|
||||
_load_token: Optional[int] = None,
|
||||
_base_local_dir: Optional[str] = None,
|
||||
|
|
@ -901,6 +913,7 @@ class VideoBackend:
|
|||
model_kind = model_kind,
|
||||
transformer_quant = transformer_quant,
|
||||
text_encoder_quant = text_encoder_quant,
|
||||
vae_quant = vae_quant,
|
||||
)
|
||||
kind = resolve_video_model_kind(gguf_filename, model_kind)
|
||||
# text_encoder_quant tri-state (mirrors the image backend + transformer_quant): UNSET
|
||||
|
|
@ -908,6 +921,10 @@ class VideoBackend:
|
|||
# "none"/"off" pins the encoder dense; a scheme forces it. So the shipped default is auto.
|
||||
if text_encoder_quant is None or str(text_encoder_quant).strip() == "":
|
||||
text_encoder_quant = TE_QUANT_AUTO
|
||||
# vae_quant tri-state, same contract. The vae_force_fp32 families (Wan) keep the VAE dense
|
||||
# regardless (quantize_vae's force_fp32 gate), so auto is safe as the shipped default.
|
||||
if vae_quant is None or str(vae_quant).strip() == "":
|
||||
vae_quant = VAE_QUANT_AUTO
|
||||
base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo)
|
||||
|
||||
with self._lock:
|
||||
|
|
@ -1206,6 +1223,19 @@ class VideoBackend:
|
|||
offload_active = plan.offload_policy != "none",
|
||||
logger = logger,
|
||||
)
|
||||
# Quantise the dense convolutional VAE (opt-in fp8 layerwise / fp8_dynamic torchao conv).
|
||||
# The vae_force_fp32 families (Wan) run the VAE in fp32 for numerical stability, so
|
||||
# force_fp32 pins them dense (quantising the fp32 VAE bands the decode); auto skips the
|
||||
# torchao mode under offload. Best-effort: a failure leaves the VAE dense.
|
||||
vae_quant_engaged = quantize_vae(
|
||||
pipe,
|
||||
target,
|
||||
mode = vae_quant,
|
||||
family = fam.name,
|
||||
offload_active = plan.offload_policy != "none",
|
||||
force_fp32 = getattr(fam, "vae_force_fp32", False),
|
||||
logger = logger,
|
||||
)
|
||||
|
||||
# ── optimisation layers, in the image backend's order: step cache FIRST
|
||||
# (compile keys its fullgraph decision off an active cache: FBCache hooks
|
||||
|
|
@ -1383,6 +1413,15 @@ class VideoBackend:
|
|||
if text_encoder_quant_engaged is not None
|
||||
else "not engaged (dense bf16 text encoder loaded)",
|
||||
),
|
||||
"vae_quant": (
|
||||
vae_quant,
|
||||
vae_quant_engaged or "off",
|
||||
"dense VAE quantised in place"
|
||||
if vae_quant_engaged is not None
|
||||
else "not engaged (fp32 VAE family / offload / disabled -> dense)"
|
||||
if getattr(fam, "vae_force_fp32", False)
|
||||
else "not engaged (dense VAE loaded)",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -1416,6 +1455,7 @@ class VideoBackend:
|
|||
cache_threshold = transformer_cache_threshold,
|
||||
transformer_quant = transformer_quant_engaged,
|
||||
text_encoder_quant = text_encoder_quant_engaged,
|
||||
vae_quant = vae_quant_engaged,
|
||||
resolved = resolved,
|
||||
)
|
||||
# Ownership of the globals transferred to _state / _teardown_state.
|
||||
|
|
@ -1770,6 +1810,7 @@ class VideoBackend:
|
|||
"transformer_cache": None,
|
||||
"transformer_quant": None,
|
||||
"text_encoder_quant": None,
|
||||
"vae_quant": None,
|
||||
"has_audio": False,
|
||||
"defaults": None,
|
||||
"resolved": None,
|
||||
|
|
@ -1798,6 +1839,7 @@ class VideoBackend:
|
|||
"transformer_cache": state.transformer_cache,
|
||||
"transformer_quant": state.transformer_quant,
|
||||
"text_encoder_quant": state.text_encoder_quant,
|
||||
"vae_quant": state.vae_quant,
|
||||
"has_audio": fam.has_audio,
|
||||
"defaults": {
|
||||
"steps": default_steps,
|
||||
|
|
|
|||
|
|
@ -1758,6 +1758,15 @@ class DiffusionLoadRequest(BaseModel):
|
|||
"cc>=8.0), or nvfp4 (~4x smaller, Blackwell sm_100+). none/off keeps it dense bf16. A "
|
||||
"memory-vs-quality tradeoff (shifts fine detail); pairs well with balanced mode.",
|
||||
)
|
||||
vae_quant: Optional[Literal["auto", "none", "off", "fp8", "fp8_dynamic"]] = Field(
|
||||
None,
|
||||
description = "Quantise the VAE (image decoder). auto (the default when unset) picks the "
|
||||
"best accurate scheme for this GPU + family: fp8_dynamic (torchao PER-TENSOR fp8 COMPUTE "
|
||||
"on the conv tensor cores, cc >= 8.9, resident only) where a conv probe passes, else "
|
||||
"layerwise fp8 (diffusers 8-bit storage, cc >= 8.9, survives offload), else dense. The "
|
||||
"3-channel RGB head + output norms always stay dense. There is no int8 (no Conv3d int8 "
|
||||
"kernel). none/off keeps the VAE dense bf16; an explicit scheme forces it.",
|
||||
)
|
||||
transformer_quant: Optional[Literal["auto", "none", "off", "int8", "fp8", "nvfp4", "mxfp8"]] = (
|
||||
Field(
|
||||
None,
|
||||
|
|
@ -2117,6 +2126,9 @@ class DiffusionStatusResponse(BaseModel):
|
|||
text_encoder_quant: Optional[str] = Field(
|
||||
None, description = "Text-encoder quantisation engaged: fp8 | nvfp4 | null"
|
||||
)
|
||||
vae_quant: Optional[str] = Field(
|
||||
None, description = "VAE quantisation engaged: fp8 | fp8_dynamic | null"
|
||||
)
|
||||
transformer_quant: Optional[str] = Field(
|
||||
None,
|
||||
description = "Transformer quant engaged on the dense fast path: int8 | fp8 | "
|
||||
|
|
@ -2364,6 +2376,15 @@ class VideoLoadRequest(BaseModel):
|
|||
"for a family without a measured schedule); nvfp4 = torchao 4-bit weight-only (Blackwell "
|
||||
"sm_100+). none/off keeps the encoder dense. Mirrors the image backend's field.",
|
||||
)
|
||||
vae_quant: Optional[Literal["auto", "none", "off", "fp8", "fp8_dynamic"]] = Field(
|
||||
None,
|
||||
description = "Quantise the VAE (video decoder). auto (the default when unset) picks the "
|
||||
"best accurate scheme for this GPU + family: fp8_dynamic (torchao PER-TENSOR fp8 COMPUTE "
|
||||
"on the Conv2d/Conv3d tensor cores, cc >= 8.9, resident only), else layerwise fp8 "
|
||||
"(diffusers 8-bit storage, cc >= 8.9, survives offload), else dense. The fp32-VAE families "
|
||||
"(Wan) always stay dense. No int8 (no Conv3d int8 kernel). none/off keeps it dense; an "
|
||||
"explicit scheme forces it. Mirrors the image backend's field.",
|
||||
)
|
||||
|
||||
@field_validator("attention_backend", mode = "before")
|
||||
@classmethod
|
||||
|
|
@ -2535,6 +2556,11 @@ class VideoStatusResponse(BaseModel):
|
|||
"(null = the dense bf16 encoder is loaded). An int8 request without a per-family "
|
||||
"keep-bf16 schedule is reported as the fp8 it fell back to.",
|
||||
)
|
||||
vae_quant: Optional[str] = Field(
|
||||
None,
|
||||
description = "VAE quant engaged: fp8 | fp8_dynamic | null (null = the dense VAE is "
|
||||
"loaded; the fp32-VAE families always report null).",
|
||||
)
|
||||
has_audio: bool = Field(
|
||||
False, description = "Whether the loaded family produces a synchronized audio track"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -12515,6 +12515,7 @@ async def load_diffusion_model(
|
|||
memory_mode = request.memory_mode,
|
||||
speed_mode = request.speed_mode,
|
||||
text_encoder_quant = request.text_encoder_quant,
|
||||
vae_quant = request.vae_quant,
|
||||
transformer_quant = request.transformer_quant,
|
||||
transformer_quant_fast_accum = request.transformer_quant_fast_accum,
|
||||
transformer_prequant_path = request.transformer_prequant_path,
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ async def load_video_model(
|
|||
model_kind = request.model_kind,
|
||||
transformer_quant = request.transformer_quant,
|
||||
text_encoder_quant = request.text_encoder_quant,
|
||||
vae_quant = request.vae_quant,
|
||||
)
|
||||
# Refuse while training is running: a multi-GB video pipeline would compete
|
||||
# with the training subprocess for VRAM. Mirrors the image-load guard.
|
||||
|
|
@ -125,6 +126,7 @@ async def load_video_model(
|
|||
transformer_cache_threshold = request.transformer_cache_threshold,
|
||||
transformer_quant = request.transformer_quant,
|
||||
text_encoder_quant = request.text_encoder_quant,
|
||||
vae_quant = request.vae_quant,
|
||||
model_kind = request.model_kind,
|
||||
)
|
||||
return VideoStatusResponse(**status_dict)
|
||||
|
|
|
|||
379
studio/backend/tests/test_diffusion_vae_quant.py
Normal file
379
studio/backend/tests/test_diffusion_vae_quant.py
Normal file
|
|
@ -0,0 +1,379 @@
|
|||
# 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]
|
||||
Loading…
Add table
Add a link
Reference in a new issue