diff --git a/src/soup_cli/commands/serve.py b/src/soup_cli/commands/serve.py index 9b68bfc..bb9d725 100644 --- a/src/soup_cli/commands/serve.py +++ b/src/soup_cli/commands/serve.py @@ -1303,9 +1303,9 @@ def _load_draft_model(speculative_model: str, device: str): # from_pretrained torch.load's them — a poisoned ~/.soup/drafts.json entry # must not become a load-time RCE. if os.path.isdir(speculative_model): - from soup_cli.utils.strict_safetensors import check_strict_safetensors + from soup_cli.utils.strict_safetensors import assert_safe_top_level_weights - check_strict_safetensors(speculative_model, strict=True) + assert_safe_top_level_weights(speculative_model) console.print(f"[dim]Loading draft model: {escape(speculative_model)}...[/]") draft = AutoModelForCausalLM.from_pretrained( diff --git a/src/soup_cli/utils/adapter_fuse.py b/src/soup_cli/utils/adapter_fuse.py index 98a4d1d..9d3f5db 100644 --- a/src/soup_cli/utils/adapter_fuse.py +++ b/src/soup_cli/utils/adapter_fuse.py @@ -64,12 +64,13 @@ def merge_adapter_to_dense( # Refuse pickle / PyTorch-classic weights in the adapter dir before PEFT # torch.load's them — the adapter dir may have been produced (or swapped) - # by an untrusted process. Mirrors the strict-safetensors policy applied on - # other weight-loading paths. - from soup_cli.utils.strict_safetensors import check_strict_safetensors + # by an untrusted process. Shallow scan: PEFT loads adapter_model.* from + # the TOP LEVEL, while a training output dir also holds the HF Trainer's + # own checkpoint-N/optimizer.pt pickles, which are not the threat (a + # recursive scan would make Soup refuse its own trainer's output). + from soup_cli.utils.strict_safetensors import assert_safe_top_level_weights - if os.path.isdir(adapter_dir): - check_strict_safetensors(adapter_dir, strict=True) + assert_safe_top_level_weights(adapter_dir) from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer diff --git a/src/soup_cli/utils/gpu.py b/src/soup_cli/utils/gpu.py index 6146a4a..da902bd 100644 --- a/src/soup_cli/utils/gpu.py +++ b/src/soup_cli/utils/gpu.py @@ -1,6 +1,7 @@ """GPU detection, memory calculation, and auto batch size.""" import math +import re def detect_device() -> tuple[str, str]: @@ -115,16 +116,31 @@ def model_size_from_name(model_name: str) -> float: if marker in name_lower: return size + # Longer markers first: "1.7b" contains "7b", so a naive scan would call a + # 1.7B model a 7B one (and over-predict its VRAM by 4x). size_markers = [ ("70b", 70), ("65b", 65), ("34b", 34), ("33b", 33), - ("13b", 13), ("8b", 8), ("7b", 7), ("3b", 3), - ("1.5b", 1.5), ("1b", 1), ("0.5b", 0.5), + ("13b", 13), ("8b", 8), ("3b", 3), + ("1.5b", 1.5), ("1.7b", 1.7), ("0.5b", 0.5), ("0.6b", 0.6), + ("7b", 7), ("1b", 1), ] for marker, size in size_markers: if marker in name_lower: return size + # Sub-billion checkpoints carry their size in MILLIONS (SmolLM2-135M, + # SmolVLM-256M, ...). Without this they fell through to the 7B default and + # the hardware-fit gate refused to train them — which blocked `soup draft` + # for exactly the tiny models drafts are made of (v0.71.33 live smoke; + # same class as the v0.71.32 whisper fix). + # + # Checked AFTER the "b" markers on purpose: `Qwen2.5-7B-Instruct-1M` is a + # 7B model with a 1M *context*, not a 1M-parameter model. + million = re.search(r"(? Tuple[str, ...]: return tuple(sorted(offenders)) +# Filenames that ``from_pretrained`` / ``PeftModel.from_pretrained`` actually +# deserialize as WEIGHTS. A training output dir is full of other pickles the HF +# Trainer wrote itself (``training_args.bin``, ``checkpoint-N/optimizer.pt``) — +# those are never loaded by ``from_pretrained``, so treating them as an attack +# just makes Soup refuse its own trainer's output. +_WEIGHT_STEM_PREFIXES = ( + "adapter_model", + "pytorch_model", + "model", + "consolidated", +) + + +def _is_loadable_weight_name(filename: str) -> bool: + stem = os.path.splitext(filename)[0].lower() + return any(stem.startswith(prefix) for prefix in _WEIGHT_STEM_PREFIXES) + + +def find_unsafe_weight_files_shallow(model_dir: str) -> Tuple[str, ...]: + """Unsafe *loadable* weight files at the TOP LEVEL of ``model_dir`` (v0.71.33). + + Narrower than the recursive :func:`find_unsafe_weight_files` in two ways, + both deliberate: + + * **Top level only** — ``from_pretrained`` loads weights from the top level; + ``checkpoint-N/optimizer.pt`` below it is the Trainer's own pickle. + * **Weight filenames only** — ``training_args.bin`` is a pickled + ``TrainingArguments``, not a tensor file, and ``from_pretrained`` never + deserializes it. + + What it still catches is the actual threat: a pickle ``adapter_model.bin`` / + ``pytorch_model.bin`` sitting exactly where the loader will unpickle it. + """ + if not isinstance(model_dir, str) or not model_dir: + raise ValueError("model_dir must be a non-empty str") + if not os.path.isdir(model_dir): + return () + + offenders: list[str] = [] + for filename in os.listdir(model_dir): + full = os.path.join(model_dir, filename) + if not os.path.isfile(full) or not _is_loadable_weight_name(filename): + continue + ext = os.path.splitext(filename)[1].lower() + if ext in UNSAFE_EXTENSIONS: + offenders.append(full) + elif ext == SAFETENSORS_EXTENSION and not is_safetensors_magic(full): + offenders.append(full) + return tuple(sorted(offenders)) + + +def assert_safe_top_level_weights(model_dir: str) -> None: + """Raise ``ValueError`` if ``model_dir`` has loadable pickle weights. + + Shallow counterpart of ``check_strict_safetensors(strict=True)``; see + :func:`find_unsafe_weight_files_shallow` for the (deliberate) scope. + """ + offenders = find_unsafe_weight_files_shallow(model_dir) + if offenders: + names = ", ".join(os.path.basename(path) for path in offenders) + raise ValueError( + "unsafe weight file (pickle / PyTorch-classic / invalid " + f"safetensors): {names}. Re-save the model as safetensors." + ) + + def check_strict_safetensors( model_dir: str, *, strict: bool = False, ) -> StrictSafetensorsReport: diff --git a/tests/test_v07133.py b/tests/test_v07133.py index 8169933..602c139 100644 --- a/tests/test_v07133.py +++ b/tests/test_v07133.py @@ -138,6 +138,68 @@ class TestAdapterFuse: base_model="org/tiny", adapter_dir=str(adapter), out_dir="out" ) + def test_trainer_checkpoint_pickles_are_not_mistaken_for_an_attack( + self, tmp_path, monkeypatch + ): + """Live-smoke regression: the HF Trainer writes checkpoint-N/optimizer.pt + (a pickle IT wrote). A recursive scan refused Soup's own distill output + and made `soup draft distill` impossible. Only TOP-LEVEL weights — the + ones PEFT actually loads — are the threat surface.""" + from soup_cli.utils.strict_safetensors import ( + assert_safe_top_level_weights, + check_strict_safetensors, + ) + + monkeypatch.chdir(tmp_path) # check_strict_safetensors is cwd-contained + adapter = tmp_path / "_adapter" + (adapter / "checkpoint-60").mkdir(parents=True) + (adapter / "adapter_config.json").write_text("{}", encoding="utf-8") + # Both pickles below are written by the HF Trainer itself, and NEITHER is + # ever deserialized by from_pretrained. Both broke the live smoke. + (adapter / "training_args.bin").write_bytes(b"\x80\x04pickle") + (adapter / "checkpoint-60" / "optimizer.pt").write_bytes(b"\x80\x04pickle") + + # The RECURSIVE, extension-only scan flags them — that is exactly what + # made `soup draft distill` refuse its own trainer's output... + with pytest.raises(ValueError, match="(?i)unsafe"): + check_strict_safetensors(str(adapter), strict=True) + # ...while the shallow weight-name scan (what the merge uses) accepts it. + assert_safe_top_level_weights(str(adapter)) + + # A top-level pickle ADAPTER — the file PEFT actually unpickles, i.e. + # the real threat — is still refused. + (adapter / "adapter_model.bin").write_bytes(b"\x80\x04pickle") + with pytest.raises(ValueError, match="(?i)unsafe"): + assert_safe_top_level_weights(str(adapter)) + + def test_shallow_scan_ignores_trainer_bookkeeping_pickles(self, tmp_path): + from soup_cli.utils.strict_safetensors import ( + assert_safe_top_level_weights, + find_unsafe_weight_files_shallow, + ) + + adapter = tmp_path / "_adapter" + (adapter / "checkpoint-60").mkdir(parents=True) + (adapter / "adapter_config.json").write_text("{}", encoding="utf-8") + (adapter / "training_args.bin").write_bytes(b"\x80\x04pickle") + (adapter / "checkpoint-60" / "optimizer.pt").write_bytes(b"\x80\x04pickle") + + assert find_unsafe_weight_files_shallow(str(adapter)) == () + assert_safe_top_level_weights(str(adapter)) # must not raise + + def test_shallow_scan_still_catches_a_pickle_model_file(self, tmp_path): + """serve --auto-spec loads a registry draft dir: pytorch_model.bin is + the file from_pretrained would unpickle.""" + from soup_cli.utils.strict_safetensors import assert_safe_top_level_weights + + model_dir = tmp_path / "draft" + model_dir.mkdir() + (model_dir / "config.json").write_text("{}", encoding="utf-8") + (model_dir / "pytorch_model.bin").write_bytes(b"\x80\x04pickle") + + with pytest.raises(ValueError, match="(?i)unsafe"): + assert_safe_top_level_weights(str(model_dir)) + def test_merge_loads_base_weights_from_the_base_model_not_out_dir( self, tmp_path, monkeypatch ): @@ -1739,6 +1801,47 @@ class TestDraftListCli: assert "org/evil" in result.output # literal, not interpreted as markup +class TestMillionParamSizeGate: + """Live-smoke regression: `soup draft distill` on SmolLM2-135M was REFUSED + by the hardware-fit gate, which predicted 14 GB of weights — because + model_size_from_name knew no "M" (millions) suffix and fell back to the 7B + default. Every draft-sized model hit this. Same class as the v0.71.32 + whisper fix.""" + + def test_million_suffix_is_parsed(self): + from soup_cli.utils.gpu import model_size_from_name + + assert model_size_from_name( + "HuggingFaceTB/SmolLM2-135M-Instruct" + ) == pytest.approx(0.135) + assert model_size_from_name( + "HuggingFaceTB/SmolLM2-360M-Instruct" + ) == pytest.approx(0.360) + assert model_size_from_name("HuggingFaceTB/SmolVLM-256M") == pytest.approx( + 0.256 + ) + + def test_billion_marker_still_wins_over_a_context_length_suffix(self): + """`Qwen2.5-7B-Instruct-1M` is a 7B model with a 1M CONTEXT — the "1M" + must not be read as 1M parameters.""" + from soup_cli.utils.gpu import model_size_from_name + + assert model_size_from_name("Qwen/Qwen2.5-7B-Instruct-1M") == 7 + + def test_one_point_seven_b_is_not_seven_b(self): + """"1.7b" contains "7b" — the marker list must match the longer one.""" + from soup_cli.utils.gpu import model_size_from_name + + assert model_size_from_name( + "HuggingFaceTB/SmolLM2-1.7B-Instruct" + ) == pytest.approx(1.7) + + def test_unknown_model_still_defaults_to_7b(self): + from soup_cli.utils.gpu import model_size_from_name + + assert model_size_from_name("some-unknown-model") == 7.0 + + class TestPromptTexts: def test_extracts_messages_and_drops_unusable_rows(self): from soup_cli.commands.draft import _prompt_texts