Studio diffusion (Phase 9): pre-quantized transformer loading

The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on
the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs
13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline
(scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton
on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True)
the quantized weights, so the dense bf16 never touches the GPU.

Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching
GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same
torchao config + min_features filter the runtime path uses, applied ahead of time.

New core/inference/diffusion_prequant.py (resolve_prequant_source +
load_prequantized_transformer, best-effort, lazy imports). diffusion.py
_load_dense_quant_pipeline tries the pre-quant source first and falls back to the
dense materialise+quantise path, then to GGUF, so the default is unchanged.
DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty
prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for
the resolver, the meta-init+assign loader, and the backend branch selection +
fallbacks; GPU verification via scripts/verify_prequant_backend.py.

studio/backend/core/inference/diffusion_prequant.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Load a *pre-quantized* transformer instead of quantising a dense one on the GPU.
+
+The opt-in fast transformer_quant path (see ``diffusion_transformer_quant.py``) loads
+the dense bf16 transformer and torchao-``quantize_``s it in place. That materialises the
+full bf16 weights on the GPU before quantising, so the load peak is ~2x the GGUF's and it
+pulls the full bf16 download. When a transformer has already been quantised once and saved
+(``scripts/build_prequant_checkpoint.py``), this module loads those weights directly:
+
+  1. build the transformer skeleton on the ``meta`` device (no storage) via
+     ``accelerate.init_empty_weights`` + ``from_config``;
+  2. ``load_state_dict(assign=True)`` the quantized state dict (the torchao weight subclass
+     tensors are assigned in, not copied), so the dense bf16 never touches the GPU;
+  3. move to the device.
+
+Measured (B200, Z-Image fp8): transformer GPU load peak 12.9 -> 6.3 GB, download 12 ->
+6.28 GB, output bit-identical (LPIPS 0.0). The checkpoint carries the exact same scheme +
+``min_features`` as the runtime path, so the result is identical to quantising on the fly.
+
+Best-effort and lazily imported throughout: a missing / mismatched / unreadable checkpoint
+returns None and the caller falls back to the dense-quantise path (and then to GGUF). All
+behaviour is gated on a configured source -- with nothing configured this module is inert.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+# torch.save dict layout this module reads (and the build script writes). Bumped if the
+# on-disk structure changes so an old/foreign artifact is rejected rather than mis-loaded.
+PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1"
+
+
+@dataclass(frozen = True)
+class PrequantSource:
+    """Where a pre-quantized transformer checkpoint lives. ``kind`` is "path" (a local
+    file) or "repo" (a Hub repo id in ``location`` + ``filename`` inside it)."""
+
+    kind: str
+    location: str
+    filename: Optional[str] = None
+
+
+def prequant_filename(scheme: str) -> str:
+    """The conventional checkpoint filename for ``scheme`` inside a Hub repo."""
+    return f"transformer_{scheme}.pt"
+
+
+def resolve_prequant_source(
+    fam: Any,
+    scheme: str,
+    *,
+    path_override: Optional[str] = None,
+) -> Optional[PrequantSource]:
+    """Resolve where the pre-quantized checkpoint for ``(fam, scheme)`` should come from.
+
+    Priority: (1) an explicit local ``path_override`` (testing / power users); (2) the
+    family's hosted repo for ``scheme``; (3) None -> no pre-quant, caller quantises dense.
+    Pure: no IO, no torch -- it only decides the source, the loader fetches it.
+    """
+    override = (path_override or "").strip()
+    if override:
+        return PrequantSource(kind = "path", location = override, filename = None)
+    try:
+        from .diffusion_families import family_prequant_repo
+
+        repo_id = family_prequant_repo(fam, scheme)
+    except Exception:  # noqa: BLE001 — a bad family object must not break the load
+        repo_id = None
+    if repo_id:
+        return PrequantSource(kind = "repo", location = repo_id, filename = prequant_filename(scheme))
+    return None
+
+
+def load_prequantized_transformer(
+    transformer_cls: Any,
+    base: str,
+    source: PrequantSource,
+    *,
+    device: str,
+    dtype: Any,
+    hf_token: Optional[str] = None,
+    scheme: str,
+    logger: Any = None,
+) -> Optional[Any]:
+    """Load the pre-quantized transformer described by ``source`` onto ``device``.
+
+    Returns the placed, already-quantized transformer, or None on any problem (missing /
+    mismatched / unreadable checkpoint, or a meta-init the class does not support) so the
+    caller falls back to the dense-quantise path. Best-effort: never raises for an
+    ordinary unavailable artifact.
+    """
+    try:
+        path = _resolve_checkpoint_path(source, hf_token)
+        if path is None:
+            return None
+
+        import torch
+
+        # torchao weight subclasses are not safetensors-serializable, so the checkpoint is
+        # a torch.save pickle. weights_only=False is required to rebuild those subclasses;
+        # only a configured family repo (first-party) or an explicit local path reaches
+        # here, which is the trust signal -- this never loads an arbitrary remote pickle.
+        ckpt = torch.load(path, weights_only = False, map_location = "cpu")
+        if not _validate_checkpoint(ckpt, scheme, base, logger):
+            return None
+        state_dict = ckpt["state_dict"]
+
+        config = transformer_cls.load_config(base, subfolder = "transformer", token = hf_token)
+        from accelerate import init_empty_weights
+
+        with init_empty_weights():
+            transformer = transformer_cls.from_config(config)
+        # assign=True swaps in the loaded (quantized) tensors rather than copying into the
+        # meta tensors (a copy into meta is a no-op); strict=True since the saved state
+        # dict is the full state dict of the same class (non-persistent buffers excluded).
+        transformer.load_state_dict(state_dict, strict = True, assign = True)
+        if _has_meta_tensors(transformer):
+            # A class with non-persistent buffers (computed in __init__, absent from the
+            # state dict) leaves those on meta. Rebuild on CPU so the buffers hold their
+            # real values, then re-assign the quantized weights. The dense bf16 lives in
+            # CPU RAM only -- the GPU still receives just the quantized footprint.
+            transformer = transformer_cls.from_config(config)
+            transformer.load_state_dict(state_dict, strict = True, assign = True)
+
+        transformer = transformer.to(device)
+        try:  # diagnostic marker, mirrors the runtime-quant path
+            transformer._unsloth_runtime_quant = scheme
+        except Exception:  # noqa: BLE001 — marker is best-effort
+            pass
+        if logger is not None:
+            logger.info(
+                "diffusion.prequant: loaded %s checkpoint (%s) onto %s",
+                scheme,
+                source.kind,
+                device,
+            )
+        return transformer
+    except Exception as exc:  # noqa: BLE001 — fall back to the dense-quantise path
+        _warn(logger, f"{scheme}:{source.kind}", exc)
+        return None
+
+
+def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) -> Optional[str]:
+    """The local file path for ``source``, downloading from the Hub if needed; None if absent."""
+    if source.kind == "path":
+        import os
+
+        return source.location if os.path.isfile(source.location) else None
+    if source.kind == "repo":
+        from huggingface_hub import hf_hub_download
+
+        return hf_hub_download(
+            repo_id = source.location, filename = source.filename, token = hf_token
+        )
+    return None
+
+
+def _validate_checkpoint(ckpt: Any, scheme: str, base: str, logger: Any) -> bool:
+    """Reject a checkpoint that is the wrong format / scheme / base model."""
+    if not isinstance(ckpt, dict) or ckpt.get("format") != PREQUANT_FORMAT:
+        _warn(logger, scheme, ValueError("unrecognised pre-quant checkpoint format"))
+        return False
+    if "state_dict" not in ckpt:
+        _warn(logger, scheme, ValueError("pre-quant checkpoint has no state_dict"))
+        return False
+    meta = ckpt.get("metadata") or {}
+    if meta.get("scheme") != scheme:
+        _warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}"))
+        return False
+    ckpt_base = meta.get("base_model_id")
+    if ckpt_base and base and ckpt_base != base:
+        _warn(logger, scheme, ValueError(f"checkpoint base {ckpt_base!r} != {base!r}"))
+        return False
+    return True
+
+
+def _has_meta_tensors(module: Any) -> bool:
+    """True if any parameter or buffer is still on the meta device after loading."""
+    try:
+        for tensor in list(module.parameters()) + list(module.buffers()):
+            if getattr(tensor, "is_meta", False):
+                return True
+    except Exception:  # noqa: BLE001
+        return False
+    return False
+
+
+def _warn(logger: Any, what: str, exc: Exception) -> None:
+    if logger is not None:
+        logger.warning("diffusion.prequant: %s failed: %s", what, exc)