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)