fix(draft): two bugs the live smoke found (v0.71.33)

Neither was caught by 15k green tests; both made 'soup draft distill'
impossible on real models.

1. model_size_from_name knew no 'M' (millions) suffix, so SmolLM2-135M fell
   through to the 7B default -> the hardware-fit gate predicted 14 GB of
   weights and REFUSED to train. That blocked every draft-sized model, i.e.
   exactly what this feature is for. Same class as the v0.71.32 whisper fix.
   Also fixed '1.7b' matching the '7b' marker (a 1.7B model was read as 7B),
   while keeping 'Qwen2.5-7B-Instruct-1M' at 7B (1M = context, not params).

2. The strict-safetensors guard added in the security round scanned
   recursively by file EXTENSION, so it flagged the HF Trainer's own pickles
   (training_args.bin, checkpoint-N/optimizer.pt) and made the merge refuse
   Soup's own distill output. Scoped to top-level files that from_pretrained
   actually deserializes as WEIGHTS (adapter_model/pytorch_model/model/...),
   via the new strict_safetensors.assert_safe_top_level_weights. A pickle
   adapter_model.bin / pytorch_model.bin is still refused.

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(

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

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"(?<![a-z0-9.])(\d+(?:\.\d+)?)m(?![a-z0-9])", name_lower)
+    if million:
+        return float(million.group(1)) / 1000.0
+
     return 7.0  # default guess
 
 

src/soup_cli/utils/strict_safetensors.py

@@ -144,6 +144,72 @@ def find_unsafe_weight_files(model_dir: str) -> 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:

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