mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-19 08:54:21 +00:00
Studio: whole-model NVFP4 for the video families with hosted pre-quantized denoisers (#10729)
* studio: whole-model NVFP4 for the video families with hosted pre-quantized denoisers Wan2.2-TI2V-5B, Wan2.2-T2V-A14B (both experts), HunyuanVideo-1.5 480p and 720p run NVFP4 on the torchao backend from hosted pre-quantized denoisers, so a load never quantises 5 to 28 GB of dense weights on the fly. Every admitted linear (min 512 features, in and out divisible by 16) is NVFP4; nothing is per-layer here. - one shared auto walk with a per-family head in front of it, and divisible_for_scheme - ZeroRowSafeLinear, a live bug: torchao's nvfp4 activation path raises on HunyuanVideo-1.5's empty image stream on every text-to-video render - loader gates on the packed-weight fingerprint, the component and require_divisible - conventional video pre-quant seeding, all components or none, with the memory plan taking the artifact's resident size rather than the dense build peak - registry rows for the four hosted NVFP4 repos nvfp4 stays an explicit opt-in here; no family gets an auto row on this branch. The tooling that builds and fingerprints these artifacts is not in this repo. It lives in the private scripts repo with the rest of the quant generation tooling. * Tighten the comments and docstrings on the video prequant path Collapse the long single-line comments this branch added to the repo's wrap width, shorten the prose docstrings to the parts that say why, and drop the test docstrings that only restate the test name. Comments, docstrings and whitespace only: no code changes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Credit the resident video pipeline's memory in the seed plan * Pin the residency check in the two-expert download plan test * Credit only the target card's resident bytes and fingerprint built checkpoints The residency plan counted every accelerator storage of the old pipeline as reclaimable on the card the new load targets, so switching GPUs could skip the dense denoiser shards on a card that never held them. Count only the target's. The builder never wrote metadata["fingerprint"], so the loader's packed-weight check had nothing to compare against for artifacts it produced. * Fingerprint torchao 0.18 Int8Tensor prequant weights torchao >= 0.18 int8 quantization produces Int8Tensor, which the packed fingerprint payload map did not list, so every int8 weight was skipped and the recorded fingerprint verified nothing. Hash its qdata, scale and optional zero point / activation slots, with a stub test and a real torchao round trip that catches a flipped byte. * Apply the GEMM alignment floor in the offline prequant builder and record it The runtime quantize_ skips Linears whose in/out features miss the scheme's tiling floor, but build_prequant_checkpoint.py did not, so a ragged layer could be baked as a quantized GEMM the runtime leaves dense. The builder now passes divisible_for_scheme(scheme) to make_filter_fn and stamps require_divisible so the loader can verify it. * Trim comments * Studio: NVFP4 kill switch for image and video, default off The hosted unsloth/*-NVFP4 checkpoints are still private, so NVFP4 is disabled in general until UNSLOTH_NVFP4_DIFFUSION is set to a truthy value (1, true, yes, on). The default lives in one constant, NVFP4_DIFFUSION_DEFAULT in core/inference/diffusion_nvfp4_flag.py, so turning it on for everyone later is a one-line change. With the switch on, behaviour is unchanged. While it is off: - auto never picks nvfp4 for the dense transformer (ladder and family prefer rows), and no smoke probe is spent on it; - an explicit transformer_quant or text_encoder_quant of nvfp4 gets a 400 "NVFP4 is disabled in this build" from every image and video load and download-plan route, ahead of the precision fallback; - no hosted NVFP4 denoiser or encoder is resolved, seeded, planned or sized, and the Hub helpers never ask about a *-NVFP4 repo; - /images/info estimates and the /api/system dense ladder omit nvfp4, and /api/system reports nvfp4_diffusion so the Precision and Text encoder selects hide NVFP4. The offline builder does not consult the switch. Existing NVFP4 test modules run with it on through a conftest fixture; the default-off behaviour, and that the enabled path is still reachable end to end, are covered in test_nvfp4_diffusion_flag.py. * Studio: keep the diffusion bench NVFP4-capable under the kill switch An explicit --transformer-quant or --text-encoder-quant nvfp4 turns the switch on for the bench itself, like the other developer tools. The switch tests ask auto the way a load does, with a hosted checkpoint where the selector requires one, so the enabled-path check stays meaningful on the branches that add that requirement. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: reset a hidden NVFP4 pick to Auto when the backend switch is off A Precision or Text encoder select reseeded to nvfp4 by an earlier load kept that value after a restart turned UNSLOTH_NVFP4_DIFFUSION off: the option was filtered out, so the select rendered blank and the next load or Reapply sent nvfp4 and got a 400. Once /api/system reports the switch off, a held nvfp4 now falls back to auto on the image and video pages. Nothing is reset before system info arrives, since the switch reads off by default until then. * Studio: keep the hosted video denoiser seed for an explicit scheme under Speed off Only an auto precision is rewritten to off under Speed off; an explicit scheme is still quantized (and upgrades the speed), so suppressing the seed staged the dense denoiser for a runtime quantize. Matches diffusion.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: update the Speed off video plan test for the kept explicit seed * Studio: refuse an NVFP4 checkpoint loaded directly as the model while NVFP4 is off With UNSLOTH_NVFP4_DIFFUSION off, a hosted NVFP4 repo (or any repo or local path whose pre-quant metadata records scheme nvfp4) picked as the model itself named no precision, so none of the scheme gates saw it. The image and video load and download-plan routes now 400 such a pick with the same disabled message style. Detection reads the recorded scheme of a local file, a local directory or a cached Hub snapshot, then the family NVFP4 repo tables, then the -NVFP4 name. Switched on, the pick proceeds as before. * Studio: format the NVFP4 checkpoint refusal tests * Studio: skip pipeline roots and the thread hop in the NVFP4 checkpoint refusal A diffusers pipeline loads from its component folders, so a stray root artifact no longer makes its directory read as an NVFP4 checkpoint. With the switch on the route helper returns before handing anything to a thread. * Studio: stamp the built denoiser component and drop an unused metadata reader * Studio: keep the NVFP4 flag and speed-off video plan tests off the runner's own GPU * Studio: trim comments in the NVFP4 video layer * Studio: check the hosted video seed rows before spending a quant smoke probe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: let the prequant builder target the second Wan A14B expert Add --component (default transformer) for the subfolder, metadata and upload name, and resolve video families so the Wan and HunyuanVideo denoisers can be built. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
095b0dee3c
commit
8ecfd30923
36 changed files with 4366 additions and 238 deletions
|
|
@ -28,6 +28,17 @@ 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]
|
||||
|
|
@ -55,6 +66,7 @@ def upload_destination(
|
|||
safetensors: bool = False,
|
||||
override: Optional[str] = None,
|
||||
upload_repo: Optional[str] = None,
|
||||
component: str = DEFAULT_COMPONENT,
|
||||
) -> str:
|
||||
"""The repo-root filename this build should publish under.
|
||||
|
||||
|
|
@ -87,6 +99,26 @@ def upload_destination(
|
|||
"would publish an artifact nothing can open. Rename the upload, or change --out."
|
||||
)
|
||||
return override
|
||||
if component and component != DEFAULT_COMPONENT:
|
||||
# A second denoiser (Wan A14B's transformer_2) resolves ONLY its task-specific row, with no
|
||||
# fallback, so any other name would publish an artifact the loader never asks for.
|
||||
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:
|
||||
|
|
@ -129,6 +161,20 @@ def upload_destination(
|
|||
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(
|
||||
|
|
@ -142,6 +188,12 @@ def main(argv = None) -> int:
|
|||
"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,
|
||||
|
|
@ -179,7 +231,6 @@ def main(argv = None) -> int:
|
|||
import torchao
|
||||
import diffusers
|
||||
|
||||
from core.inference.diffusion_families import detect_family
|
||||
from core.inference.diffusion_prequant import prequant_format_for
|
||||
|
||||
# Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant).
|
||||
|
|
@ -187,10 +238,8 @@ def main(argv = None) -> int:
|
|||
FP8_GRANULARITY,
|
||||
TQ_FP8,
|
||||
TQ_SCHEMES,
|
||||
_REQUIRE_BF16_SCHEMES,
|
||||
_make_quant_config,
|
||||
_resolve_fast_accum,
|
||||
exclude_tokens_for_scheme,
|
||||
make_filter_fn,
|
||||
)
|
||||
from torchao.quantization import quantize_
|
||||
|
|
@ -199,7 +248,8 @@ def main(argv = None) -> int:
|
|||
if scheme not in TQ_SCHEMES:
|
||||
print(f"error: --scheme must be one of {TQ_SCHEMES} (not 'auto')", flush = True)
|
||||
return 2
|
||||
fam = detect_family(args.base, override = args.family)
|
||||
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
|
||||
|
|
@ -268,29 +318,29 @@ def main(argv = None) -> int:
|
|||
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=transformer) ...", 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 = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token
|
||||
args.base, subfolder = component, torch_dtype = torch.bfloat16, token = args.hf_token
|
||||
).to("cuda")
|
||||
print(f" quantising in place ({scheme}) ...", flush = True)
|
||||
# Mirror the runtime exclusions: int8 skips the M=1 modulation projections (torch._int_mm needs M>16) plus
|
||||
# per-family ones; family=None bakes linears the runtime rejects.
|
||||
exclude_name_tokens = exclude_tokens_for_scheme(scheme, fam.name)
|
||||
# fp8 / mxfp8 need bf16 weights, so skip non-bf16 Linears; nvfp4 handles fp32. Mirrors the runtime gate.
|
||||
require_bf16 = scheme in _REQUIRE_BF16_SCHEMES
|
||||
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"],
|
||||
)
|
||||
|
||||
# ConvRot, BEFORE quantize_: rotating the weights is only worth anything if the quantizer then sees the rotated
|
||||
|
|
@ -332,14 +382,20 @@ def main(argv = None) -> int:
|
|||
# Let the loader reject a checkpoint that would not match the runtime path.
|
||||
"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,
|
||||
# The subfolder built above: the loader refuses it as another denoiser (e.g. transformer_2).
|
||||
"component": component,
|
||||
"torch_version": torch.__version__,
|
||||
"torchao_version": getattr(torchao, "__version__", "?"),
|
||||
"diffusers_version": diffusers.__version__,
|
||||
}
|
||||
from core.inference.diffusion_prequant import packed_weight_fingerprint
|
||||
|
||||
metadata["fingerprint"] = packed_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
|
||||
|
|
|
|||
|
|
@ -598,6 +598,8 @@ def _build_parser() -> argparse.ArgumentParser:
|
|||
|
||||
def main(argv: Optional[list[str]] = None) -> int:
|
||||
args = _build_parser().parse_args(argv)
|
||||
if "nvfp4" in (args.transformer_quant, args.text_encoder_quant):
|
||||
os.environ["UNSLOTH_NVFP4_DIFFUSION"] = "1"
|
||||
if isinstance(args.gguf, str):
|
||||
args.gguf = args.gguf.strip() or None
|
||||
if bool(args.write_baseline) == bool(args.compare):
|
||||
|
|
|
|||
|
|
@ -84,6 +84,7 @@ def _lpips(ref, arr):
|
|||
|
||||
def run(mode, steps, seed, res):
|
||||
sys.path.insert(0, str(BACKEND))
|
||||
os.environ["UNSLOTH_NVFP4_DIFFUSION"] = "1"
|
||||
import torch
|
||||
import diffusers
|
||||
from core.inference.diffusion_prequant import (
|
||||
|
|
|
|||
|
|
@ -1877,7 +1877,15 @@ class DiffusionBackend:
|
|||
``begin_load`` runs inside ``acquire_for``, which evicts chat under the arbiter lock before
|
||||
the register callback, and after ``select_and_activate_engine``, which unloads the resident
|
||||
model on an engine switch.
|
||||
|
||||
NVFP4 while the NVFP4 switch is off raises ``ValueError`` (the route's 400) first, ahead of
|
||||
the opt-in fallback: a disabled scheme is refused, never swapped.
|
||||
"""
|
||||
from .diffusion_nvfp4_flag import refuse_disabled_nvfp4
|
||||
|
||||
refuse_disabled_nvfp4(
|
||||
transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant
|
||||
)
|
||||
if precision_fallback_allowed():
|
||||
return
|
||||
pinned = normalize_transformer_quant(transformer_quant)
|
||||
|
|
@ -3405,6 +3413,10 @@ class DiffusionBackend:
|
|||
candidate answer, and skipping them all reports the miss, which keeps the shards."""
|
||||
if source is None or getattr(source, "kind", None) != "repo":
|
||||
return None
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked, nvfp4_repo_blocked
|
||||
|
||||
if nvfp4_blocked(scheme) or nvfp4_repo_blocked(source.location):
|
||||
return None
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
info = HfApi(token = hf_token or None).model_info(source.location, files_metadata = True)
|
||||
|
|
|
|||
|
|
@ -231,6 +231,10 @@ def estimate_dense_quant(
|
|||
) -> Optional[DenseQuantEstimate]:
|
||||
"""Estimate the candidate's footprint from the family table, or None when the
|
||||
family (or scheme factor) is unknown."""
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked
|
||||
|
||||
if nvfp4_blocked(scheme):
|
||||
return None
|
||||
components = family_bf16_components_gb(fam, base_repo)
|
||||
factor = _QUANT_STEADY_FACTOR.get(scheme)
|
||||
if components is None or factor is None:
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ from pathlib import Path, PurePosixPath
|
|||
from typing import NamedTuple, Optional, Sequence
|
||||
from utils.paths.path_utils import is_appledouble_metadata
|
||||
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked
|
||||
|
||||
|
||||
# Runtime->route contract: the /images/generate route matches these messages EXACTLY for a 409 (vs a 500), so both
|
||||
# engines raise them verbatim.
|
||||
|
|
@ -1183,6 +1185,8 @@ def family_prequant_repo(
|
|||
close enough that planning around it costs nothing, since the base_model_id validation
|
||||
refuses the artifact well after the plan was made. A base whose weights really differ belongs
|
||||
in ``prequant_excluded_bases``, which returns None here instead."""
|
||||
if nvfp4_blocked(scheme):
|
||||
return None
|
||||
# Both tables are keyed on lowercased upstream ids.
|
||||
base = canonical_base(base_repo).lower()
|
||||
if base:
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from __future__ import annotations
|
|||
from typing import Any
|
||||
|
||||
from .diffusion_auto_policy import _FAMILY_BF16_GB, _QUANT_STEADY_FACTOR
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked
|
||||
|
||||
|
||||
def _round1(value: float) -> float:
|
||||
|
|
@ -33,6 +34,8 @@ def family_inference_infos() -> list[dict[str, Any]]:
|
|||
companions_gb = text_encoders_gb + vae_gb
|
||||
estimated = {"bf16": _round1(transformer_gb + companions_gb)}
|
||||
for scheme, factor in _QUANT_STEADY_FACTOR.items():
|
||||
if nvfp4_blocked(scheme):
|
||||
continue
|
||||
estimated[scheme] = _round1(transformer_gb * factor + companions_gb)
|
||||
infos.append(
|
||||
{
|
||||
|
|
|
|||
98
studio/backend/core/inference/diffusion_nvfp4_flag.py
Normal file
98
studio/backend/core/inference/diffusion_nvfp4_flag.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""The one NVFP4 switch for Studio image/video; read at call time, never cached, so tests can monkeypatch it."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Iterable, Optional
|
||||
|
||||
NVFP4_DIFFUSION_ENV = "UNSLOTH_NVFP4_DIFFUSION"
|
||||
|
||||
# Flip to True once the hosted *-NVFP4 repos are public.
|
||||
NVFP4_DIFFUSION_DEFAULT = False
|
||||
|
||||
NVFP4_SCHEME = "nvfp4"
|
||||
|
||||
_TRUTHY = frozenset({"1", "true", "yes", "on"})
|
||||
_FALSY = frozenset({"0", "false", "no", "off"})
|
||||
|
||||
|
||||
def nvfp4_diffusion_enabled() -> bool:
|
||||
"""Whether NVFP4 may be used; unrecognised values take ``NVFP4_DIFFUSION_DEFAULT``."""
|
||||
raw = os.environ.get(NVFP4_DIFFUSION_ENV)
|
||||
if raw is None:
|
||||
return NVFP4_DIFFUSION_DEFAULT
|
||||
value = raw.strip().lower()
|
||||
if value in _TRUTHY:
|
||||
return True
|
||||
if value in _FALSY:
|
||||
return False
|
||||
return NVFP4_DIFFUSION_DEFAULT
|
||||
|
||||
|
||||
def is_nvfp4(value: Optional[str]) -> bool:
|
||||
"""Whether ``value`` names the NVFP4 scheme, in any case or spacing."""
|
||||
if value is None:
|
||||
return False
|
||||
return str(value).strip().lower().replace("-", "_") == NVFP4_SCHEME
|
||||
|
||||
|
||||
def nvfp4_blocked(scheme: Optional[str]) -> bool:
|
||||
"""True when ``scheme`` is NVFP4 and the switch is off: the one test every gate asks."""
|
||||
return is_nvfp4(scheme) and not nvfp4_diffusion_enabled()
|
||||
|
||||
|
||||
def nvfp4_repo_blocked(repo_id: Optional[str]) -> bool:
|
||||
"""True for a hosted ``*-NVFP4`` repo while the switch is off (Hub-helper backstop)."""
|
||||
if nvfp4_diffusion_enabled() or not repo_id:
|
||||
return False
|
||||
return str(repo_id).strip().rstrip("/").lower().endswith("-nvfp4")
|
||||
|
||||
|
||||
def without_nvfp4(schemes: Iterable[str]) -> tuple[str, ...]:
|
||||
"""``schemes`` in order, minus NVFP4 while the switch is off."""
|
||||
if nvfp4_diffusion_enabled():
|
||||
return tuple(schemes)
|
||||
return tuple(scheme for scheme in schemes if not is_nvfp4(scheme))
|
||||
|
||||
|
||||
def nvfp4_disabled_message(control: str = "transformer_quant") -> str:
|
||||
"""The refusal a request naming NVFP4 gets while the switch is off."""
|
||||
return (
|
||||
f"{control}='nvfp4' could not be used: NVFP4 is disabled in this build. Pick another "
|
||||
f"precision, or set {NVFP4_DIFFUSION_ENV}=1 on the Studio backend to enable it."
|
||||
)
|
||||
|
||||
|
||||
def refuse_disabled_nvfp4(
|
||||
*, transformer_quant: Optional[str] = None, text_encoder_quant: Optional[str] = None
|
||||
) -> None:
|
||||
"""Raise ``ValueError`` (400) for NVFP4 while off, before any silent fallback could swap it."""
|
||||
if nvfp4_diffusion_enabled():
|
||||
return
|
||||
if is_nvfp4(transformer_quant):
|
||||
raise ValueError(nvfp4_disabled_message("transformer_quant"))
|
||||
if is_nvfp4(text_encoder_quant):
|
||||
raise ValueError(nvfp4_disabled_message("text_encoder_quant"))
|
||||
|
||||
|
||||
def nvfp4_checkpoint_disabled_message() -> str:
|
||||
"""The refusal a load or plan of an NVFP4 checkpoint gets while the switch is off."""
|
||||
return (
|
||||
"model_path could not be used: it is an NVFP4 checkpoint and NVFP4 is disabled in this "
|
||||
f"build. Pick another model, or set {NVFP4_DIFFUSION_ENV}=1 on the Studio backend to "
|
||||
"enable it."
|
||||
)
|
||||
|
||||
|
||||
def refuse_disabled_nvfp4_checkpoint(*model_paths: Optional[str]) -> None:
|
||||
"""Raise ``ValueError`` (400) for an NVFP4 checkpoint model path while off; it bypasses the scheme gates."""
|
||||
if nvfp4_diffusion_enabled():
|
||||
return
|
||||
from .diffusion_prequant import declares_nvfp4_checkpoint
|
||||
|
||||
for path in model_paths:
|
||||
if path and declares_nvfp4_checkpoint(path):
|
||||
raise ValueError(nvfp4_checkpoint_disabled_message())
|
||||
|
|
@ -38,6 +38,8 @@ from functools import lru_cache
|
|||
|
||||
from core._torchao_stub import is_stubbed, torch_is_rocm
|
||||
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked, nvfp4_disabled_message
|
||||
|
||||
TE_QUANT_FP8 = "fp8"
|
||||
TE_QUANT_NVFP4 = "nvfp4"
|
||||
TE_QUANT_INT8 = "int8"
|
||||
|
|
@ -79,6 +81,8 @@ def normalize_te_quant(value: Optional[str]) -> Optional[str]:
|
|||
raise ValueError(
|
||||
f"Unsupported text_encoder_quant '{value}'. Use one of: {', '.join(TE_QUANT_MODES)}."
|
||||
)
|
||||
if nvfp4_blocked(normalized):
|
||||
raise ValueError(nvfp4_disabled_message("text_encoder_quant"))
|
||||
return normalized
|
||||
|
||||
|
||||
|
|
@ -110,7 +114,7 @@ def resolve_te_quant_request(
|
|||
"""
|
||||
if not te_quant_is_auto(value):
|
||||
return normalize_te_quant(value), False
|
||||
if auto_scheme is None:
|
||||
if auto_scheme is None or nvfp4_blocked(auto_scheme):
|
||||
return None, False
|
||||
# Validate the family's own field rather than trusting it: a typo here would otherwise reach
|
||||
# quantize_text_encoders as an unknown mode on every default load of that family.
|
||||
|
|
@ -166,6 +170,8 @@ def te_quant_supported(target: Any, mode: str) -> bool:
|
|||
Blackwell sm_100+ (nvfp4)."""
|
||||
if getattr(target, "device", None) != "cuda":
|
||||
return False
|
||||
if nvfp4_blocked(mode):
|
||||
return False
|
||||
# Torchao modes cannot use the Windows stub or ROCm's non-SM capability values. Plain fp8 is only a dtype cast and
|
||||
# remains supported.
|
||||
if mode in _TE_TORCHAO_MODES and (is_stubbed("torchao") or torch_is_rocm()):
|
||||
|
|
|
|||
|
|
@ -34,6 +34,8 @@ from dataclasses import dataclass
|
|||
from collections.abc import Sequence
|
||||
from typing import Any, Optional
|
||||
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked
|
||||
|
||||
# torch.save dict layout tag; bump on an on-disk change so old/foreign artifacts are rejected
|
||||
PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1"
|
||||
|
||||
|
|
@ -47,6 +49,8 @@ PREQUANT_FORMAT_ROTATED = "unsloth_prequant_transformer_state_dict_v2"
|
|||
|
||||
PREQUANT_FORMATS = (PREQUANT_FORMAT, PREQUANT_FORMAT_ROTATED)
|
||||
|
||||
DEFAULT_PREQUANT_COMPONENT = "transformer"
|
||||
|
||||
|
||||
def prequant_format_for(metadata: Any) -> str:
|
||||
"""The on-disk format tag an offline builder should stamp for ``metadata``."""
|
||||
|
|
@ -432,23 +436,31 @@ def prequant_repo_filename(
|
|||
repo_id: str,
|
||||
scheme: str,
|
||||
suffix: str = ".pt",
|
||||
*,
|
||||
component: Optional[str] = None,
|
||||
) -> str:
|
||||
"""The model-name checkpoint filename for ``scheme`` in ``repo_id``: the hosted repos are named
|
||||
<Model>-FP8 (or -INT8 / -quantized) and carry <Model>-<SCHEME>.pt files, e.g.
|
||||
unsloth/Z-Image-Turbo-FP8 -> Z-Image-Turbo-INT8.pt / Z-Image-Turbo-FP8.pt.
|
||||
"""Filename for ``scheme`` in ``repo_id``; a non-default ``component`` inserts -<component>-.
|
||||
|
||||
``suffix`` picks the container. It defaults to ``.pt`` so every existing caller keeps naming the
|
||||
artifact it already names; ``derived_prequant_filenames`` is what puts the safetensors spelling
|
||||
of the same name ahead of it."""
|
||||
model = repo_id.rsplit("/", 1)[-1]
|
||||
for drop in ("-fp8", "-int8", "-quantized"):
|
||||
for drop in ("-fp8", "-int8", "-nvfp4", "-mxfp8", "-quantized"):
|
||||
if model.lower().endswith(drop):
|
||||
model = model[: -len(drop)]
|
||||
break
|
||||
part = (component or "").strip()
|
||||
if part and part != DEFAULT_PREQUANT_COMPONENT:
|
||||
return f"{model}-{part}-{scheme.upper()}{suffix}"
|
||||
return f"{model}-{scheme.upper()}{suffix}"
|
||||
|
||||
|
||||
def derived_prequant_filenames(repo_id: str, scheme: str) -> tuple[str, ...]:
|
||||
def derived_prequant_filenames(
|
||||
repo_id: str,
|
||||
scheme: str,
|
||||
*,
|
||||
component: Optional[str] = None,
|
||||
) -> tuple[str, ...]:
|
||||
"""The names to try for ``(repo_id, scheme)``, best first, safetensors AHEAD of the pickle.
|
||||
|
||||
Preferring safetensors is a policy decision rather than a detail: it needs no constructor
|
||||
|
|
@ -458,11 +470,14 @@ def derived_prequant_filenames(repo_id: str, scheme: str) -> tuple[str, ...]:
|
|||
picked up with no code change, and a repo that never does keeps resolving exactly what it
|
||||
resolves today, because the ``.pt`` names stay in the chain behind it.
|
||||
"""
|
||||
return (
|
||||
prequant_repo_filename(repo_id, scheme, ".safetensors"),
|
||||
prequant_repo_filename(repo_id, scheme, ".pt"),
|
||||
prequant_filename(scheme),
|
||||
names = (
|
||||
prequant_repo_filename(repo_id, scheme, ".safetensors", component = component),
|
||||
prequant_repo_filename(repo_id, scheme, ".pt", component = component),
|
||||
)
|
||||
part = (component or "").strip()
|
||||
if part and part != DEFAULT_PREQUANT_COMPONENT:
|
||||
return names
|
||||
return names + (prequant_filename(scheme),)
|
||||
|
||||
|
||||
def resolve_prequant_source(
|
||||
|
|
@ -488,6 +503,8 @@ def resolve_prequant_source(
|
|||
checkpoints at the root, so there is no directory to prepend; a repo that nested them would 404
|
||||
on the primary AND on the fallback and the load would silently fall back to dense.
|
||||
"""
|
||||
if nvfp4_blocked(scheme):
|
||||
return None
|
||||
override = (path_override or "").strip()
|
||||
if override:
|
||||
return PrequantSource(kind = "path", location = override, filename = None)
|
||||
|
|
@ -591,6 +608,294 @@ def local_prequant_scheme(path: str) -> Optional[str]:
|
|||
return scheme
|
||||
|
||||
|
||||
def hosted_nvfp4_repo_ids() -> frozenset:
|
||||
"""Lowercased repo ids any image/video family registers for nvfp4, read off the family tables."""
|
||||
from dataclasses import fields, is_dataclass
|
||||
|
||||
from .diffusion_nvfp4_flag import is_nvfp4
|
||||
|
||||
families: list = []
|
||||
for module in ("diffusion_families", "video_families"):
|
||||
try:
|
||||
mod = __import__(f"{__package__}.{module}", fromlist = ["_FAMILIES"])
|
||||
families.extend(getattr(mod, "_FAMILIES", ()) or ())
|
||||
except Exception: # noqa: BLE001 - a table that fails to import names no repo
|
||||
continue
|
||||
repos = set()
|
||||
for fam in families:
|
||||
if not is_dataclass(fam):
|
||||
continue
|
||||
for field in fields(fam):
|
||||
value = getattr(fam, field.name, None)
|
||||
if not isinstance(value, tuple):
|
||||
continue
|
||||
for row in value:
|
||||
if not isinstance(row, tuple) or not row or not isinstance(row[-1], str):
|
||||
continue
|
||||
if "/" in row[-1] and any(isinstance(x, str) and is_nvfp4(x) for x in row[:-1]):
|
||||
repos.add(row[-1].strip().lower())
|
||||
return frozenset(repos)
|
||||
|
||||
|
||||
# Root only: a diffusers pipeline keeps its weights in component folders.
|
||||
_PREQUANT_PROBE_SUFFIXES = (".safetensors", ".pt", ".pth")
|
||||
_PREQUANT_PROBE_LIMIT = 16
|
||||
|
||||
|
||||
def _cached_snapshot_dirs(repo_id: str) -> list:
|
||||
"""The local snapshot folders of Hub repo ``repo_id``, newest first. No network."""
|
||||
import os
|
||||
|
||||
parts = repo_id.strip().strip("/").split("/")
|
||||
if len(parts) != 2 or not all(parts):
|
||||
return []
|
||||
roots = []
|
||||
try:
|
||||
from utils.hf_cache_settings import active_hf_hub_cache
|
||||
roots.append(active_hf_hub_cache())
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
try:
|
||||
from huggingface_hub import constants
|
||||
roots.append(constants.HF_HUB_CACHE)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
found = []
|
||||
for root in dict.fromkeys(r for r in roots if r):
|
||||
snapshots = os.path.join(root, "models--" + "--".join(parts), "snapshots")
|
||||
try:
|
||||
entries = [os.path.join(snapshots, name) for name in os.listdir(snapshots)]
|
||||
except OSError:
|
||||
continue
|
||||
entries = [e for e in entries if os.path.isdir(e)]
|
||||
entries.sort(key = lambda e: os.path.getmtime(e), reverse = True)
|
||||
found.extend(entries)
|
||||
return found
|
||||
|
||||
|
||||
def _artifacts_declare_nvfp4(folder: str) -> bool:
|
||||
"""Whether a pre-quant artifact at the root of ``folder`` records ``scheme = nvfp4``."""
|
||||
import os
|
||||
|
||||
from .diffusion_nvfp4_flag import is_nvfp4
|
||||
|
||||
try:
|
||||
names = sorted(os.listdir(folder))
|
||||
except OSError:
|
||||
return False
|
||||
if "model_index.json" in names:
|
||||
return False
|
||||
probed = 0
|
||||
for name in names:
|
||||
if not name.lower().endswith(_PREQUANT_PROBE_SUFFIXES):
|
||||
continue
|
||||
path = os.path.join(folder, name)
|
||||
if not os.path.isfile(path):
|
||||
continue
|
||||
if is_nvfp4(local_prequant_scheme(path)):
|
||||
return True
|
||||
probed += 1
|
||||
if probed >= _PREQUANT_PROBE_LIMIT:
|
||||
break
|
||||
return False
|
||||
|
||||
|
||||
def declares_nvfp4_checkpoint(model_path: Optional[str]) -> bool:
|
||||
"""Whether ``model_path`` is an NVFP4 prequant checkpoint: recorded metadata first, else registry or name. Never raises."""
|
||||
import os
|
||||
|
||||
raw = str(model_path or "").strip()
|
||||
if not raw:
|
||||
return False
|
||||
try:
|
||||
local = os.path.expanduser(raw)
|
||||
if os.path.isfile(local):
|
||||
if not local.lower().endswith(_PREQUANT_PROBE_SUFFIXES):
|
||||
return False
|
||||
from .diffusion_nvfp4_flag import is_nvfp4
|
||||
return is_nvfp4(local_prequant_scheme(local))
|
||||
if os.path.isdir(local):
|
||||
if _artifacts_declare_nvfp4(local):
|
||||
return True
|
||||
else:
|
||||
key = raw.strip("/").lower()
|
||||
if key in hosted_nvfp4_repo_ids():
|
||||
return True
|
||||
if any(_artifacts_declare_nvfp4(snap) for snap in _cached_snapshot_dirs(raw)):
|
||||
return True
|
||||
except Exception: # noqa: BLE001 - a probe that fails is "not known to be NVFP4"
|
||||
pass
|
||||
return raw.rstrip("/\\").lower().endswith("-nvfp4")
|
||||
|
||||
|
||||
# Bump on any payload-order or hash change, so two recipes are never compared under one name.
|
||||
FINGERPRINT_ALGO = "md5-packed-v1"
|
||||
|
||||
# Payload attrs per torchao class NAME, hash order; unlisted reads "not covered", never "equal".
|
||||
_FINGERPRINT_PAYLOAD: dict = {
|
||||
"NVFP4Tensor": ("qdata", "scale", "per_tensor_scale"),
|
||||
"Float8Tensor": ("qdata", "scale"),
|
||||
"Int8Tensor": (
|
||||
"qdata",
|
||||
"scale",
|
||||
"zero_point",
|
||||
"act_quant_scale",
|
||||
"act_quant_zero_point",
|
||||
"act_pre_scale",
|
||||
),
|
||||
"MXTensor": ("qdata", "scale"),
|
||||
"LinearActivationQuantizedTensor": ("original_weight_tensor",),
|
||||
"AffineQuantizedTensor": ("tensor_impl",),
|
||||
"PlainAQTTensorImpl": ("int_data", "scale", "zero_point"),
|
||||
}
|
||||
|
||||
|
||||
def _packed_bytes(tensor: Any, torch: Any) -> bytes:
|
||||
"""Raw bytes of ``tensor``; a uint8 view reinterprets, so fp8 / fp4 / bf16 hash exactly."""
|
||||
t = tensor.detach().contiguous()
|
||||
if t.dim() == 0:
|
||||
t = t.reshape(1)
|
||||
if t.dtype is not torch.uint8:
|
||||
t = t.view(torch.uint8)
|
||||
return t.cpu().numpy().tobytes()
|
||||
|
||||
|
||||
def _hash_packed_payload(tensor: Any, digest: Any, torch: Any) -> bool:
|
||||
"""Hash slot names + bytes into ``digest``; False if uncovered or no slot present."""
|
||||
names = _FINGERPRINT_PAYLOAD.get(type(tensor).__name__)
|
||||
if names is None:
|
||||
return False
|
||||
hashed = False
|
||||
for name in names:
|
||||
value = getattr(tensor, name, None)
|
||||
if value is None:
|
||||
continue
|
||||
digest.update(name.encode("utf-8"))
|
||||
if type(value).__name__ in _FINGERPRINT_PAYLOAD:
|
||||
if not _hash_packed_payload(value, digest, torch):
|
||||
return False
|
||||
hashed = True
|
||||
continue
|
||||
digest.update(_packed_bytes(value, torch))
|
||||
hashed = True
|
||||
return hashed
|
||||
|
||||
|
||||
def packed_weight_fingerprint(state_dict: Any, *, select: Any = None) -> dict:
|
||||
"""md5 of every ``.weight``'s packed bytes by fqn; not the pickle, so re-saves match."""
|
||||
import hashlib
|
||||
|
||||
import torch
|
||||
|
||||
modules: dict = {}
|
||||
skipped: list = []
|
||||
items = state_dict.items() if hasattr(state_dict, "items") else ()
|
||||
for key, tensor in items:
|
||||
if key != "weight" and not str(key).endswith(".weight"):
|
||||
continue
|
||||
if select is not None and not select(str(key)):
|
||||
continue
|
||||
digest = hashlib.md5()
|
||||
try:
|
||||
covered = _hash_packed_payload(tensor, digest, torch)
|
||||
except Exception: # noqa: BLE001 -- an unreadable payload is uncovered, never a raise
|
||||
covered = False
|
||||
if covered:
|
||||
modules[key] = digest.hexdigest()
|
||||
else:
|
||||
skipped.append(key)
|
||||
return {
|
||||
"algo": FINGERPRINT_ALGO,
|
||||
"count": len(modules),
|
||||
"modules": modules,
|
||||
"skipped": skipped,
|
||||
}
|
||||
|
||||
|
||||
FINGERPRINT_MODE_ENV = "UNSLOTH_PREQUANT_FINGERPRINT"
|
||||
FINGERPRINT_MODES = ("full", "sample", "off")
|
||||
FINGERPRINT_SAMPLE_RATE = 8
|
||||
|
||||
|
||||
def _fingerprint_mode() -> str:
|
||||
"""``full`` (default) / ``sample`` / ``off``. An unrecognised value reads as the default."""
|
||||
import os
|
||||
|
||||
mode = (os.environ.get(FINGERPRINT_MODE_ENV) or "").strip().lower()
|
||||
return mode if mode in FINGERPRINT_MODES else "full"
|
||||
|
||||
|
||||
def _fingerprint_sampled(fqn: str) -> bool:
|
||||
"""Stable 1-in-8 by md5 of the fqn; ``hash()`` is randomised per process by PYTHONHASHSEED."""
|
||||
import hashlib
|
||||
return hashlib.md5(fqn.encode("utf-8")).digest()[0] % FINGERPRINT_SAMPLE_RATE == 0
|
||||
|
||||
|
||||
def _verify_packed_fingerprint(
|
||||
state_dict: Any,
|
||||
metadata: Any,
|
||||
*,
|
||||
logger: Any = None,
|
||||
) -> bool:
|
||||
"""Refuse on a fingerprint mismatch; no block, or one this build cannot compute, passes."""
|
||||
block = (metadata or {}).get("fingerprint")
|
||||
expected = (block or {}).get("modules") if isinstance(block, dict) else None
|
||||
if not expected:
|
||||
return True
|
||||
mode = _fingerprint_mode()
|
||||
if mode == "off":
|
||||
if logger is not None:
|
||||
logger.debug(
|
||||
"diffusion.prequant: fingerprint check disabled (%s=off)", FINGERPRINT_MODE_ENV
|
||||
)
|
||||
return True
|
||||
try:
|
||||
actual = (
|
||||
packed_weight_fingerprint(
|
||||
state_dict,
|
||||
select = _fingerprint_sampled if mode == "sample" else None,
|
||||
).get("modules")
|
||||
or {}
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 -- an uncomputable fingerprint checks nothing
|
||||
_warn(logger, "fingerprint", exc)
|
||||
return True
|
||||
if not actual:
|
||||
_warn(
|
||||
logger,
|
||||
"fingerprint",
|
||||
RuntimeError(
|
||||
f"this build recognised none of the {len(expected)} quantized weights the "
|
||||
"checkpoint fingerprinted (a torchao payload rename?); loading it unverified"
|
||||
),
|
||||
)
|
||||
return True
|
||||
checked = [k for k in expected if mode != "sample" or _fingerprint_sampled(k)]
|
||||
differing = sorted(k for k in checked if actual.get(k) != expected[k])
|
||||
counted = mode != "sample" and len(actual) != len(expected)
|
||||
if not differing and not counted:
|
||||
if logger is not None:
|
||||
logger.info(
|
||||
"diffusion.prequant: fingerprint verified (%d of %d quantized weights, %s)",
|
||||
len(checked),
|
||||
len(expected),
|
||||
mode,
|
||||
)
|
||||
return True
|
||||
if logger is not None:
|
||||
logger.error(
|
||||
"diffusion.prequant: fingerprint MISMATCH (%d of %d checked weights differ, %d "
|
||||
"weights present vs %d recorded): %s. The checkpoint does not hold the bytes it was "
|
||||
"built with; falling back to the dense path",
|
||||
len(differing),
|
||||
len(checked),
|
||||
len(actual),
|
||||
len(expected),
|
||||
", ".join(differing[:5]) or "counts only",
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def usable_prequant_source(
|
||||
fam: Any,
|
||||
scheme: str,
|
||||
|
|
@ -782,6 +1087,7 @@ def load_prequantized_transformer(
|
|||
cache_dir: Optional[str] = None,
|
||||
prepare_model: Optional[Any] = None,
|
||||
config_subfolder: str = "transformer",
|
||||
component: Optional[str] = None,
|
||||
local_files_only: bool = False,
|
||||
logger: Any = None,
|
||||
) -> Optional[Any]:
|
||||
|
|
@ -791,24 +1097,8 @@ def load_prequantized_transformer(
|
|||
lands under huggingface_hub's import-time constant, so a mid-session cache change re-downloads
|
||||
into a root Unsloth no longer reads.
|
||||
|
||||
``config_subfolder`` is where the DENOISER CONFIG lives inside ``base``, defaulting to the
|
||||
universal ``transformer``. A family hosting several denoiser partitions in one repo overrides it
|
||||
with the one this checkpoint belongs to (MiniMax-H3's ``transformer_ref``): the scoped download
|
||||
stages only that partition, so reading the config from the other one would send an otherwise
|
||||
fully staged load back to the Hub.
|
||||
|
||||
``prepare_model`` (optional) is called as ``prepare_model(transformer, metadata)`` on the
|
||||
freshly built skeleton, AFTER ``from_config`` and BEFORE ``load_state_dict``. That window is the
|
||||
only one where a family can reshape the module to match how the checkpoint was baked (a swapped
|
||||
submodule, a patched attention class): earlier there is no module, and later ``strict=True`` has
|
||||
already rejected the mismatch. It gets the checkpoint's own metadata so it can key on what was
|
||||
baked rather than on today's defaults. A raising callback falls out to the outer handler below,
|
||||
i.e. a warning and a dense fallback, never a failed load.
|
||||
|
||||
A checkpoint that declares an ACTIVATION ROTATION (``diffusion_convrot``) has the matching
|
||||
online half installed here, on exactly the fqns it records. That is unconditional and central
|
||||
rather than a family opt-in, because the one failure mode worth designing against is the silent
|
||||
one: rotated weights met by unrotated activations render wrong pixels and raise nothing.
|
||||
``component`` is checked (experts share every other key); ``prepare_model`` may reshape before
|
||||
``load_state_dict``; a recorded ``diffusion_convrot`` always gets its online half installed.
|
||||
|
||||
Returns the placed transformer, or None on any problem (missing / mismatched / unreadable
|
||||
checkpoint, unsupported meta-init, or a rotation this build cannot apply exactly) so the caller
|
||||
|
|
@ -843,10 +1133,19 @@ def load_prequantized_transformer(
|
|||
# hand back the same dict, so every check below applies to them equally.
|
||||
ckpt = _load_prequant_checkpoint(path, map_location = "cpu")
|
||||
if not _validate_checkpoint(
|
||||
ckpt, scheme, base, logger, min_features = min_features, fast_accum = fast_accum
|
||||
ckpt,
|
||||
scheme,
|
||||
base,
|
||||
logger,
|
||||
min_features = min_features,
|
||||
fast_accum = fast_accum,
|
||||
component = component,
|
||||
):
|
||||
return None
|
||||
state_dict = ckpt["state_dict"]
|
||||
# The only check reading what the artifact HOLDS: corruption after build passes the rest.
|
||||
if not _verify_packed_fingerprint(state_dict, ckpt.get("metadata") or {}, logger = logger):
|
||||
return None
|
||||
_pin_kernel_preference(state_dict, logger)
|
||||
|
||||
# Read from the root that actually supplied the checkpoint: after a mid-session cache change the pinned root
|
||||
|
|
@ -892,14 +1191,19 @@ def load_prequantized_transformer(
|
|||
|
||||
apply_activation_rotation(transformer, metadata, logger = logger)
|
||||
|
||||
# assign=True shares the checkpoint tensors; a live ref keeps the CPU copy past to(device).
|
||||
del state_dict
|
||||
del ckpt
|
||||
|
||||
transformer = transformer.to(device)
|
||||
# Same small-M row padding the runtime quantise path applies, and for the same reason: a checkpoint built
|
||||
# under the current exclusion set QUANTISES the family's small-M linears, so without the wrappers they would
|
||||
# raise inside _int_mm the moment the compiled scope reaches them. After load_state_dict, since wrapping
|
||||
# reparents the Linears; after .to() so the granularity probe reads the device tensors the GEMM will see.
|
||||
from .diffusion_transformer_quant import apply_small_m_padding
|
||||
from .diffusion_transformer_quant import apply_small_m_padding, apply_zero_row_guard
|
||||
|
||||
apply_small_m_padding(transformer, scheme, metadata.get("family"), logger = logger)
|
||||
apply_zero_row_guard(transformer, scheme, metadata.get("family"), logger = logger)
|
||||
# from_config starts in TRAIN mode while the dense/GGUF paths use from_pretrained (eval()'d). Match it so
|
||||
# train/eval-sensitive layers cannot make prequant inference diverge.
|
||||
try:
|
||||
|
|
@ -1206,14 +1510,9 @@ def _validate_checkpoint(
|
|||
logger: Any,
|
||||
min_features: Optional[int] = None,
|
||||
fast_accum: Optional[bool] = None,
|
||||
component: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""Reject a checkpoint that is the wrong format / scheme / base model / filter. ``min_features``
|
||||
(when given) is the runtime Linear-feature threshold: a different ``--min-features``
|
||||
quantises a different set of Linears, so assign=True would silently install a mismatched
|
||||
model while status still reports the scheme. Reject it. ``fast_accum`` (fp8 only): when the
|
||||
caller forces it and the checkpoint baked a different value, the loaded kernels would ignore
|
||||
the request, so reject and let the dense path honor it. A checkpoint predating a metadata
|
||||
field (absent) is accepted for back-compat."""
|
||||
"""Reject a wrong format / scheme / base / filter / denoiser; absent fields predate them."""
|
||||
if not isinstance(ckpt, dict) or ckpt.get("format") not in PREQUANT_FORMATS:
|
||||
_warn(logger, scheme, ValueError("unrecognised pre-quant checkpoint format"))
|
||||
return False
|
||||
|
|
@ -1307,6 +1606,42 @@ def _validate_checkpoint(
|
|||
),
|
||||
)
|
||||
return False
|
||||
ckpt_divisible = meta.get("require_divisible")
|
||||
if ckpt_divisible is not None:
|
||||
from .diffusion_transformer_quant import divisible_for_scheme
|
||||
expected_divisible = divisible_for_scheme(scheme)
|
||||
if int(ckpt_divisible) != expected_divisible:
|
||||
_warn(
|
||||
logger,
|
||||
scheme,
|
||||
ValueError(
|
||||
f"checkpoint require_divisible {ckpt_divisible!r} != {expected_divisible!r}"
|
||||
),
|
||||
)
|
||||
return False
|
||||
elif logger is not None:
|
||||
logger.debug(
|
||||
"diffusion.prequant: checkpoint records no require_divisible (built before the "
|
||||
"field); accepting it for %s",
|
||||
scheme,
|
||||
)
|
||||
if component:
|
||||
ckpt_component = meta.get("component")
|
||||
if ckpt_component is not None and str(ckpt_component) != str(component):
|
||||
_warn(
|
||||
logger,
|
||||
scheme,
|
||||
ValueError(
|
||||
f"checkpoint component {ckpt_component!r} != {component!r}; this is another "
|
||||
"denoiser of the same family and would load clean and render wrong"
|
||||
),
|
||||
)
|
||||
return False
|
||||
if ckpt_component is None and logger is not None:
|
||||
logger.debug(
|
||||
"diffusion.prequant: checkpoint records no component; accepting it for %r",
|
||||
component,
|
||||
)
|
||||
# fp8 fast-accum is baked into the saved kernels; only enforce when the caller forces it.
|
||||
if fast_accum is not None:
|
||||
ckpt_fa = meta.get("fast_accum")
|
||||
|
|
|
|||
|
|
@ -219,6 +219,80 @@ class PadToMinM(nn.Module):
|
|||
return f"min_m = {self.min_m}, pad_to = {self.pad_to}"
|
||||
|
||||
|
||||
class ZeroRowSafeLinear(nn.Module):
|
||||
"""Answer an EMPTY activation here: torchao NVFP4's whole-input ``max()`` raises on it."""
|
||||
|
||||
def __init__(self, inner: nn.Linear) -> None:
|
||||
super().__init__()
|
||||
self.inner = inner
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.numel() == 0:
|
||||
# The bias add broadcasts over zero elements, kept so the result matches F.linear's.
|
||||
out = x.new_zeros((*x.shape[:-1], self.inner.out_features))
|
||||
bias = getattr(self.inner, "bias", None)
|
||||
return out if bias is None else out + bias
|
||||
return self.inner(x)
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
if name == "inner":
|
||||
raise
|
||||
inner = self._modules.get("inner")
|
||||
if inner is None:
|
||||
raise
|
||||
return getattr(inner, name)
|
||||
|
||||
def state_dict(self, *args: Any, **kwargs: Any) -> Any: # type: ignore[override]
|
||||
"""Emit the inner Linear's tensors under the wrapper's prefix, loadable unwrapped."""
|
||||
destination = kwargs.pop("destination", args[0] if args else None)
|
||||
prefix = kwargs.pop("prefix", args[1] if len(args) > 1 else "")
|
||||
keep_vars = kwargs.pop("keep_vars", args[2] if len(args) > 2 else False)
|
||||
if destination is None:
|
||||
return self.inner.state_dict(prefix = prefix, keep_vars = keep_vars)
|
||||
self.inner.state_dict(destination = destination, prefix = prefix, keep_vars = keep_vars)
|
||||
return destination
|
||||
|
||||
def _load_from_state_dict(
|
||||
self,
|
||||
state_dict: Any,
|
||||
prefix: str,
|
||||
local_metadata: Any,
|
||||
strict: bool,
|
||||
missing_keys: list,
|
||||
unexpected_keys: list,
|
||||
error_msgs: list,
|
||||
) -> None:
|
||||
"""Accept the unwrapped key names ``state_dict`` above writes, and hand them to ``inner``."""
|
||||
for key in [k for k in state_dict if k.startswith(prefix)]:
|
||||
leaf = key[len(prefix) :]
|
||||
if not leaf or leaf.startswith("inner."):
|
||||
continue
|
||||
state_dict[prefix + "inner." + leaf] = state_dict.pop(key)
|
||||
super()._load_from_state_dict(
|
||||
state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
|
||||
)
|
||||
|
||||
|
||||
def wrap_zero_row_linears(model: nn.Module, fqns: Iterable[str]) -> tuple[str, ...]:
|
||||
"""Wrap each quantized Linear in ``fqns`` with ``ZeroRowSafeLinear`` (idempotent)."""
|
||||
done: list[str] = []
|
||||
for fqn in sorted(set(fqns)):
|
||||
parent_name, _, leaf = fqn.rpartition(".")
|
||||
try:
|
||||
parent = model.get_submodule(parent_name) if parent_name else model
|
||||
module = getattr(parent, leaf)
|
||||
except AttributeError:
|
||||
continue
|
||||
if not is_quantized_linear(module):
|
||||
continue
|
||||
setattr(parent, leaf, ZeroRowSafeLinear(module))
|
||||
done.append(fqn)
|
||||
return tuple(done)
|
||||
|
||||
|
||||
def padding_is_bitwise_exact(
|
||||
module: Any,
|
||||
m: int,
|
||||
|
|
|
|||
|
|
@ -229,6 +229,10 @@ def family_te_prequant_repo(fam: Any, scheme: str, component: str) -> Optional[s
|
|||
Reads the family's ``te_prequant_repos`` (scheme, component, repo_id) triples; the field
|
||||
is optional on both DiffusionFamily and VideoFamily, so one resolver serves both loaders.
|
||||
"""
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked
|
||||
|
||||
if nvfp4_blocked(scheme):
|
||||
return None
|
||||
for entry in getattr(fam, "te_prequant_repos", ()) or ():
|
||||
try:
|
||||
entry_scheme, entry_component, repo_id = entry
|
||||
|
|
|
|||
|
|
@ -26,11 +26,13 @@ import os as _os
|
|||
import re as _re
|
||||
import sys as _sys
|
||||
import threading as _threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
# stdlib-only module (no torch), so this stays inside the "imported lazily" promise above.
|
||||
from core._torchao_stub import is_stubbed, torch_is_rocm
|
||||
from .diffusion_native_quant import int8_act_disabled, int8_act_requested, is_native_linear
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked, nvfp4_disabled_message, without_nvfp4
|
||||
from .diffusion_torchao_patches import install_torchao_int_mm_patch
|
||||
|
||||
# Also runs in the spawned smoke-probe child, which imports this module and nothing else of the backend.
|
||||
|
|
@ -197,6 +199,45 @@ def apply_small_m_padding(
|
|||
return wrapped
|
||||
|
||||
|
||||
# nvfp4 raises on an empty activation; HunyuanVideo-1.5's attention trim hits it every t2v render.
|
||||
_HUNYUAN15_NVFP4_ZERO_ROW_TOKENS = ("image_embedder", "context_embedder_2")
|
||||
_NVFP4_FAMILY_ZERO_ROW_NAME_TOKENS: dict[str, tuple[str, ...]] = {
|
||||
"hunyuanvideo-1.5": _HUNYUAN15_NVFP4_ZERO_ROW_TOKENS,
|
||||
"hunyuanvideo-1.5-720p": _HUNYUAN15_NVFP4_ZERO_ROW_TOKENS,
|
||||
}
|
||||
|
||||
|
||||
def zero_row_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tuple[str, ...]:
|
||||
"""Tokens needing the empty-activation guard; nvfp4 only (fp8 / mxfp8 reduce per row)."""
|
||||
if scheme != TQ_NVFP4:
|
||||
return ()
|
||||
return _NVFP4_FAMILY_ZERO_ROW_NAME_TOKENS.get(str(family or "").strip().lower(), ())
|
||||
|
||||
|
||||
def apply_zero_row_guard(
|
||||
transformer: Any,
|
||||
scheme: str,
|
||||
family: Optional[str] = None,
|
||||
*,
|
||||
logger: Any = None,
|
||||
) -> tuple[str, ...]:
|
||||
"""Wrap this family's zero-row-reachable quantized Linears after quantize; not best-effort."""
|
||||
tokens = zero_row_tokens_for_scheme(scheme, family)
|
||||
if not tokens:
|
||||
return ()
|
||||
from .diffusion_quant_pad import matching_linear_fqns, wrap_zero_row_linears
|
||||
|
||||
wrapped = wrap_zero_row_linears(transformer, matching_linear_fqns(transformer, tokens))
|
||||
if wrapped and logger is not None:
|
||||
logger.info(
|
||||
"diffusion.transformer_quant: zero-row guarded %d %s linears on %s",
|
||||
len(wrapped),
|
||||
scheme,
|
||||
family,
|
||||
)
|
||||
return wrapped
|
||||
|
||||
|
||||
def exclude_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tuple[str, ...]:
|
||||
"""Name tokens to exclude from quantisation for ``scheme`` (optionally family-specific). int8
|
||||
(M>16) skips the M=1 modulation / conditioning-embedder projections
|
||||
|
|
@ -212,6 +253,14 @@ def exclude_tokens_for_scheme(scheme: str, family: Optional[str] = None) -> tupl
|
|||
return ()
|
||||
|
||||
|
||||
_SCHEME_DIVISIBLE: dict[str, int] = {TQ_FP8: 16, TQ_NVFP4: 16, TQ_MXFP8: 32}
|
||||
|
||||
|
||||
def divisible_for_scheme(scheme: str) -> int:
|
||||
"""The in/out feature alignment ``scheme``'s GEMM requires, 0 when it has none."""
|
||||
return _SCHEME_DIVISIBLE.get(scheme, 0)
|
||||
|
||||
|
||||
# Per-arch preference for ``auto``, best first. int8 leads every tier: it runs full-rate on consumer and
|
||||
# workstation cards (which halve fp8 FP32 accumulate) and is within a few percent of fp8 elsewhere.
|
||||
# nvfp4 is explicit opt-in and kept OUT of the ladder, being slower and less accurate at DiT shapes.
|
||||
|
|
@ -239,6 +288,18 @@ _FAMILY_SCHEME_DENY: dict[str, frozenset[str]] = {
|
|||
}
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class _AutoPrefer:
|
||||
"""Family head of the ``auto`` order; off below ``floor`` and on consumer GPUs by default."""
|
||||
|
||||
floor: tuple[int, int]
|
||||
schemes: tuple[str, ...]
|
||||
consumer_ok: bool = False
|
||||
|
||||
|
||||
_FAMILY_AUTO_PREFER: dict[str, _AutoPrefer] = {}
|
||||
|
||||
|
||||
# Schemes denied for TRAINING on top of the inference table. Training holds a stricter bar because the evidence above
|
||||
# is rendering evidence: it says a frozen fp8 forward reconstructs the bf16 image, not that a LoRA converges when its
|
||||
# frozen linears are fp8. Nobody has run that, so qwen fp8 stays out of the Train UI until someone does. Delete the
|
||||
|
|
@ -279,6 +340,8 @@ def explain_unusable_scheme(family: Optional[str], scheme: str) -> str:
|
|||
has to separate them: a family the scheme is measured to break, a torchao that cannot run at
|
||||
all in this install, and a GPU without the kernels. Only the last is a hardware limit, and
|
||||
only the last is what the message used to say."""
|
||||
if nvfp4_blocked(scheme):
|
||||
return nvfp4_disabled_message()
|
||||
if family_denies_scheme(family, scheme):
|
||||
return (
|
||||
f"'{scheme}' is ruled out for family '{family}' by the measured accuracy gate (it "
|
||||
|
|
@ -428,6 +491,8 @@ def normalize_transformer_quant(value: Optional[str]) -> Optional[str]:
|
|||
raise ValueError(
|
||||
f"Unsupported transformer_quant '{value}'. Use one of: {', '.join(TQ_MODES)}."
|
||||
)
|
||||
if nvfp4_blocked(normalized):
|
||||
raise ValueError(nvfp4_disabled_message("transformer_quant"))
|
||||
return normalized
|
||||
|
||||
|
||||
|
|
@ -759,14 +824,11 @@ def select_transformer_quant_scheme(
|
|||
cap = _capability()
|
||||
if cap is None:
|
||||
return None
|
||||
for floor, schemes in _AUTO_LADDER:
|
||||
if cap >= floor:
|
||||
for scheme in schemes:
|
||||
if _family_denied(family, scheme):
|
||||
continue
|
||||
if _scheme_supported(scheme, device):
|
||||
return scheme
|
||||
return None
|
||||
for scheme in _auto_scheme_order(family, device, cap):
|
||||
if _family_denied(family, scheme):
|
||||
continue
|
||||
if _scheme_supported(scheme, device):
|
||||
return scheme
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -803,7 +865,7 @@ def dense_quant_host_capable(target: Any) -> bool:
|
|||
card = _smoke_cache_device_key(str(getattr(target, "device", "cuda")))
|
||||
for floor, schemes in _AUTO_LADDER:
|
||||
if cap >= floor:
|
||||
return any(_SMOKE_CACHE.get((scheme, card), True) for scheme in schemes)
|
||||
return any(_SMOKE_CACHE.get((scheme, card), True) for scheme in without_nvfp4(schemes))
|
||||
return False
|
||||
|
||||
|
||||
|
|
@ -820,14 +882,33 @@ def auto_scheme_candidates(target: Any, family: Optional[str] = None) -> tuple[s
|
|||
cap = _capability()
|
||||
if cap is None:
|
||||
return ()
|
||||
return tuple(
|
||||
scheme
|
||||
for scheme in _auto_scheme_order(family, device, cap)
|
||||
if not _family_denied(family, scheme) and _scheme_supported(scheme, device)
|
||||
)
|
||||
|
||||
|
||||
def _auto_scheme_order(family: Optional[str], device: Any, cap: tuple[int, int]) -> tuple[str, ...]:
|
||||
"""``auto``'s scheme order for this family and GPU, best first; empty below every tier."""
|
||||
tier: tuple[str, ...] = ()
|
||||
for floor, schemes in _AUTO_LADDER:
|
||||
if cap >= floor:
|
||||
return tuple(
|
||||
scheme
|
||||
for scheme in schemes
|
||||
if not _family_denied(family, scheme) and _scheme_supported(scheme, device)
|
||||
)
|
||||
return ()
|
||||
tier = schemes
|
||||
break
|
||||
if not tier:
|
||||
return ()
|
||||
prefer = _FAMILY_AUTO_PREFER.get(str(family or "").strip().lower())
|
||||
head: tuple[str, ...] = ()
|
||||
if prefer is not None and cap >= prefer.floor:
|
||||
if prefer.consumer_ok or not _is_consumer_gpu(device):
|
||||
head = prefer.schemes
|
||||
order: list[str] = []
|
||||
# Filtered here so the NVFP4 switch reaches every reader of the order.
|
||||
for scheme in without_nvfp4(head + tier):
|
||||
if scheme not in order:
|
||||
order.append(scheme)
|
||||
return tuple(order)
|
||||
|
||||
|
||||
def auto_scheme_candidates_cached(target: Any, family: Optional[str] = None) -> tuple[str, ...]:
|
||||
|
|
@ -842,15 +923,13 @@ def auto_scheme_candidates_cached(target: Any, family: Optional[str] = None) ->
|
|||
cap = _capability()
|
||||
if cap is None:
|
||||
return ()
|
||||
card = _smoke_cache_device_key(str(getattr(target, "device", "cuda")))
|
||||
for floor, schemes in _AUTO_LADDER:
|
||||
if cap >= floor:
|
||||
return tuple(
|
||||
scheme
|
||||
for scheme in schemes
|
||||
if not _family_denied(family, scheme) and _SMOKE_CACHE.get((scheme, card), True)
|
||||
)
|
||||
return ()
|
||||
device = str(getattr(target, "device", "cuda"))
|
||||
card = _smoke_cache_device_key(device)
|
||||
return tuple(
|
||||
scheme
|
||||
for scheme in _auto_scheme_order(family, device, cap)
|
||||
if not _family_denied(family, scheme) and _SMOKE_CACHE.get((scheme, card), True)
|
||||
)
|
||||
|
||||
|
||||
def _capability() -> Optional[tuple[int, int]]:
|
||||
|
|
@ -869,6 +948,8 @@ def _scheme_supported(
|
|||
unproven_ok: bool = False,
|
||||
) -> bool:
|
||||
"""CUDA + (for fp8) the fp8 dtype + a cached quantise+matmul smoke test for ``scheme``."""
|
||||
if nvfp4_blocked(scheme):
|
||||
return False
|
||||
try:
|
||||
import torch
|
||||
if not torch.cuda.is_available():
|
||||
|
|
@ -966,7 +1047,14 @@ def _child_probe_table(device: str) -> Optional[dict[str, Optional[bool]]]:
|
|||
# and binds the child to this process's lifetime, so a wedged probe cannot outlive the backend.
|
||||
proc = ctx.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (__name__, "_child_probe_entry", {}, device, TQ_SCHEMES, queue),
|
||||
args = (
|
||||
__name__,
|
||||
"_child_probe_entry",
|
||||
{},
|
||||
device,
|
||||
without_nvfp4(TQ_SCHEMES),
|
||||
queue,
|
||||
),
|
||||
daemon = True,
|
||||
)
|
||||
proc.start()
|
||||
|
|
@ -1463,9 +1551,7 @@ def quantize_transformer(
|
|||
# non-bf16 ones. "lora_" keeps a baked adapter's side path high precision. Runtime only: NOT part of
|
||||
# exclude_tokens_for_scheme, whose list is baked into prequant metadata.
|
||||
exclude = exclude_tokens_for_scheme(scheme, family) + ("lora_",)
|
||||
# GEMM tiling floors per scheme: scaled_mm needs 16-aligned dims, MX block scaling 32. int8's _int_mm has no
|
||||
# such floor and keeps the historical filter.
|
||||
divisible = {TQ_FP8: 16, TQ_NVFP4: 16, TQ_MXFP8: 32}.get(scheme, 0)
|
||||
divisible = divisible_for_scheme(scheme)
|
||||
quantize_(
|
||||
transformer,
|
||||
_make_quant_config(scheme, fast_accum = fast_accum),
|
||||
|
|
@ -1480,6 +1566,7 @@ def quantize_transformer(
|
|||
# here means the transformer is quantized but not safely compilable, so it falls into the except below and the
|
||||
# caller loads GGUF.
|
||||
apply_small_m_padding(transformer, scheme, family, logger = logger)
|
||||
apply_zero_row_guard(transformer, scheme, family, logger = logger)
|
||||
try:
|
||||
transformer._unsloth_runtime_quant = scheme
|
||||
except Exception: # noqa: BLE001 - marker is best-effort
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
140
studio/backend/core/inference/video_denoiser_prequant.py
Normal file
140
studio/backend/core/inference/video_denoiser_prequant.py
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Seed a conventional video pipeline with hosted pre-quantized denoisers, all-or-nothing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
MOE_SECOND_DENOISER = "transformer_2"
|
||||
|
||||
|
||||
def denoiser_components(fam: Any) -> tuple[str, ...]:
|
||||
"""Pipeline attribute name(s) of this family's denoiser(s), from the registry."""
|
||||
if getattr(fam, "is_moe", False):
|
||||
return ("transformer", MOE_SECOND_DENOISER)
|
||||
return ("transformer",)
|
||||
|
||||
|
||||
def denoiser_prequant_sources(
|
||||
fam: Any, scheme: Optional[str], base_repo: Optional[str]
|
||||
) -> Optional[dict[str, Any]]:
|
||||
"""``{component: PrequantSource}`` for EVERY denoiser, or None (never partial); never raises."""
|
||||
wanted = (scheme or "").strip().lower()
|
||||
if wanted in ("", "auto", "off", "none"):
|
||||
return None
|
||||
try:
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked
|
||||
|
||||
if nvfp4_blocked(wanted):
|
||||
return None
|
||||
from .diffusion_prequant import resolve_prequant_source
|
||||
from .video_families import video_family_prequant_task_specific
|
||||
|
||||
sources: dict[str, Any] = {}
|
||||
for component in denoiser_components(fam):
|
||||
task = None if component == "transformer" else component
|
||||
if task is not None and not video_family_prequant_task_specific(fam, wanted, task):
|
||||
# Without its own row, ``task`` resolves expert 1's file: loads fine, wrong weights.
|
||||
return None
|
||||
sources[component] = resolve_prequant_source(
|
||||
fam,
|
||||
wanted,
|
||||
base_repo = base_repo,
|
||||
task = task,
|
||||
)
|
||||
if sources[component] is None:
|
||||
return None
|
||||
return sources or None
|
||||
except Exception: # noqa: BLE001 -- an unanswerable registry keeps the dense denoiser
|
||||
return None
|
||||
|
||||
|
||||
def denoiser_prequant_pipe_kwargs(
|
||||
fam: Any,
|
||||
base_repo: str,
|
||||
*,
|
||||
scheme: Optional[str],
|
||||
dtype: Any,
|
||||
device: str,
|
||||
hf_token: Optional[str] = None,
|
||||
target: Any = None,
|
||||
local_files_only: bool = False,
|
||||
cache_dir: Optional[str] = None,
|
||||
logger: Any = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Overrides for assembly, or ``{}`` (dense bf16 + runtime quant) unless all denoisers seed."""
|
||||
try:
|
||||
import gc
|
||||
|
||||
sources = denoiser_prequant_sources(fam, scheme, base_repo)
|
||||
if not sources:
|
||||
return {}
|
||||
if target is not None:
|
||||
from .diffusion_transformer_quant import dense_transformer_supported
|
||||
if not dense_transformer_supported(target):
|
||||
return {}
|
||||
import diffusers
|
||||
|
||||
from .diffusion_prequant import load_prequantized_transformer
|
||||
from .diffusion_transformer_quant import DEFAULT_MIN_LINEAR_FEATURES
|
||||
|
||||
transformer_cls = getattr(diffusers, str(getattr(fam, "transformer_class", "")), None)
|
||||
if transformer_cls is None:
|
||||
_warn(
|
||||
logger,
|
||||
scheme or "",
|
||||
ValueError(
|
||||
f"transformer class {getattr(fam, 'transformer_class', None)!r} "
|
||||
"not found in diffusers"
|
||||
),
|
||||
)
|
||||
return {}
|
||||
seeded: dict[str, Any] = {}
|
||||
for component, source in sources.items():
|
||||
module = load_prequantized_transformer(
|
||||
transformer_cls,
|
||||
base_repo,
|
||||
source,
|
||||
device = device,
|
||||
dtype = dtype,
|
||||
hf_token = hf_token,
|
||||
scheme = scheme,
|
||||
min_features = DEFAULT_MIN_LINEAR_FEATURES,
|
||||
config_subfolder = component,
|
||||
component = component,
|
||||
cache_dir = cache_dir,
|
||||
local_files_only = local_files_only,
|
||||
logger = logger,
|
||||
)
|
||||
if module is None:
|
||||
seeded.clear()
|
||||
gc.collect()
|
||||
if logger is not None:
|
||||
logger.info(
|
||||
"video.denoiser_prequant: no usable %s checkpoint for %s %s, so the "
|
||||
"dense denoiser(s) are loaded and quantised at runtime instead",
|
||||
scheme,
|
||||
getattr(fam, "name", None),
|
||||
component,
|
||||
)
|
||||
return {}
|
||||
seeded[component] = module
|
||||
gc.collect()
|
||||
if logger is not None:
|
||||
logger.info(
|
||||
"video.denoiser_prequant: seeded %s pre-quantized denoiser(s) for %s (%s)",
|
||||
len(seeded),
|
||||
getattr(fam, "name", None),
|
||||
scheme,
|
||||
)
|
||||
return seeded
|
||||
except Exception as exc: # noqa: BLE001 -- seeding is an optimisation, never a blocker
|
||||
_warn(logger, "pipe_kwargs", exc)
|
||||
return {}
|
||||
|
||||
|
||||
def _warn(logger: Any, what: str, exc: Exception) -> None:
|
||||
if logger is not None:
|
||||
logger.warning("video.denoiser_prequant: %s failed: %s", what, exc)
|
||||
|
|
@ -21,6 +21,8 @@ import re
|
|||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
from .diffusion_nvfp4_flag import nvfp4_blocked, without_nvfp4
|
||||
|
||||
# The request model's ceiling on num_frames, declared HERE so the shape gate and the bound cannot drift: the gate's
|
||||
# refusal names the lattice point above the request, and suggesting one the request model would itself reject is a
|
||||
# dead end. VideoGenerateRequest imports this for its `le`.
|
||||
|
|
@ -100,6 +102,7 @@ class VideoFamily:
|
|||
# parameters), so 0.55 x 66.3 GB over-states it by 16 GB and a hard refusal turns away a load that fits. Measured
|
||||
# from Hub file metadata (2026-08-09): MiniMax-H3-FP8.pt 20,260,192,855 bytes, MiniMax-H3-INT8.pt 20,253,894,865.
|
||||
prequant_resident_gb: Optional[float] = None
|
||||
prequant_resident_gb_by_scheme: tuple[tuple[str, float], ...] = field(default_factory = tuple)
|
||||
# Per-variant overrides as (base_repo, scheme, repo_id), keyed on the LOWERCASED upstream base id. A pre-quantized
|
||||
# checkpoint is baked from ONE base's weights and the loader refuses it for any other base, so a variant that ships
|
||||
# its own denoiser needs its own entry; a variant without one falls through to prequant_repos and, if that
|
||||
|
|
@ -249,6 +252,9 @@ _FAMILIES: tuple[VideoFamily, ...] = (
|
|||
pipeline_class = "WanPipeline",
|
||||
transformer_class = "WanTransformer3DModel",
|
||||
base_repo = "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
prequant_repos = (("nvfp4", "unsloth/Wan2.2-TI2V-5B-NVFP4"),),
|
||||
prequant_filenames = (("nvfp4", "Wan2.2-TI2V-5B-NVFP4.pt"),),
|
||||
prequant_resident_gb_by_scheme = (("nvfp4", 2.9),),
|
||||
# "wan2.2-5b"/"wan-ti2v" are the picker/GGUF short ids; "wan2.2-ti2v" catches the repo stem
|
||||
aliases = ("wan2.2-5b", "wan-ti2v", "wan2.2-ti2v", "wan-ti2v-5b"),
|
||||
has_audio = False,
|
||||
|
|
@ -276,6 +282,13 @@ _FAMILIES: tuple[VideoFamily, ...] = (
|
|||
pipeline_class = "WanPipeline",
|
||||
transformer_class = "WanTransformer3DModel",
|
||||
base_repo = "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
prequant_repos = (("nvfp4", "unsloth/Wan2.2-T2V-A14B-NVFP4"),),
|
||||
prequant_filenames = (
|
||||
("nvfp4", "Wan2.2-T2V-A14B-NVFP4.pt"),
|
||||
("nvfp4", "transformer_2", "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt"),
|
||||
),
|
||||
# BOTH experts: the plan subtracts one denoiser term and this family builds two.
|
||||
prequant_resident_gb_by_scheme = (("nvfp4", 16.2),),
|
||||
aliases = ("wan2.2-14b", "wan-t2v", "wan2.2-t2v", "wan-t2v-a14b", "wan-a14b"),
|
||||
has_audio = False,
|
||||
# is_moe drives the dual-DiT optimisation layers; cfg2_kwarg names the pipeline kwarg for transformer_2's
|
||||
|
|
@ -306,6 +319,9 @@ _FAMILIES: tuple[VideoFamily, ...] = (
|
|||
pipeline_class = "HunyuanVideo15Pipeline",
|
||||
transformer_class = "HunyuanVideo15Transformer3DModel",
|
||||
base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
|
||||
prequant_repos = (("nvfp4", "unsloth/HunyuanVideo-1.5-NVFP4"),),
|
||||
prequant_filenames = (("nvfp4", "HunyuanVideo-1.5-Diffusers-480p_t2v-NVFP4.pt"),),
|
||||
prequant_resident_gb_by_scheme = (("nvfp4", 4.8),),
|
||||
# No bare "hunyuanvideo" alias: it would also claim the incompatible 1.0 repos.
|
||||
aliases = ("hunyuanvideo-1-5", "hunyuanvideo1.5", "hunyuanvideo1-5", "hv15"),
|
||||
has_audio = False,
|
||||
|
|
@ -331,6 +347,10 @@ _FAMILIES: tuple[VideoFamily, ...] = (
|
|||
pipeline_class = "HunyuanVideo15Pipeline",
|
||||
transformer_class = "HunyuanVideo15Transformer3DModel",
|
||||
base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v",
|
||||
# RTN, not GPTQ: the 720p transformer is separately trained; 480p corrections do not apply.
|
||||
prequant_repos = (("nvfp4", "unsloth/HunyuanVideo-1.5-NVFP4"),),
|
||||
prequant_filenames = (("nvfp4", "HunyuanVideo-1.5-Diffusers-720p_t2v-NVFP4.pt"),),
|
||||
prequant_resident_gb_by_scheme = (("nvfp4", 4.8),),
|
||||
aliases = ("hunyuanvideo-1.5-diffusers-720p_t2v", "hv15-720p"),
|
||||
has_audio = False,
|
||||
guidance_via_guider = True,
|
||||
|
|
@ -422,6 +442,8 @@ def video_family_prequant_repo(
|
|||
the refusal path of a load request, and a table typo must not turn a legitimate pick into a
|
||||
500. A family object that predates these fields simply has no hosted checkpoint.
|
||||
"""
|
||||
if nvfp4_blocked(scheme):
|
||||
return None
|
||||
base = _prequant_base_key(base_repo)
|
||||
if base:
|
||||
for entry in getattr(fam, "prequant_variant_repos", ()) or ():
|
||||
|
|
@ -439,6 +461,26 @@ def video_family_prequant_repo(
|
|||
return None
|
||||
|
||||
|
||||
def video_family_prequant_resident_gb(fam: VideoFamily, scheme: str) -> Optional[float]:
|
||||
"""The MEASURED resident size in decimal GB of this family's hosted ``scheme`` denoiser."""
|
||||
if nvfp4_blocked(scheme):
|
||||
return None
|
||||
for entry in getattr(fam, "prequant_resident_gb_by_scheme", ()) or ():
|
||||
if not isinstance(entry, (tuple, list)) or len(entry) != 2:
|
||||
continue
|
||||
entry_scheme, resident_gb = entry
|
||||
if entry_scheme == scheme and resident_gb:
|
||||
try:
|
||||
return float(resident_gb)
|
||||
except (TypeError, ValueError): # a malformed row is "not measured", never a 500
|
||||
continue
|
||||
measured = getattr(fam, "prequant_resident_gb", None)
|
||||
try:
|
||||
return float(measured) if measured else None
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def video_family_prequant_task_specific(fam: VideoFamily, scheme: str, task: str) -> bool:
|
||||
"""True when the family names an artifact for exactly this ``(scheme, task)`` pair.
|
||||
|
||||
|
|
@ -499,7 +541,7 @@ def video_family_prequant_schemes(fam: VideoFamily, task: Optional[str] = None)
|
|||
schemes.append(entry[1])
|
||||
if task:
|
||||
schemes = [s for s in schemes if video_family_prequant_available(fam, s, task = task)]
|
||||
return tuple(schemes)
|
||||
return without_nvfp4(schemes)
|
||||
|
||||
|
||||
def snap_num_frames(fam: VideoFamily, num_frames: int) -> int:
|
||||
|
|
|
|||
|
|
@ -2318,6 +2318,15 @@ def _dense_quant_supported() -> bool:
|
|||
return bool(_dense_quant_capability)
|
||||
|
||||
|
||||
def _nvfp4_diffusion_enabled() -> bool:
|
||||
"""Whether image and video generation may offer NVFP4 (``UNSLOTH_NVFP4_DIFFUSION``)."""
|
||||
try:
|
||||
from core.inference.diffusion_nvfp4_flag import nvfp4_diffusion_enabled
|
||||
return nvfp4_diffusion_enabled()
|
||||
except Exception: # noqa: BLE001 -- a capability read must never fail a status request
|
||||
return False
|
||||
|
||||
|
||||
def _dense_quant_schemes() -> list[str]:
|
||||
"""The scheme ladder for ``/api/system``, a pure read of already-resolved state: the polled route
|
||||
must never import torch, and the entry beside it refreshed both in one pass."""
|
||||
|
|
@ -2424,6 +2433,8 @@ def get_system_info(
|
|||
# pure read of that same pass.
|
||||
"dense_quant_supported": _dense_quant_supported(),
|
||||
"dense_quant_schemes": _dense_quant_schemes(),
|
||||
# Torch-free env read, safe on this polled route.
|
||||
"nvfp4_diffusion": _nvfp4_diffusion_enabled(),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -39951,6 +39951,44 @@ async def _selected_gpu_ordinal(gpu_ids, *, allow_ranking: bool = True) -> Optio
|
|||
)
|
||||
|
||||
|
||||
def _refuse_disabled_nvfp4_request(request: Any) -> None:
|
||||
"""400 a load or plan naming NVFP4 while the NVFP4 switch (``UNSLOTH_NVFP4_DIFFUSION``) is off.
|
||||
|
||||
First thing in every image and video load / plan route, ahead of the precision gates and their
|
||||
opt-in silent fallback, so the request is refused outright, never swapped for another scheme,
|
||||
and nothing is resolved, planned or fetched for it."""
|
||||
from core.inference.diffusion_nvfp4_flag import refuse_disabled_nvfp4
|
||||
try:
|
||||
refuse_disabled_nvfp4(
|
||||
transformer_quant = getattr(request, "transformer_quant", None),
|
||||
text_encoder_quant = getattr(request, "text_encoder_quant", None),
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code = 400, detail = str(exc))
|
||||
|
||||
|
||||
async def _refuse_disabled_nvfp4_checkpoint(request: Any) -> None:
|
||||
"""400 a load or plan whose model is itself an NVFP4 checkpoint while the NVFP4 switch is off.
|
||||
|
||||
The scheme gates above never see such a pick: it names no precision, it IS one. Runs after the
|
||||
account checks, since it reads the local path or cached snapshot the request names."""
|
||||
from core.inference.diffusion_nvfp4_flag import (
|
||||
nvfp4_diffusion_enabled,
|
||||
refuse_disabled_nvfp4_checkpoint,
|
||||
)
|
||||
|
||||
if nvfp4_diffusion_enabled():
|
||||
return
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
refuse_disabled_nvfp4_checkpoint,
|
||||
getattr(request, "model_path", None),
|
||||
getattr(request, "base_repo", None),
|
||||
)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code = 400, detail = str(exc))
|
||||
|
||||
|
||||
@studio_router.post("/images/download-plan", response_model = DiffusionDownloadPlanResponse)
|
||||
async def diffusion_download_plan(
|
||||
request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
|
||||
|
|
@ -39960,6 +39998,7 @@ async def diffusion_download_plan(
|
|||
|
||||
Validates the same way /images/load does, so an unloadable pick fails here rather than
|
||||
after a multi-GB download."""
|
||||
_refuse_disabled_nvfp4_request(request)
|
||||
if account_access.managed_account():
|
||||
await asyncio.to_thread(account_access.require_media_references, request)
|
||||
if account_access.managed_account():
|
||||
|
|
@ -39971,6 +40010,7 @@ async def diffusion_download_plan(
|
|||
request = request.model_copy(
|
||||
update = {"hf_token": account_access.account_hf_token(request.hf_token)}
|
||||
)
|
||||
await _refuse_disabled_nvfp4_checkpoint(request)
|
||||
from core.inference.diffusion import (
|
||||
get_diffusion_backend,
|
||||
resolve_local_single_file,
|
||||
|
|
@ -40099,9 +40139,13 @@ def _assert_native_precision_unset(
|
|||
|
||||
Raises RuntimeError, which the route maps to 409 alongside the diffusers refusals."""
|
||||
from core.inference.diffusion_auto_policy import precision_fallback_allowed
|
||||
from core.inference.diffusion_nvfp4_flag import refuse_disabled_nvfp4
|
||||
from core.inference.diffusion_precision import normalize_te_quant
|
||||
from core.inference.diffusion_transformer_quant import TQ_AUTO, normalize_transformer_quant
|
||||
|
||||
refuse_disabled_nvfp4(
|
||||
transformer_quant = transformer_quant, text_encoder_quant = text_encoder_quant
|
||||
)
|
||||
if precision_fallback_allowed():
|
||||
return
|
||||
pinned = normalize_transformer_quant(transformer_quant)
|
||||
|
|
@ -40162,6 +40206,7 @@ async def load_diffusion_model_gated(
|
|||
Media auto-switch awaits this rather than the route so the idle unload can tell an
|
||||
API-loaded pipeline from one the user picked on the Images page.
|
||||
"""
|
||||
_refuse_disabled_nvfp4_request(request)
|
||||
if account_access.managed_account():
|
||||
await asyncio.to_thread(account_access.require_media_references, request)
|
||||
account_access.require_idle_other_accounts()
|
||||
|
|
@ -40174,6 +40219,7 @@ async def load_diffusion_model_gated(
|
|||
request = request.model_copy(
|
||||
update = {"hf_token": account_access.account_hf_token(request.hf_token)}
|
||||
)
|
||||
await _refuse_disabled_nvfp4_checkpoint(request)
|
||||
# Tested at ENTRY because `begin_load` returns before the worker moves a byte, so the only
|
||||
# fact available is that a repo ALREADY cached is not one this load will fetch. The WRITE is
|
||||
# deferred to the launch below, since the validation in between 400s without moving a byte.
|
||||
|
|
|
|||
|
|
@ -144,6 +144,12 @@ async def video_download_plan(
|
|||
):
|
||||
"""The repos + files this pick needs, so the frontend stages them through the Hub
|
||||
download manager instead of the load downloading inline. Mirrors /images/download-plan."""
|
||||
from routes.inference import (
|
||||
_refuse_disabled_nvfp4_checkpoint,
|
||||
_refuse_disabled_nvfp4_request,
|
||||
)
|
||||
|
||||
_refuse_disabled_nvfp4_request(request)
|
||||
if account_access.managed_account():
|
||||
await asyncio.to_thread(account_access.require_media_references, request)
|
||||
if account_access.managed_account():
|
||||
|
|
@ -155,6 +161,7 @@ async def video_download_plan(
|
|||
request = request.model_copy(
|
||||
update = {"hf_token": account_access.account_hf_token(request.hf_token)}
|
||||
)
|
||||
await _refuse_disabled_nvfp4_checkpoint(request)
|
||||
from core.inference.diffusion import resolve_local_single_file
|
||||
from core.inference.video import (
|
||||
assert_video_precision_available,
|
||||
|
|
@ -269,6 +276,12 @@ async def load_video_model_gated(
|
|||
"""Everything ``POST /video/load`` does, plus who asked for it. Media auto-switch awaits this rather than the
|
||||
route so the idle unload can tell an API-loaded pipeline from one the user picked on the Video page.
|
||||
"""
|
||||
from routes.inference import (
|
||||
_refuse_disabled_nvfp4_checkpoint,
|
||||
_refuse_disabled_nvfp4_request,
|
||||
)
|
||||
|
||||
_refuse_disabled_nvfp4_request(request)
|
||||
if account_access.managed_account():
|
||||
await asyncio.to_thread(account_access.require_media_references, request)
|
||||
account_access.require_idle_other_accounts()
|
||||
|
|
@ -281,6 +294,7 @@ async def load_video_model_gated(
|
|||
request = request.model_copy(
|
||||
update = {"hf_token": account_access.account_hf_token(request.hf_token)}
|
||||
)
|
||||
await _refuse_disabled_nvfp4_checkpoint(request)
|
||||
# Same as the image load: tested at entry, because `begin_load` returns before the worker
|
||||
# moves a byte, and written at the launch below, because the validation in between 400s
|
||||
# without starting one and the record would be permanent.
|
||||
|
|
|
|||
|
|
@ -1239,3 +1239,43 @@ def _drop_the_idle_reload_stash_between_tests():
|
|||
yield
|
||||
finally:
|
||||
_keepwarm._set_last_unloaded(None)
|
||||
|
||||
|
||||
# Run with the NVFP4 switch on; test_nvfp4_diffusion_flag and test_build_prequant_checkpoint must stay off.
|
||||
_NVFP4_ENABLED_TEST_MODULES = frozenset(
|
||||
{
|
||||
"test_dense_quant_rocm_gate_9396",
|
||||
"test_diffusion_auto_policy",
|
||||
"test_diffusion_backend",
|
||||
"test_diffusion_inference_info",
|
||||
"test_diffusion_lora",
|
||||
"test_diffusion_native_quant",
|
||||
"test_diffusion_pipeline_prequant",
|
||||
"test_diffusion_precision",
|
||||
"test_diffusion_prequant",
|
||||
"test_diffusion_quant_pad",
|
||||
"test_diffusion_routes",
|
||||
"test_diffusion_te_prequant",
|
||||
"test_diffusion_transformer_quant",
|
||||
"test_train_precision_scheme_contract",
|
||||
"test_video_backend",
|
||||
"test_video_families",
|
||||
"test_video_h3_te_quant",
|
||||
"test_video_prequant",
|
||||
"test_video_routes",
|
||||
"test_xformers_stub_diffusion_parity",
|
||||
}
|
||||
)
|
||||
_NVFP4_ENABLED_TEST_PREFIX = "test_diffusion_nvfp4_"
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _nvfp4_diffusion_enabled_for_nvfp4_tests(request, monkeypatch):
|
||||
"""Switch NVFP4 on for the modules above; every other module sees the default (off)."""
|
||||
module = getattr(request, "module", None)
|
||||
name = getattr(module, "__name__", "").rsplit(".", 1)[-1]
|
||||
if name in _NVFP4_ENABLED_TEST_MODULES or name.startswith(_NVFP4_ENABLED_TEST_PREFIX):
|
||||
monkeypatch.setenv("UNSLOTH_NVFP4_DIFFUSION", "1")
|
||||
else:
|
||||
monkeypatch.delenv("UNSLOTH_NVFP4_DIFFUSION", raising = False)
|
||||
yield
|
||||
|
|
|
|||
|
|
@ -253,3 +253,118 @@ def test_the_recorded_base_must_be_the_canonical_id_not_just_the_same_tail(capsy
|
|||
assert fam.transformer_class in str(exc), exc
|
||||
continue
|
||||
assert rc != 2 or "is not" not in capsys.readouterr().out, accepted
|
||||
|
||||
|
||||
def _build_nvfp4(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
fam = None,
|
||||
extra_argv = (),
|
||||
loaded = None,
|
||||
):
|
||||
"""Run the build on a tiny aligned + ragged dense model; returns (quantized names, metadata)."""
|
||||
torch = pytest.importorskip("torch")
|
||||
pytest.importorskip("torchao")
|
||||
import torchao.quantization as tq
|
||||
|
||||
import core.inference.diffusion_transformer_quant as dtq
|
||||
|
||||
fam = fam or detect_family("Tongyi-MAI/Z-Image-Turbo")
|
||||
assert fam is not None
|
||||
|
||||
class _Dense(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.aligned = torch.nn.Linear(1024, 1024, dtype = torch.bfloat16)
|
||||
self.ragged = torch.nn.Linear(1024, 1000, dtype = torch.bfloat16)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args, **kwargs):
|
||||
if loaded is not None:
|
||||
loaded.append(kwargs.get("subfolder"))
|
||||
return cls()
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
return self
|
||||
|
||||
fake_diffusers = types.ModuleType("diffusers")
|
||||
fake_diffusers.__version__ = "0"
|
||||
setattr(fake_diffusers, fam.transformer_class, _Dense)
|
||||
monkeypatch.setitem(sys.modules, "diffusers", fake_diffusers)
|
||||
|
||||
quantized = []
|
||||
|
||||
def _fake_quantize(model, config, filter_fn):
|
||||
quantized.extend(n for n, m in model.named_modules() if n and filter_fn(m, n))
|
||||
|
||||
monkeypatch.setattr(tq, "quantize_", _fake_quantize)
|
||||
monkeypatch.setattr(dtq, "_make_quant_config", lambda *a, **k: object())
|
||||
|
||||
build = _script()
|
||||
out = tmp_path / "out.pt"
|
||||
argv = ["--base", "b", "--family", fam.name, "--scheme", "nvfp4", "--out", str(out)]
|
||||
assert build.main(argv + list(extra_argv)) == 0
|
||||
return quantized, torch.load(out, weights_only = False)["metadata"]
|
||||
|
||||
|
||||
def test_the_build_skips_ragged_linears_and_records_the_alignment_floor(monkeypatch, tmp_path):
|
||||
"""A build skips Linears off the GEMM tiling floor and stamps ``require_divisible``."""
|
||||
import core.inference.diffusion_transformer_quant as dtq
|
||||
|
||||
quantized, meta = _build_nvfp4(monkeypatch, tmp_path)
|
||||
assert quantized == ["aligned"]
|
||||
assert meta["require_divisible"] == dtq.divisible_for_scheme("nvfp4") == 16
|
||||
|
||||
|
||||
def test_the_build_stamps_its_denoiser_so_it_cannot_load_as_the_second_expert(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""The build reads ``subfolder="transformer"``; published under a transformer_2 name, the loader must refuse it."""
|
||||
from core.inference.diffusion_prequant import PREQUANT_FORMAT, _validate_checkpoint
|
||||
|
||||
_, meta = _build_nvfp4(monkeypatch, tmp_path)
|
||||
assert meta["component"] == "transformer"
|
||||
ckpt = {"format": PREQUANT_FORMAT, "metadata": {**meta, "base_model_id": ""}, "state_dict": {}}
|
||||
assert _validate_checkpoint(ckpt, "nvfp4", "", None, component = "transformer")
|
||||
assert not _validate_checkpoint(ckpt, "nvfp4", "", None, component = "transformer_2")
|
||||
|
||||
|
||||
def test_the_second_expert_builds_from_its_own_subfolder_and_loads_as_transformer_2(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""``--component transformer_2`` reads that subfolder, stamps it, and publishes under the name the loader asks for."""
|
||||
from core.inference.diffusion_prequant import PREQUANT_FORMAT, _validate_checkpoint
|
||||
|
||||
wan = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers", override = "wan2.2-t2v-a14b")
|
||||
assert wan is not None
|
||||
loaded = []
|
||||
_, meta = _build_nvfp4(
|
||||
monkeypatch, tmp_path, fam = wan, extra_argv = ["--component", "transformer_2"], loaded = loaded
|
||||
)
|
||||
assert loaded == ["transformer_2"]
|
||||
assert meta["component"] == "transformer_2" and meta["family"] == wan.name
|
||||
ckpt = {"format": PREQUANT_FORMAT, "metadata": {**meta, "base_model_id": ""}, "state_dict": {}}
|
||||
assert _validate_checkpoint(ckpt, "nvfp4", "", None, component = "transformer_2")
|
||||
assert not _validate_checkpoint(ckpt, "nvfp4", "", None, component = "transformer")
|
||||
|
||||
build = _script()
|
||||
assert (
|
||||
build.upload_destination(wan, "nvfp4", rotated = False, component = "transformer_2")
|
||||
== "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt"
|
||||
)
|
||||
# The default component keeps the name it has always published under.
|
||||
assert build.upload_destination(wan, "nvfp4", rotated = False) == "transformer_nvfp4.pt"
|
||||
# A component the family declares no row for would land where nothing ever looks.
|
||||
with pytest.raises(ValueError, match = "transformer_2"):
|
||||
build.upload_destination(
|
||||
detect_family("Tongyi-MAI/Z-Image-Turbo"),
|
||||
"nvfp4",
|
||||
rotated = False,
|
||||
component = "transformer_2",
|
||||
)
|
||||
|
||||
|
||||
def test_a_default_build_still_reads_the_transformer_subfolder(monkeypatch, tmp_path):
|
||||
loaded = []
|
||||
_, meta = _build_nvfp4(monkeypatch, tmp_path, loaded = loaded)
|
||||
assert loaded == ["transformer"] and meta["component"] == "transformer"
|
||||
|
|
|
|||
|
|
@ -88,6 +88,18 @@ def test_prequant_repo_filename_convention():
|
|||
)
|
||||
assert prequant_repo_filename("org/Some-Model-quantized", "fp8") == "Some-Model-FP8.pt"
|
||||
assert prequant_repo_filename("org/PlainRepo", "int8") == "PlainRepo-INT8.pt"
|
||||
assert (
|
||||
prequant_repo_filename("unsloth/Wan2.2-TI2V-5B-NVFP4", "nvfp4") == "Wan2.2-TI2V-5B-NVFP4.pt"
|
||||
)
|
||||
assert prequant_repo_filename("unsloth/Model-MXFP8", "mxfp8") == "Model-MXFP8.pt"
|
||||
assert (
|
||||
prequant_repo_filename("unsloth/Wan2.2-T2V-A14B-NVFP4", "nvfp4", component = "transformer_2")
|
||||
== "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt"
|
||||
)
|
||||
assert (
|
||||
prequant_repo_filename("unsloth/Wan2.2-T2V-A14B-NVFP4", "nvfp4", component = "transformer")
|
||||
== "Wan2.2-T2V-A14B-NVFP4.pt"
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_variant_base_picks_variant_repo():
|
||||
|
|
@ -302,6 +314,225 @@ def test_usable_source_repo_unaffected_by_allowlist(monkeypatch, restricted_load
|
|||
assert src is not None and src.kind == "repo" and src.location == "org/hosted-fp8"
|
||||
|
||||
|
||||
_UINT8 = object() # the stub torch's uint8, so the fingerprint's dtype view is a no-op here
|
||||
|
||||
|
||||
class _Bytes:
|
||||
"""Enough of a tensor for the fingerprint: it views the payload as uint8 and hashes it."""
|
||||
|
||||
dtype = _UINT8
|
||||
|
||||
def __init__(self, payload):
|
||||
self._payload = payload
|
||||
|
||||
def detach(self):
|
||||
return self
|
||||
|
||||
def contiguous(self):
|
||||
return self
|
||||
|
||||
def dim(self):
|
||||
return 1
|
||||
|
||||
def view(self, dtype):
|
||||
return self
|
||||
|
||||
def cpu(self):
|
||||
return self
|
||||
|
||||
def numpy(self):
|
||||
import numpy as np
|
||||
return np.frombuffer(self._payload, dtype = np.uint8)
|
||||
|
||||
|
||||
class Float8Tensor:
|
||||
"""A quantized weight as far as the fingerprint is concerned: the class NAME is the key."""
|
||||
|
||||
def __init__(self, qdata):
|
||||
self.qdata = _Bytes(qdata)
|
||||
self.scale = _Bytes(b"scale")
|
||||
|
||||
|
||||
class _Recorder:
|
||||
"""A logger that keeps what it was told, so a refusal can be asserted to NAME the weight."""
|
||||
|
||||
def __init__(self):
|
||||
self.lines: list = []
|
||||
|
||||
def _record(self, msg, *args):
|
||||
self.lines.append(msg % args if args else msg)
|
||||
|
||||
debug = info = warning = error = _record
|
||||
|
||||
@property
|
||||
def text(self):
|
||||
return "\n".join(self.lines)
|
||||
|
||||
|
||||
def _fingerprinted_ckpt(payloads, *, recorded = None):
|
||||
"""A checkpoint whose state dict holds ``payloads`` and whose metadata fingerprints them."""
|
||||
from core.inference.diffusion_prequant import packed_weight_fingerprint
|
||||
|
||||
state_dict = {fqn: Float8Tensor(data) for fqn, data in payloads.items()}
|
||||
ckpt = _good_ckpt()
|
||||
ckpt["state_dict"] = state_dict
|
||||
ckpt["metadata"]["fingerprint"] = recorded or packed_weight_fingerprint(state_dict)
|
||||
return ckpt
|
||||
|
||||
|
||||
def test_a_flipped_byte_in_a_packed_weight_is_refused_and_named(monkeypatch, tmp_path):
|
||||
good = _fingerprinted_ckpt(
|
||||
{"blocks.0.attn1.to_q.weight": b"q0", "blocks.1.attn1.to_q.weight": b"q1"}
|
||||
)
|
||||
logger = _Recorder()
|
||||
assert _load(monkeypatch, tmp_path, good, logger = logger) is not None
|
||||
corrupt = _fingerprinted_ckpt(
|
||||
{"blocks.0.attn1.to_q.weight": b"q0", "blocks.1.attn1.to_q.weight": b"q1"}
|
||||
)
|
||||
corrupt["state_dict"]["blocks.1.attn1.to_q.weight"].qdata = _Bytes(b"Q1")
|
||||
logger = _Recorder()
|
||||
assert _load(monkeypatch, tmp_path, corrupt, logger = logger) is None
|
||||
assert "blocks.1.attn1.to_q.weight" in logger.text
|
||||
assert "blocks.0.attn1.to_q.weight" not in logger.text
|
||||
|
||||
|
||||
def test_a_checkpoint_without_a_fingerprint_block_still_loads(monkeypatch, tmp_path):
|
||||
ckpt = _good_ckpt()
|
||||
ckpt["state_dict"] = {"blocks.0.attn1.to_q.weight": Float8Tensor(b"q0")}
|
||||
assert _load(monkeypatch, tmp_path, ckpt) is not None
|
||||
ckpt["metadata"]["fingerprint"] = {"algo": "md5-packed-v1", "count": 0, "modules": {}}
|
||||
assert _load(monkeypatch, tmp_path, ckpt) is not None
|
||||
|
||||
|
||||
def test_the_fingerprint_mode_env_picks_how_much_is_checked(monkeypatch, tmp_path):
|
||||
from core.inference.diffusion_prequant import (
|
||||
FINGERPRINT_MODE_ENV,
|
||||
_fingerprint_sampled,
|
||||
)
|
||||
|
||||
sampled = [f"blocks.{i}.attn1.to_q.weight" for i in range(64)]
|
||||
checked = [fqn for fqn in sampled if _fingerprint_sampled(fqn)]
|
||||
assert 0 < len(checked) < len(sampled)
|
||||
assert checked == [fqn for fqn in sampled if _fingerprint_sampled(fqn)]
|
||||
|
||||
payloads = {fqn: fqn.encode("utf-8") for fqn in sampled}
|
||||
corrupt = _fingerprinted_ckpt(payloads)
|
||||
missed = next(fqn for fqn in sampled if not _fingerprint_sampled(fqn))
|
||||
corrupt["state_dict"][checked[0]].qdata = _Bytes(b"flipped")
|
||||
monkeypatch.setenv(FINGERPRINT_MODE_ENV, "sample")
|
||||
assert _load(monkeypatch, tmp_path, corrupt) is None
|
||||
only_missed = _fingerprinted_ckpt(payloads)
|
||||
only_missed["state_dict"][missed].qdata = _Bytes(b"flipped")
|
||||
assert _load(monkeypatch, tmp_path, only_missed) is not None
|
||||
monkeypatch.setenv(FINGERPRINT_MODE_ENV, "full")
|
||||
assert _load(monkeypatch, tmp_path, only_missed) is None
|
||||
monkeypatch.setenv(FINGERPRINT_MODE_ENV, "off")
|
||||
assert _load(monkeypatch, tmp_path, only_missed) is not None
|
||||
monkeypatch.setenv(FINGERPRINT_MODE_ENV, "sometimes")
|
||||
assert _load(monkeypatch, tmp_path, only_missed) is None
|
||||
|
||||
|
||||
class _CountedBytes(_Bytes):
|
||||
"""A payload that records the fqn whose bytes were actually read."""
|
||||
|
||||
def __init__(self, fqn, payload, hashed):
|
||||
super().__init__(payload)
|
||||
self._fqn = fqn
|
||||
self._hashed = hashed
|
||||
|
||||
def numpy(self):
|
||||
self._hashed.append(self._fqn)
|
||||
return super().numpy()
|
||||
|
||||
|
||||
def test_a_class_whose_payload_slots_all_read_none_is_not_fingerprinted():
|
||||
from core.inference.diffusion_prequant import packed_weight_fingerprint
|
||||
|
||||
renamed = Float8Tensor(b"q0")
|
||||
renamed.qdata = None
|
||||
renamed.scale = None
|
||||
fingerprint = packed_weight_fingerprint({"blocks.0.attn1.to_q.weight": renamed})
|
||||
assert fingerprint["modules"] == {}
|
||||
assert fingerprint["skipped"] == ["blocks.0.attn1.to_q.weight"]
|
||||
assert fingerprint["count"] == 0
|
||||
|
||||
|
||||
class Int8Tensor:
|
||||
"""torchao 0.18's int8 weight, as far as the fingerprint is concerned."""
|
||||
|
||||
def __init__(self, qdata):
|
||||
self.qdata = _Bytes(qdata)
|
||||
self.scale = _Bytes(b"scale")
|
||||
self.zero_point = None
|
||||
self.act_quant_scale = None
|
||||
self.act_quant_zero_point = None
|
||||
self.act_pre_scale = None
|
||||
|
||||
|
||||
def test_torchao_018_int8_weights_are_fingerprinted():
|
||||
from core.inference.diffusion_prequant import packed_weight_fingerprint
|
||||
|
||||
fqn = "blocks.0.attn1.to_q.weight"
|
||||
good = packed_weight_fingerprint({fqn: Int8Tensor(b"q0")})
|
||||
assert good["count"] == 1 and good["skipped"] == []
|
||||
flipped = packed_weight_fingerprint({fqn: Int8Tensor(b"q1")})
|
||||
assert flipped["modules"][fqn] != good["modules"][fqn]
|
||||
|
||||
|
||||
def test_real_torchao_int8_round_trip_fingerprint_catches_corruption():
|
||||
torch = pytest.importorskip("torch")
|
||||
pytest.importorskip("torchao")
|
||||
import io
|
||||
|
||||
from torchao.quantization import quantize_
|
||||
|
||||
from core.inference.diffusion_prequant import (
|
||||
_verify_packed_fingerprint,
|
||||
packed_weight_fingerprint,
|
||||
)
|
||||
from core.inference.diffusion_transformer_quant import _make_quant_config
|
||||
|
||||
torch.manual_seed(0)
|
||||
model = torch.nn.Sequential(torch.nn.Linear(64, 64, bias = False)).to(torch.bfloat16)
|
||||
quantize_(model, _make_quant_config("int8"))
|
||||
fingerprint = packed_weight_fingerprint(model.state_dict())
|
||||
assert fingerprint["count"] == 1, fingerprint
|
||||
assert fingerprint["skipped"] == []
|
||||
|
||||
buf = io.BytesIO()
|
||||
torch.save(model.state_dict(), buf)
|
||||
buf.seek(0)
|
||||
reloaded = torch.load(buf, weights_only = False)
|
||||
assert _verify_packed_fingerprint(reloaded, {"fingerprint": fingerprint})
|
||||
|
||||
weight = reloaded["0.weight"]
|
||||
inner = getattr(weight, "qdata", None)
|
||||
if inner is None: # torchao <= 0.17: the AffineQuantizedTensor chain
|
||||
inner = weight.original_weight_tensor.tensor_impl.int_data
|
||||
inner.view(-1)[0] ^= 1
|
||||
assert not _verify_packed_fingerprint(reloaded, {"fingerprint": fingerprint})
|
||||
|
||||
|
||||
def test_sample_mode_hashes_only_the_weights_it_compares(monkeypatch, tmp_path):
|
||||
from core.inference.diffusion_prequant import FINGERPRINT_MODE_ENV, _fingerprint_sampled
|
||||
|
||||
fqns = [f"blocks.{i}.attn1.to_q.weight" for i in range(64)]
|
||||
ckpt = _fingerprinted_ckpt({fqn: fqn.encode("utf-8") for fqn in fqns})
|
||||
hashed: list = []
|
||||
for fqn, tensor in ckpt["state_dict"].items():
|
||||
tensor.qdata = _CountedBytes(fqn, fqn.encode("utf-8"), hashed)
|
||||
tensor.scale = _CountedBytes(fqn, b"scale", hashed)
|
||||
monkeypatch.setenv(FINGERPRINT_MODE_ENV, "sample")
|
||||
assert _load(monkeypatch, tmp_path, ckpt) is not None
|
||||
assert sorted(set(hashed)) == sorted(fqn for fqn in fqns if _fingerprint_sampled(fqn))
|
||||
|
||||
|
||||
def test_a_build_that_recognises_no_quantized_weight_loads_unverified(monkeypatch, tmp_path):
|
||||
ckpt = _fingerprinted_ckpt({"blocks.0.attn1.to_q.weight": b"q0"})
|
||||
ckpt["state_dict"] = {"blocks.0.attn1.to_q.weight": object()}
|
||||
assert _load(monkeypatch, tmp_path, ckpt) is not None
|
||||
|
||||
|
||||
# ── load_prequantized_transformer ────────────────────────────────────────────────
|
||||
class _FakeTransformer:
|
||||
calls: dict = {}
|
||||
|
|
@ -380,6 +611,7 @@ def _stub_torch_accelerate(
|
|||
seen["safe_globals"] = list(entries)
|
||||
|
||||
torch.load = _load
|
||||
torch.uint8 = _UINT8
|
||||
# A stub without this namespace would let a regression to an unrestricted load pass silently.
|
||||
torch.serialization = types.SimpleNamespace(add_safe_globals = _add_safe_globals)
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
|
|
@ -413,6 +645,8 @@ def _load(
|
|||
exists = True,
|
||||
allow_local = True,
|
||||
fast_accum = None,
|
||||
component = None,
|
||||
logger = None,
|
||||
):
|
||||
_FakeTransformer.calls = {}
|
||||
_stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises)
|
||||
|
|
@ -434,7 +668,8 @@ def _load(
|
|||
hf_token = None,
|
||||
scheme = scheme,
|
||||
fast_accum = fast_accum,
|
||||
logger = None,
|
||||
component = component,
|
||||
logger = logger,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -590,6 +825,38 @@ def test_load_require_bf16_nvfp4_true_is_none(monkeypatch, tmp_path):
|
|||
assert _load(monkeypatch, tmp_path, ckpt, scheme = "nvfp4") is None
|
||||
|
||||
|
||||
def test_load_require_divisible_mismatch_is_none(monkeypatch, tmp_path):
|
||||
ckpt = _good_ckpt(scheme = "fp8")
|
||||
ckpt["metadata"]["require_divisible"] = 0
|
||||
assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is None
|
||||
|
||||
|
||||
def test_load_require_divisible_match_ok(monkeypatch, tmp_path):
|
||||
ckpt = _good_ckpt(scheme = "fp8")
|
||||
ckpt["metadata"]["require_divisible"] = 16
|
||||
assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is not None
|
||||
|
||||
|
||||
def test_load_require_divisible_absent_is_accepted(monkeypatch, tmp_path):
|
||||
ckpt = _good_ckpt(scheme = "fp8")
|
||||
assert "require_divisible" not in ckpt["metadata"]
|
||||
assert _load(monkeypatch, tmp_path, ckpt, scheme = "fp8") is not None
|
||||
|
||||
|
||||
def test_load_component_mismatch_is_none(monkeypatch, tmp_path):
|
||||
ckpt = _good_ckpt()
|
||||
ckpt["metadata"]["component"] = "transformer_2"
|
||||
assert _load(monkeypatch, tmp_path, ckpt, component = "transformer") is None
|
||||
assert _load(monkeypatch, tmp_path, ckpt, component = "transformer_2") is not None
|
||||
|
||||
|
||||
def test_load_component_absent_is_accepted(monkeypatch, tmp_path):
|
||||
ckpt = _good_ckpt()
|
||||
assert _load(monkeypatch, tmp_path, ckpt, component = "transformer") is not None
|
||||
ckpt["metadata"]["component"] = "transformer_2"
|
||||
assert _load(monkeypatch, tmp_path, ckpt) is not None
|
||||
|
||||
|
||||
def test_resolve_checkpoint_path_expands_user(monkeypatch, tmp_path):
|
||||
# The allowlist gate expands ~, so the existence check must too, or a "~/..." checkpoint that passed the gate is silently skipped.
|
||||
import os
|
||||
|
|
@ -1890,6 +2157,57 @@ def test_the_floor_check_ignores_dense_and_unreadable_state_dicts():
|
|||
assert pq._fp8_activation_floor_present({}, None) is True
|
||||
|
||||
|
||||
def test_the_checkpoint_is_released_before_the_device_copy(monkeypatch, tmp_path):
|
||||
"""assign=True shares the checkpoint's tensors: unreference them before ``.to(device)``."""
|
||||
import weakref
|
||||
|
||||
seen: dict = {}
|
||||
|
||||
class _StateDict(dict):
|
||||
"""A weak-referenceable state dict: plain dicts cannot be weakly referenced."""
|
||||
|
||||
class _ReleaseProbe(_FakeTransformer):
|
||||
def load_state_dict(
|
||||
self,
|
||||
sd,
|
||||
strict = True,
|
||||
assign = False,
|
||||
):
|
||||
_FakeTransformer.calls["load_state_dict"] = {"strict": strict, "assign": assign}
|
||||
seen["state_dict"] = weakref.ref(sd)
|
||||
|
||||
def to(self, device):
|
||||
seen["alive_at_move"] = seen["state_dict"]() is not None
|
||||
return super().to(device)
|
||||
|
||||
_FakeTransformer.calls = {}
|
||||
_stub_torch_accelerate(monkeypatch, None)
|
||||
monkeypatch.setattr(
|
||||
pq,
|
||||
"_torch_load_prequant",
|
||||
lambda path, **kwargs: {
|
||||
"format": PREQUANT_FORMAT,
|
||||
"metadata": {"scheme": "int8", "base_model_id": "Tongyi-MAI/Z-Image-Turbo"},
|
||||
"state_dict": _StateDict(weight = object()),
|
||||
},
|
||||
)
|
||||
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path))
|
||||
path = tmp_path / "ckpt.pt"
|
||||
path.write_bytes(b"x")
|
||||
|
||||
out = load_prequantized_transformer(
|
||||
_ReleaseProbe,
|
||||
"Tongyi-MAI/Z-Image-Turbo",
|
||||
PrequantSource(kind = "path", location = str(path), filename = None),
|
||||
device = "cuda",
|
||||
dtype = "bfloat16",
|
||||
scheme = "int8",
|
||||
)
|
||||
|
||||
assert out is not None
|
||||
assert seen["alive_at_move"] is False
|
||||
|
||||
|
||||
def test_the_plan_only_commits_to_a_hosted_name_this_install_can_open(monkeypatch):
|
||||
"""The call that drops the released dense shards must ask both questions, not one.
|
||||
|
||||
|
|
|
|||
|
|
@ -23,11 +23,13 @@ from core.inference.diffusion_quant_pad import ( # noqa: E402
|
|||
DEFAULT_PAD_TO,
|
||||
INT_MM_MIN_M,
|
||||
PadToMinM,
|
||||
ZeroRowSafeLinear,
|
||||
activation_granularity_is_per_row,
|
||||
is_quantized_linear,
|
||||
matching_linear_fqns,
|
||||
padding_is_bitwise_exact,
|
||||
wrap_small_m_linears,
|
||||
wrap_zero_row_linears,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -306,6 +308,86 @@ def test_is_quantized_linear_only_accepts_a_torchao_weight():
|
|||
assert is_quantized_linear(nn.LayerNorm(8)) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bias", [True, False])
|
||||
@pytest.mark.parametrize("shape", [(1, 0, 8), (0, 8), (2, 0, 8)])
|
||||
def test_an_empty_activation_comes_back_at_the_projected_width(bias, shape):
|
||||
inner = _RecordingLinear(8, 6, bias = bias).to(torch.bfloat16)
|
||||
wrapped = ZeroRowSafeLinear(inner)
|
||||
x = torch.zeros(*shape, dtype = torch.bfloat16)
|
||||
with torch.no_grad():
|
||||
got, want = wrapped(x), torch.nn.functional.linear(x, inner.weight, inner.bias)
|
||||
assert got.shape == (*shape[:-1], 6) == want.shape
|
||||
assert got.dtype == want.dtype == torch.bfloat16
|
||||
assert torch.equal(got, want)
|
||||
assert not hasattr(inner, "seen"), "the inner GEMM must not be called with an empty input"
|
||||
|
||||
|
||||
def test_the_guard_is_inert_on_a_non_empty_activation():
|
||||
torch.manual_seed(0)
|
||||
inner = _RecordingLinear(8, 6)
|
||||
wrapped = ZeroRowSafeLinear(inner)
|
||||
for shape in ((1, 8), (5, 8), (2, 7, 8)):
|
||||
x = torch.randn(*shape)
|
||||
with torch.no_grad():
|
||||
assert torch.equal(wrapped(x), inner(x))
|
||||
assert inner.seen.shape == x.shape
|
||||
|
||||
|
||||
def test_the_guard_passes_attributes_and_keys_through():
|
||||
plain, guarded = _Tiny(), _Tiny()
|
||||
guarded.context_embedder = ZeroRowSafeLinear(guarded.context_embedder)
|
||||
assert guarded.context_embedder.in_features == 8
|
||||
assert guarded.context_embedder.out_features == 6
|
||||
assert guarded.context_embedder.weight is guarded.context_embedder.inner.weight
|
||||
with pytest.raises(AttributeError):
|
||||
guarded.context_embedder.definitely_not_a_linear_attribute
|
||||
|
||||
assert sorted(guarded.state_dict()) == sorted(plain.state_dict())
|
||||
assert "context_embedder.inner.weight" not in guarded.state_dict()
|
||||
assert plain.load_state_dict(dict(guarded.state_dict()), strict = True)
|
||||
assert guarded.load_state_dict(dict(plain.state_dict()), strict = True)
|
||||
assert torch.equal(guarded.context_embedder.weight, plain.context_embedder.weight)
|
||||
|
||||
|
||||
def test_wrapping_for_zero_rows_is_idempotent_and_skips_dense_linears():
|
||||
model = _Tiny()
|
||||
assert wrap_zero_row_linears(model, ["context_embedder"]) == ()
|
||||
assert isinstance(model.context_embedder, nn.Linear)
|
||||
|
||||
model.context_embedder = _fake_quantized_linear()
|
||||
assert wrap_zero_row_linears(model, ["context_embedder"]) == ("context_embedder",)
|
||||
assert isinstance(model.context_embedder, ZeroRowSafeLinear)
|
||||
assert wrap_zero_row_linears(model, ["context_embedder"]) == ()
|
||||
assert not isinstance(model.context_embedder.inner, ZeroRowSafeLinear)
|
||||
|
||||
|
||||
def test_the_two_wrappers_never_stack():
|
||||
model = _Tiny()
|
||||
model.context_embedder = _fake_quantized_linear()
|
||||
assert wrap_small_m_linears(model, ["context_embedder"]) == ("context_embedder",)
|
||||
assert wrap_zero_row_linears(model, ["context_embedder"]) == ()
|
||||
assert isinstance(model.context_embedder, PadToMinM)
|
||||
assert model.context_embedder(torch.randn(0, 8)).shape == (0, 6)
|
||||
|
||||
|
||||
def test_the_guard_list_is_read_with_the_same_substring_rule():
|
||||
class _Projection(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.linear_1 = _fake_quantized_linear(8, 8)
|
||||
self.linear_2 = _fake_quantized_linear(8, 6)
|
||||
|
||||
def forward(self, x):
|
||||
return self.linear_2(self.linear_1(x))
|
||||
|
||||
model = _Tiny()
|
||||
model.image_embedder = _Projection()
|
||||
fqns = matching_linear_fqns(model, ("image_embedder",))
|
||||
assert fqns == ("image_embedder.linear_1", "image_embedder.linear_2")
|
||||
assert wrap_zero_row_linears(model, fqns) == fqns
|
||||
assert model.image_embedder(torch.zeros(1, 0, 8)).shape == (1, 0, 6)
|
||||
|
||||
|
||||
# ── the real thing ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
|
@ -324,3 +406,34 @@ def test_int8_padding_is_bitwise_exact_on_a_real_quantized_linear():
|
|||
assert activation_granularity_is_per_row(lin) is True
|
||||
for m in (1, 10, 13, 16, 17, 19, 64):
|
||||
assert padding_is_bitwise_exact(lin, m), f"padding changed the kept rows at M = {m}"
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason = "nvfp4 dynamic quant needs Blackwell")
|
||||
def test_the_zero_row_guard_on_a_real_nvfp4_linear():
|
||||
"""A real nvfp4 Linear survives a zero-row call once wrapped."""
|
||||
pytest.importorskip("torchao")
|
||||
if torch.cuda.get_device_capability() < (10, 0):
|
||||
pytest.skip("nvfp4 needs sm_100+")
|
||||
from torchao.quantization import quantize_
|
||||
|
||||
from core.inference.diffusion_transformer_quant import _make_quant_config, make_filter_fn
|
||||
|
||||
lin = nn.Linear(1152, 3072, dtype = torch.bfloat16).cuda().eval()
|
||||
quantize_(lin, _make_quant_config("nvfp4"), filter_fn = make_filter_fn(0))
|
||||
assert is_quantized_linear(lin)
|
||||
|
||||
empty = torch.zeros(1, 0, 1152, dtype = torch.bfloat16, device = "cuda")
|
||||
with pytest.raises(RuntimeError, match = "numel"):
|
||||
lin(empty)
|
||||
|
||||
guarded = ZeroRowSafeLinear(lin)
|
||||
out = guarded(empty)
|
||||
assert out.shape == (1, 0, 3072) and out.dtype == torch.bfloat16
|
||||
|
||||
torch.manual_seed(0)
|
||||
for m in (1, 5, 7):
|
||||
x = torch.randn(1, m, 1152, dtype = torch.bfloat16, device = "cuda")
|
||||
with torch.no_grad():
|
||||
got = guarded(x)
|
||||
assert torch.equal(got, lin(x))
|
||||
assert torch.isfinite(got).all()
|
||||
|
|
|
|||
|
|
@ -1375,6 +1375,107 @@ def test_apply_small_m_padding_is_inert_without_a_pad_list(monkeypatch):
|
|||
apply_small_m_padding(object(), TQ_INT8, "minimax-h3")
|
||||
|
||||
|
||||
def test_the_zero_row_guard_is_nvfp4_only():
|
||||
from core.inference.diffusion_transformer_quant import zero_row_tokens_for_scheme
|
||||
for scheme in (TQ_FP8, TQ_INT8, TQ_MXFP8, "auto"):
|
||||
for family in ("hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"):
|
||||
assert zero_row_tokens_for_scheme(scheme, family) == (), (scheme, family)
|
||||
|
||||
|
||||
def test_both_hunyuan_tiers_guard_their_trimmable_streams():
|
||||
from core.inference.diffusion_transformer_quant import zero_row_tokens_for_scheme
|
||||
for family in ("hunyuanvideo-1.5", "hunyuanvideo-1.5-720p"):
|
||||
tokens = zero_row_tokens_for_scheme(TQ_NVFP4, family)
|
||||
assert "image_embedder" in tokens
|
||||
assert "context_embedder_2" in tokens
|
||||
|
||||
|
||||
def test_an_unlisted_family_has_no_zero_row_guard():
|
||||
"""Scoped to the families whose attention path is measured to produce an empty activation."""
|
||||
from core.inference.diffusion_transformer_quant import zero_row_tokens_for_scheme
|
||||
for family in ("z-image", "qwen-image", "wan2.2-ti2v-5b", "minimax-h3", None):
|
||||
assert zero_row_tokens_for_scheme(TQ_NVFP4, family) == ()
|
||||
|
||||
|
||||
def test_the_guard_list_does_not_double_up_with_the_pad_or_exclude_lists():
|
||||
from core.inference.diffusion_transformer_quant import (
|
||||
_NVFP4_FAMILY_ZERO_ROW_NAME_TOKENS,
|
||||
exclude_tokens_for_scheme,
|
||||
pad_tokens_for_scheme,
|
||||
zero_row_tokens_for_scheme,
|
||||
)
|
||||
for family in _NVFP4_FAMILY_ZERO_ROW_NAME_TOKENS:
|
||||
guarded = zero_row_tokens_for_scheme(TQ_NVFP4, family)
|
||||
assert guarded
|
||||
assert exclude_tokens_for_scheme(TQ_NVFP4, family) == ()
|
||||
assert pad_tokens_for_scheme(TQ_NVFP4, family) == ()
|
||||
|
||||
|
||||
def test_quantize_transformer_guards_after_padding(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_NVFP4})
|
||||
order = []
|
||||
tqz = types.ModuleType("torchao.quantization")
|
||||
tqz.quantize_ = lambda module, config, filter_fn = None: order.append("quantize_")
|
||||
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
||||
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg")
|
||||
monkeypatch.setattr(
|
||||
tq,
|
||||
"apply_small_m_padding",
|
||||
lambda transformer, scheme, family = None, logger = None: (order.append("pad") or ()),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
tq,
|
||||
"apply_zero_row_guard",
|
||||
lambda transformer, scheme, family = None, logger = None: (
|
||||
order.append(("guard", scheme, family)) or ()
|
||||
),
|
||||
)
|
||||
transformer = types.SimpleNamespace()
|
||||
pipe = types.SimpleNamespace(transformer = transformer)
|
||||
assert (
|
||||
quantize_transformer(pipe, _target(), mode = "nvfp4", family = "hunyuanvideo-1.5") == TQ_NVFP4
|
||||
)
|
||||
assert order == ["quantize_", "pad", ("guard", TQ_NVFP4, "hunyuanvideo-1.5")]
|
||||
assert transformer._unsloth_runtime_quant == TQ_NVFP4
|
||||
|
||||
|
||||
def test_a_guard_failure_fails_the_whole_quantise(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_NVFP4})
|
||||
tqz = types.ModuleType("torchao.quantization")
|
||||
tqz.quantize_ = lambda module, config, filter_fn = None: None
|
||||
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
||||
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg")
|
||||
|
||||
def _boom(
|
||||
transformer,
|
||||
scheme,
|
||||
family = None,
|
||||
logger = None,
|
||||
):
|
||||
raise RuntimeError("cannot resolve the guard list")
|
||||
|
||||
monkeypatch.setattr(tq, "apply_zero_row_guard", _boom)
|
||||
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
||||
assert quantize_transformer(pipe, _target(), mode = "nvfp4", family = "hunyuanvideo-1.5") is None
|
||||
assert not hasattr(pipe.transformer, "_unsloth_runtime_quant")
|
||||
|
||||
|
||||
def test_apply_zero_row_guard_is_inert_without_a_guard_list(monkeypatch):
|
||||
"""apply_zero_row_guard is inert, and imports nothing, without a guard list."""
|
||||
from core.inference.diffusion_transformer_quant import apply_zero_row_guard
|
||||
|
||||
stub = types.ModuleType("core.inference.diffusion_quant_pad") # no names to import
|
||||
monkeypatch.setitem(sys.modules, "core.inference.diffusion_quant_pad", stub)
|
||||
|
||||
assert apply_zero_row_guard(object(), TQ_NVFP4, "z-image") == ()
|
||||
assert apply_zero_row_guard(object(), TQ_FP8, "hunyuanvideo-1.5") == ()
|
||||
assert apply_zero_row_guard(object(), TQ_NVFP4, None) == ()
|
||||
with pytest.raises(ImportError):
|
||||
apply_zero_row_guard(object(), TQ_NVFP4, "hunyuanvideo-1.5")
|
||||
|
||||
|
||||
def test_the_training_deny_is_a_superset_of_the_inference_deny():
|
||||
# The two tables are separate because rendering evidence is not training evidence, but the
|
||||
# relationship must only ever go one way: anything inference refuses, training refuses too.
|
||||
|
|
@ -1437,6 +1538,158 @@ def test_the_candidate_list_agrees_with_the_selector_on_the_winner(monkeypatch):
|
|||
assert (candidates[0] if candidates else None) == chosen, (cc, family)
|
||||
|
||||
|
||||
def _prefer(
|
||||
monkeypatch,
|
||||
row,
|
||||
family = "fake-video",
|
||||
):
|
||||
"""Install one ``_FAMILY_AUTO_PREFER`` row."""
|
||||
monkeypatch.setattr(tq, "_FAMILY_AUTO_PREFER", {family: row})
|
||||
|
||||
|
||||
def _nvfp4_head(**kw):
|
||||
return tq._AutoPrefer(floor = (10, 0), schemes = (TQ_NVFP4,), **kw)
|
||||
|
||||
|
||||
def test_a_prefer_row_leads_the_tier_on_datacenter_blackwell(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8})
|
||||
_prefer(monkeypatch, _nvfp4_head())
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_NVFP4
|
||||
assert tq.auto_scheme_candidates(_target(), "fake-video") == (
|
||||
TQ_NVFP4,
|
||||
TQ_INT8,
|
||||
TQ_FP8,
|
||||
TQ_MXFP8,
|
||||
)
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "z-image") == TQ_INT8
|
||||
assert tq.auto_scheme_candidates(_target()) == (TQ_INT8, TQ_FP8, TQ_MXFP8)
|
||||
|
||||
|
||||
def test_a_prefer_row_is_dropped_below_its_capability_floor(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA L40S")
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_INT8})
|
||||
_prefer(monkeypatch, _nvfp4_head())
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_INT8
|
||||
assert tq.auto_scheme_candidates(_target(), "fake-video") == (TQ_INT8, TQ_FP8)
|
||||
|
||||
|
||||
def test_a_head_the_probe_rejects_falls_through_to_the_tier(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_FP8, TQ_MXFP8, TQ_INT8})
|
||||
_prefer(monkeypatch, _nvfp4_head())
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_INT8
|
||||
assert tq.auto_scheme_candidates(_target(), "fake-video") == (TQ_INT8, TQ_FP8, TQ_MXFP8)
|
||||
|
||||
|
||||
def test_the_deny_list_outranks_a_prefer_row(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8})
|
||||
_prefer(monkeypatch, _nvfp4_head(), family = "qwen-image")
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "qwen-image") == TQ_INT8
|
||||
assert TQ_NVFP4 not in tq.auto_scheme_candidates(_target(), "qwen-image")
|
||||
|
||||
|
||||
def test_a_prefer_row_needs_consumer_ok_to_apply_to_a_consumer_gpu(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (12, 0), device_name = "NVIDIA GeForce RTX 5090")
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8})
|
||||
_prefer(monkeypatch, _nvfp4_head())
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_INT8
|
||||
_prefer(monkeypatch, _nvfp4_head(consumer_ok = True))
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") == TQ_NVFP4
|
||||
assert tq.auto_scheme_candidates(_target(), "fake-video") == (
|
||||
TQ_NVFP4,
|
||||
TQ_INT8,
|
||||
TQ_FP8,
|
||||
TQ_MXFP8,
|
||||
)
|
||||
|
||||
|
||||
def test_a_head_already_in_the_tier_is_not_listed_twice(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_FP8, TQ_MXFP8, TQ_INT8})
|
||||
_prefer(monkeypatch, tq._AutoPrefer(floor = (10, 0), schemes = (TQ_INT8,)))
|
||||
assert tq.auto_scheme_candidates(_target(), "fake-video") == (TQ_INT8, TQ_FP8, TQ_MXFP8)
|
||||
|
||||
|
||||
def test_a_capability_below_every_tier_has_no_order_even_with_a_head(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (7, 5))
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_INT8})
|
||||
_prefer(monkeypatch, tq._AutoPrefer(floor = (7, 0), schemes = (TQ_NVFP4,)))
|
||||
assert tq._auto_scheme_order("fake-video", "cuda", (7, 5)) == ()
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = "fake-video") is None
|
||||
assert tq.auto_scheme_candidates(_target(), "fake-video") == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("family", [None, "fake-video", "z-image", "qwen-image"])
|
||||
@pytest.mark.parametrize(
|
||||
"allowed",
|
||||
[
|
||||
{TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8},
|
||||
{TQ_FP8, TQ_MXFP8, TQ_INT8},
|
||||
{TQ_NVFP4, TQ_INT8},
|
||||
{TQ_INT8},
|
||||
set(),
|
||||
],
|
||||
)
|
||||
def test_the_candidate_head_stays_the_selector_winner_with_a_prefer_row(
|
||||
monkeypatch, family, allowed
|
||||
):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, allowed)
|
||||
_prefer(monkeypatch, _nvfp4_head())
|
||||
candidates = tq.auto_scheme_candidates(_target(), family)
|
||||
chosen = select_transformer_quant_scheme(_target(), "auto", family = family)
|
||||
assert (candidates[0] if candidates else None) == chosen, (family, allowed)
|
||||
|
||||
|
||||
def test_the_shipped_prefer_table_keeps_the_ladder_as_it_was(monkeypatch):
|
||||
_stub_torch(monkeypatch, cc = (10, 0))
|
||||
_allow(monkeypatch, {TQ_NVFP4, TQ_FP8, TQ_MXFP8, TQ_INT8})
|
||||
assert tq._FAMILY_AUTO_PREFER == {}
|
||||
for family in (None, "hunyuanvideo-1.5", "hunyuanvideo-1.5-720p", "wan2.2-ti2v-5b"):
|
||||
assert select_transformer_quant_scheme(_target(), "auto", family = family) == TQ_INT8
|
||||
|
||||
|
||||
def test_divisible_for_scheme_matches_each_gemm():
|
||||
"""scaled_mm (fp8, and torchao's nvfp4 path) needs 16-aligned dims and MX block scaling 32."""
|
||||
from core.inference.diffusion_transformer_quant import divisible_for_scheme
|
||||
|
||||
assert divisible_for_scheme(TQ_FP8) == 16
|
||||
assert divisible_for_scheme(TQ_NVFP4) == 16
|
||||
assert divisible_for_scheme(TQ_MXFP8) == 32
|
||||
assert divisible_for_scheme(TQ_INT8) == 0
|
||||
assert divisible_for_scheme("auto") == 0
|
||||
|
||||
|
||||
def test_quantize_transformer_filters_on_the_scheme_alignment(monkeypatch):
|
||||
torch_stub = _stub_torch(monkeypatch, cc = (10, 0))
|
||||
|
||||
class _Linear:
|
||||
def __init__(self, in_features, out_features):
|
||||
self.in_features, self.out_features = in_features, out_features
|
||||
self.weight = types.SimpleNamespace(dtype = torch_stub.bfloat16)
|
||||
|
||||
torch_stub.nn = types.SimpleNamespace(Linear = _Linear)
|
||||
monkeypatch.setattr(tq, "_make_quant_config", lambda scheme, fast_accum = None: "cfg")
|
||||
tqz = types.ModuleType("torchao.quantization")
|
||||
seen: list = []
|
||||
tqz.quantize_ = lambda module, config, filter_fn = None: seen.append(filter_fn)
|
||||
monkeypatch.setitem(sys.modules, "torchao.quantization", tqz)
|
||||
|
||||
filters = {}
|
||||
for scheme in (TQ_NVFP4, TQ_MXFP8):
|
||||
_allow(monkeypatch, {scheme})
|
||||
pipe = types.SimpleNamespace(transformer = types.SimpleNamespace())
|
||||
assert quantize_transformer(pipe, _target(), mode = scheme) == scheme
|
||||
filters[scheme] = seen[-1]
|
||||
|
||||
aligned_16 = _Linear(528, 2048) # 528 % 16 == 0, 528 % 32 == 16
|
||||
assert filters[TQ_NVFP4](aligned_16, "image_embedder.linear_2") is True
|
||||
assert filters[TQ_MXFP8](aligned_16, "image_embedder.linear_2") is False
|
||||
assert filters[TQ_NVFP4](_Linear(520, 2048), "blocks.0.ff.net.0") is False
|
||||
|
||||
|
||||
def test_the_pre_eviction_gate_does_not_refuse_on_an_indeterminate_probe(monkeypatch):
|
||||
# The route-level precision gate asks the selector, not the probe, so unproven_ok has to reach
|
||||
# through select_transformer_quant_scheme for the leniency to exist where it matters.
|
||||
|
|
|
|||
621
studio/backend/tests/test_nvfp4_diffusion_flag.py
Normal file
621
studio/backend/tests/test_nvfp4_diffusion_flag.py
Normal file
|
|
@ -0,0 +1,621 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""The NVFP4 switch (``UNSLOTH_NVFP4_DIFFUSION``), default OFF."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from core.inference import diffusion_nvfp4_flag as flag
|
||||
from core.inference import diffusion_transformer_quant as tq
|
||||
|
||||
ENV = "UNSLOTH_NVFP4_DIFFUSION"
|
||||
DISABLED = "NVFP4 is disabled in this build"
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _default_off(monkeypatch):
|
||||
monkeypatch.delenv(ENV, raising = False)
|
||||
yield
|
||||
|
||||
|
||||
def _enable(monkeypatch):
|
||||
monkeypatch.setenv(ENV, "1")
|
||||
|
||||
|
||||
def test_the_default_is_off_and_lives_in_one_constant():
|
||||
assert flag.NVFP4_DIFFUSION_DEFAULT is False
|
||||
assert flag.NVFP4_DIFFUSION_ENV == ENV
|
||||
assert flag.nvfp4_diffusion_enabled() is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("requested", ["auto", "fp8", "int8"])
|
||||
def test_a_video_family_without_a_usable_hosted_seed_spends_no_smoke_probe(monkeypatch, requested):
|
||||
import torch
|
||||
|
||||
import core.inference.video as vid
|
||||
from core.inference.video_families import _FAMILIES
|
||||
|
||||
calls = []
|
||||
monkeypatch.setattr(
|
||||
tq,
|
||||
"_child_probe_table",
|
||||
lambda card: calls.append(card) or {s: None for s in tq.TQ_SCHEMES},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
tq, "_smoke_probe", lambda s, d, unproven_ok = False: calls.append(s) or unproven_ok
|
||||
)
|
||||
monkeypatch.setattr(tq, "dense_transformer_supported", lambda t: True)
|
||||
monkeypatch.setattr(tq, "_capability", lambda: (10, 0))
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
target = types.SimpleNamespace(device = "cuda", dtype = torch.bfloat16, ordinal = 0)
|
||||
for fam in _FAMILIES:
|
||||
if getattr(fam, "modular_workflow", None):
|
||||
continue
|
||||
assert (
|
||||
vid._video_auto_denoiser_scheme(
|
||||
fam, target = target, requested = requested, base_repo = fam.base_repo
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert calls == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["1", "true", "TRUE", " yes ", "on", "On"])
|
||||
def test_truthy_values_enable(monkeypatch, value):
|
||||
monkeypatch.setenv(ENV, value)
|
||||
assert flag.nvfp4_diffusion_enabled() is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "false", "no", "off", "", " ", "garbage", "2", "enable"])
|
||||
def test_everything_else_disables(monkeypatch, value):
|
||||
monkeypatch.setenv(ENV, value)
|
||||
assert flag.nvfp4_diffusion_enabled() is False
|
||||
|
||||
|
||||
def test_flipping_the_default_is_the_one_line_change(monkeypatch):
|
||||
"""Unset (and unrecognised) follow the constant; an explicit 0 still turns it off."""
|
||||
monkeypatch.setattr(flag, "NVFP4_DIFFUSION_DEFAULT", True)
|
||||
assert flag.nvfp4_diffusion_enabled() is True
|
||||
monkeypatch.setenv(ENV, "0")
|
||||
assert flag.nvfp4_diffusion_enabled() is False
|
||||
monkeypatch.setenv(ENV, "garbage")
|
||||
assert flag.nvfp4_diffusion_enabled() is True
|
||||
|
||||
|
||||
def test_the_switch_is_read_at_call_time(monkeypatch):
|
||||
assert flag.nvfp4_blocked("nvfp4")
|
||||
_enable(monkeypatch)
|
||||
assert not flag.nvfp4_blocked("nvfp4")
|
||||
monkeypatch.setenv(ENV, "0")
|
||||
assert flag.nvfp4_blocked("NVFP4")
|
||||
|
||||
|
||||
def test_only_nvfp4_is_blocked():
|
||||
for scheme in ("int8", "fp8", "mxfp8", "fp8_dynamic", "auto", None, ""):
|
||||
assert not flag.nvfp4_blocked(scheme)
|
||||
assert flag.without_nvfp4(("int8", "nvfp4", "fp8")) == ("int8", "fp8")
|
||||
assert flag.nvfp4_repo_blocked("unsloth/Wan2.2-TI2V-5B-NVFP4")
|
||||
assert not flag.nvfp4_repo_blocked("unsloth/Wan2.2-TI2V-5B-FP8")
|
||||
assert not flag.nvfp4_repo_blocked(None)
|
||||
|
||||
|
||||
def _blackwell(monkeypatch, *, supported = ("int8", "fp8", "nvfp4", "mxfp8")):
|
||||
"""A datacenter sm_100 host whose smoke probe passes ``supported``, nothing allocated."""
|
||||
import torch
|
||||
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
monkeypatch.setattr(tq, "_capability", lambda: (10, 0))
|
||||
monkeypatch.setattr(tq, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(tq, "_is_consumer_gpu", lambda device = None: False)
|
||||
monkeypatch.setattr(tq, "_child_probe_table", lambda device: None)
|
||||
monkeypatch.setattr(
|
||||
tq, "_smoke_probe", lambda scheme, device, unproven_ok = False: scheme in supported
|
||||
)
|
||||
monkeypatch.setattr(tq, "_SMOKE_CACHE", {})
|
||||
monkeypatch.setattr(
|
||||
tq,
|
||||
"_AUTO_LADDER",
|
||||
(((10, 0), (tq.TQ_INT8, tq.TQ_FP8, tq.TQ_NVFP4, tq.TQ_MXFP8)),),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
tq,
|
||||
"_FAMILY_AUTO_PREFER",
|
||||
{"wan2.2-t2v-a14b": tq._AutoPrefer(floor = (10, 0), schemes = (tq.TQ_NVFP4,))},
|
||||
)
|
||||
return types.SimpleNamespace(device = "cuda", dtype = None)
|
||||
|
||||
|
||||
def _auto(target, family):
|
||||
"""``auto`` on ``family`` with a hosted checkpoint for every scheme, so only the switch blocks nvfp4."""
|
||||
kwargs = {}
|
||||
if "has_prequant" in inspect.signature(tq.select_transformer_quant_scheme).parameters:
|
||||
kwargs["has_prequant"] = lambda scheme: True
|
||||
return tq.select_transformer_quant_scheme(target, "auto", family = family, **kwargs)
|
||||
|
||||
|
||||
def test_auto_never_offers_nvfp4(monkeypatch):
|
||||
target = _blackwell(monkeypatch)
|
||||
fam = "wan2.2-t2v-a14b"
|
||||
assert tq.TQ_NVFP4 not in tq._auto_scheme_order(fam, "cuda", (10, 0))
|
||||
assert tq.TQ_NVFP4 not in tq.auto_scheme_candidates(target, fam)
|
||||
assert tq.TQ_NVFP4 not in tq.auto_scheme_candidates_cached(target, fam)
|
||||
assert _auto(target, fam) == tq.TQ_INT8
|
||||
|
||||
|
||||
def test_a_host_that_only_runs_nvfp4_has_nothing_automatic(monkeypatch):
|
||||
target = _blackwell(monkeypatch, supported = ("nvfp4",))
|
||||
assert _auto(target, "wan2.2-t2v-a14b") is None
|
||||
probed = []
|
||||
monkeypatch.setattr(tq, "_smoke_probe", lambda scheme, *a, **k: probed.append(scheme) or True)
|
||||
assert tq._scheme_supported(tq.TQ_NVFP4, "cuda") is False
|
||||
assert probed == []
|
||||
|
||||
|
||||
def test_the_child_probe_is_not_asked_about_nvfp4(monkeypatch):
|
||||
asked = []
|
||||
|
||||
class _Stop(Exception):
|
||||
pass
|
||||
|
||||
def _spawn(*args, **kwargs):
|
||||
asked.append(kwargs.get("args") or args)
|
||||
raise _Stop
|
||||
|
||||
import multiprocessing as mp
|
||||
|
||||
class _Ctx:
|
||||
def Queue(self):
|
||||
return types.SimpleNamespace(close = lambda: None, join_thread = lambda: None)
|
||||
|
||||
def Process(self, *args, **kwargs):
|
||||
return _spawn(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(mp, "get_context", lambda kind: _Ctx())
|
||||
monkeypatch.setattr(tq, "_CHILD_PROBE_UNAVAILABLE", False)
|
||||
monkeypatch.setattr(tq, "_close_probe_child", lambda proc, queue: True)
|
||||
assert tq._child_probe_table("cuda") is None
|
||||
assert asked, "the spawn was never attempted"
|
||||
schemes = asked[0][4]
|
||||
assert tq.TQ_NVFP4 not in schemes and tq.TQ_INT8 in schemes
|
||||
|
||||
|
||||
def test_an_explicit_nvfp4_transformer_quant_is_refused_not_swapped(monkeypatch):
|
||||
target = _blackwell(monkeypatch)
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
tq.normalize_transformer_quant("nvfp4")
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
tq.normalize_transformer_quant(" NVFP4 ")
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
tq.select_transformer_quant_scheme(target, "nvfp4", family = "z-image")
|
||||
assert DISABLED in tq.explain_unusable_scheme("z-image", "nvfp4")
|
||||
assert tq.normalize_transformer_quant("fp8") == tq.TQ_FP8
|
||||
assert tq.select_transformer_quant_scheme(target, "mxfp8", family = "z-image") == tq.TQ_MXFP8
|
||||
|
||||
|
||||
def test_the_precision_gates_refuse_nvfp4_even_under_the_silent_fallback(monkeypatch):
|
||||
"""The opt-in fallback swaps a declined scheme for another one; a disabled scheme is refused."""
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
from core.inference.video import assert_video_precision_available
|
||||
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ALLOW_PRECISION_FALLBACK", "1")
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
DiffusionBackend.assert_precision_available(
|
||||
object(), None, model_kind = "pipeline", transformer_quant = "nvfp4"
|
||||
)
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
DiffusionBackend.assert_precision_available(
|
||||
object(), None, model_kind = "pipeline", text_encoder_quant = "nvfp4"
|
||||
)
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
assert_video_precision_available(None, model_kind = "pipeline", transformer_quant = "nvfp4")
|
||||
from routes.inference import _assert_native_precision_unset
|
||||
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
_assert_native_precision_unset(transformer_quant = "nvfp4")
|
||||
|
||||
|
||||
def _run(coro):
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field,control",
|
||||
[("transformer_quant", "transformer_quant"), ("text_encoder_quant", "text_encoder_quant")],
|
||||
)
|
||||
def test_every_load_and_plan_route_400s_an_nvfp4_request(monkeypatch, field, control):
|
||||
from fastapi import HTTPException
|
||||
|
||||
from models.inference import DiffusionLoadRequest, VideoLoadRequest
|
||||
from routes import inference as image_routes
|
||||
from routes import video as video_routes
|
||||
|
||||
touched = []
|
||||
monkeypatch.setattr(
|
||||
image_routes.account_access, "managed_account", lambda: touched.append("acct") or False
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
video_routes.account_access, "managed_account", lambda: touched.append("acct") or False
|
||||
)
|
||||
image = DiffusionLoadRequest(model_path = "Tongyi-MAI/Z-Image-Turbo", **{field: "nvfp4"})
|
||||
video = VideoLoadRequest(model_path = "Wan-AI/Wan2.2-TI2V-5B-Diffusers", **{field: "nvfp4"})
|
||||
calls = [
|
||||
image_routes.diffusion_download_plan(image, current_subject = "u"),
|
||||
image_routes.load_diffusion_model_gated(image, "u"),
|
||||
video_routes.video_download_plan(video, current_subject = "u"),
|
||||
video_routes.load_video_model_gated(video, "u"),
|
||||
]
|
||||
for call in calls:
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_run(call)
|
||||
assert exc.value.status_code == 400
|
||||
assert DISABLED in exc.value.detail
|
||||
assert f"{control}='nvfp4'" in exc.value.detail
|
||||
assert touched == []
|
||||
|
||||
|
||||
def test_text_encoder_nvfp4_is_refused():
|
||||
from core.inference import diffusion_precision as dp
|
||||
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
dp.normalize_te_quant("nvfp4")
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
dp.resolve_te_quant_request("nvfp4", None)
|
||||
assert dp.resolve_te_quant_request(None, "nvfp4") == (None, False)
|
||||
assert dp.resolve_te_quant_request("auto", "fp8") == ("fp8", True)
|
||||
target = types.SimpleNamespace(device = "cuda", dtype = None)
|
||||
assert dp.te_quant_supported(target, dp.TE_QUANT_NVFP4) is False
|
||||
assert dp.normalize_te_quant("fp8_dynamic") == dp.TE_QUANT_FP8_DYNAMIC
|
||||
|
||||
|
||||
def test_no_hosted_nvfp4_text_encoder_is_resolved():
|
||||
from core.inference.diffusion_te_prequant import family_te_prequant_repo
|
||||
fam = types.SimpleNamespace(
|
||||
te_prequant_repos = (("nvfp4", "text_encoder", "unsloth/Some-Model-NVFP4"),)
|
||||
)
|
||||
assert family_te_prequant_repo(fam, "nvfp4", "text_encoder") is None
|
||||
|
||||
|
||||
_WAN_T2V = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
|
||||
|
||||
|
||||
def _wan_a14b():
|
||||
from core.inference.video_families import detect_video_family
|
||||
return detect_video_family(_WAN_T2V)
|
||||
|
||||
|
||||
def test_no_hosted_nvfp4_repo_is_resolved_for_any_family():
|
||||
from core.inference.diffusion_families import family_prequant_repo
|
||||
from core.inference.diffusion_prequant import resolve_prequant_source
|
||||
from core.inference.video_families import (
|
||||
_FAMILIES as VIDEO_FAMILIES,
|
||||
video_family_prequant_available,
|
||||
video_family_prequant_repo,
|
||||
video_family_prequant_resident_gb,
|
||||
video_family_prequant_schemes,
|
||||
)
|
||||
from core.inference.video_denoiser_prequant import denoiser_prequant_sources
|
||||
|
||||
hosted = [f for f in VIDEO_FAMILIES if any(s == "nvfp4" for s, _r in f.prequant_repos)]
|
||||
assert hosted, "the registry no longer hosts an NVFP4 denoiser; retarget this test"
|
||||
for fam in hosted:
|
||||
assert video_family_prequant_repo(fam, "nvfp4") is None
|
||||
assert family_prequant_repo(fam, "nvfp4") is None
|
||||
assert resolve_prequant_source(fam, "nvfp4") is None
|
||||
assert resolve_prequant_source(fam, "nvfp4", path_override = "/x/local-NVFP4.pt") is None
|
||||
assert denoiser_prequant_sources(fam, "nvfp4", fam.base_repo) is None
|
||||
assert not video_family_prequant_available(fam, "nvfp4")
|
||||
assert video_family_prequant_resident_gb(fam, "nvfp4") is None
|
||||
assert "nvfp4" not in video_family_prequant_schemes(fam)
|
||||
|
||||
|
||||
class _RecordingApi:
|
||||
def __init__(self, repos):
|
||||
self.repos = repos
|
||||
self.asked: list[str] = []
|
||||
|
||||
def model_info(
|
||||
self,
|
||||
repo_id,
|
||||
files_metadata = False,
|
||||
token = None,
|
||||
):
|
||||
self.asked.append(repo_id)
|
||||
return types.SimpleNamespace(siblings = self.repos[repo_id])
|
||||
|
||||
|
||||
def _sib(name, size):
|
||||
return types.SimpleNamespace(rfilename = name, size = size)
|
||||
|
||||
|
||||
_WAN_SIBLINGS = {
|
||||
_WAN_T2V: [
|
||||
_sib("model_index.json", 1000),
|
||||
_sib("transformer/config.json", 1000),
|
||||
_sib("transformer/diffusion_pytorch_model.safetensors", 28_000_000_000),
|
||||
_sib("transformer_2/config.json", 1000),
|
||||
_sib("transformer_2/diffusion_pytorch_model.safetensors", 28_000_000_000),
|
||||
_sib("text_encoder/model-00001-of-00001.safetensors", 11_000_000_000),
|
||||
_sib("vae/diffusion_pytorch_model.safetensors", 500_000_000),
|
||||
],
|
||||
"unsloth/Wan2.2-T2V-A14B-NVFP4": [
|
||||
_sib("Wan2.2-T2V-A14B-NVFP4.pt", 8_000_000_000),
|
||||
_sib("Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", 8_100_000_000),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _wan_plan(monkeypatch):
|
||||
"""A Wan2.2-A14B pipeline plan on a host whose auto selector lands on nvfp4, Hub mocked."""
|
||||
import core.inference.video as video_mod
|
||||
from core.inference import diffusion_prequant
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
api = _RecordingApi(_WAN_SIBLINGS)
|
||||
monkeypatch.setattr("huggingface_hub.HfApi", lambda *a, **k: api)
|
||||
monkeypatch.setattr(
|
||||
DiffusionBackend,
|
||||
"_hub_file_is_cached",
|
||||
staticmethod(lambda repo_id, filename, revision = None, expected_size = None, **kw: False),
|
||||
)
|
||||
monkeypatch.setattr(video_mod, "select_transformer_quant_scheme", lambda *a, **k: "nvfp4")
|
||||
monkeypatch.setattr(
|
||||
diffusion_prequant, "restricted_prequant_load_supported", lambda *a, **k: True
|
||||
)
|
||||
monkeypatch.setattr(video_mod, "_video_seed_stays_resident", lambda fam, **kw: True)
|
||||
monkeypatch.setattr(
|
||||
video_mod,
|
||||
"resolve_diffusion_device_target",
|
||||
lambda *a, **k: types.SimpleNamespace(backend = "cuda", device = "cuda", dtype = None),
|
||||
)
|
||||
from core.inference.video import VideoBackend
|
||||
|
||||
plan = VideoBackend().download_plan(_WAN_T2V, model_kind = "pipeline")
|
||||
staged = {f for e in plan["entries"] for f in e["files"]}
|
||||
return plan, staged, api
|
||||
|
||||
|
||||
def test_the_video_download_plan_asks_no_nvfp4_repo_and_keeps_the_dense_shards(monkeypatch):
|
||||
plan, staged, api = _wan_plan(monkeypatch)
|
||||
assert not any(repo.lower().endswith("-nvfp4") for repo in api.asked), api.asked
|
||||
assert not any("NVFP4" in f for f in staged)
|
||||
assert "transformer/diffusion_pytorch_model.safetensors" in staged
|
||||
assert "transformer_2/diffusion_pytorch_model.safetensors" in staged
|
||||
|
||||
|
||||
def test_the_video_seed_hub_probe_never_reaches_an_nvfp4_repo():
|
||||
from core.inference.video import VideoBackend
|
||||
|
||||
api = _RecordingApi(_WAN_SIBLINGS)
|
||||
repo, files = VideoBackend._denoiser_prequant_hub_files(_wan_a14b(), "nvfp4", _WAN_T2V, api)
|
||||
assert (repo, files) == (None, [])
|
||||
assert api.asked == []
|
||||
|
||||
|
||||
def test_the_image_hub_entry_never_asks_about_an_nvfp4_repo(monkeypatch):
|
||||
from core.inference.diffusion import DiffusionBackend
|
||||
|
||||
asked = []
|
||||
monkeypatch.setattr(
|
||||
"huggingface_hub.HfApi", lambda *a, **k: types.SimpleNamespace(model_info = asked.append)
|
||||
)
|
||||
source = types.SimpleNamespace(
|
||||
kind = "repo", location = "unsloth/Z-Image-Turbo-NVFP4", filename = "x-NVFP4.pt"
|
||||
)
|
||||
assert DiffusionBackend._prequant_source_hub_entry(source, None, scheme = "nvfp4") is None
|
||||
assert DiffusionBackend._prequant_source_hub_entry(source, None) is None
|
||||
assert asked == []
|
||||
|
||||
|
||||
def test_the_info_payload_and_estimates_omit_nvfp4(monkeypatch):
|
||||
from core.inference.diffusion_auto_policy import estimate_dense_quant
|
||||
from core.inference.diffusion_families import detect_family
|
||||
from core.inference.diffusion_inference_info import family_inference_infos
|
||||
|
||||
infos = family_inference_infos()
|
||||
assert infos
|
||||
for info in infos:
|
||||
assert "nvfp4" not in info["estimated_resident_gb"]
|
||||
assert "int8" in info["estimated_resident_gb"]
|
||||
assert estimate_dense_quant(detect_family("Tongyi-MAI/Z-Image-Turbo"), "nvfp4") is None
|
||||
assert estimate_dense_quant(detect_family("Tongyi-MAI/Z-Image-Turbo"), "int8") is not None
|
||||
|
||||
|
||||
def test_api_system_reports_the_switch(monkeypatch):
|
||||
import main
|
||||
|
||||
assert main._nvfp4_diffusion_enabled() is False
|
||||
_enable(monkeypatch)
|
||||
assert main._nvfp4_diffusion_enabled() is True
|
||||
|
||||
|
||||
def test_enabled_the_nvfp4_path_is_reachable_end_to_end(monkeypatch):
|
||||
_enable(monkeypatch)
|
||||
from core.inference import diffusion_precision as dp
|
||||
from core.inference.diffusion_inference_info import family_inference_infos
|
||||
from core.inference.video_denoiser_prequant import denoiser_prequant_sources
|
||||
from core.inference.video_families import video_family_prequant_repo
|
||||
|
||||
target = _blackwell(monkeypatch)
|
||||
assert _auto(target, "wan2.2-t2v-a14b") == "nvfp4"
|
||||
assert tq.normalize_transformer_quant("nvfp4") == "nvfp4"
|
||||
assert tq.select_transformer_quant_scheme(target, "nvfp4", family = "z-image") == "nvfp4"
|
||||
assert dp.normalize_te_quant("nvfp4") == "nvfp4"
|
||||
fam = _wan_a14b()
|
||||
assert video_family_prequant_repo(fam, "nvfp4") == "unsloth/Wan2.2-T2V-A14B-NVFP4"
|
||||
sources = denoiser_prequant_sources(fam, "nvfp4", _WAN_T2V)
|
||||
assert sources and set(sources) == {"transformer", "transformer_2"}
|
||||
assert all(s.location == "unsloth/Wan2.2-T2V-A14B-NVFP4" for s in sources.values())
|
||||
assert all("nvfp4" in info["estimated_resident_gb"] for info in family_inference_infos())
|
||||
|
||||
|
||||
def test_enabled_the_video_plan_stages_the_hosted_nvfp4_denoisers(monkeypatch):
|
||||
_enable(monkeypatch)
|
||||
plan, staged, api = _wan_plan(monkeypatch)
|
||||
assert "unsloth/Wan2.2-T2V-A14B-NVFP4" in api.asked
|
||||
assert {"Wan2.2-T2V-A14B-NVFP4.pt", "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt"} <= staged
|
||||
assert "transformer/diffusion_pytorch_model.safetensors" not in staged
|
||||
|
||||
|
||||
def test_enabled_the_routes_let_nvfp4_through_the_switch(monkeypatch):
|
||||
_enable(monkeypatch)
|
||||
from routes.inference import _refuse_disabled_nvfp4_request
|
||||
_refuse_disabled_nvfp4_request(
|
||||
types.SimpleNamespace(transformer_quant = "nvfp4", text_encoder_quant = "nvfp4")
|
||||
)
|
||||
|
||||
|
||||
def _prequant_dir(root, name, scheme):
|
||||
"""A directory holding one tiny safetensors pre-quant artifact that records ``scheme``."""
|
||||
import json
|
||||
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
|
||||
from core.inference.diffusion_prequant import PREQUANT_FORMAT
|
||||
from core.inference.prequant_safetensors import UNSLOTH_FORMAT_KEY, UNSLOTH_METADATA_KEY
|
||||
|
||||
folder = root / name
|
||||
folder.mkdir()
|
||||
save_file(
|
||||
{"w": torch.zeros(2)},
|
||||
str(folder / "model.safetensors"),
|
||||
metadata = {
|
||||
UNSLOTH_FORMAT_KEY: PREQUANT_FORMAT,
|
||||
UNSLOTH_METADATA_KEY: json.dumps({"scheme": scheme}),
|
||||
},
|
||||
)
|
||||
return folder
|
||||
|
||||
|
||||
def _checkpoint_route_calls(image_path, video_path, monkeypatch):
|
||||
from models.inference import DiffusionLoadRequest, VideoLoadRequest
|
||||
from routes import inference as image_routes
|
||||
from routes import video as video_routes
|
||||
|
||||
monkeypatch.setattr(image_routes.account_access, "managed_account", lambda: False)
|
||||
monkeypatch.setattr(video_routes.account_access, "managed_account", lambda: False)
|
||||
monkeypatch.setattr(image_routes.account_access, "require_idle_other_accounts", lambda: None)
|
||||
monkeypatch.setattr(video_routes.account_access, "require_idle_other_accounts", lambda: None)
|
||||
image = DiffusionLoadRequest(model_path = image_path)
|
||||
video = VideoLoadRequest(model_path = video_path)
|
||||
return [
|
||||
lambda: image_routes.diffusion_download_plan(image, current_subject = "u"),
|
||||
lambda: image_routes.load_diffusion_model_gated(image, "u"),
|
||||
lambda: video_routes.video_download_plan(video, current_subject = "u"),
|
||||
lambda: video_routes.load_video_model_gated(video, "u"),
|
||||
]
|
||||
|
||||
|
||||
HOSTED_NVFP4 = [
|
||||
"unsloth/Z-Image-Turbo-NVFP4",
|
||||
"unsloth/FLUX.1-schnell-NVFP4",
|
||||
"unsloth/Qwen-Image-2512-NVFP4",
|
||||
"unsloth/Wan2.2-TI2V-5B-NVFP4",
|
||||
"unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
"unsloth/HunyuanVideo-1.5-NVFP4",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("repo", HOSTED_NVFP4)
|
||||
def test_a_hosted_nvfp4_repo_as_the_model_is_refused(repo):
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
flag.refuse_disabled_nvfp4_checkpoint(repo)
|
||||
|
||||
|
||||
def test_every_family_registered_nvfp4_repo_is_recognised_by_the_table():
|
||||
from core.inference.diffusion_prequant import hosted_nvfp4_repo_ids
|
||||
from core.inference.video_families import _FAMILIES
|
||||
|
||||
registered = {
|
||||
repo.lower()
|
||||
for fam in _FAMILIES
|
||||
for scheme, repo in fam.prequant_repos
|
||||
if scheme == "nvfp4"
|
||||
}
|
||||
assert registered and registered <= hosted_nvfp4_repo_ids()
|
||||
|
||||
|
||||
def test_a_cached_repo_is_judged_by_its_metadata_not_its_name(tmp_path, monkeypatch):
|
||||
from core.inference import diffusion_prequant as dpq
|
||||
|
||||
snap = _prequant_dir(tmp_path, "snap", "nvfp4")
|
||||
monkeypatch.setattr(dpq, "_cached_snapshot_dirs", lambda repo: [str(snap)])
|
||||
assert dpq.declares_nvfp4_checkpoint("someone/renamed-checkpoint") is True
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
flag.refuse_disabled_nvfp4_checkpoint("someone/renamed-checkpoint")
|
||||
|
||||
|
||||
def test_a_local_dir_whose_metadata_declares_nvfp4_is_refused(tmp_path):
|
||||
folder = _prequant_dir(tmp_path, "my-checkpoint", "nvfp4")
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
flag.refuse_disabled_nvfp4_checkpoint(str(folder))
|
||||
with pytest.raises(ValueError, match = DISABLED):
|
||||
flag.refuse_disabled_nvfp4_checkpoint(str(folder / "model.safetensors"))
|
||||
|
||||
|
||||
def test_the_routes_400_an_nvfp4_checkpoint_as_the_model(tmp_path, monkeypatch):
|
||||
from fastapi import HTTPException
|
||||
folder = _prequant_dir(tmp_path, "local-z-image", "nvfp4")
|
||||
for image_path, video_path in (
|
||||
("unsloth/Z-Image-Turbo-NVFP4", "unsloth/Wan2.2-TI2V-5B-NVFP4"),
|
||||
(str(folder), str(folder)),
|
||||
):
|
||||
for call in _checkpoint_route_calls(image_path, video_path, monkeypatch):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_run(call())
|
||||
assert exc.value.status_code == 400
|
||||
assert DISABLED in exc.value.detail
|
||||
assert "NVFP4 checkpoint" in exc.value.detail
|
||||
assert str(tmp_path) not in exc.value.detail
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"path",
|
||||
[
|
||||
"Tongyi-MAI/Z-Image-Turbo",
|
||||
"unsloth/Z-Image-Turbo-FP8",
|
||||
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"unsloth/Qwen-Image-2512-GGUF",
|
||||
"",
|
||||
None,
|
||||
],
|
||||
)
|
||||
def test_non_nvfp4_models_are_not_refused(path):
|
||||
flag.refuse_disabled_nvfp4_checkpoint(path)
|
||||
|
||||
|
||||
def test_a_local_dir_with_another_scheme_is_not_refused(tmp_path):
|
||||
folder = _prequant_dir(tmp_path, "my-fp8-checkpoint", "fp8")
|
||||
flag.refuse_disabled_nvfp4_checkpoint(str(folder))
|
||||
plain = tmp_path / "plain-pipeline"
|
||||
plain.mkdir()
|
||||
(plain / "model_index.json").write_text("{}")
|
||||
flag.refuse_disabled_nvfp4_checkpoint(str(plain))
|
||||
pipeline = _prequant_dir(tmp_path, "pipeline-with-extra", "nvfp4")
|
||||
(pipeline / "model_index.json").write_text("{}")
|
||||
flag.refuse_disabled_nvfp4_checkpoint(str(pipeline))
|
||||
|
||||
|
||||
def test_enabled_an_nvfp4_checkpoint_as_the_model_proceeds(tmp_path, monkeypatch):
|
||||
_enable(monkeypatch)
|
||||
folder = _prequant_dir(tmp_path, "my-checkpoint", "nvfp4")
|
||||
for path in (*HOSTED_NVFP4, str(folder)):
|
||||
flag.refuse_disabled_nvfp4_checkpoint(path)
|
||||
|
||||
reached = []
|
||||
|
||||
def _stop(*args, **kwargs):
|
||||
reached.append(True)
|
||||
raise RuntimeError("reached the backend")
|
||||
|
||||
monkeypatch.setattr("core.inference.diffusion.get_diffusion_backend", _stop)
|
||||
calls = _checkpoint_route_calls("unsloth/Z-Image-Turbo-NVFP4", str(folder), monkeypatch)
|
||||
with pytest.raises(RuntimeError, match = "reached the backend"):
|
||||
_run(calls[0]())
|
||||
assert reached == [True]
|
||||
|
|
@ -9245,6 +9245,572 @@ def test_h3_generate_non_oom_error_leaves_the_graphs_alone(fake_runtime):
|
|||
assert handle.resets == 0
|
||||
|
||||
|
||||
def _stub_denoiser_seed(
|
||||
monkeypatch,
|
||||
*,
|
||||
scheme = "nvfp4",
|
||||
components = ("transformer", "transformer_2"),
|
||||
seeded = True,
|
||||
repo = "unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
plan_scheme = "nvfp4",
|
||||
):
|
||||
"""Stub the whole seeding route: plan-time scheme, registry lookup and checkpoint load."""
|
||||
import core.inference.video as video_mod
|
||||
import core.inference.video_denoiser_prequant as dq
|
||||
|
||||
if plan_scheme is not None:
|
||||
monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: plan_scheme)
|
||||
source = types.SimpleNamespace(
|
||||
kind = "repo", location = repo, filename = "x.pt", fallback_filename = None
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dq, "denoiser_prequant_sources", lambda fam, s, base: {c: source for c in components}
|
||||
)
|
||||
modules = {c: object() for c in components}
|
||||
calls = []
|
||||
|
||||
def _fake_pipe_kwargs(fam, base, **kwargs):
|
||||
calls.append({"base": base, **kwargs})
|
||||
return dict(modules) if seeded else {}
|
||||
|
||||
monkeypatch.setattr(dq, "denoiser_prequant_pipe_kwargs", _fake_pipe_kwargs)
|
||||
return calls, modules
|
||||
|
||||
|
||||
def test_wan_a14b_prequant_seeds_both_dits_instead_of_quantising_them(fake_runtime, monkeypatch):
|
||||
"""The prequant twin of ``test_wan_a14b_dense_quant_applies_to_both_dits``."""
|
||||
import core.inference.video as video_mod
|
||||
|
||||
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
||||
quantised = []
|
||||
|
||||
def _fake_quant(
|
||||
view,
|
||||
target,
|
||||
*,
|
||||
mode,
|
||||
family,
|
||||
logger = None,
|
||||
):
|
||||
quantised.append(view.transformer)
|
||||
return "int8"
|
||||
|
||||
monkeypatch.setattr(video_mod, "quantize_transformer", _fake_quant)
|
||||
calls, modules = _stub_denoiser_seed(monkeypatch)
|
||||
|
||||
backend = VideoBackend()
|
||||
status = backend.load_pipeline(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
|
||||
assembled = _FakeWanPipelineSingle.last
|
||||
assert assembled["transformer"] is modules["transformer"]
|
||||
assert assembled["transformer_2"] is modules["transformer_2"]
|
||||
assert quantised == [], "a seeded DiT must not be quantised again"
|
||||
assert status["transformer_quant"] == "nvfp4"
|
||||
resolved = status["resolved"]["transformer_quant"]
|
||||
assert "unsloth/Wan2.2-T2V-A14B-NVFP4" in resolved["reason"]
|
||||
assert resolved["requested"] == "nvfp4" and resolved["value"] == "nvfp4"
|
||||
assert resolved["status"] == "applied"
|
||||
assert calls and calls[0]["scheme"] == "nvfp4"
|
||||
|
||||
|
||||
def test_a_checkpoint_that_will_not_load_falls_back_to_the_dense_quant(fake_runtime, monkeypatch):
|
||||
import core.inference.video as video_mod
|
||||
|
||||
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
||||
quantised = []
|
||||
|
||||
def _fake_quant(
|
||||
view,
|
||||
target,
|
||||
*,
|
||||
mode,
|
||||
family,
|
||||
logger = None,
|
||||
):
|
||||
quantised.append(view.transformer)
|
||||
return "nvfp4"
|
||||
|
||||
monkeypatch.setattr(video_mod, "quantize_transformer", _fake_quant)
|
||||
_stub_denoiser_seed(monkeypatch, seeded = False)
|
||||
|
||||
backend = VideoBackend()
|
||||
status = backend.load_pipeline(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
pipe = backend._state.pipe
|
||||
assert quantised == [pipe.transformer, pipe.transformer_2]
|
||||
assert status["transformer_quant"] == "nvfp4"
|
||||
assert "unsloth/" not in status["resolved"]["transformer_quant"]["reason"]
|
||||
|
||||
|
||||
def test_seeding_is_skipped_entirely_under_offload(fake_runtime, monkeypatch):
|
||||
import core.inference.video as video_mod
|
||||
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ALLOW_PRECISION_FALLBACK", "1")
|
||||
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: None)
|
||||
calls, _modules = _stub_denoiser_seed(monkeypatch)
|
||||
real_plan = video_mod.plan_diffusion_memory
|
||||
monkeypatch.setattr(
|
||||
video_mod,
|
||||
"plan_diffusion_memory",
|
||||
lambda **kwargs: dataclasses.replace(real_plan(**kwargs), offload_policy = "model"),
|
||||
)
|
||||
_stub_apply_memory_plan(monkeypatch, video_mod)
|
||||
|
||||
backend = VideoBackend()
|
||||
status = backend.load_pipeline(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
assert calls == [], "no checkpoint may be fetched for a load that cannot keep it resident"
|
||||
assert status["transformer_quant"] is None
|
||||
assert "moves the DiT" in status["resolved"]["transformer_quant"]["reason"]
|
||||
|
||||
|
||||
def _plans_for(monkeypatch, video_mod):
|
||||
"""Every ``model_dense_mib`` the memory planner was asked to judge, in order."""
|
||||
real_plan = video_mod.plan_diffusion_memory
|
||||
seen: list = []
|
||||
|
||||
def _record(**kwargs):
|
||||
seen.append(kwargs.get("model_dense_mib"))
|
||||
return real_plan(**kwargs)
|
||||
|
||||
monkeypatch.setattr(video_mod, "plan_diffusion_memory", _record)
|
||||
return seen
|
||||
|
||||
|
||||
def test_the_memory_plan_prices_a_seeded_denoiser_at_the_measured_row(fake_runtime, monkeypatch):
|
||||
import core.inference.video as video_mod
|
||||
|
||||
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: None)
|
||||
monkeypatch.setattr(video_mod, "video_family_prequant_resident_gb", lambda fam, scheme: 8.0)
|
||||
_stub_denoiser_seed(monkeypatch)
|
||||
seen = _plans_for(monkeypatch, video_mod)
|
||||
|
||||
backend = VideoBackend()
|
||||
backend.load_pipeline(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
fam = _detect_load_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers", None, None)
|
||||
mib_per_gb = 1000.0**3 / (1024.0 * 1024.0)
|
||||
companions = fam.bf16_components_gb[1] + fam.bf16_components_gb[2]
|
||||
assert seen == [int((8.0 + companions) * mib_per_gb)]
|
||||
assert seen[0] < int((fam.bf16_components_gb[0] + companions) * mib_per_gb)
|
||||
|
||||
|
||||
def test_a_failed_seed_replans_at_bf16_and_refuses_again(fake_runtime, monkeypatch):
|
||||
"""The artifact-sized budget is valid only once the checkpoint is in hand."""
|
||||
import core.inference.video as video_mod
|
||||
|
||||
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: "nvfp4")
|
||||
monkeypatch.setattr(video_mod, "video_family_prequant_resident_gb", lambda fam, scheme: 8.0)
|
||||
_stub_denoiser_seed(monkeypatch, seeded = False)
|
||||
seen = _plans_for(monkeypatch, video_mod)
|
||||
refusals: list = []
|
||||
monkeypatch.setattr(
|
||||
video_mod,
|
||||
"raise_on_unified_memory_shortfall",
|
||||
lambda plan, family = None, logger = None: refusals.append(
|
||||
plan.model_dense_mib if hasattr(plan, "model_dense_mib") else None
|
||||
),
|
||||
)
|
||||
|
||||
backend = VideoBackend()
|
||||
backend.load_pipeline(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
fam = _detect_load_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers", None, None)
|
||||
mib_per_gb = 1000.0**3 / (1024.0 * 1024.0)
|
||||
companions = fam.bf16_components_gb[1] + fam.bf16_components_gb[2]
|
||||
assert seen == [
|
||||
int((8.0 + companions) * mib_per_gb),
|
||||
int((fam.bf16_components_gb[0] + companions) * mib_per_gb),
|
||||
]
|
||||
assert len(refusals) == 2
|
||||
|
||||
|
||||
def test_the_planned_scheme_is_what_the_load_seeds(fake_runtime, monkeypatch):
|
||||
import core.inference.video as video_mod
|
||||
|
||||
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: None)
|
||||
calls, _modules = _stub_denoiser_seed(monkeypatch, plan_scheme = None)
|
||||
monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: None)
|
||||
|
||||
backend = VideoBackend()
|
||||
status = backend.load_pipeline(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
_video_auto_denoiser_planned = "nvfp4",
|
||||
)
|
||||
assert [c["scheme"] for c in calls] == ["nvfp4"]
|
||||
assert status["transformer_quant"] == "nvfp4"
|
||||
|
||||
|
||||
def test_a_seed_the_plan_declined_is_not_re_taken_by_the_load(fake_runtime, monkeypatch):
|
||||
"""The load honours the plan's decline: re-deciding here would fetch the artifact inline."""
|
||||
import core.inference.video as video_mod
|
||||
|
||||
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: "nvfp4")
|
||||
calls, _modules = _stub_denoiser_seed(monkeypatch, plan_scheme = None)
|
||||
monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4")
|
||||
|
||||
backend = VideoBackend()
|
||||
status = backend.load_pipeline(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
_video_auto_denoiser_planned = video_mod.DENOISER_SEED_DECLINED,
|
||||
)
|
||||
assert calls == [], "a seed the plan declined may not be fetched inline by the load"
|
||||
assert status["transformer_quant"] == "nvfp4"
|
||||
assert "unsloth/" not in status["resolved"]["transformer_quant"]["reason"]
|
||||
|
||||
|
||||
_A14B_SIBLINGS = [
|
||||
_sibling("model_index.json", 10),
|
||||
_sibling("transformer/config.json", 1),
|
||||
_sibling("transformer/diffusion_pytorch_model-00001-of-00002.safetensors", 14),
|
||||
_sibling("transformer/diffusion_pytorch_model-00002-of-00002.safetensors", 14),
|
||||
_sibling("transformer_2/config.json", 1),
|
||||
_sibling("transformer_2/diffusion_pytorch_model-00001-of-00002.safetensors", 14),
|
||||
_sibling("transformer_2/diffusion_pytorch_model-00002-of-00002.safetensors", 14),
|
||||
_sibling("text_encoder/model-00001-of-00001.safetensors", 11),
|
||||
_sibling("vae/diffusion_pytorch_model.safetensors", 3),
|
||||
_sibling("scheduler/scheduler_config.json", 1),
|
||||
]
|
||||
|
||||
|
||||
def test_base_download_files_drops_both_experts_and_keeps_both_configs():
|
||||
info = types.SimpleNamespace(siblings = _A14B_SIBLINGS)
|
||||
|
||||
dense = dict(VideoBackend._base_download_files(info, "pipeline"))
|
||||
seeded = dict(
|
||||
VideoBackend._base_download_files(
|
||||
info,
|
||||
"pipeline",
|
||||
skip_transformer_weights = True,
|
||||
skip_transformer_components = ("transformer", "transformer_2"),
|
||||
)
|
||||
)
|
||||
assert not any(n.endswith(".safetensors") and n.startswith("transformer") for n in seeded)
|
||||
assert seeded["transformer/config.json"] == 1
|
||||
assert seeded["transformer_2/config.json"] == 1
|
||||
assert seeded["text_encoder/model-00001-of-00001.safetensors"] == 11
|
||||
assert sum(dense.values()) - sum(seeded.values()) == 56
|
||||
|
||||
|
||||
def test_base_download_files_keeps_the_h3_partition_default():
|
||||
info = types.SimpleNamespace(siblings = _H3_SIBLINGS)
|
||||
references = dict(
|
||||
VideoBackend._base_download_files(
|
||||
info, "pipeline", skip_transformer_weights = True, h3_task = "ref2va"
|
||||
)
|
||||
)
|
||||
assert not any(n.startswith("transformer_ref/diffusion_pytorch_model") for n in references)
|
||||
assert "transformer_ref/config.json" in references
|
||||
|
||||
|
||||
def test_the_download_plan_stages_both_experts_artifacts(monkeypatch):
|
||||
import core.inference.video as video_mod
|
||||
import core.inference.video_denoiser_prequant as dq
|
||||
|
||||
_plan_api(
|
||||
monkeypatch,
|
||||
{
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": [
|
||||
_PlanSibling("model_index.json", 1000),
|
||||
_PlanSibling("transformer/config.json", 1000),
|
||||
_PlanSibling("transformer/diffusion_pytorch_model.safetensors", 28_000_000_000),
|
||||
_PlanSibling("transformer_2/config.json", 1000),
|
||||
_PlanSibling("transformer_2/diffusion_pytorch_model.safetensors", 28_000_000_000),
|
||||
_PlanSibling("text_encoder/model-00001-of-00001.safetensors", 11_000_000_000),
|
||||
_PlanSibling("vae/diffusion_pytorch_model.safetensors", 500_000_000),
|
||||
],
|
||||
"unsloth/Wan2.2-T2V-A14B-NVFP4": [
|
||||
_PlanSibling("Wan2.2-T2V-A14B-NVFP4.pt", 8_000_000_000),
|
||||
_PlanSibling("Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", 8_100_000_000),
|
||||
],
|
||||
},
|
||||
)
|
||||
sources = {
|
||||
"transformer": types.SimpleNamespace(
|
||||
kind = "repo",
|
||||
location = "unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
filename = "Wan2.2-T2V-A14B-NVFP4.pt",
|
||||
fallback_filename = None,
|
||||
),
|
||||
"transformer_2": types.SimpleNamespace(
|
||||
kind = "repo",
|
||||
location = "unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
filename = "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt",
|
||||
fallback_filename = None,
|
||||
),
|
||||
}
|
||||
monkeypatch.setattr(dq, "denoiser_prequant_sources", lambda fam, scheme, base: sources)
|
||||
monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4")
|
||||
monkeypatch.setattr(video_mod, "_video_seed_stays_resident", lambda fam, **kw: True)
|
||||
|
||||
plan = VideoBackend().download_plan(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
|
||||
staged = {f for e in plan["entries"] for f in e["files"]}
|
||||
assert "Wan2.2-T2V-A14B-NVFP4.pt" in staged
|
||||
assert "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt" in staged
|
||||
assert not any(
|
||||
f.endswith("diffusion_pytorch_model.safetensors") and f.startswith("transformer")
|
||||
for f in staged
|
||||
)
|
||||
assert "transformer/config.json" in staged and "transformer_2/config.json" in staged
|
||||
assert (
|
||||
plan["required_bytes"]
|
||||
== 8_000_000_000 + 8_100_000_000 + 1000 * 3 + 11_000_000_000 + 500_000_000
|
||||
)
|
||||
|
||||
|
||||
def test_an_explicit_scheme_under_speed_off_stages_the_hosted_experts(monkeypatch):
|
||||
"""An EXPLICIT scheme under speed_mode="off" is still quantized, so the plan keeps its seed."""
|
||||
import core.inference.video as video_mod
|
||||
import core.inference.video_denoiser_prequant as dq
|
||||
|
||||
_plan_api(
|
||||
monkeypatch,
|
||||
{
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": [
|
||||
_PlanSibling("model_index.json", 1000),
|
||||
_PlanSibling("transformer/config.json", 1000),
|
||||
_PlanSibling("transformer/diffusion_pytorch_model.safetensors", 28_000_000_000),
|
||||
_PlanSibling("transformer_2/config.json", 1000),
|
||||
_PlanSibling("transformer_2/diffusion_pytorch_model.safetensors", 28_000_000_000),
|
||||
_PlanSibling("text_encoder/model-00001-of-00001.safetensors", 11_000_000_000),
|
||||
_PlanSibling("vae/diffusion_pytorch_model.safetensors", 500_000_000),
|
||||
],
|
||||
"unsloth/Wan2.2-T2V-A14B-NVFP4": [
|
||||
_PlanSibling("Wan2.2-T2V-A14B-NVFP4.pt", 8_000_000_000),
|
||||
_PlanSibling("Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", 8_100_000_000),
|
||||
],
|
||||
},
|
||||
)
|
||||
sources = {
|
||||
"transformer": types.SimpleNamespace(
|
||||
kind = "repo",
|
||||
location = "unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
filename = "Wan2.2-T2V-A14B-NVFP4.pt",
|
||||
fallback_filename = None,
|
||||
),
|
||||
"transformer_2": types.SimpleNamespace(
|
||||
kind = "repo",
|
||||
location = "unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
filename = "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt",
|
||||
fallback_filename = None,
|
||||
),
|
||||
}
|
||||
monkeypatch.setattr(dq, "denoiser_prequant_sources", lambda fam, scheme, base: sources)
|
||||
monkeypatch.setattr(
|
||||
video_mod, "select_transformer_quant_scheme", lambda target, requested, **kw: requested
|
||||
)
|
||||
monkeypatch.setattr(video_mod, "_video_seed_stays_resident", lambda fam, **kw: True)
|
||||
|
||||
plan = VideoBackend().download_plan(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
speed_mode = "off",
|
||||
)
|
||||
|
||||
staged = {f for e in plan["entries"] for f in e["files"]}
|
||||
assert "transformer/diffusion_pytorch_model.safetensors" not in staged
|
||||
assert "transformer_2/diffusion_pytorch_model.safetensors" not in staged
|
||||
assert {"Wan2.2-T2V-A14B-NVFP4.pt", "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt"} <= staged
|
||||
|
||||
|
||||
def _cuda_plan_target(monkeypatch, video_mod, *, free_gib):
|
||||
"""Point the planning path at a cuda card of ``free_gib``, off the test host's own hardware."""
|
||||
import torch
|
||||
|
||||
from core.inference.diffusion_device import DiffusionDeviceTarget
|
||||
from core.inference.diffusion_memory import DeviceMemory
|
||||
|
||||
target = DiffusionDeviceTarget(
|
||||
device = "cuda",
|
||||
dtype = torch.bfloat16,
|
||||
backend = "cuda",
|
||||
vendor = "nvidia",
|
||||
supports_model_cpu_offload = True,
|
||||
supports_default_torch_compile = True,
|
||||
supports_pinned_transfer = True,
|
||||
)
|
||||
monkeypatch.setattr(video_mod, "resolve_diffusion_device_target", lambda **kw: target)
|
||||
monkeypatch.setattr(
|
||||
video_mod,
|
||||
"settled_snapshot_device_memory",
|
||||
lambda t, *a, **k: DeviceMemory(
|
||||
backend = "cuda",
|
||||
device = "cuda",
|
||||
memory_kind = "discrete_vram",
|
||||
free_mib = int(free_gib * 1024),
|
||||
total_mib = int(free_gib * 1024),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _a14b_plan(monkeypatch):
|
||||
"""The A14B repo pair the seeded-plan tests resolve against."""
|
||||
import core.inference.video_denoiser_prequant as dq
|
||||
|
||||
_plan_api(
|
||||
monkeypatch,
|
||||
{
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers": [
|
||||
_PlanSibling("model_index.json", 1000),
|
||||
_PlanSibling("transformer/config.json", 1000),
|
||||
_PlanSibling("transformer/diffusion_pytorch_model.safetensors", 28_000_000_000),
|
||||
_PlanSibling("transformer_2/config.json", 1000),
|
||||
_PlanSibling("transformer_2/diffusion_pytorch_model.safetensors", 28_000_000_000),
|
||||
_PlanSibling("text_encoder/model-00001-of-00001.safetensors", 11_000_000_000),
|
||||
_PlanSibling("vae/diffusion_pytorch_model.safetensors", 500_000_000),
|
||||
],
|
||||
"unsloth/Wan2.2-T2V-A14B-NVFP4": [
|
||||
_PlanSibling("Wan2.2-T2V-A14B-NVFP4.pt", 8_000_000_000),
|
||||
_PlanSibling("Wan2.2-T2V-A14B-transformer_2-NVFP4.pt", 8_100_000_000),
|
||||
],
|
||||
},
|
||||
)
|
||||
sources = {
|
||||
"transformer": types.SimpleNamespace(
|
||||
kind = "repo",
|
||||
location = "unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
filename = "Wan2.2-T2V-A14B-NVFP4.pt",
|
||||
fallback_filename = None,
|
||||
),
|
||||
"transformer_2": types.SimpleNamespace(
|
||||
kind = "repo",
|
||||
location = "unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
filename = "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt",
|
||||
fallback_filename = None,
|
||||
),
|
||||
}
|
||||
monkeypatch.setattr(dq, "denoiser_prequant_sources", lambda fam, scheme, base: sources)
|
||||
|
||||
|
||||
def test_a_plan_that_still_offloads_at_artifact_size_stages_the_dense_experts(monkeypatch):
|
||||
"""A card the artifact-sized plan still offloads on cannot seed: stage the dense experts."""
|
||||
import core.inference.video as video_mod
|
||||
|
||||
_a14b_plan(monkeypatch)
|
||||
monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4")
|
||||
_cuda_plan_target(monkeypatch, video_mod, free_gib = 24)
|
||||
|
||||
plan = VideoBackend().download_plan(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
|
||||
staged = {f for e in plan["entries"] for f in e["files"]}
|
||||
assert "transformer/diffusion_pytorch_model.safetensors" in staged
|
||||
assert "transformer_2/diffusion_pytorch_model.safetensors" in staged
|
||||
assert not any(f.endswith(".pt") for f in staged)
|
||||
|
||||
|
||||
def test_a_card_the_artifact_fits_on_still_stages_the_artifacts(monkeypatch):
|
||||
"""Where the artifact-sized plan stays resident the load seeds, so dense shards stay out."""
|
||||
import core.inference.video as video_mod
|
||||
|
||||
_a14b_plan(monkeypatch)
|
||||
monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4")
|
||||
_cuda_plan_target(monkeypatch, video_mod, free_gib = 180)
|
||||
|
||||
plan = VideoBackend().download_plan(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
|
||||
staged = {f for e in plan["entries"] for f in e["files"]}
|
||||
assert "Wan2.2-T2V-A14B-NVFP4.pt" in staged
|
||||
assert "Wan2.2-T2V-A14B-transformer_2-NVFP4.pt" in staged
|
||||
assert not any(
|
||||
f.endswith("diffusion_pytorch_model.safetensors") and f.startswith("transformer")
|
||||
for f in staged
|
||||
)
|
||||
|
||||
|
||||
def test_an_offloading_memory_mode_stages_the_dense_experts_on_any_card(monkeypatch):
|
||||
"""An explicit offload memory_mode stages the dense experts even on a roomy card."""
|
||||
import core.inference.video as video_mod
|
||||
|
||||
_a14b_plan(monkeypatch)
|
||||
monkeypatch.setattr(video_mod, "_video_auto_denoiser_scheme", lambda fam, **kw: "nvfp4")
|
||||
_cuda_plan_target(monkeypatch, video_mod, free_gib = 180)
|
||||
|
||||
plan = VideoBackend().download_plan(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
memory_mode = "low_vram",
|
||||
)
|
||||
|
||||
staged = {f for e in plan["entries"] for f in e["files"]}
|
||||
assert "transformer/diffusion_pytorch_model.safetensors" in staged
|
||||
assert not any(f.endswith(".pt") for f in staged)
|
||||
|
||||
|
||||
def test_a_dense_encoder_fallback_that_forces_offload_also_drops_the_seed(
|
||||
fake_runtime, monkeypatch
|
||||
):
|
||||
"""A failed pre-cast encoder re-plans at bf16, which can offload: re-decide the seed there."""
|
||||
import core.inference.diffusion_te_prequant as te
|
||||
import core.inference.video as video_mod
|
||||
|
||||
monkeypatch.setenv("UNSLOTH_DIFFUSION_ALLOW_PRECISION_FALLBACK", "1")
|
||||
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
|
||||
monkeypatch.setattr(video_mod, "quantize_transformer", lambda *a, **k: None)
|
||||
calls, _modules = _stub_denoiser_seed(monkeypatch)
|
||||
monkeypatch.setattr(te, "te_prequant_budget_scale", lambda fam, **kwargs: 0.5)
|
||||
monkeypatch.setattr(te, "te_prequant_pipe_kwargs", lambda fam, base, **kwargs: {})
|
||||
real_plan = video_mod.plan_diffusion_memory
|
||||
seen: list = []
|
||||
|
||||
def _plan(**kwargs):
|
||||
seen.append(kwargs.get("model_dense_mib"))
|
||||
planned = real_plan(**kwargs)
|
||||
return planned if len(seen) == 1 else dataclasses.replace(planned, offload_policy = "model")
|
||||
|
||||
monkeypatch.setattr(video_mod, "plan_diffusion_memory", _plan)
|
||||
_stub_apply_memory_plan(monkeypatch, video_mod)
|
||||
|
||||
backend = VideoBackend()
|
||||
status = backend.load_pipeline(
|
||||
"Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
model_kind = "pipeline",
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
|
||||
assert calls == [], "no checkpoint may be seeded into a load that will offload the DiT"
|
||||
assert status["transformer_quant"] is None
|
||||
|
||||
|
||||
# ── Progress truth: the bar must track the GPU, not the host queue ─────────────
|
||||
# Measured on MiniMax-H3 (960x544x124, 30 steps, int8 DiT + int8 TE, speed=default with a captured
|
||||
# CUDA graph, B200): the 29 host-side scheduler.step calls land in 0.9 s and the GPU is still
|
||||
|
|
|
|||
|
|
@ -412,3 +412,106 @@ def test_minimax_h3_offers_every_advertised_aspect_ratio():
|
|||
assert width * height <= H3_CANVAS_MAX_PIXELS, (width, height)
|
||||
rule_width, rule_height = h3_canvas_for_aspect(width, height)
|
||||
assert width <= rule_width and height <= rule_height, (width, height)
|
||||
|
||||
|
||||
def _resident_fam(**kwargs):
|
||||
from core.inference.video_families import VideoFamily
|
||||
|
||||
base = dict(
|
||||
name = "test-video",
|
||||
pipeline_class = "TestPipeline",
|
||||
transformer_class = "TestTransformer3DModel",
|
||||
base_repo = "org/test-video",
|
||||
)
|
||||
base.update(kwargs)
|
||||
return VideoFamily(**base)
|
||||
|
||||
|
||||
def test_the_per_scheme_row_wins_and_the_float_is_the_fallback():
|
||||
from core.inference.video_families import video_family_prequant_resident_gb
|
||||
|
||||
fam = _resident_fam(
|
||||
prequant_resident_gb = 20.3,
|
||||
prequant_resident_gb_by_scheme = (("nvfp4", 8.1), ("fp8", 13.6)),
|
||||
)
|
||||
assert video_family_prequant_resident_gb(fam, "nvfp4") == pytest.approx(8.1)
|
||||
assert video_family_prequant_resident_gb(fam, "fp8") == pytest.approx(13.6)
|
||||
assert video_family_prequant_resident_gb(fam, "int8") == pytest.approx(20.3)
|
||||
|
||||
|
||||
def test_an_unmeasured_family_reports_nothing_rather_than_guessing():
|
||||
import types
|
||||
|
||||
from core.inference.video_families import video_family_prequant_resident_gb
|
||||
|
||||
assert video_family_prequant_resident_gb(_resident_fam(), "nvfp4") is None
|
||||
assert video_family_prequant_resident_gb(types.SimpleNamespace(), "nvfp4") is None
|
||||
fam = _resident_fam(prequant_resident_gb_by_scheme = (("nvfp4",), ("nvfp4", "big")))
|
||||
assert video_family_prequant_resident_gb(fam, "nvfp4") is None
|
||||
|
||||
|
||||
def test_the_h3_measurement_still_answers_through_the_helper():
|
||||
from core.inference.video_families import video_family_prequant_resident_gb
|
||||
|
||||
fam = detect_video_family("MiniMaxAI/MiniMax-H3")
|
||||
assert video_family_prequant_resident_gb(fam, "int8") == pytest.approx(fam.prequant_resident_gb)
|
||||
assert video_family_prequant_resident_gb(fam, "fp8") == pytest.approx(fam.prequant_resident_gb)
|
||||
|
||||
|
||||
def test_every_hosted_nvfp4_denoiser_carries_its_measured_resident_size():
|
||||
from core.inference.video_families import (
|
||||
_FAMILIES,
|
||||
video_family_prequant_resident_gb,
|
||||
)
|
||||
|
||||
expected = {
|
||||
"wan2.2-ti2v-5b": 2.9,
|
||||
"wan2.2-t2v-a14b": 16.2,
|
||||
"hunyuanvideo-1.5": 4.8,
|
||||
"hunyuanvideo-1.5-720p": 4.8,
|
||||
}
|
||||
hosting = {
|
||||
fam.name
|
||||
for fam in _FAMILIES
|
||||
if any(scheme == "nvfp4" for scheme, _repo in fam.prequant_repos)
|
||||
}
|
||||
assert hosting == set(expected)
|
||||
for fam in _FAMILIES:
|
||||
if fam.name not in expected:
|
||||
continue
|
||||
assert video_family_prequant_resident_gb(fam, "nvfp4") == pytest.approx(expected[fam.name])
|
||||
|
||||
|
||||
def test_the_measured_nvfp4_size_is_a_4_bit_fraction_of_the_term_it_replaces():
|
||||
from core.inference.video_families import (
|
||||
_FAMILIES,
|
||||
video_family_prequant_resident_gb,
|
||||
)
|
||||
|
||||
seen = 0
|
||||
for fam in _FAMILIES:
|
||||
if not any(scheme == "nvfp4" for scheme, _repo in fam.prequant_repos):
|
||||
continue
|
||||
measured = video_family_prequant_resident_gb(fam, "nvfp4")
|
||||
if measured is None or not fam.bf16_components_gb:
|
||||
continue
|
||||
seen += 1
|
||||
assert (
|
||||
0.2 * fam.bf16_components_gb[0] < measured < 0.4 * fam.bf16_components_gb[0]
|
||||
), fam.name
|
||||
assert seen == 4
|
||||
|
||||
|
||||
def test_the_a14b_row_prices_both_experts_not_one():
|
||||
from core.inference.video_families import (
|
||||
detect_video_family,
|
||||
video_family_prequant_resident_gb,
|
||||
)
|
||||
|
||||
fam = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers")
|
||||
assert fam.is_moe
|
||||
names = {entry[-1] for entry in fam.prequant_filenames}
|
||||
assert len(names) == 2
|
||||
measured = video_family_prequant_resident_gb(fam, "nvfp4")
|
||||
assert measured == pytest.approx(16.2)
|
||||
assert measured > 0.25 * fam.bf16_components_gb[0]
|
||||
|
|
|
|||
|
|
@ -11,6 +11,8 @@ still run (and still get tested) in an environment where diffusers cannot be imp
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from core.inference.diffusion_prequant import (
|
||||
|
|
@ -1118,7 +1120,7 @@ def test_the_fallback_is_resolved_before_the_download_is_planned(monkeypatch):
|
|||
verified = src.index("_denoiser_prequant_verified")
|
||||
predownload = src.index("_predownload_base")
|
||||
assert planned < verified < predownload
|
||||
assert 'h3_auto_denoiser or kwargs.get("transformer_quant")' in src
|
||||
assert "h3_auto_denoiser or video_auto_denoiser or requested_denoiser" in src
|
||||
|
||||
|
||||
def _h3_placement_probe(
|
||||
|
|
@ -1176,6 +1178,496 @@ def test_an_unreadable_card_keeps_the_rotation(monkeypatch):
|
|||
assert vid._h3_dense_denoiser_fits(vid._h3_dense_denoiser_resident_bytes(None), None) is False
|
||||
|
||||
|
||||
def _moe_fam(**kwargs):
|
||||
"""A dual-expert MoE family (Wan2.2-A14B's shape) hosting both experts in one repo."""
|
||||
return _fam(
|
||||
is_moe = True,
|
||||
prequant_repos = (("nvfp4", "org/test-quantized"),),
|
||||
prequant_filenames = (("nvfp4", "transformer_2", "test-transformer_2-NVFP4.pt"),),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def test_the_components_a_family_seeds_are_its_denoisers():
|
||||
from core.inference.video_denoiser_prequant import denoiser_components
|
||||
assert denoiser_components(_fam()) == ("transformer",)
|
||||
assert denoiser_components(_fam(is_moe = True)) == ("transformer", "transformer_2")
|
||||
|
||||
|
||||
def test_the_second_expert_is_addressed_through_the_task_slot():
|
||||
"""The second expert is addressed through the task slot, with no filename fallback."""
|
||||
from core.inference.video_denoiser_prequant import denoiser_prequant_sources
|
||||
|
||||
sources = denoiser_prequant_sources(_moe_fam(), "nvfp4", "org/test-video")
|
||||
assert set(sources) == {"transformer", "transformer_2"}
|
||||
from core.inference.diffusion_prequant import candidate_filenames_of
|
||||
|
||||
assert candidate_filenames_of(sources["transformer"]) == (
|
||||
"test-NVFP4.safetensors",
|
||||
"test-NVFP4.pt",
|
||||
"transformer_nvfp4.pt",
|
||||
)
|
||||
assert candidate_filenames_of(sources["transformer_2"]) == ("test-transformer_2-NVFP4.pt",)
|
||||
|
||||
|
||||
def test_every_component_or_none():
|
||||
"""Partial coverage is not coverage: every component or none."""
|
||||
from core.inference.video_denoiser_prequant import denoiser_prequant_sources
|
||||
|
||||
single = _fam(prequant_repos = (("nvfp4", "org/test-NVFP4"),))
|
||||
assert set(denoiser_prequant_sources(single, "nvfp4", "org/test-video")) == {"transformer"}
|
||||
assert (
|
||||
denoiser_prequant_sources(
|
||||
_fam(is_moe = True, prequant_repos = (("nvfp4", "org/x"),)), "nvfp4", "org/test-video"
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert denoiser_prequant_sources(_moe_fam(), "int8", "org/test-video") is None
|
||||
for scheme in (None, "", "auto", "off", "none"):
|
||||
assert denoiser_prequant_sources(_moe_fam(), scheme, "org/test-video") is None
|
||||
|
||||
|
||||
def _stub_prequant_loader(monkeypatch, outcomes):
|
||||
"""Record every ``load_prequantized_transformer`` call and return ``outcomes`` in order."""
|
||||
import gc
|
||||
import sys
|
||||
import types
|
||||
|
||||
import core.inference.diffusion_prequant as pq
|
||||
|
||||
diffusers = types.ModuleType("diffusers")
|
||||
diffusers.TestTransformer3DModel = object
|
||||
monkeypatch.setitem(sys.modules, "diffusers", diffusers)
|
||||
|
||||
events: list = []
|
||||
calls: list = []
|
||||
|
||||
def _fake_load(transformer_cls, base, source, **kwargs):
|
||||
calls.append({"base": base, "source": source, **kwargs})
|
||||
events.append(("load", kwargs.get("component")))
|
||||
return outcomes[len(calls) - 1]
|
||||
|
||||
monkeypatch.setattr(pq, "load_prequantized_transformer", _fake_load)
|
||||
monkeypatch.setattr(gc, "collect", lambda *a, **k: events.append(("collect", None)))
|
||||
return calls, events
|
||||
|
||||
|
||||
def test_seeding_loads_every_expert_into_its_own_component(monkeypatch):
|
||||
"""Seeding loads every expert into its own component, from its own config."""
|
||||
from core.inference.video_denoiser_prequant import denoiser_prequant_pipe_kwargs
|
||||
|
||||
first, second = object(), object()
|
||||
calls, events = _stub_prequant_loader(monkeypatch, [first, second])
|
||||
|
||||
seeded = denoiser_prequant_pipe_kwargs(
|
||||
_moe_fam(),
|
||||
"org/test-video",
|
||||
scheme = "nvfp4",
|
||||
dtype = "bfloat16",
|
||||
device = "cuda:0",
|
||||
hf_token = "tok",
|
||||
cache_dir = "/cache",
|
||||
)
|
||||
assert seeded == {"transformer": first, "transformer_2": second}
|
||||
assert [c["config_subfolder"] for c in calls] == ["transformer", "transformer_2"]
|
||||
assert [c["component"] for c in calls] == ["transformer", "transformer_2"]
|
||||
assert [c["source"].filename for c in calls] == [
|
||||
"test-NVFP4.safetensors",
|
||||
"test-transformer_2-NVFP4.pt",
|
||||
]
|
||||
from core.inference.diffusion_transformer_quant import DEFAULT_MIN_LINEAR_FEATURES
|
||||
|
||||
assert {c["min_features"] for c in calls} == {DEFAULT_MIN_LINEAR_FEATURES}
|
||||
assert events == [
|
||||
("load", "transformer"),
|
||||
("collect", None),
|
||||
("load", "transformer_2"),
|
||||
("collect", None),
|
||||
]
|
||||
|
||||
|
||||
def test_a_second_expert_that_will_not_load_seeds_nothing(monkeypatch):
|
||||
"""A second expert that will not load seeds nothing and releases what was loaded."""
|
||||
from core.inference.video_denoiser_prequant import denoiser_prequant_pipe_kwargs
|
||||
|
||||
calls, events = _stub_prequant_loader(monkeypatch, [object(), None])
|
||||
|
||||
assert (
|
||||
denoiser_prequant_pipe_kwargs(
|
||||
_moe_fam(),
|
||||
"org/test-video",
|
||||
scheme = "nvfp4",
|
||||
dtype = "bfloat16",
|
||||
device = "cuda:0",
|
||||
)
|
||||
== {}
|
||||
)
|
||||
assert len(calls) == 2
|
||||
assert events[-1] == ("collect", None)
|
||||
|
||||
|
||||
def test_seeding_declines_a_host_that_cannot_run_the_quant(monkeypatch):
|
||||
import types
|
||||
|
||||
import core.inference.diffusion_transformer_quant as tq
|
||||
from core.inference.video_denoiser_prequant import denoiser_prequant_pipe_kwargs
|
||||
|
||||
calls, _events = _stub_prequant_loader(monkeypatch, [object(), object()])
|
||||
monkeypatch.setattr(tq, "dense_transformer_supported", lambda target: False)
|
||||
assert (
|
||||
denoiser_prequant_pipe_kwargs(
|
||||
_moe_fam(),
|
||||
"org/test-video",
|
||||
scheme = "nvfp4",
|
||||
dtype = "bfloat16",
|
||||
device = "cuda:0",
|
||||
target = types.SimpleNamespace(device = "cpu"),
|
||||
)
|
||||
== {}
|
||||
)
|
||||
assert calls == []
|
||||
|
||||
|
||||
def _video_auto(
|
||||
monkeypatch,
|
||||
*,
|
||||
scheme,
|
||||
fam = None,
|
||||
**over,
|
||||
):
|
||||
"""``_video_auto_denoiser_scheme`` with the device selector stubbed to ``scheme``."""
|
||||
from core.inference import video as vid
|
||||
|
||||
monkeypatch.setattr(
|
||||
vid, "select_transformer_quant_scheme", lambda target, requested, family = None: scheme
|
||||
)
|
||||
kw = dict(
|
||||
target = None,
|
||||
requested = "auto",
|
||||
base_repo = "org/test-video",
|
||||
speed_mode = None,
|
||||
)
|
||||
kw.update(over)
|
||||
return vid._video_auto_denoiser_scheme(fam if fam is not None else _moe_fam(), **kw)
|
||||
|
||||
|
||||
def test_the_conventional_auto_scheme_needs_an_artifact_for_every_expert(monkeypatch):
|
||||
assert _video_auto(monkeypatch, scheme = "nvfp4") == "nvfp4"
|
||||
assert _video_auto(monkeypatch, scheme = "int8") is None
|
||||
assert (
|
||||
_video_auto(
|
||||
monkeypatch, scheme = "nvfp4", fam = _fam(is_moe = True, prequant_repos = (("nvfp4", "org/x"),))
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert _video_auto(monkeypatch, scheme = None) is None
|
||||
|
||||
|
||||
def test_speed_off_and_the_modular_workflow_are_never_seeded_here(monkeypatch):
|
||||
assert _video_auto(monkeypatch, scheme = "nvfp4", speed_mode = "off") is None
|
||||
assert _video_auto(monkeypatch, scheme = "nvfp4", speed_mode = " OFF ") is None
|
||||
assert _video_auto(monkeypatch, scheme = "nvfp4", speed_mode = "off", requested = None) is None
|
||||
assert _video_auto(monkeypatch, scheme = "nvfp4", fam = _moe_fam(modular_workflow = "fl2va")) is None
|
||||
|
||||
|
||||
def test_an_explicit_scheme_keeps_its_seed_under_speed_off(monkeypatch):
|
||||
"""An explicit scheme under Speed="off" is honored, so the plan keeps the hosted checkpoint."""
|
||||
assert _video_auto(monkeypatch, scheme = "nvfp4", speed_mode = "off", requested = "nvfp4") == "nvfp4"
|
||||
from core.inference.diffusion_transformer_quant import select_transformer_quant_scheme
|
||||
|
||||
assert select_transformer_quant_scheme(object(), "off") is None
|
||||
|
||||
|
||||
def test_an_install_that_cannot_open_a_checkpoint_keeps_the_dense_denoiser(monkeypatch):
|
||||
import core.inference.diffusion_prequant as pq
|
||||
monkeypatch.setattr(pq, "restricted_prequant_load_supported", lambda scheme = None: False)
|
||||
assert _video_auto(monkeypatch, scheme = "nvfp4") is None
|
||||
|
||||
|
||||
def test_the_conventional_coverage_probe_reads_every_component():
|
||||
from core.inference.video import VideoBackend
|
||||
|
||||
fam = _moe_fam()
|
||||
assert VideoBackend._denoiser_prequant_covered(fam, "nvfp4", "org/test-video")
|
||||
assert not VideoBackend._denoiser_prequant_covered(fam, "int8", "org/test-video")
|
||||
assert not VideoBackend._denoiser_prequant_covered(fam, "auto", "org/test-video")
|
||||
assert not VideoBackend._denoiser_prequant_covered(fam, None, "org/test-video")
|
||||
assert not VideoBackend._denoiser_prequant_covered(
|
||||
fam, "nvfp4", "org/test-video", None, kind = "gguf"
|
||||
)
|
||||
half = _fam(is_moe = True, prequant_repos = (("nvfp4", "org/x"),))
|
||||
assert not VideoBackend._denoiser_prequant_covered(half, "nvfp4", "org/test-video")
|
||||
|
||||
|
||||
def _planned_denoiser_request(monkeypatch, fam, **load_kwargs):
|
||||
"""The scheme the planner hands ``_denoiser_prequant_verified`` (stubbed, as is the build)."""
|
||||
from core.inference import video as vid
|
||||
|
||||
backend = vid.VideoBackend()
|
||||
backend._load_token = 1
|
||||
backend._loading = vid._VideoLoadingState(repo_id = fam.base_repo, base_repo = fam.base_repo)
|
||||
monkeypatch.setattr(vid, "_detect_load_family", lambda *a, **k: fam)
|
||||
monkeypatch.setattr(vid, "_assert_pick_is_not_speech", lambda *a, **k: None)
|
||||
monkeypatch.setattr(backend, "_h3_planned_auto_denoiser_scheme", lambda *a, **k: None)
|
||||
monkeypatch.setattr(backend, "_video_planned_auto_denoiser_scheme", lambda *a, **k: None)
|
||||
monkeypatch.setattr(backend, "load_pipeline", lambda **kwargs: None)
|
||||
monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: None)
|
||||
seen: list = []
|
||||
|
||||
def _verified(_fam, transformer_quant, *a, **k):
|
||||
seen.append(transformer_quant)
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(backend, "_denoiser_prequant_verified", _verified)
|
||||
backend._run_load(
|
||||
repo_id = fam.base_repo,
|
||||
local_files_only = True,
|
||||
_load_token = 1,
|
||||
**load_kwargs,
|
||||
)
|
||||
assert seen, "the plan never reached the denoiser probe"
|
||||
return seen[0]
|
||||
|
||||
|
||||
def test_a_conventional_plan_drops_no_shard_the_load_will_not_seed(monkeypatch):
|
||||
"""Only the planned seed drops shards: a conventional load never forwards the raw request."""
|
||||
from core.inference.video_families import detect_video_family
|
||||
|
||||
wan = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
|
||||
assert wan is not None and not wan.modular_workflow
|
||||
assert (
|
||||
_planned_denoiser_request(monkeypatch, wan, transformer_quant = "nvfp4", speed_mode = "off")
|
||||
is None
|
||||
)
|
||||
h3 = detect_video_family("MiniMaxAI/MiniMax-H3")
|
||||
assert h3 is not None and h3.modular_workflow
|
||||
assert (
|
||||
_planned_denoiser_request(
|
||||
monkeypatch, h3, transformer_quant = "int8", speed_mode = "off", h3_task = "fl2va"
|
||||
)
|
||||
== "int8"
|
||||
)
|
||||
|
||||
|
||||
def test_a_seed_the_plan_declined_is_pinned_into_the_load(monkeypatch):
|
||||
"""The plan's decline travels as its own value: the load cannot re-decide on a roomier card."""
|
||||
from core.inference import video as vid
|
||||
from core.inference.video_families import detect_video_family
|
||||
|
||||
wan = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
|
||||
assert wan is not None and not wan.modular_workflow
|
||||
backend = vid.VideoBackend()
|
||||
backend._load_token = 1
|
||||
backend._loading = vid._VideoLoadingState(repo_id = wan.base_repo, base_repo = wan.base_repo)
|
||||
monkeypatch.setattr(vid, "_detect_load_family", lambda *a, **k: wan)
|
||||
monkeypatch.setattr(vid, "_assert_pick_is_not_speech", lambda *a, **k: None)
|
||||
monkeypatch.setattr(backend, "_h3_planned_auto_denoiser_scheme", lambda *a, **k: None)
|
||||
monkeypatch.setattr(
|
||||
backend,
|
||||
"_video_planned_auto_denoiser_scheme",
|
||||
lambda *a, **k: vid.DENOISER_SEED_DECLINED,
|
||||
)
|
||||
monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: None)
|
||||
probed: list = []
|
||||
|
||||
def _verified(_fam, transformer_quant, *a, **k):
|
||||
probed.append(transformer_quant)
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(backend, "_denoiser_prequant_verified", _verified)
|
||||
loaded: list = []
|
||||
monkeypatch.setattr(backend, "load_pipeline", lambda **kwargs: loaded.append(kwargs))
|
||||
backend._run_load(
|
||||
repo_id = wan.base_repo,
|
||||
local_files_only = True,
|
||||
_load_token = 1,
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
assert probed == [None], "a declined seed drops no dense shard from the pull"
|
||||
assert loaded and loaded[0]["_video_auto_denoiser_planned"] == vid.DENOISER_SEED_DECLINED
|
||||
|
||||
|
||||
def test_the_seeded_denoiser_repo_is_claimed_against_a_concurrent_delete(monkeypatch):
|
||||
"""The seeded denoiser's repo (neither repo_id nor base_repo) joins the in-flight claim."""
|
||||
from core.inference import video as vid
|
||||
from core.inference.video_families import detect_video_family
|
||||
|
||||
wan = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
|
||||
assert wan is not None and not wan.modular_workflow
|
||||
backend = vid.VideoBackend()
|
||||
backend._load_token = 1
|
||||
backend._loading = vid._VideoLoadingState(repo_id = wan.base_repo, base_repo = wan.base_repo)
|
||||
monkeypatch.setattr(vid, "_detect_load_family", lambda *a, **k: wan)
|
||||
monkeypatch.setattr(vid, "_assert_pick_is_not_speech", lambda *a, **k: None)
|
||||
monkeypatch.setattr(backend, "_h3_planned_auto_denoiser_scheme", lambda *a, **k: None)
|
||||
monkeypatch.setattr(backend, "_video_planned_auto_denoiser_scheme", lambda *a, **k: "nvfp4")
|
||||
monkeypatch.setattr(backend, "_denoiser_prequant_verified", lambda *a, **k: True)
|
||||
monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: None)
|
||||
claimed: list = []
|
||||
monkeypatch.setattr(
|
||||
backend, "load_pipeline", lambda **kwargs: claimed.extend(backend.loading_repo_ids())
|
||||
)
|
||||
backend._run_load(
|
||||
repo_id = wan.base_repo,
|
||||
local_files_only = True,
|
||||
_load_token = 1,
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
assert "unsloth/Wan2.2-TI2V-5B-NVFP4" in claimed
|
||||
|
||||
|
||||
def test_both_experts_of_an_moe_resolve_to_one_claimed_repo():
|
||||
"""The two A14B experts are two files in ONE repo, so the claim names it once."""
|
||||
from core.inference.video_families import detect_video_family
|
||||
|
||||
a14b = detect_video_family("Wan-AI/Wan2.2-T2V-A14B-Diffusers")
|
||||
assert a14b is not None and a14b.is_moe
|
||||
assert VideoBackend._denoiser_prequant_repo_ids(a14b, "nvfp4", a14b.base_repo) == (
|
||||
"unsloth/Wan2.2-T2V-A14B-NVFP4",
|
||||
)
|
||||
assert VideoBackend._denoiser_prequant_repo_ids(a14b, "int8", a14b.base_repo) == ()
|
||||
assert VideoBackend._denoiser_prequant_repo_ids(a14b, "auto", a14b.base_repo) == ()
|
||||
|
||||
|
||||
def test_the_seeded_denoiser_artifact_is_fetched_under_the_load_cancel_event(monkeypatch):
|
||||
"""The seeded artifact is prefetched under the load's cancel event, not by the injection."""
|
||||
import threading
|
||||
|
||||
from core.inference import video as vid
|
||||
from core.inference.video_families import detect_video_family
|
||||
|
||||
wan = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
|
||||
assert wan is not None and not wan.modular_workflow
|
||||
backend = vid.VideoBackend()
|
||||
backend._load_token = 1
|
||||
backend._loading = vid._VideoLoadingState(repo_id = wan.base_repo, base_repo = wan.base_repo)
|
||||
monkeypatch.setattr(vid, "_detect_load_family", lambda *a, **k: wan)
|
||||
monkeypatch.setattr(vid, "_assert_pick_is_not_speech", lambda *a, **k: None)
|
||||
monkeypatch.setattr(backend, "_h3_planned_auto_denoiser_scheme", lambda *a, **k: None)
|
||||
monkeypatch.setattr(backend, "_video_planned_auto_denoiser_scheme", lambda *a, **k: "nvfp4")
|
||||
monkeypatch.setattr(backend, "_denoiser_prequant_verified", lambda *a, **k: True)
|
||||
monkeypatch.setattr(backend, "_run_load_h3_native", lambda **kwargs: None)
|
||||
monkeypatch.setattr(backend, "load_pipeline", lambda **kwargs: None)
|
||||
fetched: list = []
|
||||
import utils.hf_xet_fallback as xet
|
||||
|
||||
monkeypatch.setattr(
|
||||
xet,
|
||||
"hf_hub_download_with_xet_fallback",
|
||||
lambda repo, filename, token = None, **kwargs: fetched.append(
|
||||
(repo, filename, kwargs.get("cancel_event"))
|
||||
),
|
||||
)
|
||||
cancel = threading.Event()
|
||||
|
||||
backend._run_load(
|
||||
repo_id = wan.base_repo,
|
||||
local_files_only = True,
|
||||
_load_token = 1,
|
||||
_cancel_event = cancel,
|
||||
transformer_quant = "nvfp4",
|
||||
)
|
||||
|
||||
assert [(r, f) for r, f, _ in fetched] == [
|
||||
("unsloth/Wan2.2-TI2V-5B-NVFP4", "Wan2.2-TI2V-5B-NVFP4.pt")
|
||||
]
|
||||
assert fetched[0][2] is cancel
|
||||
|
||||
|
||||
def test_the_seed_plan_credits_the_resident_pipelines_memory(monkeypatch):
|
||||
"""The plan credits the resident pipeline's bytes, not other tenants', capped by the total."""
|
||||
import core.inference.video as vid
|
||||
from core.inference.diffusion_memory import DeviceMemory
|
||||
from core.inference.video_families import detect_video_family
|
||||
|
||||
fam = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
|
||||
seen: list = []
|
||||
monkeypatch.setattr(
|
||||
vid,
|
||||
"settled_snapshot_device_memory",
|
||||
lambda _t: DeviceMemory("cuda", "cuda", "discrete_vram", 10_000, 32_000),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
vid,
|
||||
"plan_diffusion_memory",
|
||||
lambda **kw: seen.append(kw["device_memory"].free_mib)
|
||||
or types.SimpleNamespace(offload_policy = "none"),
|
||||
)
|
||||
target = types.SimpleNamespace(device = "cuda", dtype = None)
|
||||
for credit in (0, 15_000, 40_000):
|
||||
vid._video_seed_stays_resident(
|
||||
fam,
|
||||
target = target,
|
||||
scheme = "fp8",
|
||||
memory_mode = None,
|
||||
text_encoder_quant = None,
|
||||
base_repo = None,
|
||||
reclaimable_mib = credit,
|
||||
)
|
||||
assert seen == [10_000, 25_000, 32_000]
|
||||
|
||||
|
||||
def test_the_planner_hands_the_resident_pipelines_bytes_to_the_seed_plan(monkeypatch):
|
||||
import core.inference.video as vid
|
||||
from core.inference.video_families import detect_video_family
|
||||
|
||||
fam = detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers")
|
||||
backend = vid.VideoBackend()
|
||||
resident_pipe = object()
|
||||
backend._state = types.SimpleNamespace(pipe = resident_pipe)
|
||||
monkeypatch.setattr(vid, "_video_auto_denoiser_scheme", lambda fam, **kw: "fp8")
|
||||
monkeypatch.setattr(
|
||||
vid,
|
||||
"_pipeline_device_mib",
|
||||
lambda pipe, ordinal = None: (
|
||||
captured.setdefault("ordinal", ordinal),
|
||||
12_345 if pipe is resident_pipe else 0,
|
||||
)[1],
|
||||
)
|
||||
captured: dict = {}
|
||||
monkeypatch.setattr(
|
||||
vid, "_video_seed_stays_resident", lambda fam, **kw: captured.update(kw) or True
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
vid,
|
||||
"resolve_diffusion_device_target",
|
||||
lambda **_k: types.SimpleNamespace(device = "cuda", dtype = None, ordinal = None),
|
||||
)
|
||||
assert (
|
||||
backend._video_planned_auto_denoiser_scheme(
|
||||
fam, base = None, kind = "pipeline", transformer_quant = "auto", speed_mode = None
|
||||
)
|
||||
== "fp8"
|
||||
)
|
||||
assert captured["reclaimable_mib"] == 12_345
|
||||
|
||||
|
||||
def test_only_the_resident_bytes_on_the_target_card_are_credited():
|
||||
"""Only bytes on the target card are credited; a storage shared by two tensors counts once."""
|
||||
import core.inference.video as vid
|
||||
|
||||
def tensor(index, ptr, nbytes):
|
||||
storage = types.SimpleNamespace(data_ptr = lambda: ptr, nbytes = lambda: nbytes)
|
||||
return types.SimpleNamespace(
|
||||
device = types.SimpleNamespace(type = "cuda", index = index),
|
||||
untyped_storage = lambda: storage,
|
||||
)
|
||||
|
||||
tensors = [tensor(0, 1, 3 << 20), tensor(0, 1, 3 << 20), tensor(1, 2, 5 << 20)]
|
||||
module = types.SimpleNamespace(parameters = lambda: iter(tensors), buffers = lambda: iter(()))
|
||||
pipe = types.SimpleNamespace(components = {"transformer": module})
|
||||
assert vid._pipeline_device_mib(pipe) == 8
|
||||
assert vid._pipeline_device_mib(pipe, ordinal = 0) == 3
|
||||
assert vid._pipeline_device_mib(pipe, ordinal = 1) == 5
|
||||
assert vid._pipeline_device_mib(pipe, ordinal = 2) == 0
|
||||
|
||||
|
||||
def test_the_credit_is_read_on_the_card_the_load_targets():
|
||||
import core.inference.video as vid
|
||||
assert vid._target_ordinal(types.SimpleNamespace(device = "cuda", ordinal = 3)) == 3
|
||||
assert vid._target_ordinal(types.SimpleNamespace(device = "mps", ordinal = None)) is None
|
||||
|
||||
|
||||
def test_unreadability_is_reported_only_when_it_is_what_kept_bfloat16(monkeypatch):
|
||||
try:
|
||||
import importlib
|
||||
|
|
|
|||
|
|
@ -78,7 +78,13 @@ import {
|
|||
curatedArtifactTakesDenseQuant,
|
||||
loadSpecFor,
|
||||
} from "@/features/model-picker/components/model-selector/model-catalog";
|
||||
import { useDenseQuantSchemes, useHostClass } from "@/hooks/use-host-class";
|
||||
import {
|
||||
useDenseQuantSchemes,
|
||||
useHostClass,
|
||||
useNvfp4Diffusion,
|
||||
useNvfp4DiffusionKnown,
|
||||
} from "@/hooks/use-host-class";
|
||||
import { nvfp4SelectionFallback, withNvfp4Option } from "@/lib/nvfp4-options";
|
||||
import type {
|
||||
ModelOption,
|
||||
ModelSelectorChangeMeta,
|
||||
|
|
@ -1275,6 +1281,8 @@ export function ImagesPage({
|
|||
const { isMobile, pinned } = useSidebar();
|
||||
const hostClass = useHostClass();
|
||||
const denseQuantSchemes = useDenseQuantSchemes();
|
||||
const nvfp4Diffusion = useNvfp4Diffusion();
|
||||
const nvfp4DiffusionKnown = useNvfp4DiffusionKnown();
|
||||
const imageModels = useImageModels(hostClass, denseQuantSchemes);
|
||||
const { rootStyle: railRootStyle } = useMediaRailWidth("images");
|
||||
const [quant, setQuant] = useState<string | null>(galleryCache.quant);
|
||||
|
|
@ -1435,6 +1443,10 @@ export function ImagesPage({
|
|||
const [attentionBackend, setAttentionBackend] = useState<"auto" | "native" | "cudnn" | "flash3" | "sage">(
|
||||
"auto",
|
||||
);
|
||||
useEffect(() => {
|
||||
setTransformerQuant((v) => nvfp4SelectionFallback(v, nvfp4DiffusionKnown, nvfp4Diffusion));
|
||||
setTextEncoderQuant((v) => nvfp4SelectionFallback(v, nvfp4DiffusionKnown, nvfp4Diffusion));
|
||||
}, [nvfp4Diffusion, nvfp4DiffusionKnown, transformerQuant, textEncoderQuant]);
|
||||
const [memoryMode, setMemoryMode] = useState<"auto" | "fast" | "balanced" | "low_vram">("auto");
|
||||
// "auto", or the physical index to pin this load to; offered only on a multi-card CUDA/ROCm
|
||||
// host. Persisted, unlike the selects around it: status carries the device a pipeline is on
|
||||
|
|
@ -4409,12 +4421,15 @@ export function ImagesPage({
|
|||
// The explicit low-precision schemes need the dense tensor-core path, which a Mac or
|
||||
// CPU-only host cannot run, so the picker does not list what the loader would refuse.
|
||||
...(hostOffersDensePrecision(hostClass)
|
||||
? ([
|
||||
["fp8", "FP8"],
|
||||
["int8", "INT8"],
|
||||
["nvfp4", "NVFP4 (Blackwell)"],
|
||||
["mxfp8", "MXFP8 (Blackwell)"],
|
||||
] as [string, string][])
|
||||
? withNvfp4Option(
|
||||
[
|
||||
["fp8", "FP8"],
|
||||
["int8", "INT8"],
|
||||
["nvfp4", "NVFP4 (Blackwell)"],
|
||||
["mxfp8", "MXFP8 (Blackwell)"],
|
||||
] as [string, string][],
|
||||
nvfp4Diffusion,
|
||||
)
|
||||
: []),
|
||||
]}
|
||||
/>
|
||||
|
|
@ -4434,16 +4449,18 @@ export function ImagesPage({
|
|||
badge={<ResolvedBadge status={status} controlKey="text_encoder_quant" />}
|
||||
value={textEncoderQuant}
|
||||
onValueChange={(v) => setTextEncoderQuant(v as typeof textEncoderQuant)}
|
||||
options={[
|
||||
["auto", "Default"],
|
||||
// The opt-out. Reachable only since a family default can pick a scheme on its own: with
|
||||
// "Default" meaning bf16 everywhere, omitting the field WAS the dense request.
|
||||
["none", "Dense (bf16)"],
|
||||
["fp8", "FP8 (storage)"],
|
||||
["fp8_dynamic", "FP8 (compute)"],
|
||||
["int8", "INT8"],
|
||||
["nvfp4", "NVFP4 (Blackwell)"],
|
||||
]}
|
||||
options={withNvfp4Option(
|
||||
[
|
||||
["auto", "Default"],
|
||||
// Opt-out: now that a family default can pick a scheme, omitting the field is no longer the dense request.
|
||||
["none", "Dense (bf16)"],
|
||||
["fp8", "FP8 (storage)"],
|
||||
["fp8_dynamic", "FP8 (compute)"],
|
||||
["int8", "INT8"],
|
||||
["nvfp4", "NVFP4 (Blackwell)"],
|
||||
] as [string, string][],
|
||||
nvfp4Diffusion,
|
||||
)}
|
||||
/>
|
||||
<AdvancedSelect
|
||||
label="Attention"
|
||||
|
|
|
|||
|
|
@ -116,7 +116,13 @@ import {
|
|||
catalogToModelOptions,
|
||||
loadSpecFor,
|
||||
} from "@/features/model-picker/components/model-selector/model-catalog";
|
||||
import { useDenseQuantSchemes, useHostClass } from "@/hooks/use-host-class";
|
||||
import {
|
||||
useDenseQuantSchemes,
|
||||
useHostClass,
|
||||
useNvfp4Diffusion,
|
||||
useNvfp4DiffusionKnown,
|
||||
} from "@/hooks/use-host-class";
|
||||
import { nvfp4SelectionFallback, withNvfp4Option } from "@/lib/nvfp4-options";
|
||||
import type {
|
||||
ModelOption,
|
||||
ModelSelectorChangeMeta,
|
||||
|
|
@ -930,6 +936,8 @@ function VideoGenerator({
|
|||
const isMobileShell = useIsMobileShell();
|
||||
const hostClass = useHostClass();
|
||||
const denseQuantSchemes = useDenseQuantSchemes();
|
||||
const nvfp4Diffusion = useNvfp4Diffusion();
|
||||
const nvfp4DiffusionKnown = useNvfp4DiffusionKnown();
|
||||
const videoModels = useVideoModels(hostClass, denseQuantSchemes);
|
||||
const [quant, setQuant] = useState<string | null>(galleryCache.quant);
|
||||
// Starts from the last prompt generated with; the example is only a placeholder.
|
||||
|
|
@ -1023,6 +1031,9 @@ function VideoGenerator({
|
|||
const [transformerQuant, setTransformerQuant] = useState<
|
||||
"auto" | "none" | "fp8" | "int8" | "nvfp4" | "mxfp8"
|
||||
>("auto");
|
||||
useEffect(() => {
|
||||
setTransformerQuant((v) => nvfp4SelectionFallback(v, nvfp4DiffusionKnown, nvfp4Diffusion));
|
||||
}, [nvfp4Diffusion, nvfp4DiffusionKnown, transformerQuant]);
|
||||
// The last load descriptor, so "Reapply" can reload the same model with new advanced options.
|
||||
const lastLoad = useRef<({ repoId: string } & VideoLoadOptions) | null>(null);
|
||||
// Render-safe mirror of whether a page-initiated load supplied a complete Reapply target.
|
||||
|
|
@ -3499,12 +3510,15 @@ function VideoGenerator({
|
|||
// The explicit low-precision schemes need the dense tensor-core path, which a Mac or
|
||||
// CPU-only host cannot run, so the picker does not list what the loader would refuse.
|
||||
...(hostOffersDensePrecision(hostClass)
|
||||
? ([
|
||||
["fp8", "FP8"],
|
||||
["int8", "INT8"],
|
||||
["nvfp4", "NVFP4 (Blackwell)"],
|
||||
["mxfp8", "MXFP8 (Blackwell)"],
|
||||
] as [string, string][])
|
||||
? withNvfp4Option(
|
||||
[
|
||||
["fp8", "FP8"],
|
||||
["int8", "INT8"],
|
||||
["nvfp4", "NVFP4 (Blackwell)"],
|
||||
["mxfp8", "MXFP8 (Blackwell)"],
|
||||
] as [string, string][],
|
||||
nvfp4Diffusion,
|
||||
)
|
||||
: []),
|
||||
]}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -52,6 +52,8 @@ export interface GpuInfo {
|
|||
/** The dense quant schemes the backend says this host can run, best first ("fp8", "int8"). Empty
|
||||
* until system info arrives, and on a backend too old to report the field. */
|
||||
denseQuantSchemes: readonly string[];
|
||||
/** False until system info arrives, and on backends that do not report it. */
|
||||
nvfp4Diffusion: boolean;
|
||||
name: string;
|
||||
memoryTotalGb: number;
|
||||
memorySharedGb: number;
|
||||
|
|
@ -89,6 +91,7 @@ const DEFAULT_GPU: GpuInfo = {
|
|||
backend: "",
|
||||
denseQuantSupported: false,
|
||||
denseQuantSchemes: [],
|
||||
nvfp4Diffusion: false,
|
||||
name: "Unknown",
|
||||
memoryTotalGb: 0,
|
||||
memorySharedGb: 0,
|
||||
|
|
@ -116,6 +119,7 @@ function toGpuInfo(
|
|||
backend: data?.device_backend ?? "",
|
||||
denseQuantSupported: data?.dense_quant_supported === true,
|
||||
denseQuantSchemes: normalizeDenseQuantSchemes(data?.dense_quant_schemes),
|
||||
nvfp4Diffusion: data?.nvfp4_diffusion === true,
|
||||
cpuCore: data?.cpu?.physical_count ?? 0,
|
||||
cpuThread: data?.cpu?.logical_count ?? 0,
|
||||
systemRamAvailableGb: data?.memory?.available_gb ?? 0,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,15 @@ export function useHostClass(): HostClass {
|
|||
);
|
||||
}
|
||||
|
||||
export function useNvfp4Diffusion(): boolean {
|
||||
return useGpuInfo().nvfp4Diffusion;
|
||||
}
|
||||
|
||||
/** Whether `/api/system` answered, so a false `useNvfp4Diffusion()` is not just the unloaded default. */
|
||||
export function useNvfp4DiffusionKnown(): boolean {
|
||||
return useGpuInfo().budgetKnown;
|
||||
}
|
||||
|
||||
/** The dense quant schemes this host can run, best first. */
|
||||
export function useDenseQuantSchemes(): readonly string[] {
|
||||
return useGpuInfo().denseQuantSchemes;
|
||||
|
|
|
|||
|
|
@ -59,6 +59,8 @@ export interface SystemInfoResponse {
|
|||
/** The dense quant schemes this host can run, best first. Absent on older backends, where readers
|
||||
* default to [] and name no precision. */
|
||||
dense_quant_schemes?: string[];
|
||||
/** Absent on older backends, where readers treat it as off. */
|
||||
nvfp4_diffusion?: boolean;
|
||||
uptime_seconds: number | null;
|
||||
cpu: {
|
||||
logical_count: number;
|
||||
|
|
|
|||
20
studio/frontend/src/lib/nvfp4-options.ts
Normal file
20
studio/frontend/src/lib/nvfp4-options.ts
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
export function withNvfp4Option<T extends readonly [string, string]>(
|
||||
options: readonly T[],
|
||||
nvfp4Enabled: boolean,
|
||||
): T[] {
|
||||
if (nvfp4Enabled) return [...options];
|
||||
return options.filter(([value]) => value.trim().toLowerCase() !== "nvfp4");
|
||||
}
|
||||
|
||||
/** A held `nvfp4` falls back to `auto` once the backend says it is off; before `known`, nothing resets. */
|
||||
export function nvfp4SelectionFallback<T extends string>(
|
||||
value: T,
|
||||
nvfp4Known: boolean,
|
||||
nvfp4Enabled: boolean,
|
||||
): T | "auto" {
|
||||
if (!nvfp4Known || nvfp4Enabled) return value;
|
||||
return value.trim().toLowerCase() === "nvfp4" ? "auto" : value;
|
||||
}
|
||||
|
|
@ -42,6 +42,7 @@ function snapshot(
|
|||
backend: "cuda",
|
||||
denseQuantSupported: schemes.length > 0,
|
||||
denseQuantSchemes: schemes,
|
||||
nvfp4Diffusion: false,
|
||||
name: "NVIDIA GeForce RTX 4090",
|
||||
memoryTotalGb: 24,
|
||||
memorySharedGb: 0,
|
||||
|
|
|
|||
58
studio/frontend/tests/nvfp4-options.test.ts
Normal file
58
studio/frontend/tests/nvfp4-options.test.ts
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import assert from "node:assert/strict";
|
||||
import { test } from "node:test";
|
||||
import { nvfp4SelectionFallback, withNvfp4Option } from "../src/lib/nvfp4-options.ts";
|
||||
|
||||
const OPTIONS: [string, string][] = [
|
||||
["fp8", "FP8"],
|
||||
["int8", "INT8"],
|
||||
["nvfp4", "NVFP4 (Blackwell)"],
|
||||
["mxfp8", "MXFP8 (Blackwell)"],
|
||||
];
|
||||
|
||||
test("NVFP4 is hidden while the backend switch is off", () => {
|
||||
assert.deepEqual(
|
||||
withNvfp4Option(OPTIONS, false).map(([value]) => value),
|
||||
["fp8", "int8", "mxfp8"],
|
||||
);
|
||||
});
|
||||
|
||||
test("NVFP4 is offered, in place, when the backend enables it", () => {
|
||||
assert.deepEqual(withNvfp4Option(OPTIONS, true), OPTIONS);
|
||||
});
|
||||
|
||||
test("the input list is never mutated", () => {
|
||||
const copy = OPTIONS.map((o) => [...o]);
|
||||
withNvfp4Option(OPTIONS, false);
|
||||
assert.deepEqual(OPTIONS, copy);
|
||||
});
|
||||
|
||||
test("a held NVFP4 pick falls back to auto once the backend reports the switch off", () => {
|
||||
assert.equal(nvfp4SelectionFallback("nvfp4", true, false), "auto");
|
||||
assert.equal(nvfp4SelectionFallback("NVFP4", true, false), "auto");
|
||||
});
|
||||
|
||||
test("other picks, an enabled switch, and a not-yet-known switch keep the selection", () => {
|
||||
assert.equal(nvfp4SelectionFallback("fp8", true, false), "fp8");
|
||||
assert.equal(nvfp4SelectionFallback("none", true, false), "none");
|
||||
assert.equal(nvfp4SelectionFallback("nvfp4", true, true), "nvfp4");
|
||||
assert.equal(nvfp4SelectionFallback("nvfp4", false, false), "nvfp4");
|
||||
});
|
||||
|
||||
test("both pages reset the hidden NVFP4 selects through the fallback", async () => {
|
||||
const { readFile } = await import("node:fs/promises");
|
||||
const images = await readFile(new URL("../src/features/images/images-page.tsx", import.meta.url), "utf8");
|
||||
const video = await readFile(new URL("../src/features/video/video-page.tsx", import.meta.url), "utf8");
|
||||
for (const setter of ["setTransformerQuant", "setTextEncoderQuant"]) {
|
||||
assert.match(
|
||||
images,
|
||||
new RegExp(`${setter}\\(\\(v\\) => nvfp4SelectionFallback\\(v, nvfp4DiffusionKnown, nvfp4Diffusion\\)\\)`),
|
||||
);
|
||||
}
|
||||
assert.match(
|
||||
video,
|
||||
/setTransformerQuant\(\(v\) => nvfp4SelectionFallback\(v, nvfp4DiffusionKnown, nvfp4Diffusion\)\)/,
|
||||
);
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue