diff --git a/src/soup_cli/commands/shrink.py b/src/soup_cli/commands/shrink.py index c73783e..83175ff 100644 --- a/src/soup_cli/commands/shrink.py +++ b/src/soup_cli/commands/shrink.py @@ -407,6 +407,7 @@ def _shrink_impl( out_dir=str(adapter_dir), heal_rows=heal_rows, trc=trc, + device=dev, ) healed = True @@ -465,7 +466,10 @@ def _attach_to_registry(registry_id: str, report_path: str) -> None: # --------------------------------------------------------------------------- # Distill-heal (subprocess) + fuse # --------------------------------------------------------------------------- -_HEAL_BATCH_SIZE = 4 +# batch 1 + gradient checkpointing + a bounded max_length keep the heal within a +# consumer-GPU / CPU budget (teacher + student are both resident during distill). +_HEAL_BATCH_SIZE = 1 +_HEAL_MAX_LENGTH = 1024 _HEAL_TIMEOUT_SECONDS = 24 * 60 * 60 @@ -518,10 +522,12 @@ def _build_heal_config_yaml( "data:\n" " train: {data}\n" " format: auto\n" + " max_length: {max_length}\n" "training:\n" " teacher_model: {teacher}\n" " epochs: {epochs}\n" " batch_size: {batch}\n" + " gradient_checkpointing: true\n" " lora:\n" " r: 16\n" " alpha: 32\n" @@ -532,6 +538,7 @@ def _build_heal_config_yaml( teacher=json.dumps(teacher), epochs=epochs, batch=_HEAL_BATCH_SIZE, + max_length=_HEAL_MAX_LENGTH, ) @@ -544,13 +551,16 @@ def _run_heal( out_dir: str, heal_rows: int, trc: bool = False, + device: Optional[str] = None, ) -> 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. + artifact stays a single dense model. When ``device == "cpu"`` the subprocess + runs with ``CUDA_VISIBLE_DEVICES=""`` so the heal honours the requested + device (and sidesteps the GPU hardware-fit gate on a small card). """ import subprocess import sys @@ -569,6 +579,12 @@ def _run_heal( config_path = Path(pruned_dir).parent / "heal_config.yaml" atomic_write_text(yaml_text, str(config_path), field="heal config") + env = dict(os.environ) + if device is not None and device.lower() == "cpu": + # -1 is the canonical "hide every GPU" value; an empty string trips an + # "Invalid device id" assertion in some torch/accelerate paths. + env["CUDA_VISIBLE_DEVICES"] = "-1" + argv = [ sys.executable, "-m", @@ -580,17 +596,21 @@ def _run_heal( ] try: result = subprocess.run( # noqa: S603 — argv list, no shell. - argv, capture_output=True, check=False, timeout=_HEAL_TIMEOUT_SECONDS + argv, capture_output=True, check=False, timeout=_HEAL_TIMEOUT_SECONDS, env=env ) except subprocess.TimeoutExpired as exc: raise RuntimeError( f"heal distill exceeded {_HEAL_TIMEOUT_SECONDS}s timeout" ) from exc if result.returncode != 0: - # Strip control bytes: the child's stderr is attacker-influenceable and + # Strip control bytes: the child's output is attacker-influenceable and # reaches the terminal via the friendly error handler (escape() does not - # neutralise raw ESC/OSC sequences). - tail = _for_terminal((result.stderr or b"").decode("utf-8", "replace")[-500:]) + # neutralise raw ESC/OSC sequences). Include stdout since Rich panels + # (e.g. the hardware-fit gate) print there, not to stderr. + combined = (result.stderr or b"").decode("utf-8", "replace") + ( + result.stdout or b"" + ).decode("utf-8", "replace") + tail = _for_terminal(combined[-800:]) raise RuntimeError(f"heal distill failed (rc={result.returncode}): {tail}") _fuse_adapter(base_dir=pruned_dir, adapter_dir=out_dir, trc=trc) diff --git a/tests/test_v07129.py b/tests/test_v07129.py index 398a454..e973417 100644 --- a/tests/test_v07129.py +++ b/tests/test_v07129.py @@ -861,9 +861,10 @@ class TestReviewFixes: # TDD-review gap closure (tdd agent findings #3-#17) # --------------------------------------------------------------------------- class _StubProc: - def __init__(self, returncode=0, stderr=b""): + def __init__(self, returncode=0, stderr=b"", stdout=b""): self.returncode = returncode self.stderr = stderr + self.stdout = stdout class TestRunHeal: @@ -923,6 +924,37 @@ class TestRunHeal: sc._run_heal(pruned_dir="./model", teacher="t", heal_data="./h.jsonl", steps=5, out_dir="./adapter", heal_rows=10) + def test_device_cpu_hides_gpu_in_subprocess_env(self, tmp_path, monkeypatch): + """device='cpu' must run the heal with CUDA_VISIBLE_DEVICES=-1 so the + distill honours CPU (and dodges the GPU hardware-fit gate).""" + import soup_cli.commands.shrink as sc + + monkeypatch.chdir(tmp_path) + (tmp_path / "model").mkdir() + seen = {} + + def _capture(*a, **k): + seen["env"] = k.get("env") + return _StubProc(returncode=0) + + monkeypatch.setattr("subprocess.run", _capture) + monkeypatch.setattr(sc, "_fuse_adapter", lambda **kw: None) + sc._run_heal(pruned_dir="./model", teacher="t", heal_data="./h.jsonl", + steps=5, out_dir="./adapter", heal_rows=10, device="cpu") + assert seen["env"]["CUDA_VISIBLE_DEVICES"] == "-1" + + def test_config_has_gradient_checkpointing_and_batch1(self): + from soup_cli.commands.shrink import _build_heal_config_yaml + from soup_cli.config.loader import load_config_from_string + + cfg = load_config_from_string( + _build_heal_config_yaml(pruned_dir="./m", teacher="t", + heal_data="./h.jsonl", steps=8, out_dir="./o", + heal_rows=8) + ) + assert cfg.training.batch_size == 1 + assert cfg.training.gradient_checkpointing is True + class TestDropCountEdges: def test_ratio_rounds_to_zero_rejected(self):