From a3942d9924f775389fcc9d3cbfb5d947161f88cb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 18 Jul 2026 07:55:24 +0000 Subject: [PATCH] Extend the fp8 TE quant to HiDream's Llama text_encoder_4 The generic quantize_text_encoders pass only covers text_encoder.._3, so HiDream's HEAVIEST encoder (Llama-3.1-8B TE4, 16.1 GB bf16) always stayed dense. TE4 is assembled separately (hidream_te4_kwargs), so the fp8 path now lives there: when the requested TE quant is layerwise fp8 and the device/family qualify, TE4 prefers the hosted pre-cast checkpoint (unsloth/HiDream-I1-Full-FP8, 8.6 GB) and falls back to dense-load-then- cast; a mid-pass cast failure reloads a fresh dense encoder instead of shipping partial state. The pre-cast loader and builder gain config_subfolder/config_overrides for standalone encoder repos whose config sits at the root and whose pipeline needs forward flags (output_hidden_states/attentions). Verified on B200: bit-identity 291 tensors (225 fp8, 0 mismatches), hosted checkpoint engages through the real backend (marker + status fp8), load 24.3 s vs 48.0 s dense, LPIPS 0.133 mean over 3 same-seed pairs vs the dense-TE render (gate 0.25), non-black frames. --- scripts/build_te_prequant_checkpoint.py | 21 ++- studio/backend/core/inference/diffusion.py | 30 +++- .../core/inference/diffusion_families.py | 3 + .../core/inference/diffusion_hidream.py | 89 +++++++++++- .../core/inference/diffusion_te_prequant.py | 20 ++- .../tests/test_diffusion_te_prequant.py | 130 ++++++++++++++++++ 6 files changed, 277 insertions(+), 16 deletions(-) diff --git a/scripts/build_te_prequant_checkpoint.py b/scripts/build_te_prequant_checkpoint.py index 83483d6129..94a7b713b1 100644 --- a/scripts/build_te_prequant_checkpoint.py +++ b/scripts/build_te_prequant_checkpoint.py @@ -37,6 +37,13 @@ def main(argv = None) -> int: default = "text_encoder", help = "pipeline component attribute (also the repo subfolder)", ) + p.add_argument( + "--config-subfolder", + default = None, + help = "where the encoder lives inside --base (default: the component name; " + "pass '' for a standalone encoder repo whose config sits at the root, " + "e.g. HiDream's Llama text_encoder_4)", + ) p.add_argument("--scheme", default = "fp8", choices = ["fp8"]) p.add_argument("--out", required = True, help = "output .pt path for the checkpoint") p.add_argument("--dtype", default = "bfloat16", choices = ["bfloat16"]) @@ -54,12 +61,15 @@ def main(argv = None) -> int: # branch (diffusion_families vs video_families), so resolve best-effort by name. family = args.family.strip().lower() + subfolder = args.component if args.config_subfolder is None else args.config_subfolder + from_pretrained_kwargs = {"token": args.hf_token} + if subfolder: + from_pretrained_kwargs["subfolder"] = subfolder + print(f"== build TE prequant ({family}/{args.component}/{args.scheme}) ==", flush = True) - print(f" loading dense encoder from {args.base} (subfolder={args.component}) ...", flush = True) + print(f" loading dense encoder from {args.base} (subfolder={subfolder!r}) ...", flush = True) t0 = time.time() - config = transformers.AutoConfig.from_pretrained( - args.base, subfolder = args.component, token = args.hf_token - ) + config = transformers.AutoConfig.from_pretrained(args.base, **from_pretrained_kwargs) # Prefer the checkpoint's own architecture (what the diffusers pipeline instantiates, # e.g. Gemma3ForConditionalGeneration); AutoModel.from_config would give the bare base # class and record a te_class whose state dict the pipeline cannot use. @@ -72,9 +82,8 @@ def main(argv = None) -> int: del encoder encoder = getattr(transformers, encoder_cls_name).from_pretrained( args.base, - subfolder = args.component, torch_dtype = torch.bfloat16, - token = args.hf_token, + **from_pretrained_kwargs, ) print(f" casting in place (layerwise {args.scheme}) ...", flush = True) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8ffae17831..7b1db85362 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1463,7 +1463,15 @@ class DiffusionBackend: if fam.name == HIDREAM_FAMILY_NAME: # The repo names a Llama text_encoder_4 it does not ship; # supply it from the open mirror (diffusion_hidream.py). - pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token)) + pipe_kwargs.update( + hidream_te4_kwargs( + dtype, + hf_token, + fam = fam, + te_quant_mode = text_encoder_quant, + target = target, + ) + ) # A hosted pre-cast fp8 text encoder (when the family ships one and # the runtime cast would engage) skips the dense TE download; the # later quantize_text_encoders re-applies the cast idempotently. @@ -1520,7 +1528,15 @@ class DiffusionBackend: pipe_kwargs["token"] = hf_token if fam.name == HIDREAM_FAMILY_NAME: # Same Llama TE4 assembly as the full-pipeline branch above. - pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token)) + pipe_kwargs.update( + hidream_te4_kwargs( + dtype, + hf_token, + fam = fam, + te_quant_mode = text_encoder_quant, + target = target, + ) + ) # Same pre-cast TE injection as the full-pipeline branch: the GGUF # supplies the transformer, so the companion TE is the big download. pipe_kwargs.update( @@ -2008,7 +2024,15 @@ class DiffusionBackend: if getattr(fam, "name", None) == HIDREAM_FAMILY_NAME: # 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)) + pipe_kwargs.update( + hidream_te4_kwargs( + dtype, + hf_token, + fam = fam, + te_quant_mode = te_quant_mode, + target = target, + ) + ) # 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: diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index eec7585a4b..82975544d7 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -381,6 +381,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( ("int8", "unsloth/HiDream-I1-Full-FP8"), ("fp8", "unsloth/HiDream-I1-Full-FP8"), ), + # Pre-cast Llama-3.1-8B TE4 (16.1 GB bf16 -> 8.1 GB). The generic TE pass only covers + # text_encoder.._3, so TE4 engages via hidream_te4_kwargs, not te_prequant_pipe_kwargs. + te_prequant_repos = (("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),), pipeline_class = "HiDreamImagePipeline", transformer_class = "HiDreamImageTransformer2DModel", base_repo = "HiDream-ai/HiDream-I1-Full", diff --git a/studio/backend/core/inference/diffusion_hidream.py b/studio/backend/core/inference/diffusion_hidream.py index 82ac27f7d5..b24c621d76 100644 --- a/studio/backend/core/inference/diffusion_hidream.py +++ b/studio/backend/core/inference/diffusion_hidream.py @@ -27,16 +27,77 @@ HIDREAM_FAMILY_NAME = "hidream-i1" HIDREAM_LLAMA_REPO = "unsloth/Meta-Llama-3.1-8B-Instruct" -def hidream_te4_kwargs(dtype: Any, hf_token: Optional[str] = None) -> dict[str, Any]: +def hidream_te4_kwargs( + dtype: Any, + hf_token: Optional[str] = None, + *, + fam: Any = None, + te_quant_mode: Optional[str] = None, + target: Any = None, +) -> dict[str, Any]: """``{text_encoder_4, tokenizer_4}`` kwargs for a HiDream pipeline ``from_pretrained``. Loaded eagerly (~16 GB bf16) before the pipeline call so a failure surfaces as a - clear error instead of a half-built pipeline.""" + clear error instead of a half-built pipeline. + + The generic ``quantize_text_encoders`` pass only covers ``text_encoder``..``_3``, so + TE4 -- HiDream's HEAVIEST encoder -- is handled here: when the requested TE quant is + layerwise fp8 (and the device/family qualify, same gates as the runtime cast), TE4 is + fp8-cast too, preferring the hosted pre-cast checkpoint (~half the download) and + falling back to dense-load-then-cast. Any other mode keeps today's dense bf16 TE4.""" import torch # noqa: F401 -- dtype values are torch dtypes; import keeps parity with callers from transformers import AutoTokenizer, LlamaForCausalLM - logger.info("diffusion.hidream: loading Llama TE4 from %s", HIDREAM_LLAMA_REPO) tokenizer_4 = AutoTokenizer.from_pretrained(HIDREAM_LLAMA_REPO, token = hf_token) + + fp8_engages = False + if target is not None: + try: + from . import diffusion_precision as precision + from .diffusion_precision import ( + TE_QUANT_FP8, + normalize_te_quant, + te_quant_supported, + ) + + mode = normalize_te_quant(te_quant_mode) + denied = getattr(precision, "_te_family_denied", None) + fp8_engages = ( + mode == TE_QUANT_FP8 + and te_quant_supported(target, mode) + and not (callable(denied) and denied(getattr(fam, "name", None), mode)) + ) + except Exception: # noqa: BLE001 -- quant probe failure keeps the dense bf16 path + fp8_engages = False + + if fp8_engages and fam is not None: + from .diffusion_te_prequant import ( + load_prequant_text_encoder, + resolve_te_prequant_source, + ) + + source = resolve_te_prequant_source(fam, "text_encoder_4", "fp8") + if source is not None: + encoder = load_prequant_text_encoder( + HIDREAM_LLAMA_REPO, + "text_encoder_4", + source, + dtype = dtype, + hf_token = hf_token, + scheme = "fp8", + logger = logger, + # The Llama TE4 lives in its own standalone repo (config at the root), and + # the pipeline needs hidden states/attentions from its forward. + config_subfolder = "", + config_overrides = { + "output_hidden_states": True, + "output_attentions": True, + }, + ) + if encoder is not None: + return {"text_encoder_4": encoder, "tokenizer_4": tokenizer_4} + + logger.info("diffusion.hidream: loading Llama TE4 from %s", HIDREAM_LLAMA_REPO) text_encoder_4 = LlamaForCausalLM.from_pretrained( HIDREAM_LLAMA_REPO, output_hidden_states = True, @@ -44,4 +105,26 @@ def hidream_te4_kwargs(dtype: Any, hf_token: Optional[str] = None) -> dict[str, torch_dtype = dtype, token = hf_token, ) + if fp8_engages: + try: + from .diffusion_precision import _cast_fp8 + + class _Target: + pass + + cast_target = _Target() + cast_target.dtype = dtype + _cast_fp8(text_encoder_4, cast_target) + logger.info("diffusion.hidream: TE4 layerwise fp8 cast engaged") + except Exception as exc: # noqa: BLE001 -- best-effort like the generic TE pass + # A mid-pass failure can leave fp8 storage / upcast hooks behind; a half-cast + # encoder cannot run as dense, so rebuild it fresh instead of shipping partial state. + logger.warning("diffusion.hidream: TE4 fp8 cast failed, reloading dense: %s", exc) + text_encoder_4 = LlamaForCausalLM.from_pretrained( + HIDREAM_LLAMA_REPO, + output_hidden_states = True, + output_attentions = True, + torch_dtype = dtype, + token = hf_token, + ) return {"text_encoder_4": text_encoder_4, "tokenizer_4": tokenizer_4} diff --git a/studio/backend/core/inference/diffusion_te_prequant.py b/studio/backend/core/inference/diffusion_te_prequant.py index 035a0dbad8..4bd2a4797a 100644 --- a/studio/backend/core/inference/diffusion_te_prequant.py +++ b/studio/backend/core/inference/diffusion_te_prequant.py @@ -115,13 +115,21 @@ def load_prequant_text_encoder( hf_token: Optional[str] = None, scheme: str = "fp8", logger: Any = None, + config_subfolder: Optional[str] = None, + config_overrides: Optional[dict] = None, ) -> Optional[Any]: """Load the pre-cast text encoder described by ``source`` (on CPU, for pipeline assembly to place), with the layerwise upcast hooks already installed. Returns the encoder, or None on any problem (missing / mismatched / unreadable checkpoint) so the caller falls back to the dense download + cast. Best-effort: - never raises for an unavailable artifact.""" + never raises for an unavailable artifact. + + ``config_subfolder`` overrides where the encoder config lives in ``base`` (default: + the component name; "" means the repo root, for encoders assembled from a separate + standalone repo like HiDream's Llama TE4). ``config_overrides`` sets config fields + the pipeline's assembly normally passes to ``from_pretrained`` (forward-behaviour + flags only; the state dict is unaffected by them).""" try: if source.kind == "path" and not _local_prequant_path_allowed(source.location): _warn( @@ -161,9 +169,13 @@ def load_prequant_text_encoder( ValueError(f"checkpoint te_class {te_class!r} not found in transformers"), ) return None - config = transformers.AutoConfig.from_pretrained( - base, subfolder = component, token = hf_token - ) + subfolder = component if config_subfolder is None else config_subfolder + config_kwargs: dict[str, Any] = {"token": hf_token} + if subfolder: + config_kwargs["subfolder"] = subfolder + config = transformers.AutoConfig.from_pretrained(base, **config_kwargs) + for key, value in (config_overrides or {}).items(): + setattr(config, key, value) from accelerate import init_empty_weights with init_empty_weights(): diff --git a/studio/backend/tests/test_diffusion_te_prequant.py b/studio/backend/tests/test_diffusion_te_prequant.py index 4ddb7af59b..52aaeed83b 100644 --- a/studio/backend/tests/test_diffusion_te_prequant.py +++ b/studio/backend/tests/test_diffusion_te_prequant.py @@ -261,6 +261,136 @@ def test_hosted_te_prequant_entries(): assert te_prequant_repo_filename( "unsloth/LTX-2-FP8", "text_encoder", "fp8" ) == "LTX-2-text_encoder-FP8.pt" + # HiDream's heavyweight is TE4 (Llama-3.1-8B), engaged via hidream_te4_kwargs because + # the generic quantize_text_encoders pass only covers text_encoder.._3. + assert detect_family("HiDream-ai/HiDream-I1-Full").te_prequant_repos == ( + ("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"), + ) + assert te_prequant_repo_filename( + "unsloth/HiDream-I1-Full-FP8", "text_encoder_4", "fp8" + ) == "HiDream-I1-Full-text_encoder_4-FP8.pt" + + +def _hidream_transformers_stub(monkeypatch, recorder): + """Fake transformers surface for hidream_te4_kwargs: records from_pretrained calls.""" + import sys + + class _FakeLlama: + def __init__(self, tag): + self.tag = tag + + class _LlamaCls: + @staticmethod + def from_pretrained(repo, **kwargs): + recorder.append(("llama_from_pretrained", repo)) + return _FakeLlama(f"dense{len(recorder)}") + + class _TokCls: + @staticmethod + def from_pretrained(repo, **kwargs): + recorder.append(("tokenizer", repo)) + return "tok4" + + fake = types.ModuleType("transformers") + fake.AutoTokenizer = _TokCls + fake.LlamaForCausalLM = _LlamaCls + monkeypatch.setitem(sys.modules, "transformers", fake) + return _FakeLlama + + +def test_hidream_te4_stays_dense_without_fp8(monkeypatch): + from core.inference.diffusion_hidream import hidream_te4_kwargs + + recorder: list = [] + _hidream_transformers_stub(monkeypatch, recorder) + out = hidream_te4_kwargs( + None, None, fam = _fam(name = "hidream-i1"), te_quant_mode = None, target = _target() + ) + assert out["tokenizer_4"] == "tok4" + assert getattr(out["text_encoder_4"], "tag", "").startswith("dense") + # No cast attempted: mode None normalises to no TE quant. + assert ("llama_from_pretrained", "unsloth/Meta-Llama-3.1-8B-Instruct") in recorder + + +def test_hidream_te4_prefers_precast_checkpoint(monkeypatch): + import core.inference.diffusion_hidream as dh + import core.inference.diffusion_precision as precision + + recorder: list = [] + _hidream_transformers_stub(monkeypatch, recorder) + monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True) + precast = object() + calls: dict = {} + + def _fake_load(base, component, source, **kwargs): + calls["base"] = base + calls["component"] = component + calls["config_subfolder"] = kwargs.get("config_subfolder") + calls["config_overrides"] = kwargs.get("config_overrides") + return precast + + monkeypatch.setattr(tpq, "load_prequant_text_encoder", _fake_load) + fam = _fam( + te_prequant_repos = (("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),), + name = "hidream-i1", + ) + out = dh.hidream_te4_kwargs( + None, None, fam = fam, te_quant_mode = "fp8", target = _target() + ) + assert out["text_encoder_4"] is precast + assert calls["base"] == "unsloth/Meta-Llama-3.1-8B-Instruct" + assert calls["component"] == "text_encoder_4" + # Standalone repo: config at the root, forward flags the pipeline needs applied. + assert calls["config_subfolder"] == "" + assert calls["config_overrides"] == { + "output_hidden_states": True, + "output_attentions": True, + } + # The dense Llama download never ran. + assert ("llama_from_pretrained", "unsloth/Meta-Llama-3.1-8B-Instruct") not in recorder + + +def test_hidream_te4_falls_back_to_dense_cast(monkeypatch): + import core.inference.diffusion_hidream as dh + import core.inference.diffusion_precision as precision + + recorder: list = [] + _hidream_transformers_stub(monkeypatch, recorder) + monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True) + monkeypatch.setattr(tpq, "load_prequant_text_encoder", lambda *a, **k: None) + cast: list = [] + monkeypatch.setattr(precision, "_cast_fp8", lambda enc, tgt: cast.append(enc)) + fam = _fam( + te_prequant_repos = (("fp8", "text_encoder_4", "unsloth/HiDream-I1-Full-FP8"),), + name = "hidream-i1", + ) + out = dh.hidream_te4_kwargs( + None, None, fam = fam, te_quant_mode = "fp8", target = _target() + ) + assert cast == [out["text_encoder_4"]] + assert ("llama_from_pretrained", "unsloth/Meta-Llama-3.1-8B-Instruct") in recorder + + +def test_hidream_te4_partial_cast_reloads_dense(monkeypatch): + """A mid-pass TE4 cast failure must ship a FRESH dense encoder, not partial fp8 state.""" + import core.inference.diffusion_hidream as dh + import core.inference.diffusion_precision as precision + + recorder: list = [] + _hidream_transformers_stub(monkeypatch, recorder) + monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True) + + def _boom(enc, tgt): + raise RuntimeError("cast failed mid-pass") + + monkeypatch.setattr(precision, "_cast_fp8", _boom) + fam = _fam(name = "hidream-i1") # no hosted entry -> dense + cast path + out = dh.hidream_te4_kwargs( + None, None, fam = fam, te_quant_mode = "fp8", target = _target() + ) + dense_loads = [r for r in recorder if r[0] == "llama_from_pretrained"] + assert len(dense_loads) == 2 # initial load + the fail-safe reload + assert getattr(out["text_encoder_4"], "tag", "").startswith("dense") def test_assemble_pipe_injects_precast_te(monkeypatch):