be04ba00f4
Address the Codex review round on the video/quant work: - Companion auto-quant now honors an explicit Speed=off. Both loaders already pin the DiT dense under an explicit off (bit-exact reference), but the unset text-encoder / VAE quant still promoted to auto and silently fp8/int8'd the companions, breaking the bit-exact request. An UNSET speed still auto-quantises; an explicit companion scheme still forces it. - The HunyuanVideo joint-attention trim is a speed lever (it swaps to the fused SDPA kernel), so gate it on a non-off speed tier exactly like the adjacent attention-backend selection -- the off path keeps the stock dense-mask attention. - Explicit torchao text-encoder modes (int8 / fp8_dynamic / nvfp4) now run the same kernel smoke test the auto ladder uses. They could 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 only covers the cast, not the first forward, so the load would report engaged then crash at generation. Now it falls back to dense. Layerwise fp8 has no torchao GEMM, so the probe is a no-op for it. - The trim pre-hook's fallback restores the caller's original kwargs (it may have emptied the image stream / trimmed a text stream before failing), so the stock dense-mask path runs on exactly what it expects, matching the empty-prompt guard. - video_speedmem_bench mirrors the loader: installs the Hunyuan trim before the backend set (gated on an active tier) and skips the auto int8 quant when it is the fp8-denied memory fallback and dense fits resident, so the shipped/auto rows measure what the loader actually runs. Tests: TE explicit-mode kernel probe (+ layerwise-fp8 bypass), trim mid-trim restore, and loader-level speed=off companion suppression + trim skip for both backends. 262 backend tests pass; ruff clean.
438 lines
21 KiB
Python
438 lines
21 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) from the base repo and is often the largest resident component (a Qwen3
|
|
/ T5-XXL / Mistral encoder runs to many GB). This shrinks it in place, with four
|
|
backends:
|
|
|
|
fp8 - diffusers layerwise casting: 8-bit (e4m3) storage, upcast per layer to
|
|
the compute dtype. ~2x smaller. Works on any fp8-capable CUDA card (cc >= 8.9).
|
|
fp8_dynamic - torchao dynamic fp8 COMPUTE (per-row): keeps the matmul in fp8 on the
|
|
fp8 tensor cores (torch._scaled_mm) instead of upcasting each forward.
|
|
~2x smaller plus a tensor-core speedup; needs fp8-GEMM silicon (cc >= 8.9).
|
|
int8 - torchao dynamic int8 COMPUTE (per-token act + per-channel weight ->
|
|
torch._int_mm), with per-layer keep-bf16 selection. int8 degrades on large
|
|
encoders unless the most quant-sensitive decoder blocks stay bf16, so it is
|
|
applied only for families with a measured keep-bf16 schedule (else it falls
|
|
back to fp8). ~2x smaller; needs int8 tensor cores (cc >= 8.0).
|
|
nvfp4 - torchao NVFP4 weight-only: 4-bit float with two-level microscaling, run on
|
|
Blackwell's (sm_100+) FP4 tensor cores. ~4x smaller and the lowest-VRAM
|
|
option, but a steeper quality cost than fp8.
|
|
|
|
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 and picks the
|
|
best accurate scheme (fp8_dynamic / int8 / layerwise fp8), falling back to dense when nothing
|
|
qualifies; ``none``/``off`` keeps the encoder dense, and an explicit scheme forces it. They pair
|
|
especially well with streamed (group) offload, where the text encoder stays resident -- this is
|
|
where the companion footprint dominates. Quantify the quality cost per model with the quality
|
|
harness (scripts/diffusion_quality.py) and the hidden-state gate (scripts/diffusion_quant_builder.py).
|
|
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
|
|
|
|
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"): "auto" resolves to one of these via select_te_quant_scheme,
|
|
# and these are the values 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 (torch._int_mm) degrades on large text encoders unless the most quant-sensitive decoder
|
|
# blocks stay bf16. Per-family (skip_first, skip_last) decoder blocks to keep dense, from measured
|
|
# hidden-state fidelity (mean per-token cosine vs the bf16 reference, at the layer each pipeline
|
|
# consumes): keeping the first blocks stops early-layer error seeding, keeping the last blocks
|
|
# protects the read layer. Families absent here have no int8 schedule that clears the bar, so an
|
|
# int8 request for them falls back to fp8.
|
|
# qwen-image (Qwen2.5-VL-7B): first+last 6 -> ~0.997 cosine (both ends needed; outlier-bound).
|
|
# flux.2-dev (Mistral-Small-24B): first 3 -> ~0.98 cosine (pure 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 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 == 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 device with a bf16 compute dtype, plus
|
|
the tensor-core class each backend needs -- fp8 dtype (fp8 layerwise), fp8 GEMM sm_89+
|
|
(fp8_dynamic), int8 tensor cores sm_80+ (int8), or 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:
|
|
# Compute fp8 (torch._scaled_mm) needs fp8-GEMM silicon: Ada sm_89+ / Hopper / Blackwell.
|
|
return hasattr(torch, "float8_e4m3fn") and torch.cuda.get_device_capability() >= (8, 9)
|
|
if mode == TE_QUANT_INT8:
|
|
# int8 tensor cores (torch._int_mm) need Ampere sm_80+.
|
|
return torch.cuda.get_device_capability()[0] >= 8
|
|
if mode == TE_QUANT_NVFP4:
|
|
# NVFP4 tensor cores need Blackwell (compute capability major >= 10).
|
|
return torch.cuda.get_device_capability()[0] >= 10
|
|
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 on the tensor cores) leads on fp8-GEMM silicon; int8
|
|
# (torch._int_mm) sits second but only engages for a family with a measured keep-bf16 schedule;
|
|
# layerwise ``fp8`` (storage cast, no torchao) is the universal fallback -- it needs only the fp8
|
|
# dtype and is the sole scheme that survives group offload. nvfp4 stays explicit-only (weight-only
|
|
# 4-bit is a steeper quality cost), never an auto pick, consistent with the DiT ladder.
|
|
_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, via scripts/diffusion_quant_builder.py). Populated from the accuracy sweep; a
|
|
# denied scheme is skipped by ``auto`` and refused when requested explicitly. Empty by default:
|
|
# int8 already gates on a per-family keep-bf16 schedule (``_TE_INT8_SKIP``), so this is for the
|
|
# rarer case where even keep-bf16 int8 (or fp8) misses the bar for a specific encoder.
|
|
_TE_FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {}
|
|
|
|
# Map a TE torchao scheme to the transformer smoke-probe scheme (same torchao 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())
|
|
|
|
|
|
def _te_scheme_probe(scheme: str, device: str) -> bool:
|
|
"""True iff ``scheme`` actually runs on this build. Layerwise fp8 (no torchao GEMM) always
|
|
passes; the torchao TE modes reuse the transformer module's cached quantise+matmul smoke test."""
|
|
tq = _TE_SMOKE_SCHEME.get(scheme)
|
|
if tq is None:
|
|
return True
|
|
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 (``quantize_text_encoders`` re-gates it). ``auto`` walks
|
|
``_TE_AUTO_LADDER`` for this GPU's capability and returns the first scheme that: survives the
|
|
active offload policy (torchao modes need pinned residency -> only layerwise fp8 under offload),
|
|
has a keep-bf16 schedule if int8, is not family-denied, is hardware-supported, and passes a
|
|
real kernel smoke test. Returns None when nothing qualifies (e.g. no CUDA / pre-Ampere)."""
|
|
requested = normalize_te_quant(requested)
|
|
if requested is None or requested != TE_QUANT_AUTO:
|
|
return requested
|
|
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 tensor subclasses reject the Module.to() an offload hook uses; only the
|
|
# layerwise fp8 cast streams, so under offload skip the torchao modes.
|
|
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 mode actually applied, or None when disabled, unsupported, or no encoder was
|
|
cast. ``auto`` resolves to the best scheme for the GPU + family via ``select_te_quant_scheme``.
|
|
``int8`` needs a per-family keep-bf16 schedule (``_TE_INT8_SKIP``); a family without one falls
|
|
back to ``fp8``. When ``offload_active`` the torchao modes are skipped (their tensor subclasses
|
|
reject the ``Module.to()`` an offload hook uses); 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
|
|
# The torchao modes (int8 with a schedule, fp8_dynamic, nvfp4) produce tensor subclasses that
|
|
# reject Module.to(); an offload placement moves the encoder that way and hard-crashes -- the
|
|
# DiT path skips torchao quant under offload for exactly this reason. Layerwise fp8 is not
|
|
# torchao and streams fine, so it still engages. Skip the torchao modes under offload.
|
|
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
|
|
# Mirror select_te_quant_scheme's auto path: an EXPLICIT torchao TE mode (int8 / fp8_dynamic /
|
|
# nvfp4) can clear the capability gate above yet fail the real GEMM on a torchao/torch build
|
|
# where quantize_ wraps the encoder but the kernel is broken -- and the caster's try/except only
|
|
# catches the cast, not the first prompt-encoder forward. Run the same kernel smoke test the auto
|
|
# ladder uses so a failing kernel falls back to dense here instead of crashing at generation.
|
|
# Layerwise fp8 (no torchao GEMM) has no smoke scheme, so the probe is a no-op (True) for it.
|
|
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
|
|
_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 text-encoder quant: the VLM vision tower
|
|
and the unused lm_head (not used for prompt encoding), plus the encoder's own fp32-kept
|
|
modules (T5 ``wo``, which the gated feed-forward reads the dtype of and 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 the decoder blocks to keep bf16: the first ``skip_first`` and last ``skip_last`` of
|
|
each top-level ``nn.ModuleList`` stack (a T5 ``encoder.block`` / a decoder ``...layers``).
|
|
Structural, so it needs 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 (per-token act + per-channel weight -> torch._int_mm) on the FLOP-heavy
|
|
# Linears, but keeping the first/last decoder blocks (and the vision tower / lm_head / T5 wo)
|
|
# in bf16. Reuses the committed transformer-quant factory so the 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 _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
|
|
# torchao dynamic fp8 COMPUTE, per-row (per-token activation + per-output-channel weight ->
|
|
# torch._scaled_mm on the fp8 tensor cores). Unlike the layerwise `fp8` backend this keeps the
|
|
# matmul in fp8 instead of upcasting each forward. fp8 is robust across encoder sizes, so no
|
|
# per-layer keep-bf16 is needed; 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 the encoder
|
|
# keeps (belt-and-suspenders over the named T5 wo exclusion) rather than aborting the pass.
|
|
filter_fn = make_filter_fn(
|
|
DEFAULT_MIN_LINEAR_FEATURES, _te_exclude_tokens(encoder), require_bf16 = True
|
|
)
|
|
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
|
|
|
|
# diffusers' layerwise casting stores each supported leaf module's weights in fp8 and
|
|
# upcasts them per forward. Two things on a transformers text encoder can push an fp8
|
|
# weight or activation into an op that can't handle it, and both crash only at
|
|
# generation (the load-time guard can't see them), so skip the offending modules:
|
|
skip = tuple(DEFAULT_SKIP_MODULES_PATTERN)
|
|
|
|
# (1) dtype-sensitive modules the encoder itself flags. T5 keeps "wo" in fp32: its
|
|
# gated feed-forward reads self.wo.weight.dtype and casts the activations to match
|
|
# BEFORE calling wo (transformers#20287), racing the forward-time upcast hook so
|
|
# F.linear sees an fp8 input against a bf16 weight. Names are 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; lm_head is an nn.Linear so it
|
|
# gets cast to fp8 and, sharing one tensor, drags the embedding to fp8 with it. The
|
|
# embedding 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 (T5 "shared", Qwen "embed_tokens", etc.) full
|
|
# precision: the diffusers default pattern only skips vision pos/patch
|
|
# embeds, not nn.Embedding lookups, and fp8'ing those quantizes every prompt
|
|
# token straight to the coarse fp8 grid, hurting prompt fidelity.
|
|
skip_modules_classes = (torch.nn.Embedding,),
|
|
)
|
|
|
|
|
|
def _cast_nvfp4(encoder: Any, target: Any) -> None:
|
|
# Weight-only NVFP4: linear weights become 4-bit (packed) NVFP4 tensors and run
|
|
# on Blackwell FP4 tensor cores; norms / embeddings (not nn.Linear) are untouched.
|
|
# Exclude the VLM vision tower / lm_head / T5 wo and the sub-512 projections, exactly like
|
|
# the int8 / fp8 torchao TE modes -- 4-bit-ing a VLM encoder's image tower (qwen-image /
|
|
# qwen-image-edit's Qwen2.5-VL) degrades the image/edit conditioning the sibling schemes
|
|
# deliberately protect, and require_bf16 skips any non-bf16 Linear the encoder keeps so the
|
|
# NVFP4 (scaled_mm-family) cast engages on the bf16 linears instead of aborting the pass.
|
|
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 _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)
|