From 8b69f22809b899739b7448318f01076fc43484e6 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Tue, 30 Jun 2026 23:22:13 -0300 Subject: [PATCH] Fix fp8 text-encoder quant crashing generation on T5 and tied-embedding encoders --- .../core/inference/diffusion_precision.py | 28 +++++++- .../backend/tests/test_diffusion_precision.py | 64 +++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index b229489755..38a3434490 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -99,15 +99,41 @@ def quantize_text_encoders( 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 = DEFAULT_SKIP_MODULES_PATTERN, + 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 diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 6b4f518625..f879ad99fc 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -67,6 +67,20 @@ def _stub_casters(monkeypatch, recorder): 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 ───────────────────────────────────────────────────────────── @@ -119,6 +133,56 @@ def test_quantize_fp8_casts_all_encoders(monkeypatch): 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 = []