scripts/build_prequant_checkpoint.py
Rama main · réplica de lectura del repositorio original.
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Quantizes the DiT once so runtime loads it with assign=True, skipping the dense bf16 GPU load peak."""
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
from typing import Any, Optional, Sequence
BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"
# Mirrors core.inference.diffusion_prequant.DEFAULT_PREQUANT_COMPONENT; the backend joins sys.path only in main().
DEFAULT_COMPONENT = "transformer"
def resolve_build_family(base: str, override: Optional[str] = None) -> Optional[Any]:
"""The image family for ``base`` / ``override``, else the video family (Wan, HunyuanVideo), else None."""
from core.inference.diffusion_families import detect_family
from core.inference.video_families import detect_video_family
return detect_family(base, override = override) or detect_video_family(base, override = override)
def convrot_refusal(
group: int, rotatable: Sequence[str], not_divisible: Sequence[str]
) -> Optional[str]:
"""Refuse a ConvRot build with no rotatable Linear before the hours of quantising, not after:
its empty fqn list is refused at load time, leaving an unopenable artifact"""
if rotatable:
return None
return (
f"ConvRot group {group} divides the in_features of none of the {len(not_divisible)} "
"quantized linears, so the checkpoint would record an empty rotation and be refused at "
"load time. Pick a smaller power-of-4 group, or drop --convrot-groupsize."
)
def upload_destination(
fam: Any,
scheme: str,
*,
rotated: bool,
safetensors: bool = False,
override: Optional[str] = None,
upload_repo: Optional[str] = None,
component: str = DEFAULT_COMPONENT,
convrot_group: Optional[int] = None,
) -> str:
"""The repo-root filename this build should publish under.
The loader asks for the family's declared ``prequant_filenames`` name first and the derived
``<Model>-<SCHEME>.pt`` second, so a ROTATED artifact published under the legacy
``transformer_<scheme>.pt`` is either never resolved at all, or resolved as the fallback by a
build too old to honour the rotation, which then refuses the v2 tag and drops to the dense
download. A rotated build therefore goes to the declared name or nowhere. Plain builds keep
the legacy name they have always used.
A SAFETENSORS build is in the same position for a different reason: both derived names end in
``.pt``, so no build ever asks the Hub for a safetensors artifact unless the family names it.
Uploading one under a derived name produces a file that is reachable by nothing and a repo that
looks like it has a checkpoint when it does not, so it is refused here rather than discovered
as a silent dense fallback later.
An ``override`` skips the family table, because naming the artifact by hand is the escape hatch
for a repo the table does not describe yet. It does NOT skip the container check: every loader
dispatches on the extension alone, so a safetensors build published as ``.pt`` is read as a
pickle and a pickle published as ``.safetensors`` is read from a header it does not have. Either
way the upload succeeds and the artifact is unopenable, after the hours the quantization took.
nvfp4 publishes safetensors only: Studio never resolves or opens an nvfp4 pickle.
"""
if (scheme or "").strip().lower() == "nvfp4" and not safetensors:
raise ValueError(
"nvfp4 checkpoints are published as .safetensors only (Studio never reads an nvfp4 "
"pickle); write --out as a .safetensors file"
)
if override:
wanted = ".safetensors" if safetensors else ".pt"
if not override.lower().endswith(wanted):
container = "safetensors" if safetensors else "torch.save"
raise ValueError(
f"--upload-filename {override!r} does not end in {wanted!r}, but --out writes the "
f"{container} container. The loader dispatches on the extension alone, so this "
"would publish an artifact nothing can open. Rename the upload, or change --out."
)
return override
if component and component != DEFAULT_COMPONENT:
# A second denoiser resolves only its task-specific row, so other names would never be loaded.
from core.inference.diffusion_families import family_prequant_filename
specific = family_prequant_filename(fam, scheme, task = component)
if specific is None or specific == family_prequant_filename(fam, scheme):
raise ValueError(
f"family {getattr(fam, 'name', fam)!r} declares no prequant_filenames entry for "
f"({scheme!r}, {component!r}), so the loader would never ask for this component. "
"Add the entry to the family table, or pass --upload-filename."
)
wanted = ".safetensors" if safetensors else ".pt"
if not specific.lower().endswith(wanted):
raise ValueError(
f"family {getattr(fam, 'name', fam)!r} declares {specific!r} for "
f"({scheme!r}, {component!r}), which does not end in {wanted!r} like --out. "
"Rename --out, or pass --upload-filename."
)
return specific
from core.inference.diffusion_prequant import prequant_filename
if not rotated and not safetensors:
return prequant_filename(scheme)
from core.inference.diffusion_families import family_prequant_filename
from core.inference.diffusion_transformer_quant import (
convrot_prequant_filename,
convrot_spec_for_scheme,
)
family = getattr(fam, "name", None)
rotated_name = convrot_prequant_filename(scheme, family) if rotated else None
if rotated_name:
spec_group = convrot_spec_for_scheme(scheme, family)[0]
if convrot_group is not None and int(convrot_group) != spec_group:
raise ValueError(
f"{rotated_name} is reserved for the ConvRot group {spec_group} build that matches "
f"the runtime path; a group {convrot_group} build needs --upload-filename."
)
preferred = rotated_name or family_prequant_filename(fam, scheme)
why = "a rotated checkpoint" if rotated else "a safetensors checkpoint"
if not preferred:
# Plain safetensors builds have a derived name; rotated builds still need a declared one.
if safetensors and not rotated and upload_repo:
from core.inference.diffusion_prequant import prequant_repo_filename
return prequant_repo_filename(upload_repo, scheme, ".safetensors")
raise ValueError(
f"family {getattr(fam, 'name', fam)!r} declares no prequant_filenames entry for "
f"{scheme!r}, so {why} has no name the loader would ask for. Add the "
"entry to the family table, or pass --upload-filename."
)
# Check both directions: the declared name's extension must match the container.
wanted = ".safetensors" if safetensors else ".pt"
if not preferred.lower().endswith(wanted):
reads_as = "a pickle" if safetensors else "safetensors"
writes = "safetensors" if safetensors else "torch.save"
raise ValueError(
f"family {getattr(fam, 'name', fam)!r} declares {preferred!r} for {scheme!r}, which "
f"does not end in {wanted!r}, but --out writes the {writes} container. The loader "
f"dispatches on the extension alone, so this would be published under a name it reads "
f"as {reads_as}. Point the prequant_filenames entry at the matching artifact, or pass "
"--upload-filename."
)
return preferred
def quant_filter_settings(scheme: str, family: Optional[str]) -> dict:
"""Runtime ``quantize_`` filter inputs for ``scheme``; one dict feeds filter and metadata."""
from core.inference.diffusion_transformer_quant import (
_REQUIRE_BF16_SCHEMES,
divisible_for_scheme,
exclude_tokens_for_scheme,
)
return {
"exclude_name_tokens": list(exclude_tokens_for_scheme(scheme, family)),
"require_bf16": scheme in _REQUIRE_BF16_SCHEMES,
"require_divisible": divisible_for_scheme(scheme),
}
def main(argv = None) -> int:
p = argparse.ArgumentParser()
p.add_argument(
"--base", required = True, help = "diffusers base repo (carries the transformer subfolder)"
)
p.add_argument("--family", required = True, help = "diffusion family name/alias (e.g. z-image)")
p.add_argument(
"--base-model-id",
default = None,
help = "the base id to RECORD in the checkpoint, when --base is a local mirror whose "
"directory name differs from the Hub repo. Must still name this family's base model.",
)
p.add_argument("--scheme", required = True, help = "quant scheme: int8 | fp8 | nvfp4 | mxfp8")
p.add_argument(
"--component",
default = DEFAULT_COMPONENT,
help = "denoiser subfolder to quantise, recorded in the checkpoint (e.g. transformer_2 for "
"Wan2.2 A14B's second expert)",
)
p.add_argument(
"--out",
required = True,
help = "output path; a .safetensors extension writes the safetensors container, anything "
"else writes the torch.save one",
)
p.add_argument("--min-features", type = int, default = 512)
p.add_argument("--dtype", default = "bfloat16", choices = ["bfloat16"])
p.add_argument("--hf-token", default = None)
p.add_argument(
"--convrot-groupsize",
type = int,
default = 0,
help = "bake a ConvRot block-Hadamard activation rotation at this group size (a power of "
"4; 0 = off). A family with an int8 ConvRot spec at this group (Qwen-Image-2.1: "
"256) rotates exactly the Linears that spec names, so the artifact matches the opt-in "
"runtime path (UNSLOTH_DIFFUSION_INT8_CONVROT=1); otherwise every quantized Linear whose in_features the group divides has its "
"weight rotated before quantize_ so the quantizer sees a flatter distribution; the "
"exact fqn list is recorded in the checkpoint and the loader rotates the "
"activations of that list and nothing else. Writes the v2 format tag.",
)
p.add_argument(
"--upload-repo", default = None, help = "optional HF repo id to upload the checkpoint to"
)
p.add_argument("--upload-revision", default = None)
p.add_argument(
"--upload-filename",
default = None,
help = "repo-root filename to publish under; defaults to the family's declared "
"prequant_filenames entry for a rotated build and the legacy transformer_<scheme>.pt "
"otherwise",
)
args = p.parse_args(argv)
if (args.scheme or "").strip().lower() == "nvfp4" and not str(args.out).lower().endswith(
".safetensors"
):
print(
"error: --scheme nvfp4 writes .safetensors only (Studio never reads an NVFP4 pickle); "
"give --out a .safetensors path",
flush = True,
)
return 2
sys.path.insert(0, str(BACKEND))
import torch
import torchao
import diffusers
from core.inference.diffusion_prequant import prequant_format_for
# Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant).
from core.inference.diffusion_transformer_quant import (
FP8_GRANULARITY,
TQ_FP8,
TQ_SCHEMES,
_make_quant_config,
_resolve_fast_accum,
convrot_fqns,
convrot_spec_for_scheme,
make_filter_fn,
)
from torchao.quantization import quantize_
scheme = args.scheme.strip().lower()
if scheme not in TQ_SCHEMES:
print(f"error: --scheme must be one of {TQ_SCHEMES} (not 'auto')", flush = True)
return 2
component = (args.component or "").strip() or DEFAULT_COMPONENT
fam = resolve_build_family(args.base, override = args.family)
if fam is None:
print(f"error: unknown family '{args.family}'", flush = True)
return 2
convrot_group = int(args.convrot_groupsize)
spec_group, spec_suffixes = convrot_spec_for_scheme(scheme, fam.name)
convrot_suffixes: tuple = spec_suffixes if convrot_group and convrot_group == spec_group else ()
# The recorded base may differ from the one read: a local staging dir name would fail the
# loader's _same_base_model tail comparison, so --base-model-id (pinned to base_repo) overrides it.
recorded_base = args.base_model_id or args.base
if args.base_model_id:
# Exact match, not _same_base_model, so a published file always records the canonical id.
if recorded_base.strip() != fam.base_repo:
print(
f"error: --base-model-id {recorded_base!r} is not {fam.name}'s base "
f"({fam.base_repo!r}); it would label this checkpoint as a different model",
flush = True,
)
return 2
transformer_cls = getattr(diffusers, fam.transformer_class)
is_safetensors_out = str(args.out).lower().endswith(".safetensors")
if is_safetensors_out:
from core.inference.prequant_safetensors import (
safetensors_prequant_supported,
scheme_is_flattenable,
)
if not safetensors_prequant_supported():
print(
"error: --out names a .safetensors checkpoint but this install cannot write one "
"(needs torchao >= 0.16 for torchao.prototype.safetensors.safetensors_support, "
"plus the safetensors package)",
flush = True,
)
return 2
# Importable helpers do not mean this scheme flattens (int8 does not through torchao 0.17).
# Probe on a tiny CPU Linear before the long download; None means the probe could not run, so proceed.
from core.inference.diffusion_transformer_quant import _make_quant_config
if scheme_is_flattenable(_make_quant_config(scheme)) is False:
print(
f"error: --out names a .safetensors checkpoint but this torchao quantises "
f"{scheme!r} to a legacy tensor subclass that cannot be written to safetensors. "
"torchao >= 0.18 produces the flattenable subclasses for every scheme Unsloth "
"ships. Upgrade torchao, or write this build as a .pt checkpoint.",
flush = True,
)
return 2
# Resolved before the load so a bad upload target fails fast.
upload_dest = None
if args.upload_repo:
try:
upload_dest = upload_destination(
fam,
scheme,
rotated = bool(convrot_group),
convrot_group = convrot_group,
safetensors = is_safetensors_out,
override = args.upload_filename,
upload_repo = args.upload_repo,
component = component,
)
except ValueError as exc:
print(f"error: {exc}", flush = True)
return 2
print(f"== build prequant ({fam.name}/{scheme}, min_feat={args.min_features}) ==", flush = True)
print(f" loading dense transformer from {args.base} (subfolder={component}) ...", flush = True)
t0 = time.time()
transformer = transformer_cls.from_pretrained(
args.base, subfolder = component, torch_dtype = torch.bfloat16, token = args.hf_token
).to("cuda")
print(f" quantising in place ({scheme}) ...", flush = True)
filter_settings = quant_filter_settings(scheme, fam.name)
exclude_name_tokens = tuple(filter_settings["exclude_name_tokens"])
require_bf16 = filter_settings["require_bf16"]
# fp8 bakes the accumulate mode in; record it so the loader can reject a contradicting request.
fast_accum = _resolve_fast_accum(None) if scheme == TQ_FP8 else None
filter_fn = make_filter_fn(
args.min_features,
exclude_name_tokens = exclude_name_tokens,
require_bf16 = require_bf16,
require_divisible = filter_settings["require_divisible"],
)
# Rotate before quantize_ so the quantizer sees the rotated distribution.
rotation: dict = {}
if convrot_group:
from core.inference.diffusion_convrot import (
rotatable_fqns,
rotate_linears_,
rotation_metadata,
)
group = int(convrot_group)
rotatable, not_divisible = rotatable_fqns(transformer, filter_fn, group)
if convrot_suffixes:
rotatable = convrot_fqns(transformer, filter_fn, group, convrot_suffixes)
refusal = convrot_refusal(group, rotatable, not_divisible)
if refusal:
print(f"error: {refusal}", flush = True)
return 2
rotate_linears_(transformer, rotatable, group)
rotation = rotation_metadata(group, rotatable)
print(
f" rotated {len(rotatable)} linears at ConvRot group {group}; "
f"{len(not_divisible)} quantized linears left plain (in_features not divisible)"
+ (f", e.g. {not_divisible[0]}" if not_divisible else ""),
flush = True,
)
quantize_(transformer, _make_quant_config(scheme), filter_fn = filter_fn)
state_dict = {
k: (v.detach().to("cpu") if hasattr(v, "detach") else v)
for k, v in transformer.state_dict().items()
}
metadata = {
"base_model_id": recorded_base,
"family": fam.name,
"scheme": scheme,
"min_features": args.min_features,
"exclude_name_tokens": list(exclude_name_tokens),
"require_bf16": require_bf16,
"require_divisible": filter_settings["require_divisible"],
"fast_accum": fast_accum,
"torch_dtype": args.dtype,
"quant_backend": "torchao",
"transformer_class": fam.transformer_class,
"component": component,
"torch_version": torch.__version__,
"torchao_version": getattr(torchao, "__version__", "?"),
"diffusers_version": diffusers.__version__,
}
from core.inference.diffusion_prequant import canonical_weight_fingerprint
# Version-neutral, so a file built on one torchao generation verifies on every other.
metadata["fingerprint"] = canonical_weight_fingerprint(state_dict)
# fp8 granularity: lets the loader reject a stale per-tensor checkpoint (runtime needs per-row).
if scheme == TQ_FP8:
metadata["fp8_granularity"] = FP8_GRANULARITY
metadata.update(rotation)
ckpt = {
# v2 when rotated, so an older Unsloth refuses the file instead of mismatching activations.
"format": prequant_format_for(metadata),
"metadata": metadata,
"state_dict": state_dict,
}
out = Path(args.out)
out.parent.mkdir(parents = True, exist_ok = True)
if is_safetensors_out:
from core.inference.prequant_safetensors import save_prequant_safetensors
save_prequant_safetensors(
str(out),
fmt = ckpt["format"],
state_dict = state_dict,
metadata = metadata,
)
else:
torch.save(ckpt, out)
size_gb = out.stat().st_size / 1e9
print(f" saved {out} ({size_gb:.2f} GB) in {time.time() - t0:.0f}s", flush = True)
print(f" metadata: {ckpt['metadata']}", flush = True)
if args.upload_repo:
from huggingface_hub import HfApi
dest = upload_dest
print(f" uploading -> {args.upload_repo}:{dest} ...", flush = True)
api = HfApi(token = args.hf_token)
api.create_repo(args.upload_repo, exist_ok = True)
api.upload_file(
path_or_fileobj = str(out),
path_in_repo = dest,
repo_id = args.upload_repo,
revision = args.upload_revision,
)
print(f" uploaded {dest} to {args.upload_repo}", flush = True)
print("BUILD-PREQUANT-DONE", flush = True)
return 0
if __name__ == "__main__":
sys.exit(main())