RomanNum3ral
/
soup
mirror of https://github.com/razor-ai/soup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat(shrink): prune orchestration + CLI registration (v0.71.29)

This commit is contained in:
Alpamys 2026-07-05 10:34:09 +05:00
parent 76456743cd
commit 368d02be7a
3 changed files with 516 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/soup_cli/cli.py
View File

@ -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.

352
src/soup_cli/commands/shrink.py Normal file
View File

@ -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 <id|path> --drop-ratio 0.25 --calib calib.jsonl -o shrunk
soup shrink --model <id|path> --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))}")

156
tests/test_v07129.py
View File

@ -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()