ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio diffusion (Phase 9) review fixes: prequant safety + validation

Browse source 
- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
  resolves inside an operator-configured ALLOWLIST of directories
  (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
  once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
  any path a load request named (arbitrary code execution). realpath() blocks symlink
  escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
  checkpoint that quantised a different layer set is rejected instead of silently
  loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
  path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
  author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.
This commit is contained in:
Daniel Han 2026-06-29 05:27:19 +00:00
commit 44770ce2d9
5 changed files with 177 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

2
scripts/prequant_probe.py
View file

@ -26,7 +26,7 @@ import numpy as np
BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
ROOT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research")
ROOT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research"
CKPT = ROOT / "prequant_fp8" / "transformer_fp8_state.pt"
OUT = ROOT / "prequant_images"
MIN_FEAT = 512

9
scripts/verify_prequant_backend.py
View file

@ -19,17 +19,20 @@ from __future__ import annotations
import argparse
import logging
import os
import sys
import time
from pathlib import Path
import numpy as np
BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"
_REPO = Path(__file__).resolve().parent.parent
_RESEARCH = _REPO / "outputs" / "quant_research"
BACKEND = _REPO / "studio" / "backend"
BASE = "Tongyi-MAI/Z-Image-Turbo"
CKPT = "/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/prequant_fp8/transformer_fp8.pt"
CKPT = os.environ.get("PREQUANT_CKPT", str(_RESEARCH / "prequant_fp8" / "transformer_fp8.pt"))
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/prequant_verify_images")
OUT = Path(os.environ.get("PREQUANT_OUT_DIR", str(_RESEARCH / "prequant_verify_images")))
logging.basicConfig(level = logging.INFO, format = "%(message)s")
LOGGER = logging.getLogger("verify_prequant")

4
studio/backend/core/inference/diffusion.py
View file

@ -57,6 +57,7 @@ from .diffusion_prequant import (
resolve_prequant_source,
)
from .diffusion_transformer_quant import (
DEFAULT_MIN_LINEAR_FEATURES,
dense_transformer_supported,
normalize_transformer_quant,
quantize_transformer,
@ -648,6 +649,9 @@ class DiffusionBackend:
dtype = dtype,
hf_token = hf_token,
scheme = scheme,
# Reject a checkpoint built with a different Linear filter than the
# dense path uses, so the prequant and runtime-quant models match.
min_features = DEFAULT_MIN_LINEAR_FEATURES,
logger = logger,
)
if transformer is not None:

105
studio/backend/core/inference/diffusion_prequant.py
View file

@ -37,20 +37,51 @@ PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1"
# code embedded in the pickle. A hosted family *repo* checkpoint is first-party and trusted,
# but a ``source.kind == "path"`` can originate from the ``transformer_prequant_path`` field
# of a load request -- i.e. an authenticated API caller naming an arbitrary local file.
# Unpickling that is remote code execution, so the local-path branch is refused unless an
# operator explicitly opts in via this env var. The trusted hosted-repo path is unaffected.
# Unpickling that is remote code execution, so a request-supplied path is unpickled ONLY when
# it resolves inside an operator-configured ALLOWLIST of directories. A bare on/off toggle is
# deliberately NOT accepted as a wildcard: enabling local checkpoints for one trusted
# directory must never also permit unpickling any other path a request happens to name. The
# trusted hosted-repo path is unaffected.
ALLOW_LOCAL_PREQUANT_PATH_ENV = "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH"
_PREQUANT_TOGGLE_TOKENS = {"1", "true", "yes", "on", "0", "false", "no", "off"}
def _local_prequant_path_allowed() -> bool:
"""Whether a request-supplied local pre-quant *path* may be unpickled (operator opt-in)."""
def _allowed_prequant_roots() -> list:
"""Operator-allowlisted directories whose pre-quant checkpoints may be unpickled.
Set ``UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH`` to one or more directories (separated by
``os.pathsep``). A bare truthy/falsey toggle is ignored on purpose -- it must name a
directory, so there is no "allow everything" mode."""
import os
return (os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV) or "").strip().lower() in {
"1",
"true",
"yes",
"on",
}
raw = (os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV) or "").strip()
if not raw:
return []
roots = []
for part in raw.split(os.pathsep):
part = part.strip()
if not part or part.lower() in _PREQUANT_TOGGLE_TOKENS:
continue # a bare on/off value is not a directory -> never a wildcard allow
try:
roots.append(os.path.realpath(os.path.expanduser(part)))
except Exception: # noqa: BLE001 — a bad entry is simply not allowlisted
continue
return roots
def _local_prequant_path_allowed(path: str) -> bool:
"""True only when ``path`` resolves inside an operator-allowlisted directory; an
arbitrary request-supplied path is never unpickled. ``realpath`` first so a symlink
cannot point an allowlisted name at a file outside the allowed roots."""
import os
roots = _allowed_prequant_roots()
if not roots:
return False
try:
real = os.path.realpath(os.path.expanduser(path))
except Exception: # noqa: BLE001
return False
return any(real == r or real.startswith(r + os.sep) for r in roots)
@dataclass(frozen = True)
@ -102,6 +133,7 @@ def load_prequantized_transformer(
dtype: Any,
hf_token: Optional[str] = None,
scheme: str,
min_features: Optional[int] = None,
logger: Any = None,
) -> Optional[Any]:
"""Load the pre-quantized transformer described by ``source`` onto ``device``.
@ -113,16 +145,16 @@ def load_prequantized_transformer(
"""
try:
# weights_only=False (required below) executes pickle code, so a caller-supplied
# local path is refused unless an operator opted in. The hosted family repo is
# first-party and always allowed.
if source.kind == "path" and not _local_prequant_path_allowed():
# local path is unpickled ONLY when it resolves inside an operator-allowlisted
# directory. The hosted family repo is first-party and always allowed.
if source.kind == "path" and not _local_prequant_path_allowed(source.location):
_warn(
logger,
f"{scheme}:path",
RuntimeError(
"request-supplied local pre-quant path refused (unpickling an "
f"arbitrary file is unsafe); set {ALLOW_LOCAL_PREQUANT_PATH_ENV}=1 "
"to allow trusted local checkpoints",
"request-supplied local pre-quant path refused (unpickling an arbitrary "
f"file is unsafe); set {ALLOW_LOCAL_PREQUANT_PATH_ENV} to an allowlisted "
"directory containing trusted checkpoints to permit it",
),
)
return None
@ -137,7 +169,7 @@ def load_prequantized_transformer(
# a torch.save pickle. weights_only=False is required to rebuild those subclasses.
# The local-path branch is gated above; the repo branch is a first-party artifact.
ckpt = torch.load(path, weights_only = False, map_location = "cpu")
if not _validate_checkpoint(ckpt, scheme, base, logger):
if not _validate_checkpoint(ckpt, scheme, base, logger, min_features = min_features):
return None
state_dict = ckpt["state_dict"]
@ -187,8 +219,15 @@ def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) ->
return None
def _validate_checkpoint(ckpt: Any, scheme: str, base: str, logger: Any) -> bool:
"""Reject a checkpoint that is the wrong format / scheme / base model."""
def _validate_checkpoint(
ckpt: Any, scheme: str, base: str, logger: Any, min_features: Optional[int] = 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 checkpoint
built with a different ``--min-features`` quantises a different set of Linear layers,
so ``load_state_dict(assign=True)`` would silently install a model that does not match
what the dense path produces while status still reports the requested scheme. Reject it."""
if not isinstance(ckpt, dict) or ckpt.get("format") != PREQUANT_FORMAT:
_warn(logger, scheme, ValueError("unrecognised pre-quant checkpoint format"))
return False
@ -200,21 +239,39 @@ def _validate_checkpoint(ckpt: Any, scheme: str, base: str, logger: Any) -> bool
_warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}"))
return False
ckpt_base = meta.get("base_model_id")
if ckpt_base and base and ckpt_base != base:
if ckpt_base and base and not _same_base_model(ckpt_base, base):
_warn(logger, scheme, ValueError(f"checkpoint base {ckpt_base!r} != {base!r}"))
return False
if min_features is not None:
ckpt_min = meta.get("min_features")
if ckpt_min is not None and int(ckpt_min) != int(min_features):
_warn(
logger,
scheme,
ValueError(f"checkpoint min_features {ckpt_min!r} != runtime {min_features!r}"),
)
return False
return True
def _same_base_model(a: str, b: str) -> bool:
"""Tolerant compare of two base-model ids: an exact match, or the same final
path/repo segment (so a local path or a fork id matches the canonical repo, e.g.
``/models/Z-Image-Turbo`` vs ``Tongyi-MAI/Z-Image-Turbo``)."""
def _tail(x: str) -> str:
return x.replace("\\", "/").rstrip("/").split("/")[-1].lower()
return a == b or _tail(a) == _tail(b)
def _has_meta_tensors(module: Any) -> bool:
"""True if any parameter or buffer is still on the meta device after loading."""
from itertools import chain
try:
for tensor in list(module.parameters()) + list(module.buffers()):
if getattr(tensor, "is_meta", False):
return True
return any(
getattr(t, "is_meta", False) for t in chain(module.parameters(), module.buffers())
)
except Exception: # noqa: BLE001
return False
return False
def _warn(logger: Any, what: str, exc: Exception) -> None:

88
studio/backend/tests/test_diffusion_prequant.py
View file

@ -147,10 +147,11 @@ def _load(
):
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises)
# The local-path branch is opt-in (it unpickles an arbitrary file); these tests
# exercise the load mechanics, so enable it unless a test is checking the gate.
# The local-path branch is opt-in via a directory ALLOWLIST (it unpickles an arbitrary
# file); these tests exercise the load mechanics, so allowlist tmp_path (where ckpt.pt
# lives) unless a test is checking the gate.
if allow_local:
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, "1")
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path))
else:
monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
path = tmp_path / "ckpt.pt"
@ -264,3 +265,84 @@ def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path):
logger = None,
)
assert result is not None
def test_load_local_path_outside_allowlist_refused(monkeypatch, tmp_path):
# Even with the opt-in set, a path OUTSIDE every allowlisted directory must not be
# unpickled: enabling one trusted dir is not a wildcard for arbitrary request paths.
called = {"load": False}
def _explode(*a, **k):
called["load"] = True
raise AssertionError("torch.load must not run on a path outside the allowlist")
torch = types.ModuleType("torch")
torch.load = _explode
monkeypatch.setitem(sys.modules, "torch", torch)
allowed = tmp_path / "allowed"
allowed.mkdir()
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(allowed))
outside = tmp_path / "evil.pt" # a real file, but outside the allowlisted dir
outside.write_bytes(b"x")
source = PrequantSource(kind = "path", location = str(outside), filename = None)
result = load_prequantized_transformer(
_FakeTransformer,
"Tongyi-MAI/Z-Image-Turbo",
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
logger = None,
)
assert result is None
assert called["load"] is False
def test_load_min_features_mismatch_is_none(monkeypatch, tmp_path):
# A checkpoint built with a different --min-features quantises a different Linear set,
# so it must be rejected when the runtime threshold is supplied.
ckpt = _good_ckpt()
ckpt["metadata"]["min_features"] = 256 # built with 256, runtime asks for 512
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, ckpt)
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path))
path = tmp_path / "ckpt.pt"
path.write_bytes(b"x")
source = PrequantSource(kind = "path", location = str(path), filename = None)
result = load_prequantized_transformer(
_FakeTransformer,
"Tongyi-MAI/Z-Image-Turbo",
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
min_features = 512,
logger = None,
)
assert result is None
def test_load_base_fork_tail_matches(monkeypatch, tmp_path):
# A local path / fork id with the same final segment as the canonical base is accepted.
ckpt = _good_ckpt(base = "Tongyi-MAI/Z-Image-Turbo")
_FakeTransformer.calls = {}
_stub_torch_accelerate(monkeypatch, ckpt)
monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path))
path = tmp_path / "ckpt.pt"
path.write_bytes(b"x")
source = PrequantSource(kind = "path", location = str(path), filename = None)
result = load_prequantized_transformer(
_FakeTransformer,
"/local/models/Z-Image-Turbo", # different prefix, same tail
source,
device = "cuda",
dtype = "bfloat16",
hf_token = None,
scheme = "fp8",
logger = None,
)
assert result is not None