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.
This commit is contained in:
Daniel Han 2026-07-18 06:52:46 +00:00
commit 5249a52356
3 changed files with 21 additions and 3 deletions

View file

@ -95,8 +95,10 @@ def main(argv = None) -> int:
"te_class": encoder_cls_name,
"torch_dtype": args.dtype,
"cast_backend": "diffusers_layerwise",
"torch_version": torch.__version__,
"transformers_version": transformers.__version__,
# str(): torch.__version__ is a TorchVersion object; pickling it into the
# checkpoint makes torch.load(weights_only=True) reject the whole artifact.
"torch_version": str(torch.__version__),
"transformers_version": str(transformers.__version__),
}
ckpt = {
"format": TE_PREQUANT_FORMAT,