From bdf676cbe5d931e6d341d224d1cc09891c9e23ee Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 18 Jul 2026 07:19:03 +0000 Subject: [PATCH] Report the compute dtype on fp8-cast encoders and inject the pre-cast TE on the dense fast path Two more findings from the hosted-TE GPU smokes: - Module.dtype reports the first floating parameter, which after the layerwise fp8 cast is the fp8 STORAGE dtype. Flux2 derives its prompt embed and latent dtypes from encoder.dtype and feeds them to randn_tensor, which has no fp8 kernel, so ANY flux.2 load with text_encoder_quant=fp8 crashed at generation (pre-existing, runtime cast included). The cast now swaps in a subclass whose dtype property reports the compute dtype; forward behaviour is unchanged. - The dense transformer_quant fast path assembles companions through _assemble_pipe, which never received the pre-cast TE injection, so the hosted encoder engaged on full-pipeline and GGUF builds but not on the fast path. Threaded through like the other two branches. Verified live on B200: qwen-image (full pipeline), flux.2-dev (GGUF picker with int8 DiT prequant), ltx-2 (video backend) all engage the hosted TE, render non-black, and report text_encoder_quant=fp8 truthfully. --- studio/backend/core/inference/diffusion.py | 23 ++++++++++- .../core/inference/diffusion_precision.py | 16 ++++++++ .../tests/test_diffusion_te_prequant.py | 38 +++++++++++++++++++ 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index d5a80b667f..2c531eda1e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1380,6 +1380,7 @@ class DiffusionBackend: prequant_path = transformer_prequant_path, allow_dense_fallback = dense_fallback_allowed, lora_specs = loras, + text_encoder_quant = text_encoder_quant, ) except Exception as exc: # noqa: BLE001 — fall back to the GGUF build logger.warning( @@ -1812,6 +1813,7 @@ class DiffusionBackend: base_local_dir: Optional[str] = None, allow_dense_fallback: bool = True, lora_specs: Optional[list[tuple[str, float]]] = None, + text_encoder_quant: Optional[str] = None, ) -> tuple[Any, str]: """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``. @@ -1867,7 +1869,7 @@ class DiffusionBackend: if transformer is not None: pipe = self._assemble_pipe( pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir, - fam = fam, + fam = fam, te_quant_mode = text_encoder_quant, target = target, ) return pipe, scheme @@ -1882,7 +1884,8 @@ class DiffusionBackend: base, subfolder = "transformer", torch_dtype = dtype, token = hf_token ) pipe = self._assemble_pipe( - pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir, fam = fam + pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir, + fam = fam, te_quant_mode = text_encoder_quant, target = target, ) if lora_specs: # Bake the adapters BEFORE quantize_: peft injects its wrappers on the dense @@ -1929,6 +1932,8 @@ class DiffusionBackend: device: str, base_local_dir: Optional[str] = None, fam: Optional[DiffusionFamily] = None, + te_quant_mode: Optional[str] = None, + target: Any = None, ) -> Any: """Assemble the diffusers pipeline around ``transformer`` and place it on ``device`` (a no-op for an already-placed pre-quantized transformer; it moves the companions).""" @@ -1948,6 +1953,20 @@ class DiffusionBackend: # The repo ships no Llama text_encoder_4; assemble it from the open mirror # (diffusion_hidream.py) exactly like the full-pipeline load branch. pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token)) + # Same pre-cast TE injection as the full-pipeline and GGUF branches: the dense + # fast path supplies only the transformer, so the companion TE is the big download. + if target is not None: + pipe_kwargs.update( + te_prequant_pipe_kwargs( + fam, + base, + te_quant_mode = te_quant_mode, + target = target, + dtype = dtype, + hf_token = hf_token, + logger = logger, + ) + ) pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs) pipe.to(device) return pipe diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index f0919a179e..7b19e9fa98 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -288,6 +288,22 @@ def _cast_fp8(encoder: Any, target: Any) -> None: 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. + compute_dtype = getattr(target, "dtype", None) + if compute_dtype is not None and not getattr(encoder, "_unsloth_te_dtype_override", False): + cls = type(encoder) + encoder.__class__ = type( + cls.__name__, + (cls,), + { + "dtype": property(lambda self, _d = compute_dtype: _d), + "_unsloth_te_dtype_override": True, + }, + ) + def _has_layerwise_hooks(encoder: Any) -> bool: """True when any submodule already carries the diffusers layerwise-casting hook.""" diff --git a/studio/backend/tests/test_diffusion_te_prequant.py b/studio/backend/tests/test_diffusion_te_prequant.py index 6a479a7703..4ddb7af59b 100644 --- a/studio/backend/tests/test_diffusion_te_prequant.py +++ b/studio/backend/tests/test_diffusion_te_prequant.py @@ -263,6 +263,39 @@ def test_hosted_te_prequant_entries(): ) == "LTX-2-text_encoder-FP8.pt" +def test_assemble_pipe_injects_precast_te(monkeypatch): + """The dense transformer_quant fast path assembles companions through _assemble_pipe, + which must inject the hosted pre-cast TE like the full-pipeline and GGUF branches.""" + import core.inference.diffusion as dif + + seen: dict = {} + + class FakePipe: + def to(self, device): + return self + + class FakePipelineCls: + @staticmethod + def from_pretrained(base, **kw): + seen.update(kw) + return FakePipe() + + monkeypatch.setattr( + dif, "te_prequant_pipe_kwargs", lambda *a, **k: {"text_encoder": "PRECAST"} + ) + dif.DiffusionBackend._assemble_pipe( + FakePipelineCls, "org/base", "TR", None, None, "cpu", None, + fam = None, te_quant_mode = "fp8", target = object(), + ) + assert seen["text_encoder"] == "PRECAST" + seen.clear() + # No target (defensive default) keeps the assembly unchanged. + dif.DiffusionBackend._assemble_pipe( + FakePipelineCls, "org/base", "TR", None, None, "cpu", None, fam = None, + ) + assert "text_encoder" not in seen + + def test_cast_fp8_is_idempotent_on_precast_encoder(): """A pre-cast encoder arrives with the layerwise hooks installed; the runtime re-apply in quantize_text_encoders must be a no-op (re-registering the hook name raises, which made @@ -275,8 +308,13 @@ def test_cast_fp8_is_idempotent_on_precast_encoder(): enc = torch.nn.Sequential(torch.nn.Linear(64, 64), torch.nn.LayerNorm(64)) _cast_fp8(enc, target) assert enc[0].weight.dtype == torch.float8_e4m3fn + # Module.dtype must report the COMPUTE dtype: pipelines derive tensor dtypes from it + # (Flux2 feeds it to randn_tensor, which has no fp8 kernel). + assert enc.dtype == torch.bfloat16 + assert isinstance(enc, torch.nn.Sequential) _cast_fp8(enc, target) # must not raise assert enc[0].weight.dtype == torch.float8_e4m3fn + assert enc.dtype == torch.bfloat16 def test_builder_metadata_survives_weights_only_load(tmp_path):