Fix fp8 text-encoder quant crashing generation on T5 and tied-embedding encoders

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

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 = []