From 44770ce2d9263f7194e63e9b1a70d3ede0dc119f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 29 Jun 2026 05:27:19 +0000 Subject: [PATCH] Studio diffusion (Phase 9) review fixes: prequant safety + validation - 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. --- scripts/prequant_probe.py | 2 +- scripts/verify_prequant_backend.py | 9 +- studio/backend/core/inference/diffusion.py | 4 + .../core/inference/diffusion_prequant.py | 105 ++++++++++++++---- .../backend/tests/test_diffusion_prequant.py | 88 ++++++++++++++- 5 files changed, 177 insertions(+), 31 deletions(-) diff --git a/scripts/prequant_probe.py b/scripts/prequant_probe.py index 3e80e3d2a3..cbcbc245fe 100644 --- a/scripts/prequant_probe.py +++ b/scripts/prequant_probe.py @@ -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 diff --git a/scripts/verify_prequant_backend.py b/scripts/verify_prequant_backend.py index 55e39aa134..e160597599 100644 --- a/scripts/verify_prequant_backend.py +++ b/scripts/verify_prequant_backend.py @@ -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") diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index cbb5abd2ab..8cbbed4da0 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -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: diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py index b29f658f9f..c924c705b9 100644 --- a/studio/backend/core/inference/diffusion_prequant.py +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -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: diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py index 0a6a6dc67f..a233f6225f 100644 --- a/studio/backend/tests/test_diffusion_prequant.py +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -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