fix(shrink): code-review findings (CUDA free before heal, dtype=auto, atomic fuse-swap, fuse+exit2 coverage) (v0.71.29)

src/soup_cli/commands/shrink.py

@@ -128,6 +128,20 @@ def _count_params(model: object) -> int:
     return sum(p.numel() for p in model.parameters())  # type: ignore[attr-defined]
 
 
+def _release_cuda() -> None:
+    """Return the CUDA caching-allocator pool to the driver (best-effort)."""
+    import gc
+
+    gc.collect()
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except ImportError:
+        pass
+
+
 def _perplexity(
     model: object, tokenizer: object, prompts: Sequence[str], device: str
 ) -> float:
@@ -166,21 +180,33 @@ def _perplexity(
     return math.exp(sum(losses) / len(losses))
 
 
-def _load_for_shrink(
-    model_id: str, device: Optional[str], trust_remote_code: bool
-) -> tuple[Any, Any, str]:
-    """Load a model + tokenizer for shrinking (trust_remote_code probe + warn)."""
-    from soup_cli.utils.live_eval import load_model_and_tokenizer
+def _resolve_trc(model_id: str, requested: bool) -> bool:
+    """Resolve trust_remote_code once (probe + warn), reused across loads."""
     from soup_cli.utils.trust_remote import (
         model_requires_trust_remote_code,
         resolve_trust_remote_code,
     )
 
     requires = model_requires_trust_remote_code(model_id) or False
-    trc = resolve_trust_remote_code(
-        model_id, requested=trust_remote_code, console=console, requires_remote_code=requires
+    return resolve_trust_remote_code(
+        model_id, requested=requested, console=console, requires_remote_code=requires
+    )
+
+
+def _load_for_shrink(
+    model_id: str, device: Optional[str], trc: bool
+) -> tuple[Any, Any, str]:
+    """Load a model + tokenizer for shrinking, preserving the checkpoint dtype.
+
+    ``dtype="auto"`` keeps the model at its native precision so the shipped
+    smaller model is not silently upcast to fp32 (which would shrink the layer
+    count while widening the bytes-per-parameter).
+    """
+    from soup_cli.utils.live_eval import load_model_and_tokenizer
+
+    return load_model_and_tokenizer(
+        model_id, device=device, trust_remote_code=trc, dtype="auto"
     )
-    return load_model_and_tokenizer(model_id, device=device, trust_remote_code=trc)
 
 
 def _render_importance_table(
@@ -290,6 +316,7 @@ def _shrink_impl(
     # Fail fast on the flag combination BEFORE loading a multi-GB model.
     if (drop_ratio is None) == (drop_layers is None):
         raise typer.BadParameter("set exactly one of --drop-ratio / --drop-layers")
+    heal_rows = 0
     if heal is not None:
         if not isinstance(heal_steps, int) or isinstance(heal_steps, bool):
             raise typer.BadParameter("--heal-steps must be an int")
@@ -298,13 +325,23 @@ def _shrink_impl(
         heal_rows = _count_jsonl_rows(heal)  # validates cwd containment + O_NOFOLLOW
     prompts = _load_calib(calib)
 
-    console.print(f"[dim]Loading {escape(model)} ...[/]")
-    mdl, tokenizer, dev = _load_for_shrink(model, device, trust_remote_code)
-    # Reject an unsupported architecture up front (before the importance scan).
-    shrink_arch_of(mdl)
-    n_layers = int(mdl.config.num_hidden_layers)
+    trc = _resolve_trc(model, trust_remote_code)
+    # Fail fast on arch + drop-count from the CONFIG before loading weights.
+    from transformers import AutoConfig
+
+    from soup_cli.utils.shrink import arch_family_of_config
+
+    pre_config = AutoConfig.from_pretrained(model, trust_remote_code=trc)
+    arch_family_of_config(pre_config)
+    n_layers = int(pre_config.num_hidden_layers)
     count = resolve_drop_count(n_layers, drop_ratio=drop_ratio, drop_layers=drop_layers)
 
+    console.print(f"[dim]Loading {escape(model)} ...[/]")
+    mdl, tokenizer, dev = _load_for_shrink(model, device, trc)
+    # Defence-in-depth: re-check the loaded model's arch (layer_list also
+    # re-guards independently before any slice).
+    shrink_arch_of(mdl)
+
     console.print(f"[dim]Scoring importance over {len(prompts)} calib prompts ...[/]")
     importances = compute_layer_importance(
         mdl, tokenizer, prompts, block_size=count, device=dev
@@ -335,7 +372,11 @@ def _shrink_impl(
     model_out.mkdir(parents=True, exist_ok=True)
     mdl.save_pretrained(str(model_out))
     tokenizer.save_pretrained(str(model_out))
+    # Free the parent's VRAM before spawning the heal subprocess (teacher +
+    # student both load in the child; del alone leaves the caching-allocator
+    # pool resident — mirrors utils/interference_live.py).
     del mdl
+    _release_cuda()
 
     healed = False
     if heal is not None:
@@ -350,10 +391,11 @@ def _shrink_impl(
             steps=heal_steps,
             out_dir=str(adapter_dir),
             heal_rows=heal_rows,
+            trc=trc,
         )
         healed = True
 
-    reloaded, tok2, dev2 = _load_for_shrink(str(model_out), device, trust_remote_code)
+    reloaded, tok2, dev2 = _load_for_shrink(str(model_out), device, trc)
     layers_after = int(reloaded.config.num_hidden_layers)
     params_after = _count_params(reloaded)
     ppl_final = _perplexity(reloaded, tok2, prompts, dev2)
@@ -486,6 +528,7 @@ def _run_heal(
     steps: int,
     out_dir: str,
     heal_rows: int,
+    trc: bool = False,
 ) -> None:
     """Distill the teacher into the pruned student, then fuse the adapter.
 
@@ -532,15 +575,41 @@ def _run_heal(
         tail = (result.stderr or b"").decode("utf-8", "replace")[-500:]
         raise RuntimeError(f"heal distill failed (rc={result.returncode}): {tail}")
 
-    _fuse_adapter(base_dir=pruned_dir, adapter_dir=out_dir)
+    _fuse_adapter(base_dir=pruned_dir, adapter_dir=out_dir, trc=trc)
 
 
-def _fuse_adapter(*, base_dir: str, adapter_dir: str) -> None:
-    """Merge a LoRA adapter into ``base_dir`` in place (dense healed model)."""
+def _fuse_adapter(*, base_dir: str, adapter_dir: str, trc: bool = False) -> None:
+    """Merge a LoRA adapter into ``base_dir`` (dense healed model), atomically.
+
+    The merged model is written to a sibling temp dir and only then swapped in
+    for ``base_dir``. An in-place ``save_pretrained`` over the just-loaded
+    ``base_dir`` fails on Windows (error 1224 — the source ``.safetensors`` is
+    still memory-mapped by the loaded weights), so the temp-dir swap is the
+    cross-platform-safe path.
+    """
+    import gc
+    import shutil
+    import tempfile
+
     from peft import PeftModel
     from transformers import AutoModelForCausalLM, AutoTokenizer
 
-    base = AutoModelForCausalLM.from_pretrained(base_dir)
+    base = AutoModelForCausalLM.from_pretrained(
+        base_dir, trust_remote_code=trc, torch_dtype="auto"
+    )
     merged = PeftModel.from_pretrained(base, adapter_dir).merge_and_unload()
-    merged.save_pretrained(base_dir)
-    AutoTokenizer.from_pretrained(base_dir).save_pretrained(base_dir)
+    tokenizer = AutoTokenizer.from_pretrained(base_dir, trust_remote_code=trc)
+
+    parent = os.path.dirname(os.path.abspath(base_dir)) or "."
+    staging = tempfile.mkdtemp(prefix=".fuse_", dir=parent)
+    try:
+        merged.save_pretrained(staging)
+        tokenizer.save_pretrained(staging)
+    finally:
+        # Drop every reference so Windows releases the base_dir mmap before we
+        # remove it; otherwise rmtree(base_dir) also hits error 1224.
+        del merged, base, tokenizer
+        gc.collect()
+        _release_cuda()
+    shutil.rmtree(base_dir)
+    os.replace(staging, base_dir)

src/soup_cli/utils/live_eval.py

@@ -118,11 +118,15 @@ def load_model_and_tokenizer(
     adapter: Optional[str] = None,
     device: Optional[str] = None,
     trust_remote_code: bool = False,
+    dtype: Optional[str] = None,
 ):
     """Load an ``AutoModelForCausalLM`` + tokenizer, optionally with a LoRA adapter.
 
     Returns ``(model, tokenizer, device)``. ``model`` is ``.eval()``-ed and
-    moved to the resolved device. Heavy imports are local.
+    moved to the resolved device. Heavy imports are local. ``dtype`` (e.g.
+    ``"auto"``) is forwarded as ``torch_dtype`` so a caller can preserve the
+    checkpoint's native precision instead of upcasting to fp32 (``soup shrink``
+    needs this so the shipped smaller model is not silently re-widened).
     """
     if not isinstance(model_id, str) or not model_id.strip():
         raise ValueError("model_id must be a non-empty string")
@@ -133,7 +137,10 @@ def load_model_and_tokenizer(
     tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=trust_remote_code)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=trust_remote_code)
+    model_kwargs = {"trust_remote_code": trust_remote_code}
+    if dtype is not None:
+        model_kwargs["torch_dtype"] = dtype
+    model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
     if adapter is not None:
         if not isinstance(adapter, str) or not adapter.strip():
             raise ValueError("adapter must be a non-empty string or None")

src/soup_cli/utils/shrink.py

@@ -149,16 +149,16 @@ _ARCH_PATTERNS = {
 SUPPORTED_SHRINK_ARCHS = tuple(_ARCH_PATTERNS)
 
 
-def shrink_arch_of(model: object) -> str:
-    """Return the supported family name for ``model`` or raise ``ValueError``.
+def arch_family_of_config(config: object) -> str:
+    """Return the supported family name for an HF ``config`` or raise ``ValueError``.
 
     Detection is over ``config.model_type`` + ``config.architectures`` with
     regex word-family matching (mirrors ``longlora.is_*_model``). Only the v1
     families in :data:`SUPPORTED_SHRINK_ARCHS` (Llama / Qwen / SmolLM — all of
     which expose ``model.model.layers`` + ``config.num_hidden_layers``) are
-    accepted; anything else is a friendly reject.
+    accepted; anything else is a friendly reject. Taking a bare config lets the
+    CLI fail fast before loading weights.
     """
-    config = getattr(model, "config", None)
     model_type = getattr(config, "model_type", "") or ""
     architectures = list(getattr(config, "architectures", []) or [])
     haystack = " ".join([str(model_type), *[str(a) for a in architectures]])
@@ -171,6 +171,11 @@ def shrink_arch_of(model: object) -> str:
     )
 
 
+def shrink_arch_of(model: object) -> str:
+    """Return the supported family name for a loaded ``model`` or raise."""
+    return arch_family_of_config(getattr(model, "config", None))
+
+
 def layer_list(model: object) -> Any:
     """Return ``model.model.layers`` (the decoder ``ModuleList``), arch-guarded."""
     shrink_arch_of(model)  # raises on unsupported arch

tests/test_v07129.py

@@ -688,3 +688,56 @@ class TestReviewFixes:
         src = pathlib.Path("src/soup_cli/commands/shrink.py").read_text(encoding="utf-8")
         assert "loss == loss" not in src
         assert "math.isnan(loss)" in src
+
+    def test_fuse_adapter_produces_dense_model(self, tmp_path):
+        """_fuse_adapter merges a LoRA adapter back into the base in place, so
+        the shipped dir is a single dense model (no adapter_config.json)."""
+        from peft import LoraConfig, TaskType, get_peft_model
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        from soup_cli.commands.shrink import _fuse_adapter
+
+        base_dir = tmp_path / "base"
+        tok = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct")
+        _tiny_llama(4, vocab_size=len(tok)).save_pretrained(str(base_dir))
+        tok.save_pretrained(str(base_dir))
+
+        adapter_dir = tmp_path / "adapter"
+        base = AutoModelForCausalLM.from_pretrained(str(base_dir))
+        peft_model = get_peft_model(
+            base,
+            LoraConfig(r=4, lora_alpha=8, target_modules=["q_proj", "v_proj"],
+                       task_type=TaskType.CAUSAL_LM),
+        )
+        peft_model.save_pretrained(str(adapter_dir))
+
+        _fuse_adapter(base_dir=str(base_dir), adapter_dir=str(adapter_dir))
+
+        # In-place overwrite yields a dense model — no adapter marker survives.
+        assert not (base_dir / "adapter_config.json").exists()
+        fused = AutoModelForCausalLM.from_pretrained(str(base_dir))
+        assert fused.config.num_hidden_layers == 4
+
+    def test_dont_ship_exit_code_2(self, tmp_path, monkeypatch):
+        """A genuine perplexity regression past tolerance exits 2 (DON'T SHIP)."""
+        from typer.testing import CliRunner
+
+        from soup_cli.cli import app
+        from soup_cli.commands import shrink as shrink_cmd
+
+        # Force a regression: original 10.0 -> pruned 20.0 (ratio 2.0 >> tol).
+        seq = iter([10.0, 20.0])
+        monkeypatch.setattr(shrink_cmd, "_perplexity", lambda *a, **k: next(seq))
+
+        monkeypatch.chdir(tmp_path)
+        model_dir = _write_tiny_model(tmp_path / "m_ds", layers=6)
+        calib = tmp_path / "calib.jsonl"
+        calib.write_text('{"text":"the quick brown fox"}\n', encoding="utf-8")
+        out_dir = tmp_path / "shrunk_ds"
+        r = CliRunner().invoke(
+            app,
+            ["shrink", "--model", model_dir, "--drop-layers", "2",
+             "--calib", "calib.jsonl", "--device", "cpu",
+             "--output-dir", str(out_dir), "--tolerance", "0.10"],
+        )
+        assert r.exit_code == 2, (r.output, repr(r.exception))