ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Extend the fp8 TE quant to HiDream's Llama text_encoder_4

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-18 07:55:24 +00:00
commit a3942d9924
6 changed files with 277 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

21
scripts/build_te_prequant_checkpoint.py
View file

@ -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)

30
studio/backend/core/inference/diffusion.py
View file

@ -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:

3
studio/backend/core/inference/diffusion_families.py
View file

@ -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",

89
studio/backend/core/inference/diffusion_hidream.py
View file

@ -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}

20
studio/backend/core/inference/diffusion_te_prequant.py
View file

@ -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():

130
studio/backend/tests/test_diffusion_te_prequant.py
View file

@ -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):