06691e1858
The dtype override swapped encoder.__class__ to a dynamic subclass, which breaks transformers' kwargs-based output recording: a fp8-cast Qwen3VLModel stopped returning hidden_states and every krea-2 generation with text_encoder_quant=fp8 crashed at encode_prompt (regression from the HiDream TE4 change; caught by the krea hosted-TE live smoke). The override is now a property shadowed on the ORIGINAL class that prefers a per-instance compute-dtype attribute, so class identity is preserved and uncast instances keep the stock behaviour. The idempotency test now pins exact class identity and the uncast-sibling fallback.
589 lines
28 KiB
Python
589 lines
28 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
|
|
|
|
"""Opt-in low-precision casting of the diffusion pipeline's text encoder(s).
|
|
|
|
The transformer arrives quantised in the GGUF, but the companion text encoder loads dense
|
|
(bf16) and is often the largest resident component (Qwen3 / T5-XXL / Mistral run to many GB).
|
|
This shrinks it in place, with four backends:
|
|
|
|
fp8 - diffusers layerwise casting: 8-bit (e4m3) storage, upcast per layer. ~2x
|
|
smaller. Any fp8-capable CUDA card (cc >= 8.9).
|
|
fp8_dynamic - torchao dynamic fp8 COMPUTE (per-row): keeps the matmul in fp8 on the tensor
|
|
cores (torch._scaled_mm) instead of upcasting. ~2x smaller + speedup; cc >= 8.9.
|
|
int8 - torchao dynamic int8 COMPUTE (per-token act + per-channel weight, _int_mm),
|
|
with per-layer keep-bf16 selection. Degrades on large encoders unless the
|
|
sensitive decoder blocks stay bf16, so applied only for families with a
|
|
measured schedule (else falls back to fp8). ~2x smaller; cc >= 8.0.
|
|
nvfp4 - torchao NVFP4 weight-only: 4-bit float, two-level microscaling, Blackwell
|
|
sm_100+ FP4 cores. ~4x smaller (lowest VRAM) but a steeper quality cost.
|
|
|
|
All keep normalisations / embeddings full precision and are a memory-vs-quality tradeoff.
|
|
``auto`` (the loader default) walks ``select_te_quant_scheme``'s per-GPU ladder for the best
|
|
accurate scheme (fp8_dynamic / int8 / layerwise fp8), falling back to dense; ``none``/``off``
|
|
stays dense, an explicit scheme forces it. They pair well with streamed (group) offload, where
|
|
the resident text encoder dominates the companion footprint. Quantify the quality cost per model
|
|
with scripts/diffusion_quality.py and scripts/diffusion_quant_builder.py. torch / diffusers /
|
|
torchao imported lazily.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
TE_QUANT_FP8 = "fp8"
|
|
TE_QUANT_NVFP4 = "nvfp4"
|
|
TE_QUANT_INT8 = "int8"
|
|
TE_QUANT_FP8_DYNAMIC = "fp8_dynamic"
|
|
TE_QUANT_AUTO = "auto"
|
|
# Concrete schemes (excludes "auto") the casters dispatch on.
|
|
TE_QUANT_MODES = (TE_QUANT_FP8, TE_QUANT_NVFP4, TE_QUANT_INT8, TE_QUANT_FP8_DYNAMIC)
|
|
|
|
# Pipeline attributes that hold a text encoder, in order.
|
|
_TEXT_ENCODER_ATTRS = ("text_encoder", "text_encoder_2", "text_encoder_3")
|
|
|
|
# int8 degrades on large text encoders unless the quant-sensitive decoder blocks stay bf16.
|
|
# Per-family (skip_first, skip_last) blocks to keep dense, from measured hidden-state fidelity
|
|
# (per-token cosine vs bf16 at the consumed layer): keeping first blocks stops early-layer error
|
|
# seeding, last blocks protect the read layer. Families absent have no schedule clearing the bar,
|
|
# so int8 falls back to fp8.
|
|
# qwen-image (Qwen2.5-VL-7B): first+last 6 -> ~0.997 cosine (both ends; outlier-bound).
|
|
# flux.2-dev (Mistral-Small-24B): first 3 -> ~0.98 cosine (early-layer seeding).
|
|
_TE_INT8_SKIP: dict[str, tuple[int, int]] = {
|
|
"qwen-image": (6, 6),
|
|
"qwen-image-edit": (6, 6),
|
|
"flux.2-dev": (3, 0),
|
|
}
|
|
|
|
|
|
def normalize_te_quant(value: Optional[str]) -> Optional[str]:
|
|
"""Lower/strip a requested text-encoder quant; None / "" / "none" / "off" -> None, "auto" ->
|
|
"auto" (resolved later by select_te_quant_scheme). Raises ValueError for an unsupported value."""
|
|
if value is None:
|
|
return None
|
|
normalized = str(value).strip().lower().replace("-", "_")
|
|
if not normalized or normalized in ("none", "off"):
|
|
return None
|
|
if normalized == TE_QUANT_AUTO:
|
|
return TE_QUANT_AUTO
|
|
if normalized not in TE_QUANT_MODES:
|
|
raise ValueError(
|
|
f"Unsupported text_encoder_quant '{value}'. Use one of: "
|
|
f"{', '.join((TE_QUANT_AUTO,) + TE_QUANT_MODES)}, none/off."
|
|
)
|
|
return normalized
|
|
|
|
|
|
def te_quant_supported(target: Any, mode: str) -> bool:
|
|
"""Whether ``mode`` is usable for ``target``: a CUDA bf16 device plus the tensor-core class
|
|
each backend needs -- fp8 dtype (fp8), fp8 GEMM sm_89+ (fp8_dynamic), int8 sm_80+ (int8),
|
|
Blackwell sm_100+ (nvfp4)."""
|
|
if getattr(target, "device", None) != "cuda":
|
|
return False
|
|
try:
|
|
import torch
|
|
|
|
if getattr(target, "dtype", None) is not torch.bfloat16:
|
|
return False
|
|
if mode == TE_QUANT_FP8:
|
|
return hasattr(torch, "float8_e4m3fn")
|
|
if mode == TE_QUANT_FP8_DYNAMIC:
|
|
# fp8 GEMM needs Ada sm_89+ / Hopper / Blackwell.
|
|
return hasattr(torch, "float8_e4m3fn") and torch.cuda.get_device_capability() >= (8, 9)
|
|
if mode == TE_QUANT_INT8:
|
|
return torch.cuda.get_device_capability()[0] >= 8 # int8 cores: Ampere sm_80+
|
|
if mode == TE_QUANT_NVFP4:
|
|
return torch.cuda.get_device_capability()[0] >= 10 # NVFP4 cores: Blackwell sm_100+
|
|
except Exception:
|
|
return False
|
|
return False
|
|
|
|
|
|
# Per-arch preference order for text-encoder ``auto`` (best-first), mirroring the transformer's
|
|
# ``_AUTO_LADDER``. fp8_dynamic (compute fp8) leads on fp8-GEMM silicon; int8 sits second but only
|
|
# engages for a family with a measured keep-bf16 schedule; layerwise ``fp8`` (storage cast) is the
|
|
# universal fallback and the only scheme that survives group offload. nvfp4 stays explicit-only.
|
|
_TE_AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = (
|
|
((8, 9), (TE_QUANT_FP8_DYNAMIC, TE_QUANT_INT8, TE_QUANT_FP8)), # Ada sm_89 / Hopper / Blackwell
|
|
(
|
|
(8, 0),
|
|
(TE_QUANT_INT8, TE_QUANT_FP8),
|
|
), # Ampere sm_80/86: no fp8 GEMM -> int8 or layerwise fp8
|
|
)
|
|
|
|
# Text encoders whose activation ranges break a scheme at the MODEL level (measured hidden-state
|
|
# cosine vs bf16). A denied scheme is skipped by ``auto`` and refused when requested explicitly.
|
|
# int8 already gates on a per-family keep-bf16 schedule (``_TE_INT8_SKIP``), so this covers the
|
|
# rarer case where a scheme breaks the encoder outright.
|
|
# ltx-2 / fp8_dynamic: compute-fp8 on the Gemma3-27B encoder BLACK-FRAMES the whole clip (B200,
|
|
# pairwise vs the dense encoder: mean luma 137.9 -> 0.0, LPIPS 0.78). Layerwise fp8 on the same
|
|
# encoder is near-lossless (LPIPS 0.0043) at the same ~2x shrink, so auto falls through to it.
|
|
_TE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {
|
|
"ltx-2": frozenset({TE_QUANT_FP8_DYNAMIC}),
|
|
}
|
|
|
|
# Families whose AUTO text-encoder quant resolves dense: measured out-of-bar trajectory drift for
|
|
# zero speed win. Unlike the deny table this only steers the AUTO default; an explicit scheme is
|
|
# still honored.
|
|
# wan2.2-t2v-a14b: TE fp8_dynamic ALONE moves the clip to pairwise LPIPS 0.1195 vs the dense-TE
|
|
# stack (B200, 1280x720/33f/50 steps) for 146.7 -> 142.7 s e2e -- the UMT5 encoder runs once per
|
|
# generation, so the 1.03x is noise next to being the dominant accuracy cost. The MoE trajectory
|
|
# amplifies the perturbation ~3x harder than on wan2.2-ti2v-5b (0.0396, kept quantized there for
|
|
# a real 1.09x on its faster DiT).
|
|
_TE_AUTO_DENSE_FAMILIES: frozenset[str] = frozenset(
|
|
# wan2.2-i2v-a14b inherits the T2V entry: same UMT5 encoder and MoE trajectory.
|
|
{"hunyuanvideo-1.5", "hunyuanvideo-1.5-720p", "wan2.2-t2v-a14b", "wan2.2-i2v-a14b"}
|
|
)
|
|
|
|
# Map a TE torchao scheme to the transformer smoke-probe scheme (same GEMM), so ``auto`` degrades
|
|
# gracefully when a build lacks a kernel. Layerwise fp8 has no torchao GEMM to probe.
|
|
_TE_SMOKE_SCHEME = {TE_QUANT_FP8_DYNAMIC: "fp8", TE_QUANT_INT8: "int8", TE_QUANT_NVFP4: "nvfp4"}
|
|
|
|
|
|
def _te_family_denied(family: Optional[str], scheme: str) -> bool:
|
|
return scheme in _TE_FAMILY_SCHEME_DENY.get((family or "").strip().lower(), frozenset())
|
|
|
|
|
|
# nvfp4 TE casts WEIGHT-ONLY (see _cast_nvfp4), a different kernel from the transformer's
|
|
# dynamic-activation NVFP4 GEMM, so it gets its own cached smoke probe.
|
|
_TE_NVFP4_PROBE_CACHE: dict[str, bool] = {}
|
|
|
|
|
|
def _te_nvfp4_weightonly_probe(device: str) -> bool:
|
|
"""True iff weight-only NVFP4 (the ``_cast_nvfp4`` config) runs one forward on this build. The
|
|
transformer ``_smoke_probe`` tests the DYNAMIC-activation NVFP4 GEMM, a different kernel: a
|
|
Blackwell build can carry the weight-only path without the dynamic one, so an explicit TE
|
|
``nvfp4`` needs this dedicated probe. Cached per device."""
|
|
if device in _TE_NVFP4_PROBE_CACHE:
|
|
return _TE_NVFP4_PROBE_CACHE[device]
|
|
ok = False
|
|
try:
|
|
import torch
|
|
from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig
|
|
from torchao.quantization import quantize_
|
|
|
|
from .diffusion_transformer_quant import make_filter_fn
|
|
|
|
lin = torch.nn.Linear(512, 512, bias = False).to(device = device, dtype = torch.bfloat16)
|
|
quantize_(lin, NVFP4WeightOnlyConfig(), filter_fn = make_filter_fn(0))
|
|
x = torch.randn(32, 512, device = device, dtype = torch.bfloat16)
|
|
with torch.no_grad():
|
|
lin(x)
|
|
torch.cuda.synchronize()
|
|
ok = True
|
|
except Exception: # noqa: BLE001 -- an unavailable kernel just means stay dense
|
|
ok = False
|
|
_TE_NVFP4_PROBE_CACHE[device] = ok
|
|
return ok
|
|
|
|
|
|
def _te_scheme_probe(scheme: str, device: str) -> bool:
|
|
"""True iff ``scheme`` runs on this build. Layerwise fp8 (no torchao GEMM) always passes; nvfp4
|
|
probes its weight-only kernel; the other torchao modes reuse the transformer's cached smoke
|
|
test."""
|
|
tq = _TE_SMOKE_SCHEME.get(scheme)
|
|
if tq is None:
|
|
return True
|
|
# nvfp4 casts weight-only, so probe that kernel rather than the transformer's dynamic GEMM.
|
|
if scheme == TE_QUANT_NVFP4:
|
|
return _te_nvfp4_weightonly_probe(device)
|
|
try:
|
|
from .diffusion_transformer_quant import _smoke_probe
|
|
return _smoke_probe(tq, device)
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def select_te_quant_scheme(
|
|
target: Any,
|
|
requested: Optional[str],
|
|
*,
|
|
family: Optional[str] = None,
|
|
offload_active: bool = False,
|
|
) -> Optional[str]:
|
|
"""Resolve the concrete text-encoder scheme to apply, or None to stay dense bf16.
|
|
|
|
An explicit scheme is returned as-is. ``auto`` walks ``_TE_AUTO_LADDER`` for this GPU and
|
|
returns the first scheme that survives offload (torchao needs pinned residency -> only fp8),
|
|
has a keep-bf16 schedule if int8, is not family-denied, is hardware-supported, and passes a
|
|
kernel smoke test. Returns None when nothing qualifies."""
|
|
requested = normalize_te_quant(requested)
|
|
if requested is None or requested != TE_QUANT_AUTO:
|
|
return requested
|
|
# AUTO resolves dense for these families regardless of hardware. TE quant perturbs the
|
|
# CONDITIONING and a multi-step video trajectory amplifies that chaotically: on
|
|
# HunyuanVideo-1.5-720p (B200, 720p/33f/30 steps) TE fp8_dynamic ALONE moves the clip to LPIPS
|
|
# 0.236 vs the bit-exact reference while the rest of the stack sits at 0.052-0.053, for ZERO
|
|
# speed win (35.48 vs 35.36 s e2e). The ~6.7 GB saved isn't worth being the dominant accuracy
|
|
# cost. (VAE fp8 stays in auto: 0.053, decode-only, no trajectory to amplify.)
|
|
if (family or "").strip().lower() in _TE_AUTO_DENSE_FAMILIES:
|
|
return None
|
|
from .diffusion_transformer_quant import _capability, _is_consumer_gpu
|
|
|
|
cap = _capability()
|
|
if cap is None:
|
|
return None
|
|
device = str(getattr(target, "device", "cuda"))
|
|
for floor, schemes in _TE_AUTO_LADDER:
|
|
if cap >= floor:
|
|
# Consumer GDDR parts run int8 full-rate but halve fp8 FP32-accumulate: prefer int8.
|
|
ordered = (
|
|
(TE_QUANT_INT8,) + tuple(s for s in schemes if s != TE_QUANT_INT8)
|
|
if TE_QUANT_INT8 in schemes
|
|
and schemes[0] != TE_QUANT_INT8
|
|
and _is_consumer_gpu(device)
|
|
else schemes
|
|
)
|
|
for scheme in ordered:
|
|
# torchao tensors reject Module.to(); only layerwise fp8 streams, so skip the
|
|
# torchao modes under offload.
|
|
if offload_active and scheme in (
|
|
TE_QUANT_INT8,
|
|
TE_QUANT_FP8_DYNAMIC,
|
|
TE_QUANT_NVFP4,
|
|
):
|
|
continue
|
|
# int8 only clears the bar on a family with a measured keep-bf16 schedule.
|
|
if scheme == TE_QUANT_INT8 and (family or "").lower() not in _TE_INT8_SKIP:
|
|
continue
|
|
if _te_family_denied(family, scheme):
|
|
continue
|
|
if not te_quant_supported(target, scheme):
|
|
continue
|
|
if not _te_scheme_probe(scheme, device):
|
|
continue
|
|
return scheme
|
|
return None
|
|
return None
|
|
|
|
|
|
def quantize_text_encoders(
|
|
pipe: Any,
|
|
target: Any,
|
|
*,
|
|
mode: Optional[str],
|
|
family: Optional[str] = None,
|
|
offload_active: bool = False,
|
|
logger: Any = None,
|
|
) -> Optional[str]:
|
|
"""Quantise each present text encoder in place with ``mode`` (auto / fp8 / fp8_dynamic / int8 /
|
|
nvfp4). Returns the applied mode, or None when disabled, unsupported, or nothing was cast.
|
|
``auto`` resolves via ``select_te_quant_scheme``. ``int8`` needs a per-family keep-bf16 schedule
|
|
(``_TE_INT8_SKIP``); a family without one falls back to ``fp8``. Under offload the torchao modes
|
|
are skipped (their tensors reject Module.to()); layerwise ``fp8`` still engages. Best-effort:
|
|
any failure leaves the encoder dense."""
|
|
mode = normalize_te_quant(mode)
|
|
if mode is None:
|
|
return None
|
|
if mode == TE_QUANT_AUTO:
|
|
mode = select_te_quant_scheme(
|
|
target, TE_QUANT_AUTO, family = family, offload_active = offload_active
|
|
)
|
|
if mode is None:
|
|
return None
|
|
skip: Optional[tuple[int, int]] = None
|
|
if mode == TE_QUANT_INT8:
|
|
skip = _TE_INT8_SKIP.get((family or "").lower())
|
|
if skip is None:
|
|
_note(logger, f"int8 has no keep-bf16 schedule for family '{family}'; using fp8")
|
|
mode = TE_QUANT_FP8
|
|
# A denied scheme is refused even when requested explicitly. Gate the FINAL concrete mode so
|
|
# an int8 -> fp8 fallback is re-checked too (auto already filtered in select_te_quant_scheme).
|
|
if _te_family_denied(family, mode):
|
|
_note(
|
|
logger,
|
|
f"text-encoder '{mode}' denied for family '{family}' (out-of-bar; staying dense)",
|
|
)
|
|
return None
|
|
# The torchao modes (int8, fp8_dynamic, nvfp4) produce tensors that reject Module.to(), which
|
|
# an offload placement crashes on; layerwise fp8 streams fine. Skip the torchao modes here.
|
|
if offload_active and mode in (TE_QUANT_INT8, TE_QUANT_FP8_DYNAMIC, TE_QUANT_NVFP4):
|
|
_note(
|
|
logger,
|
|
f"text-encoder '{mode}' skipped under offload (torchao tensors reject Module.to()); "
|
|
"pin a resident memory mode or use fp8",
|
|
)
|
|
return None
|
|
if not te_quant_supported(target, mode):
|
|
return None
|
|
# An EXPLICIT torchao mode can clear the capability gate yet fail the real GEMM on a build
|
|
# where quantize_ wraps the encoder but the kernel is broken (the caster's try/except catches
|
|
# the cast, not the first forward). Run the auto ladder's kernel smoke test so a failing kernel
|
|
# falls back to dense here instead of crashing at generation. No-op for layerwise fp8.
|
|
device = str(getattr(target, "device", "cuda"))
|
|
if not _te_scheme_probe(mode, device):
|
|
_note(logger, f"text-encoder '{mode}' failed the kernel smoke test; staying dense")
|
|
return None
|
|
if mode == TE_QUANT_INT8:
|
|
first, last = skip # type: ignore[misc]
|
|
|
|
def caster(enc: Any, tgt: Any) -> None:
|
|
_cast_int8_selective(enc, tgt, first, last)
|
|
elif mode == TE_QUANT_FP8_DYNAMIC:
|
|
caster = _cast_fp8_dynamic
|
|
elif mode == TE_QUANT_NVFP4:
|
|
caster = _cast_nvfp4
|
|
else:
|
|
caster = _cast_fp8
|
|
cast: list[str] = []
|
|
for attr in _TEXT_ENCODER_ATTRS:
|
|
encoder = getattr(pipe, attr, None)
|
|
if encoder is None:
|
|
continue
|
|
try:
|
|
caster(encoder, target)
|
|
cast.append(attr)
|
|
except Exception as exc: # noqa: BLE001 — leave this encoder dense
|
|
# A mid-pass caster failure may have left the encoder PARTIALLY quantized (can't
|
|
# run as dense), so fail the load for that; a clean miss stays best-effort dense.
|
|
# raise_if_partially_quantized only recognises torchao parameter subclasses, so it
|
|
# cannot see a partial layerwise fp8 mutation (diffusers apply_layerwise_casting installs
|
|
# upcast hooks + fp8 storage in place, leaving no torchao params). Detect a leftover
|
|
# layerwise hook directly and fail closed there too; a clean failure stays dense.
|
|
from .diffusion_transformer_quant import raise_if_partially_quantized
|
|
|
|
if mode == TE_QUANT_FP8 and _has_layerwise_casting(encoder):
|
|
raise RuntimeError(
|
|
f"text_encoder_quant fp8:{attr} failed after partially installing layerwise "
|
|
"casting (leftover fp8 hooks); reload the model instead of a dense fallback "
|
|
f"(original error: {exc})"
|
|
) from exc
|
|
raise_if_partially_quantized(encoder, what = f"text_encoder_quant {mode}:{attr}", exc = exc)
|
|
_warn(logger, f"{mode}:{attr}", exc)
|
|
return mode if cast else None
|
|
|
|
|
|
def _te_exclude_tokens(encoder: Any) -> tuple[str, ...]:
|
|
"""fqn tokens whose Linears stay bf16 in a torchao TE quant: the VLM vision tower, the unused
|
|
lm_head, and the encoder's own fp32-kept modules (T5 ``wo``, which explodes in low precision)."""
|
|
tokens = ["visual", "vision_tower", "lm_head"]
|
|
tokens += [str(m).lower() for m in (getattr(encoder, "_keep_in_fp32_modules", None) or ())]
|
|
return tuple(dict.fromkeys(tokens))
|
|
|
|
|
|
def _keep_bf16_block_fqns(encoder: Any, skip_first: int, skip_last: int) -> set[str]:
|
|
"""FQNs of decoder blocks to keep bf16: the first ``skip_first`` and last ``skip_last`` of
|
|
each top-level ``nn.ModuleList`` stack. Structural, so no per-architecture table."""
|
|
import torch
|
|
|
|
keep: set[str] = set()
|
|
for name, module in encoder.named_modules():
|
|
if not isinstance(module, torch.nn.ModuleList):
|
|
continue
|
|
n = len(module)
|
|
if n <= skip_first + skip_last:
|
|
continue
|
|
for i in list(range(skip_first)) + list(range(n - skip_last, n)):
|
|
keep.add(f"{name}.{i}" if name else str(i))
|
|
return keep
|
|
|
|
|
|
def _cast_int8_selective(encoder: Any, target: Any, skip_first: int, skip_last: int) -> None:
|
|
# torchao dynamic int8 on the FLOP-heavy Linears, keeping the first/last decoder blocks (and
|
|
# vision tower / lm_head / T5 wo) bf16. Reuses the transformer-quant factory so config never drifts.
|
|
from torchao.quantization import quantize_
|
|
from .diffusion_transformer_quant import (
|
|
TQ_INT8,
|
|
DEFAULT_MIN_LINEAR_FEATURES,
|
|
_make_quant_config,
|
|
make_filter_fn,
|
|
exclude_tokens_for_scheme,
|
|
)
|
|
|
|
base = make_filter_fn(
|
|
DEFAULT_MIN_LINEAR_FEATURES,
|
|
exclude_tokens_for_scheme(TQ_INT8) + _te_exclude_tokens(encoder),
|
|
)
|
|
keep = _keep_bf16_block_fqns(encoder, skip_first, skip_last)
|
|
|
|
def filter_fn(module: Any, fqn: str = "") -> bool:
|
|
if not base(module, fqn):
|
|
return False
|
|
return not any(fqn == k or fqn.startswith(k + ".") for k in keep)
|
|
|
|
quantize_(encoder, _make_quant_config(TQ_INT8), filter_fn = filter_fn)
|
|
|
|
|
|
def _weight_has_zero_output_row(module: Any) -> bool:
|
|
"""True when a Linear's weight has an all-zero OUTPUT row. torchao per-row fp8 derives a
|
|
per-channel scale from that row's amax, so a dead row gives scale 0 -> 0/0 = NaN through the
|
|
forward. Real checkpoints ship such rows: SDXL's text_encoder_2 (OpenCLIP ViT-bigG) has one in
|
|
``text_model.encoder.layers.2.self_attn.out_proj`` -- B200: every fp8_dynamic SDXL render came
|
|
out black until this Linear is left dense. Cheap (one amax per Linear); False on any error."""
|
|
try:
|
|
weight = getattr(module, "weight", None)
|
|
if weight is None or weight.ndim != 2:
|
|
return False
|
|
return bool((weight.abs().amax(dim = -1) == 0).any().item())
|
|
except Exception: # noqa: BLE001 -- unreadable weight: let quantize_ decide
|
|
return False
|
|
|
|
|
|
def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
|
|
# torchao dynamic fp8 COMPUTE, per-row (torch._scaled_mm on the fp8 cores). Unlike layerwise
|
|
# `fp8` this keeps the matmul in fp8 instead of upcasting. Robust across encoder sizes, so no
|
|
# per-layer keep-bf16; only the vision tower / lm_head / T5 wo are excluded.
|
|
from torchao.quantization import quantize_
|
|
from .diffusion_transformer_quant import (
|
|
TQ_FP8,
|
|
DEFAULT_MIN_LINEAR_FEATURES,
|
|
_make_quant_config,
|
|
make_filter_fn,
|
|
)
|
|
|
|
# require_bf16: scaled_mm asserts a bf16 weight, so skip any stray non-bf16 Linear rather than
|
|
# aborting the pass (belt-and-suspenders over the named T5 wo exclusion).
|
|
base = make_filter_fn(
|
|
DEFAULT_MIN_LINEAR_FEATURES, _te_exclude_tokens(encoder), require_bf16 = True
|
|
)
|
|
|
|
# An all-zero output row NaNs under per-row scaling (scale 0 -> 0/0); keep those dense.
|
|
def filter_fn(module: Any, fqn: str = "") -> bool:
|
|
return base(module, fqn) and not _weight_has_zero_output_row(module)
|
|
|
|
quantize_(encoder, _make_quant_config(TQ_FP8), filter_fn = filter_fn)
|
|
|
|
|
|
def _cast_fp8(encoder: 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
|
|
|
|
# Idempotent: a pre-cast encoder (diffusion_te_prequant) arrives with the layerwise hooks
|
|
# already installed, and re-registering the same hook name raises -- which would make
|
|
# quantize_text_encoders report the (actually engaged) cast as failed. Keyed on the explicit
|
|
# completion marker this function sets, NOT on hook presence alone: leftover hooks from a
|
|
# cast that failed mid-pass must still fail closed, not read as "already cast".
|
|
if getattr(encoder, "_unsloth_te_cast_complete", False) and _has_layerwise_hooks(encoder):
|
|
return
|
|
|
|
# Layerwise casting stores each leaf's weights in fp8 and upcasts per forward. Two things on a
|
|
# transformers encoder push an fp8 weight/activation into an op that can't handle it, both
|
|
# crashing only at generation, so skip the offending modules:
|
|
skip = tuple(DEFAULT_SKIP_MODULES_PATTERN)
|
|
|
|
# (1) dtype-sensitive modules the encoder flags. T5 keeps "wo" in fp32: its gated FF reads
|
|
# self.wo.weight.dtype and casts activations to match BEFORE calling wo (transformers#20287),
|
|
# racing the upcast hook so F.linear sees fp8 input vs bf16 weight. Literal substrings.
|
|
skip += tuple(re.escape(m) for m in (getattr(encoder, "_keep_in_fp32_modules", None) or ()))
|
|
|
|
# (2) an output projection tied to the input embedding. A CausalLM encoder (FLUX.2's Qwen3)
|
|
# ties lm_head.weight to embed_tokens.weight; casting lm_head (an nn.Linear) to fp8 drags the
|
|
# shared embedding down, which then emits fp8 activations that crash the first RMSNorm. Skip
|
|
# the tied projection so the shared tensor stays dense (lm_head is unused for prompt encoding).
|
|
get_out, get_in = (
|
|
getattr(encoder, "get_output_embeddings", None),
|
|
getattr(encoder, "get_input_embeddings", None),
|
|
)
|
|
out_emb = get_out() if callable(get_out) else None
|
|
in_emb = get_in() if callable(get_in) else None
|
|
if out_emb is not None and in_emb is not None and out_emb.weight is in_emb.weight:
|
|
tied_name = next((n for n, m in encoder.named_modules() if m is out_emb), None)
|
|
if tied_name:
|
|
skip += (rf"^{re.escape(tied_name)}$",)
|
|
|
|
apply_layerwise_casting(
|
|
encoder,
|
|
storage_dtype = torch.float8_e4m3fn,
|
|
compute_dtype = target.dtype,
|
|
skip_modules_pattern = skip,
|
|
# Keep token-embedding tables full precision: the diffusers default only skips vision
|
|
# pos/patch embeds, and fp8'ing nn.Embedding quantizes every prompt token to the coarse
|
|
# fp8 grid, hurting fidelity.
|
|
skip_modules_classes = (torch.nn.Embedding,),
|
|
)
|
|
|
|
# Module.dtype reports the first floating parameter, which is now fp8 STORAGE; pipelines
|
|
# derive tensor dtypes from encoder.dtype (Flux2 casts prompt embeds to it and feeds the
|
|
# result to randn_tensor, which has no fp8 kernel; VLM pipelines cast pixel_values to it,
|
|
# racing the upcast hooks). The encoder computes in target.dtype, so report that -- via a
|
|
# property shadowed on the ORIGINAL class reading a per-instance override. Swapping
|
|
# __class__ to a dynamic subclass instead breaks transformers' kwargs-based output
|
|
# recording (Qwen3VLModel returned hidden_states=None and krea-2 crashed at encode).
|
|
compute_dtype = getattr(target, "dtype", None)
|
|
try:
|
|
if compute_dtype is not None:
|
|
_install_dtype_override(type(encoder))
|
|
encoder._unsloth_te_compute_dtype = compute_dtype
|
|
# Marks the cast COMPLETE (hooks fully installed), enabling the idempotent early return
|
|
# above. Best-effort like the dtype override: a non-Module double without settable
|
|
# attributes still counts as cast, it just re-casts on a repeat call.
|
|
encoder._unsloth_te_cast_complete = True
|
|
except Exception: # noqa: BLE001 — real HF encoders are heap-type nn.Modules; only doubles fail
|
|
pass
|
|
|
|
|
|
def _install_dtype_override(cls: type) -> None:
|
|
"""Shadow ``cls.dtype`` with a property preferring the per-instance compute-dtype
|
|
override ``_cast_fp8`` sets; instances without it keep the original behaviour. Class
|
|
identity is untouched, applied once per class."""
|
|
existing = cls.__dict__.get("dtype")
|
|
if getattr(getattr(existing, "fget", None), "_unsloth_te_dtype_override", False):
|
|
return
|
|
# The property object itself when accessed through the class (property.__get__(None, cls)).
|
|
original_fget = getattr(getattr(cls, "dtype", None), "fget", None)
|
|
|
|
def _dtype(self):
|
|
override = self.__dict__.get("_unsloth_te_compute_dtype")
|
|
if override is not None:
|
|
return override
|
|
if original_fget is not None:
|
|
return original_fget(self)
|
|
raise AttributeError("dtype")
|
|
|
|
_dtype._unsloth_te_dtype_override = True
|
|
cls.dtype = property(_dtype)
|
|
|
|
|
|
def _has_layerwise_hooks(encoder: Any) -> bool:
|
|
"""True when any submodule already carries the diffusers layerwise-casting hook."""
|
|
return _has_layerwise_casting(encoder)
|
|
|
|
|
|
def _cast_nvfp4(encoder: Any, target: Any) -> None:
|
|
# Weight-only NVFP4: linear weights become 4-bit NVFP4 on Blackwell FP4 cores; norms /
|
|
# embeddings untouched. Exclude the VLM vision tower / lm_head / T5 wo and sub-512 projections
|
|
# like the int8/fp8 TE modes (4-bit-ing a VLM image tower degrades the edit conditioning);
|
|
# require_bf16 skips non-bf16 Linears so the cast engages instead of aborting.
|
|
from torchao.quantization import quantize_
|
|
from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig
|
|
from .diffusion_transformer_quant import DEFAULT_MIN_LINEAR_FEATURES, make_filter_fn
|
|
|
|
filter_fn = make_filter_fn(
|
|
DEFAULT_MIN_LINEAR_FEATURES, _te_exclude_tokens(encoder), require_bf16 = True
|
|
)
|
|
quantize_(encoder, NVFP4WeightOnlyConfig(), filter_fn = filter_fn)
|
|
|
|
|
|
def _has_layerwise_casting(module: Any) -> bool:
|
|
"""True when any submodule still carries a diffusers layerwise-casting hook -- i.e. an
|
|
``apply_layerwise_casting`` pass installed an fp8-storage upcast hook before failing. torchao's
|
|
partial-quant detector cannot see these, so a mid-pass layerwise failure would otherwise report
|
|
a dense fallback over a half-cast encoder. Best-effort: a module without ``.modules()`` or a
|
|
moved diffusers internal returns False (defer to the torchao check)."""
|
|
try:
|
|
hook_name = "layerwise_casting"
|
|
try:
|
|
from diffusers.hooks.layerwise_casting import _LAYERWISE_CASTING_HOOK
|
|
hook_name = _LAYERWISE_CASTING_HOOK
|
|
except Exception: # noqa: BLE001 -- const moved: fall back to the stable literal
|
|
pass
|
|
for sub in module.modules():
|
|
registry = getattr(sub, "_diffusers_hook", None)
|
|
get_hook = getattr(registry, "get_hook", None)
|
|
if callable(get_hook) and get_hook(hook_name) is not None:
|
|
return True
|
|
except Exception: # noqa: BLE001 -- unqueryable module: defer to the torchao check
|
|
return False
|
|
return False
|
|
|
|
|
|
def _warn(logger: Any, what: str, exc: Exception) -> None:
|
|
if logger is not None:
|
|
logger.warning("diffusion.precision: text-encoder quant (%s) failed: %s", what, exc)
|
|
|
|
|
|
def _note(logger: Any, msg: str) -> None:
|
|
if logger is not None:
|
|
logger.info("diffusion.precision: %s", msg)
|