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

feat(shrink): subprocess distill-heal + fuse (v0.71.29)

This commit is contained in:
Alpamys 2026-07-05 10:38:41 +05:00
parent 368d02be7a
commit 74ba642b40
2 changed files with 266 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

171
src/soup_cli/commands/shrink.py
View File

@ -201,6 +201,14 @@ def shrink(
calib: str = typer.Option(
..., "--calib", help="Calibration JSONL (prompts) — must stay under cwd."
),
heal: Optional[str] = typer.Option(
None,
"--heal",
help="Heal JSONL (chat rows) — distill the original into the pruned model.",
),
heal_steps: int = typer.Option(
200, "--heal-steps", help="Approx. optimiser steps for the distill heal."
),
tolerance: float = typer.Option(
0.10, "--tolerance", help="Perplexity-regression tolerance for the verdict."
),
@ -227,6 +235,8 @@ def shrink(
drop_ratio=drop_ratio,
drop_layers=drop_layers,
calib=calib,
heal=heal,
heal_steps=heal_steps,
tolerance=tolerance,
output_dir=output_dir,
device=device,
@ -247,6 +257,8 @@ def _shrink_impl(
drop_ratio: Optional[float],
drop_layers: Optional[int],
calib: str,
heal: Optional[str],
heal_steps: int,
tolerance: float,
output_dir: str,
device: Optional[str],
@ -259,6 +271,12 @@ 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")
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")
if not (1 <= heal_steps <= _MAX_HEAL_STEPS):
raise typer.BadParameter(f"--heal-steps must be in [1, {_MAX_HEAL_STEPS}]")
heal_rows = _count_jsonl_rows(heal) # validates cwd containment + O_NOFOLLOW
prompts = _load_calib(calib)
console.print(f"[dim]Loading {escape(model)} ...[/]")
@ -302,6 +320,22 @@ def _shrink_impl(
tokenizer.save_pretrained(str(model_out))
del mdl
healed = False
if heal is not None:
adapter_dir = out_root / "heal_adapter"
console.print(
f"[dim]Healing (distill original -> pruned, ~{heal_steps} steps) ...[/]"
)
_run_heal(
pruned_dir=str(model_out),
teacher=model,
heal_data=heal,
steps=heal_steps,
out_dir=str(adapter_dir),
heal_rows=heal_rows,
)
healed = True
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)
@ -317,7 +351,7 @@ def _shrink_impl(
layers_before=n_layers,
layers_after=layers_after,
params_saved_pct=params_saved_pct,
healed=False,
healed=healed,
)
console.print(render_shrink_panel(verdict))
@ -350,3 +384,138 @@ def _attach_to_registry(registry_id: str, report_path: str) -> None:
)
except Exception as exc: # noqa: BLE001
console.print(f"[yellow]Warning:[/] could not attach to registry: {escape(str(exc))}")
# ---------------------------------------------------------------------------
# Distill-heal (subprocess) + fuse
# ---------------------------------------------------------------------------
_HEAL_BATCH_SIZE = 4
_HEAL_TIMEOUT_SECONDS = 24 * 60 * 60
def _count_jsonl_rows(path: str) -> int:
"""Count non-empty lines in a JSONL file (cwd-contained, O_NOFOLLOW, capped)."""
_under_cwd(path, "heal path")
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
try:
fd = os.open(path, flags)
except OSError as exc:
raise typer.BadParameter(f"heal 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"heal file exceeds {_MAX_CALIB_BYTES} bytes")
rows = sum(1 for line in handle if line.strip())
if rows == 0:
raise typer.BadParameter("heal file has no rows")
return rows
def _build_heal_config_yaml(
*,
pruned_dir: str,
teacher: str,
heal_data: str,
steps: int,
out_dir: str,
heal_rows: int,
) -> str:
"""Render a distill ``soup.yaml`` that heals the pruned student.
``--heal-steps`` maps to epochs (there is no ``max_steps`` knob): epochs =
ceil(steps * batch / rows), clamped to >= 1, so ~``steps`` optimiser steps
run over the heal set. LoRA student (DistillTrainer always LoRA-trains) +
logit-KL distillation from the full-depth teacher.
"""
import math
epochs = max(1, math.ceil(steps * _HEAL_BATCH_SIZE / max(1, heal_rows)))
return (
"base: {pruned}\n"
"task: distill\n"
"output: {out}\n"
"data:\n"
" train: {data}\n"
" format: auto\n"
"training:\n"
" teacher_model: {teacher}\n"
" epochs: {epochs}\n"
" batch_size: {batch}\n"
" lora:\n"
" r: 16\n"
" alpha: 32\n"
).format(
pruned=json.dumps(pruned_dir),
out=json.dumps(out_dir),
data=json.dumps(heal_data),
teacher=json.dumps(teacher),
epochs=epochs,
batch=_HEAL_BATCH_SIZE,
)
def _run_heal(
*,
pruned_dir: str,
teacher: str,
heal_data: str,
steps: int,
out_dir: str,
heal_rows: int,
) -> None:
"""Distill the teacher into the pruned student, then fuse the adapter.
Writes a validated distill config, runs ``soup train`` as a subprocess
(argv list, no shell — mirrors ``ra_dit_run._run_train_subprocess``), and
merges the resulting LoRA adapter back into ``pruned_dir`` so the shipped
artifact stays a single dense model.
"""
import subprocess
import sys
from soup_cli.config.loader import load_config_from_string
yaml_text = _build_heal_config_yaml(
pruned_dir=pruned_dir,
teacher=teacher,
heal_data=heal_data,
steps=steps,
out_dir=out_dir,
heal_rows=heal_rows,
)
load_config_from_string(yaml_text) # validate before spending a subprocess
config_path = Path(pruned_dir).parent / "heal_config.yaml"
config_path.write_text(yaml_text, encoding="utf-8")
argv = [
sys.executable,
"-m",
"soup_cli.cli",
"train",
"--config",
str(config_path),
"--yes",
]
try:
result = subprocess.run( # noqa: S603 — argv list, no shell.
argv, capture_output=True, check=False, timeout=_HEAL_TIMEOUT_SECONDS
)
except subprocess.TimeoutExpired as exc:
raise RuntimeError(
f"heal distill exceeded {_HEAL_TIMEOUT_SECONDS}s timeout"
) from exc
if result.returncode != 0:
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)
def _fuse_adapter(*, base_dir: str, adapter_dir: str) -> None:
"""Merge a LoRA adapter into ``base_dir`` in place (dense healed model)."""
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
base = AutoModelForCausalLM.from_pretrained(base_dir)
merged = PeftModel.from_pretrained(base, adapter_dir).merge_and_unload()
merged.save_pretrained(base_dir)
AutoTokenizer.from_pretrained(base_dir).save_pretrained(base_dir)

96
tests/test_v07129.py
View File

@ -505,3 +505,99 @@ class TestShrinkCli:
)
assert r.exit_code != 0
assert "support" in r.output.lower() or "support" in str(r.exception).lower()
# ---------------------------------------------------------------------------
# Task 5 — subprocess distill-heal + fuse
# ---------------------------------------------------------------------------
class TestHeal:
def test_build_heal_config_parses(self):
from soup_cli.commands.shrink import _build_heal_config_yaml
from soup_cli.config.loader import load_config_from_string
y = _build_heal_config_yaml(
pruned_dir="./out/model",
teacher="orig/model",
heal_data="./heal.jsonl",
steps=200,
out_dir="./out/heal_adapter",
heal_rows=50,
)
cfg = load_config_from_string(y)
assert cfg.task == "distill"
assert cfg.training.teacher_model == "orig/model"
assert cfg.base == "./out/model"
assert cfg.output == "./out/heal_adapter"
assert cfg.training.epochs >= 1
def test_build_heal_config_epochs_scale_with_steps(self):
from soup_cli.commands.shrink import _build_heal_config_yaml
from soup_cli.config.loader import load_config_from_string
few = load_config_from_string(
_build_heal_config_yaml(
pruned_dir="./m", teacher="t", heal_data="./h.jsonl",
steps=10, out_dir="./o", heal_rows=100,
)
)
many = load_config_from_string(
_build_heal_config_yaml(
pruned_dir="./m", teacher="t", heal_data="./h.jsonl",
steps=800, out_dir="./o", heal_rows=100,
)
)
assert many.training.epochs > few.training.epochs
def test_heal_path_sets_healed(self, tmp_path, monkeypatch):
"""With --heal, the report records healed=True (subprocess + fuse are
stubbed so no real training runs)."""
import json
from typer.testing import CliRunner
from soup_cli.cli import app
from soup_cli.commands import shrink as shrink_cmd
# Stub the heavy heal step: no subprocess, no fuse (pruned model stays).
monkeypatch.setattr(shrink_cmd, "_run_heal", lambda *a, **k: None)
monkeypatch.chdir(tmp_path)
model_dir = _write_tiny_model(tmp_path / "src_model_h", layers=6)
calib = tmp_path / "calib.jsonl"
calib.write_text('{"text":"the quick brown fox jumps over"}\n', encoding="utf-8")
heal = tmp_path / "heal.jsonl"
heal.write_text(
'{"messages":[{"role":"user","content":"hi"},'
'{"role":"assistant","content":"hello"}]}\n',
encoding="utf-8",
)
out_dir = tmp_path / "shrunk_h"
r = CliRunner().invoke(
app,
["shrink", "--model", model_dir, "--drop-layers", "2",
"--calib", "calib.jsonl", "--heal", "heal.jsonl", "--heal-steps", "5",
"--device", "cpu", "--output-dir", str(out_dir), "--tolerance", "5.0"],
)
assert r.exit_code == 0, (r.output, repr(r.exception))
report = json.loads((out_dir / "shrink_report.json").read_text(encoding="utf-8"))
assert report["healed"] is True
def test_heal_outside_cwd_rejected(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_heal.jsonl"
outside.write_text('{"text":"hi"}\n', encoding="utf-8")
monkeypatch.chdir(work)
model_dir = _write_tiny_model(work / "m", layers=6)
calib = work / "calib.jsonl"
calib.write_text('{"text":"hi there friend"}\n', encoding="utf-8")
r = CliRunner().invoke(
app,
["shrink", "--model", model_dir, "--drop-layers", "2",
"--calib", "calib.jsonl", "--heal", str(outside), "--device", "cpu"],
)
assert r.exit_code != 0