From 74ba642b4099f5ff5ba3c64785b6382eea6399e5 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Sun, 5 Jul 2026 10:38:41 +0500 Subject: [PATCH] feat(shrink): subprocess distill-heal + fuse (v0.71.29) --- src/soup_cli/commands/shrink.py | 171 +++++++++++++++++++++++++++++++- tests/test_v07129.py | 96 ++++++++++++++++++ 2 files changed, 266 insertions(+), 1 deletion(-) diff --git a/src/soup_cli/commands/shrink.py b/src/soup_cli/commands/shrink.py index 64e352c..d03c8d3 100644 --- a/src/soup_cli/commands/shrink.py +++ b/src/soup_cli/commands/shrink.py @@ -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) diff --git a/tests/test_v07129.py b/tests/test_v07129.py index d094b14..dd0d3d8 100644 --- a/tests/test_v07129.py +++ b/tests/test_v07129.py @@ -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