Commit graph

4 commits

Author SHA1 Message Date
Daniel Han
36df317293 Trim the comments across the diffusion backend
Comment-only pass over the Python this PR touches: drop what the code already
says, collapse multi-line explanations that still read on one line, and keep
the reasoning that is not recoverable from the code. No code, docstring
semantics or behaviour changes; verified with an AST comparison against the
previous revision, and the backend suite is unchanged (same 37 environment
failures as before: the API integration tests that need a live keyed server,
the flash-attn install hooks, and the GPU memory fields).
2026-07-26 20:31:19 +00:00
Daniel Han
a440ee748a 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.
2026-07-18 07:55:24 +00:00
Daniel Han
645f3858d1 Fix pre-cast TE checkpoint loading and engagement reporting
Two bugs found while building the hosted checkpoints:

- The builder recorded torch.__version__ (a TorchVersion object) in the
  checkpoint metadata, so torch.load(weights_only=True) rejected every
  artifact and the loader silently fell back to the dense download.
  Record plain strings.
- Re-applying the layerwise fp8 cast to an injected pre-cast encoder
  raised on the duplicate hook registration, making quantize_text_encoders
  report the engaged cast as failed (status showed no TE quant while the
  encoder ran fp8). _cast_fp8 now returns early when the hooks are
  already installed.

Also corrects the LTX TE size note: Gemma3-12B stored fp32 (~49 GB), not 27B.
2026-07-18 06:52:46 +00:00
Daniel Han
4e9aab60ae Add the pre-cast text-encoder checkpoint builder
Applies the runtime layerwise fp8 storage cast to a model's dense text
encoder once and saves the cast state dict with baked metadata (format
tag, base_model_id, family, scheme, component, te_class, versions) in
the layout diffusion_te_prequant.py validates. Resolves the encoder
class from the checkpoint's config.architectures so the recorded
te_class matches what the pipeline instantiates. CPU-runnable: the cast
touches storage dtypes only.
2026-07-18 06:24:48 +00:00