From 368d02be7aade5877a9b7d184f9a6fa2d562a95f Mon Sep 17 00:00:00 2001 From: Alpamys Date: Sun, 5 Jul 2026 10:34:09 +0500 Subject: [PATCH] feat(shrink): prune orchestration + CLI registration (v0.71.29) --- src/soup_cli/cli.py | 11 + src/soup_cli/commands/shrink.py | 352 ++++++++++++++++++++++++++++++++ tests/test_v07129.py | 156 +++++++++++++- 3 files changed, 516 insertions(+), 3 deletions(-) create mode 100644 src/soup_cli/commands/shrink.py diff --git a/src/soup_cli/cli.py b/src/soup_cli/cli.py index bbf0cb3..b25a61c 100644 --- a/src/soup_cli/cli.py +++ b/src/soup_cli/cli.py @@ -527,6 +527,17 @@ app.add_typer( help="Model Context Protocol server - drive Soup from any MCP client (v0.71.28).", ) +# v0.71.29 — `soup shrink` depth-prune + distill-heal. +from soup_cli.commands import shrink as _shrink_cmd # noqa: E402 + +app.command( + name="shrink", + help=( + "Depth-prune a model (drop the least-important contiguous layer block " + "by residual angular distance) + optional distill-heal (v0.71.29)." + ), +)(_shrink_cmd.shrink) + def _rewrite_advise_argv(argv: list) -> list: """Inject `run` between `advise` and a non-subcommand first argument. diff --git a/src/soup_cli/commands/shrink.py b/src/soup_cli/commands/shrink.py new file mode 100644 index 0000000..64e352c --- /dev/null +++ b/src/soup_cli/commands/shrink.py @@ -0,0 +1,352 @@ +"""soup shrink — depth-prune + distill-heal (v0.71.29, arXiv:2403.17887). + +Top-level CLI command (NOT a sub-group). Ranks a model's decoder layers by the +angular distance of the residual stream across a contiguous block over a +calibration set, drops the least-important block, optionally distill-heals, and +emits a single dense smaller model with a before/after perplexity verdict:: + + soup shrink --model --drop-ratio 0.25 --calib calib.jsonl -o shrunk + soup shrink --model --drop-layers 6 --calib calib.jsonl \ + --heal heal.jsonl --heal-steps 200 -o shrunk + +Exit codes: 0 = SHIP, 2 = DON'T SHIP, 1 = runtime error (mirrors soup ship / +soup diagnose). Heavy imports (torch/transformers) are lazy inside functions. +""" +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Optional + +import typer +from rich.console import Console +from rich.markup import escape +from rich.panel import Panel +from rich.table import Table + +from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink +from soup_cli.utils.shrink import ( + DECISION_SHIP, + compute_layer_importance, + decide_shrink, + prune_model_layers, + render_shrink_panel, + resolve_drop_count, + select_drop_block, + shrink_verdict_to_dict, +) + +console = Console() + +# 64 MiB cap on the calibration JSONL (symlink-pointed-to-/dev/zero DoS guard). +_MAX_CALIB_BYTES = 64 * 1024 * 1024 +_MAX_CALIB_ROWS = 10_000 +_MAX_HEAL_STEPS = 1_000_000 +_PPL_MAX_LENGTH = 512 + + +# --------------------------------------------------------------------------- +# Path + data helpers +# --------------------------------------------------------------------------- +def _under_cwd(path: str, label: str) -> str: + """Validate ``path`` stays under cwd and is not a symlink; return it.""" + enforce_under_cwd_and_no_symlink(path, label) + return path + + +def _extract_text(row: object) -> str: + """Best-effort prompt text from a calib row (text / prompt / messages).""" + if isinstance(row, str): + return row + if isinstance(row, dict): + for key in ("text", "prompt", "content", "instruction"): + val = row.get(key) + if isinstance(val, str) and val.strip(): + return val + messages = row.get("messages") + if isinstance(messages, list): + parts = [ + m.get("content", "") + for m in messages + if isinstance(m, dict) and isinstance(m.get("content"), str) + ] + joined = "\n".join(p for p in parts if p.strip()) + if joined.strip(): + return joined + return "" + + +def _load_calib(path: str) -> list[str]: + """Load calibration prompts from a JSONL file (cwd-contained, size-capped). + + Opens with ``O_NOFOLLOW`` and fstats the open fd (TOCTOU defence, mirrors + ``commands/diagnose.py::_load_evidence``). Each non-empty line is a JSON + object; the prompt text is extracted via :func:`_extract_text`. + """ + _under_cwd(path, "calib path") + flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) + try: + fd = os.open(path, flags) + except OSError as exc: + raise typer.BadParameter(f"calib path unreadable: {exc}") from exc + with os.fdopen(fd, "r", encoding="utf-8") as handle: + if os.fstat(handle.fileno()).st_size > _MAX_CALIB_BYTES: + raise typer.BadParameter( + f"calib file exceeds {_MAX_CALIB_BYTES} bytes" + ) + prompts: list[str] = [] + for line in handle: + line = line.strip() + if not line: + continue + try: + row = json.loads(line) + except json.JSONDecodeError: + # Tolerate a raw-text line (not JSON) as a plain prompt. + row = line + text = _extract_text(row) + if text.strip(): + prompts.append(text) + if len(prompts) >= _MAX_CALIB_ROWS: + break + if not prompts: + raise typer.BadParameter("calib file yielded no usable prompt text") + return prompts + + +# --------------------------------------------------------------------------- +# Model helpers (torch-lazy) +# --------------------------------------------------------------------------- +def _count_params(model: object) -> int: + return sum(p.numel() for p in model.parameters()) # type: ignore[attr-defined] + + +def _perplexity(model: object, tokenizer: object, prompts: list[str], device: str) -> float: + """Mean unconditional-LM perplexity of ``model`` over ``prompts``. + + ``exp(mean per-example cross-entropy)`` with ``labels = input_ids`` (the + whole sequence is the target). Returns ``inf`` when no example is usable. + """ + import math + + import torch + + losses: list[float] = [] + model.eval() # type: ignore[attr-defined] + with torch.no_grad(): + for text in prompts: + enc = tokenizer( # type: ignore[operator] + text, + return_tensors="pt", + truncation=True, + max_length=_PPL_MAX_LENGTH, + ) + input_ids = enc["input_ids"].to(device) + if input_ids.shape[1] < 2: + continue + out = model(input_ids=input_ids, labels=input_ids) # type: ignore[operator] + loss = float(out.loss.item()) + if loss == loss: # not NaN + losses.append(loss) + if not losses: + return float("inf") + return math.exp(sum(losses) / len(losses)) + + +def _load_for_shrink(model_id: str, device: Optional[str], trust_remote_code: bool): + """Load a model + tokenizer for shrinking (trust_remote_code probe + warn).""" + from soup_cli.utils.live_eval import load_model_and_tokenizer + 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 load_model_and_tokenizer(model_id, device=device, trust_remote_code=trc) + + +def _render_importance_table(importances, chosen) -> None: + table = Table(title="soup shrink — layer importance (lower = safer to drop)") + table.add_column("Rank", justify="right") + table.add_column("Block (start..end)") + table.add_column("Angular distance", justify="right") + table.add_column("Chosen") + for rank, imp in enumerate(importances, start=1): + end = imp.start + imp.block_size + is_chosen = imp.start == chosen.start + table.add_row( + str(rank), + f"{imp.start}..{end - 1}", + f"{imp.angular_distance:.4f}", + "<-- drop" if is_chosen else "", + ) + console.print(table) + + +# --------------------------------------------------------------------------- +# The command +# --------------------------------------------------------------------------- +def shrink( + model: str = typer.Option(..., "--model", help="Model id or local path to shrink."), + drop_ratio: Optional[float] = typer.Option( + None, "--drop-ratio", help="Fraction of layers to drop (0-1); e.g. 0.25." + ), + drop_layers: Optional[int] = typer.Option( + None, "--drop-layers", help="Explicit number of contiguous layers to drop." + ), + calib: str = typer.Option( + ..., "--calib", help="Calibration JSONL (prompts) — must stay under cwd." + ), + tolerance: float = typer.Option( + 0.10, "--tolerance", help="Perplexity-regression tolerance for the verdict." + ), + output_dir: str = typer.Option( + "./shrunk", "--output-dir", "-o", help="Directory for the shrunk model + report." + ), + device: Optional[str] = typer.Option( + None, "--device", help="Device for the importance/ppl passes (cuda / cpu)." + ), + trust_remote_code: bool = typer.Option( + False, "--trust-remote-code", help="Allow custom modeling code (auto_map)." + ), + attach_to_registry: Optional[str] = typer.Option( + None, "--attach-to-registry", help="Attach the shrink report to a registry entry id." + ), + plan_only: bool = typer.Option( + False, "--plan-only", help="Print the importance table + chosen block and exit." + ), +) -> None: + """Depth-prune a model (least-important contiguous block) + verdict.""" + try: + _shrink_impl( + model=model, + drop_ratio=drop_ratio, + drop_layers=drop_layers, + calib=calib, + tolerance=tolerance, + output_dir=output_dir, + device=device, + trust_remote_code=trust_remote_code, + attach_to_registry=attach_to_registry, + plan_only=plan_only, + ) + except typer.Exit: + raise + except (typer.BadParameter, ValueError) as exc: + console.print(f"[red]Error:[/] {escape(str(exc))}") + raise typer.Exit(1) from exc + + +def _shrink_impl( + *, + model: str, + drop_ratio: Optional[float], + drop_layers: Optional[int], + calib: str, + tolerance: float, + output_dir: str, + device: Optional[str], + trust_remote_code: bool, + attach_to_registry: Optional[str], + plan_only: bool, +) -> None: + if not (0.0 <= tolerance <= 5.0): + raise typer.BadParameter("--tolerance must be in [0.0, 5.0]") + # 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") + 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). + from soup_cli.utils.shrink import shrink_arch_of + + shrink_arch_of(mdl) + n_layers = int(mdl.config.num_hidden_layers) + count = resolve_drop_count(n_layers, drop_ratio=drop_ratio, drop_layers=drop_layers) + + console.print(f"[dim]Scoring importance over {len(prompts)} calib prompts ...[/]") + importances = compute_layer_importance( + mdl, tokenizer, prompts, block_size=count, device=dev + ) + chosen = select_drop_block(importances) + _render_importance_table(importances, chosen) + + if plan_only: + end = chosen.start + chosen.block_size + console.print( + Panel.fit( + f"Would drop layers [bold]{chosen.start}..{end - 1}[/] " + f"({count} of {n_layers}); angular distance " + f"{chosen.angular_distance:.4f}", + title="plan only", + ) + ) + raise typer.Exit(0) + + params_before = _count_params(mdl) + ppl_original = _perplexity(mdl, tokenizer, prompts, dev) + + # Prune in memory, save, then RELOAD (slicing leaves layer_idx stale; + # from_pretrained rebuilds them contiguously — measure on the shipped dir). + prune_model_layers(mdl, chosen.start, chosen.block_size) + out_root = Path(output_dir) + model_out = out_root / "model" + model_out.mkdir(parents=True, exist_ok=True) + mdl.save_pretrained(str(model_out)) + tokenizer.save_pretrained(str(model_out)) + del mdl + + reloaded, tok2, dev2 = _load_for_shrink(str(model_out), device, trust_remote_code) + layers_after = int(reloaded.config.num_hidden_layers) + params_after = _count_params(reloaded) + ppl_final = _perplexity(reloaded, tok2, prompts, dev2) + params_saved_pct = ( + 100.0 * (params_before - params_after) / params_before if params_before else 0.0 + ) + + verdict = decide_shrink( + ppl_original, + ppl_final, + tolerance=tolerance, + layers_before=n_layers, + layers_after=layers_after, + params_saved_pct=params_saved_pct, + healed=False, + ) + console.print(render_shrink_panel(verdict)) + + report_path = out_root / "shrink_report.json" + report = shrink_verdict_to_dict(verdict) + report["model"] = model + report["dropped_block"] = [chosen.start, chosen.start + chosen.block_size - 1] + report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") + console.print(f"[green]Wrote[/] {escape(str(report_path))}") + + if attach_to_registry: + _attach_to_registry(attach_to_registry, str(report_path)) + + raise typer.Exit(0 if verdict.decision == DECISION_SHIP else 2) + + +def _attach_to_registry(registry_id: str, report_path: str) -> None: + """Attach the shrink report JSON as a registry artifact (best-effort).""" + try: + from soup_cli.registry.attach import attach_artifact + except Exception as exc: # noqa: BLE001 — registry is optional + console.print( + f"[yellow]Warning:[/] could not import registry attach helper: {escape(str(exc))}" + ) + return + try: + attach_artifact(registry_id, "shrink_report", report_path) + console.print( + f"[green]Attached[/] shrink_report to registry entry [bold]{escape(registry_id)}[/]" + ) + except Exception as exc: # noqa: BLE001 + console.print(f"[yellow]Warning:[/] could not attach to registry: {escape(str(exc))}") diff --git a/tests/test_v07129.py b/tests/test_v07129.py index 50b61bd..d094b14 100644 --- a/tests/test_v07129.py +++ b/tests/test_v07129.py @@ -121,7 +121,7 @@ class TestNoTopLevelTorch: # --------------------------------------------------------------------------- # Task 2 — arch allowlist + prune_model_layers (torch, tiny CPU model) # --------------------------------------------------------------------------- -def _tiny_llama(layers: int = 6): +def _tiny_llama(layers: int = 6, vocab_size: int = 128): from transformers import LlamaConfig, LlamaForCausalLM cfg = LlamaConfig( @@ -130,8 +130,8 @@ def _tiny_llama(layers: int = 6): num_hidden_layers=layers, num_attention_heads=4, num_key_value_heads=4, - vocab_size=128, - max_position_embeddings=64, + vocab_size=vocab_size, + max_position_embeddings=512, ) return LlamaForCausalLM(cfg) @@ -355,3 +355,153 @@ class TestImportance: shrink.compute_layer_importance( _Model(), _Tok(), ["hi"], block_size=3, device="cpu" ) + + +# --------------------------------------------------------------------------- +# Task 4 — commands/shrink.py prune orchestration + CLI registration +# --------------------------------------------------------------------------- +def _write_tiny_model(dir_path, layers: int = 6): + """Save a tiny CPU Llama + tokenizer to ``dir_path`` for CLI smoke.""" + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct") + m = _tiny_llama(layers, vocab_size=len(tok)) + m.save_pretrained(str(dir_path)) + tok.save_pretrained(str(dir_path)) + return str(dir_path) + + +class TestShrinkCli: + def test_registered_and_help(self): + from typer.testing import CliRunner + + from soup_cli.cli import app + + r = CliRunner().invoke(app, ["shrink", "--help"]) + assert r.exit_code == 0, (r.output, repr(r.exception)) + assert "drop-ratio" in r.output + assert "calib" in r.output + + def test_rejects_both_drop_flags(self, tmp_path, monkeypatch): + from typer.testing import CliRunner + + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + calib = tmp_path / "c.jsonl" + calib.write_text('{"text":"hello world"}\n', encoding="utf-8") + r = CliRunner().invoke( + app, + ["shrink", "--model", "x", "--drop-ratio", "0.25", "--drop-layers", + "2", "--calib", "c.jsonl"], + ) + assert r.exit_code != 0 + assert "exactly one" in r.output.lower() or "exactly one" in str(r.exception).lower() + + def test_rejects_calib_outside_cwd(self, tmp_path, monkeypatch): + from typer.testing import CliRunner + + from soup_cli.cli import app + + work = tmp_path / "work" + work.mkdir() + outside = tmp_path / "outside.jsonl" + outside.write_text('{"text":"hi"}\n', encoding="utf-8") + monkeypatch.chdir(work) + r = CliRunner().invoke( + app, + ["shrink", "--model", "x", "--drop-layers", "2", "--calib", str(outside)], + ) + assert r.exit_code != 0 + + def test_rejects_bad_tolerance(self, tmp_path, monkeypatch): + from typer.testing import CliRunner + + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + calib = tmp_path / "c.jsonl" + calib.write_text('{"text":"hi"}\n', encoding="utf-8") + r = CliRunner().invoke( + app, + ["shrink", "--model", "x", "--drop-layers", "2", "--calib", "c.jsonl", + "--tolerance", "9.0"], + ) + assert r.exit_code != 0 + + def test_prune_happy_path_cpu(self, tmp_path, monkeypatch): + """End-to-end prune (no heal) on a tiny CPU Llama: pruned config has + fewer layers, report JSON written, exit 0 (SHIP) — tolerance wide.""" + import json + + from typer.testing import CliRunner + + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + model_dir = _write_tiny_model(tmp_path / "src_model", layers=6) + calib = tmp_path / "calib.jsonl" + calib.write_text( + "\n".join('{"text":"the quick brown fox jumps over the lazy dog"}' + for _ in range(4)), + encoding="utf-8", + ) + out_dir = tmp_path / "shrunk" + r = CliRunner().invoke( + app, + ["shrink", "--model", model_dir, "--drop-layers", "2", + "--calib", "calib.jsonl", "--device", "cpu", + "--output-dir", str(out_dir), "--tolerance", "5.0"], + ) + assert r.exit_code == 0, (r.output, repr(r.exception)) + cfg = json.loads((out_dir / "model" / "config.json").read_text(encoding="utf-8")) + assert cfg["num_hidden_layers"] == 4 + report = json.loads((out_dir / "shrink_report.json").read_text(encoding="utf-8")) + assert report["layers_before"] == 6 and report["layers_after"] == 4 + assert report["healed"] is False + assert "ppl_original" in report + + def test_plan_only_writes_nothing(self, tmp_path, monkeypatch): + from typer.testing import CliRunner + + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + model_dir = _write_tiny_model(tmp_path / "src_model2", layers=6) + calib = tmp_path / "calib.jsonl" + calib.write_text('{"text":"the quick brown fox jumps"}\n', encoding="utf-8") + out_dir = tmp_path / "shrunk2" + r = CliRunner().invoke( + app, + ["shrink", "--model", model_dir, "--drop-layers", "2", + "--calib", "calib.jsonl", "--device", "cpu", + "--output-dir", str(out_dir), "--plan-only"], + ) + assert r.exit_code == 0, (r.output, repr(r.exception)) + assert not (out_dir / "model").exists() + + def test_reject_unsupported_arch(self, tmp_path, monkeypatch): + """A GPT-NeoX-family tiny model is a friendly reject (arch allowlist).""" + from transformers import AutoTokenizer, GPTNeoXConfig, GPTNeoXForCausalLM + from typer.testing import CliRunner + + from soup_cli.cli import app + + monkeypatch.chdir(tmp_path) + tok = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct") + cfg = GPTNeoXConfig( + hidden_size=32, intermediate_size=64, num_hidden_layers=6, + num_attention_heads=4, vocab_size=len(tok), max_position_embeddings=512, + ) + mdir = tmp_path / "neox" + GPTNeoXForCausalLM(cfg).save_pretrained(str(mdir)) + tok.save_pretrained(str(mdir)) + calib = tmp_path / "calib.jsonl" + calib.write_text('{"text":"hi there"}\n', encoding="utf-8") + r = CliRunner().invoke( + app, + ["shrink", "--model", str(mdir), "--drop-layers", "2", + "--calib", "calib.jsonl", "--device", "cpu"], + ) + assert r.exit_code != 0 + assert "support" in r.output.lower() or "support" in str(r.exception).lower()