220 lines
8.9 KiB
Python
220 lines
8.9 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Unit tests for text-encoder quantisation (``diffusion_precision.py``).
|
|
|
|
Hermetic: torch + the diffusers / torchao casters are stubbed via ``sys.modules`` so
|
|
gating and the apply path run without a GPU, real diffusers, or real torchao.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
|
|
import pytest
|
|
|
|
from core.inference.diffusion_precision import (
|
|
TE_QUANT_FP8,
|
|
TE_QUANT_NVFP4,
|
|
normalize_te_quant,
|
|
quantize_text_encoders,
|
|
te_quant_supported,
|
|
)
|
|
|
|
|
|
def _target(
|
|
*,
|
|
device = "cuda",
|
|
dtype = "bfloat16",
|
|
cc = (10, 0),
|
|
):
|
|
return types.SimpleNamespace(device = device, dtype = dtype, _cc = cc)
|
|
|
|
|
|
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"
|
|
torch.cuda = types.SimpleNamespace(get_device_capability = lambda *a: cc)
|
|
# _cast_fp8 skips nn.Embedding modules from layerwise casting.
|
|
torch.nn = types.SimpleNamespace(Embedding = type("Embedding", (), {}))
|
|
monkeypatch.setitem(sys.modules, "torch", torch)
|
|
return torch
|
|
|
|
|
|
def _stub_casters(monkeypatch, recorder):
|
|
# diffusers fp8 layerwise casting
|
|
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))
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting)
|
|
# torchao nvfp4
|
|
tq = types.ModuleType("torchao.quantization")
|
|
tq.quantize_ = lambda module, config: recorder.append(("nvfp4", module))
|
|
mx = types.ModuleType("torchao.prototype.mx_formats")
|
|
mx.NVFP4WeightOnlyConfig = lambda: "nvfp4cfg"
|
|
monkeypatch.setitem(sys.modules, "torchao.quantization", tq)
|
|
monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", mx)
|
|
|
|
|
|
def _stub_fp8_capture(monkeypatch):
|
|
# Stub the fp8 caster to record the skip_modules_pattern passed for each encoder.
|
|
captured: dict = {}
|
|
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: captured.__setitem__(
|
|
id(module), kw["skip_modules_pattern"]
|
|
)
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting)
|
|
return captured
|
|
|
|
|
|
# ── normalisation ─────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_normalize_te_quant():
|
|
assert normalize_te_quant(None) is None
|
|
assert normalize_te_quant("") is None
|
|
assert normalize_te_quant("none") is None
|
|
assert normalize_te_quant("FP8") == TE_QUANT_FP8
|
|
assert normalize_te_quant("NVFP4") == TE_QUANT_NVFP4
|
|
with pytest.raises(ValueError):
|
|
normalize_te_quant("int2")
|
|
|
|
|
|
# ── gating ────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_fp8_supported_requires_cuda_bf16_and_fp8(monkeypatch):
|
|
_stub_torch(monkeypatch, with_fp8 = True)
|
|
assert te_quant_supported(_target(), TE_QUANT_FP8) is True
|
|
assert te_quant_supported(_target(device = "cpu"), TE_QUANT_FP8) is False
|
|
assert te_quant_supported(_target(dtype = "float16"), TE_QUANT_FP8) is False
|
|
|
|
|
|
def test_nvfp4_supported_requires_blackwell(monkeypatch):
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
assert te_quant_supported(_target(), TE_QUANT_NVFP4) is True
|
|
# Hopper (cc 9.0) has no NVFP4 tensor cores.
|
|
_stub_torch(monkeypatch, cc = (9, 0))
|
|
assert te_quant_supported(_target(), TE_QUANT_NVFP4) is False
|
|
|
|
|
|
# ── apply ─────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def test_quantize_disabled_returns_none(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
pipe = types.SimpleNamespace(text_encoder = object())
|
|
assert quantize_text_encoders(pipe, _target(), mode = None) is None
|
|
assert quantize_text_encoders(pipe, _target(), mode = "none") is None
|
|
|
|
|
|
def test_quantize_fp8_casts_all_encoders(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
recorder: list = []
|
|
_stub_casters(monkeypatch, recorder)
|
|
te1, te3 = object(), object()
|
|
pipe = types.SimpleNamespace(text_encoder = te1, text_encoder_2 = None, text_encoder_3 = te3)
|
|
mode = quantize_text_encoders(pipe, _target(), mode = "fp8")
|
|
assert mode == TE_QUANT_FP8
|
|
assert recorder == [("fp8", te1), ("fp8", te3)]
|
|
|
|
|
|
def test_fp8_skips_encoder_keep_in_fp32_modules(monkeypatch):
|
|
# T5 keeps "wo" in fp32: its gated FF reads wo.weight.dtype and casts activations to
|
|
# match BEFORE calling wo, which races diffusers' forward-time upcast hook and crashes
|
|
# generation (fp8 input vs bf16 weight). _cast_fp8 must add the encoder's own
|
|
# _keep_in_fp32_modules to the layerwise-casting skip patterns; encoders without such a
|
|
# list (CLIP, Qwen) get nothing extra skipped.
|
|
_stub_torch(monkeypatch)
|
|
captured = _stub_fp8_capture(monkeypatch)
|
|
|
|
t5 = types.SimpleNamespace(_keep_in_fp32_modules = ["wo"])
|
|
qwen = types.SimpleNamespace(_keep_in_fp32_modules = None)
|
|
pipe = types.SimpleNamespace(text_encoder = t5, text_encoder_2 = qwen)
|
|
quantize_text_encoders(pipe, _target(), mode = "fp8")
|
|
|
|
assert "wo" in captured[id(t5)] and "norm" in captured[id(t5)]
|
|
assert "wo" not in captured[id(qwen)] and "norm" in captured[id(qwen)]
|
|
|
|
|
|
def test_fp8_skips_tied_output_embedding(monkeypatch):
|
|
# A CausalLM encoder (FLUX.2's Qwen3) ties lm_head.weight to the input embedding.
|
|
# lm_head is nn.Linear, so layerwise casting would fp8 it and drag the shared
|
|
# embedding tensor to fp8, making the embedding emit fp8 activations that crash the
|
|
# first norm. _cast_fp8 must skip the tied output projection by name; an untied one
|
|
# (distinct weight tensors) is left to quantise normally.
|
|
_stub_torch(monkeypatch)
|
|
captured = _stub_fp8_capture(monkeypatch)
|
|
|
|
shared = object() # the one tensor lm_head and embed_tokens share
|
|
emb = types.SimpleNamespace(weight = shared)
|
|
head = types.SimpleNamespace(weight = shared)
|
|
tied = types.SimpleNamespace(
|
|
_keep_in_fp32_modules = None,
|
|
get_input_embeddings = lambda: emb,
|
|
get_output_embeddings = lambda: head,
|
|
named_modules = lambda: [("model.embed_tokens", emb), ("lm_head", head)],
|
|
)
|
|
u_emb, u_head = types.SimpleNamespace(weight = object()), types.SimpleNamespace(weight = object())
|
|
untied = types.SimpleNamespace(
|
|
_keep_in_fp32_modules = None,
|
|
get_input_embeddings = lambda: u_emb,
|
|
get_output_embeddings = lambda: u_head,
|
|
named_modules = lambda: [("lm_head", u_head)],
|
|
)
|
|
pipe = types.SimpleNamespace(text_encoder = tied, text_encoder_2 = untied)
|
|
quantize_text_encoders(pipe, _target(), mode = "fp8")
|
|
|
|
assert r"^lm_head$" in captured[id(tied)]
|
|
assert r"^lm_head$" not in captured[id(untied)]
|
|
|
|
|
|
def test_quantize_nvfp4_uses_torchao(monkeypatch):
|
|
_stub_torch(monkeypatch, cc = (10, 0))
|
|
recorder: list = []
|
|
_stub_casters(monkeypatch, recorder)
|
|
te = object()
|
|
pipe = types.SimpleNamespace(text_encoder = te)
|
|
mode = quantize_text_encoders(pipe, _target(), mode = "nvfp4")
|
|
assert mode == TE_QUANT_NVFP4
|
|
assert recorder == [("nvfp4", te)]
|
|
|
|
|
|
def test_quantize_nvfp4_unsupported_on_hopper_is_noop(monkeypatch):
|
|
_stub_torch(monkeypatch, cc = (9, 0))
|
|
recorder: list = []
|
|
_stub_casters(monkeypatch, recorder)
|
|
pipe = types.SimpleNamespace(text_encoder = object())
|
|
assert quantize_text_encoders(pipe, _target(cc = (9, 0)), mode = "nvfp4") is None
|
|
assert recorder == []
|
|
|
|
|
|
def test_quantize_tolerates_caster_failure(monkeypatch):
|
|
_stub_torch(monkeypatch)
|
|
hooks = types.ModuleType("diffusers.hooks")
|
|
casting = types.ModuleType("diffusers.hooks.layerwise_casting")
|
|
casting.DEFAULT_SKIP_MODULES_PATTERN = ("norm",)
|
|
|
|
def _boom(module, **kwargs):
|
|
raise RuntimeError("fp8 unsupported for this layer")
|
|
|
|
hooks.apply_layerwise_casting = _boom
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
|
|
monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting)
|
|
pipe = types.SimpleNamespace(text_encoder = object())
|
|
# The only encoder fails to cast -> nothing applied -> None.
|
|
assert quantize_text_encoders(pipe, _target(), mode = "fp8") is None
|