diff --git a/src/soup_cli/commands/shrink.py b/src/soup_cli/commands/shrink.py index 0911d26..1da1d3c 100644 --- a/src/soup_cli/commands/shrink.py +++ b/src/soup_cli/commands/shrink.py @@ -128,6 +128,20 @@ def _count_params(model: object) -> int: return sum(p.numel() for p in model.parameters()) # type: ignore[attr-defined] +def _release_cuda() -> None: + """Return the CUDA caching-allocator pool to the driver (best-effort).""" + import gc + + gc.collect() + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except ImportError: + pass + + def _perplexity( model: object, tokenizer: object, prompts: Sequence[str], device: str ) -> float: @@ -166,21 +180,33 @@ def _perplexity( return math.exp(sum(losses) / len(losses)) -def _load_for_shrink( - model_id: str, device: Optional[str], trust_remote_code: bool -) -> tuple[Any, Any, str]: - """Load a model + tokenizer for shrinking (trust_remote_code probe + warn).""" - from soup_cli.utils.live_eval import load_model_and_tokenizer +def _resolve_trc(model_id: str, requested: bool) -> bool: + """Resolve trust_remote_code once (probe + warn), reused across loads.""" 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 resolve_trust_remote_code( + model_id, requested=requested, console=console, requires_remote_code=requires + ) + + +def _load_for_shrink( + model_id: str, device: Optional[str], trc: bool +) -> tuple[Any, Any, str]: + """Load a model + tokenizer for shrinking, preserving the checkpoint dtype. + + ``dtype="auto"`` keeps the model at its native precision so the shipped + smaller model is not silently upcast to fp32 (which would shrink the layer + count while widening the bytes-per-parameter). + """ + from soup_cli.utils.live_eval import load_model_and_tokenizer + + return load_model_and_tokenizer( + model_id, device=device, trust_remote_code=trc, dtype="auto" ) - return load_model_and_tokenizer(model_id, device=device, trust_remote_code=trc) def _render_importance_table( @@ -290,6 +316,7 @@ 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") + heal_rows = 0 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") @@ -298,13 +325,23 @@ def _shrink_impl( heal_rows = _count_jsonl_rows(heal) # validates cwd containment + O_NOFOLLOW 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). - shrink_arch_of(mdl) - n_layers = int(mdl.config.num_hidden_layers) + trc = _resolve_trc(model, trust_remote_code) + # Fail fast on arch + drop-count from the CONFIG before loading weights. + from transformers import AutoConfig + + from soup_cli.utils.shrink import arch_family_of_config + + pre_config = AutoConfig.from_pretrained(model, trust_remote_code=trc) + arch_family_of_config(pre_config) + n_layers = int(pre_config.num_hidden_layers) count = resolve_drop_count(n_layers, drop_ratio=drop_ratio, drop_layers=drop_layers) + console.print(f"[dim]Loading {escape(model)} ...[/]") + mdl, tokenizer, dev = _load_for_shrink(model, device, trc) + # Defence-in-depth: re-check the loaded model's arch (layer_list also + # re-guards independently before any slice). + shrink_arch_of(mdl) + console.print(f"[dim]Scoring importance over {len(prompts)} calib prompts ...[/]") importances = compute_layer_importance( mdl, tokenizer, prompts, block_size=count, device=dev @@ -335,7 +372,11 @@ def _shrink_impl( model_out.mkdir(parents=True, exist_ok=True) mdl.save_pretrained(str(model_out)) tokenizer.save_pretrained(str(model_out)) + # Free the parent's VRAM before spawning the heal subprocess (teacher + + # student both load in the child; del alone leaves the caching-allocator + # pool resident — mirrors utils/interference_live.py). del mdl + _release_cuda() healed = False if heal is not None: @@ -350,10 +391,11 @@ def _shrink_impl( steps=heal_steps, out_dir=str(adapter_dir), heal_rows=heal_rows, + trc=trc, ) healed = True - reloaded, tok2, dev2 = _load_for_shrink(str(model_out), device, trust_remote_code) + reloaded, tok2, dev2 = _load_for_shrink(str(model_out), device, trc) layers_after = int(reloaded.config.num_hidden_layers) params_after = _count_params(reloaded) ppl_final = _perplexity(reloaded, tok2, prompts, dev2) @@ -486,6 +528,7 @@ def _run_heal( steps: int, out_dir: str, heal_rows: int, + trc: bool = False, ) -> None: """Distill the teacher into the pruned student, then fuse the adapter. @@ -532,15 +575,41 @@ def _run_heal( 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) + _fuse_adapter(base_dir=pruned_dir, adapter_dir=out_dir, trc=trc) -def _fuse_adapter(*, base_dir: str, adapter_dir: str) -> None: - """Merge a LoRA adapter into ``base_dir`` in place (dense healed model).""" +def _fuse_adapter(*, base_dir: str, adapter_dir: str, trc: bool = False) -> None: + """Merge a LoRA adapter into ``base_dir`` (dense healed model), atomically. + + The merged model is written to a sibling temp dir and only then swapped in + for ``base_dir``. An in-place ``save_pretrained`` over the just-loaded + ``base_dir`` fails on Windows (error 1224 — the source ``.safetensors`` is + still memory-mapped by the loaded weights), so the temp-dir swap is the + cross-platform-safe path. + """ + import gc + import shutil + import tempfile + from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer - base = AutoModelForCausalLM.from_pretrained(base_dir) + base = AutoModelForCausalLM.from_pretrained( + base_dir, trust_remote_code=trc, torch_dtype="auto" + ) merged = PeftModel.from_pretrained(base, adapter_dir).merge_and_unload() - merged.save_pretrained(base_dir) - AutoTokenizer.from_pretrained(base_dir).save_pretrained(base_dir) + tokenizer = AutoTokenizer.from_pretrained(base_dir, trust_remote_code=trc) + + parent = os.path.dirname(os.path.abspath(base_dir)) or "." + staging = tempfile.mkdtemp(prefix=".fuse_", dir=parent) + try: + merged.save_pretrained(staging) + tokenizer.save_pretrained(staging) + finally: + # Drop every reference so Windows releases the base_dir mmap before we + # remove it; otherwise rmtree(base_dir) also hits error 1224. + del merged, base, tokenizer + gc.collect() + _release_cuda() + shutil.rmtree(base_dir) + os.replace(staging, base_dir) diff --git a/src/soup_cli/utils/live_eval.py b/src/soup_cli/utils/live_eval.py index 3ea0482..03ebfdc 100644 --- a/src/soup_cli/utils/live_eval.py +++ b/src/soup_cli/utils/live_eval.py @@ -118,11 +118,15 @@ def load_model_and_tokenizer( adapter: Optional[str] = None, device: Optional[str] = None, trust_remote_code: bool = False, + dtype: Optional[str] = None, ): """Load an ``AutoModelForCausalLM`` + tokenizer, optionally with a LoRA adapter. Returns ``(model, tokenizer, device)``. ``model`` is ``.eval()``-ed and - moved to the resolved device. Heavy imports are local. + moved to the resolved device. Heavy imports are local. ``dtype`` (e.g. + ``"auto"``) is forwarded as ``torch_dtype`` so a caller can preserve the + checkpoint's native precision instead of upcasting to fp32 (``soup shrink`` + needs this so the shipped smaller model is not silently re-widened). """ if not isinstance(model_id, str) or not model_id.strip(): raise ValueError("model_id must be a non-empty string") @@ -133,7 +137,10 @@ def load_model_and_tokenizer( tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=trust_remote_code) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=trust_remote_code) + model_kwargs = {"trust_remote_code": trust_remote_code} + if dtype is not None: + model_kwargs["torch_dtype"] = dtype + model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs) if adapter is not None: if not isinstance(adapter, str) or not adapter.strip(): raise ValueError("adapter must be a non-empty string or None") diff --git a/src/soup_cli/utils/shrink.py b/src/soup_cli/utils/shrink.py index 9f7231f..85be5cf 100644 --- a/src/soup_cli/utils/shrink.py +++ b/src/soup_cli/utils/shrink.py @@ -149,16 +149,16 @@ _ARCH_PATTERNS = { SUPPORTED_SHRINK_ARCHS = tuple(_ARCH_PATTERNS) -def shrink_arch_of(model: object) -> str: - """Return the supported family name for ``model`` or raise ``ValueError``. +def arch_family_of_config(config: object) -> str: + """Return the supported family name for an HF ``config`` or raise ``ValueError``. Detection is over ``config.model_type`` + ``config.architectures`` with regex word-family matching (mirrors ``longlora.is_*_model``). Only the v1 families in :data:`SUPPORTED_SHRINK_ARCHS` (Llama / Qwen / SmolLM — all of which expose ``model.model.layers`` + ``config.num_hidden_layers``) are - accepted; anything else is a friendly reject. + accepted; anything else is a friendly reject. Taking a bare config lets the + CLI fail fast before loading weights. """ - config = getattr(model, "config", None) model_type = getattr(config, "model_type", "") or "" architectures = list(getattr(config, "architectures", []) or []) haystack = " ".join([str(model_type), *[str(a) for a in architectures]]) @@ -171,6 +171,11 @@ def shrink_arch_of(model: object) -> str: ) +def shrink_arch_of(model: object) -> str: + """Return the supported family name for a loaded ``model`` or raise.""" + return arch_family_of_config(getattr(model, "config", None)) + + def layer_list(model: object) -> Any: """Return ``model.model.layers`` (the decoder ``ModuleList``), arch-guarded.""" shrink_arch_of(model) # raises on unsupported arch diff --git a/tests/test_v07129.py b/tests/test_v07129.py index 683c290..9617d5d 100644 --- a/tests/test_v07129.py +++ b/tests/test_v07129.py @@ -688,3 +688,56 @@ class TestReviewFixes: src = pathlib.Path("src/soup_cli/commands/shrink.py").read_text(encoding="utf-8") assert "loss == loss" not in src assert "math.isnan(loss)" in src + + def test_fuse_adapter_produces_dense_model(self, tmp_path): + """_fuse_adapter merges a LoRA adapter back into the base in place, so + the shipped dir is a single dense model (no adapter_config.json).""" + from peft import LoraConfig, TaskType, get_peft_model + from transformers import AutoModelForCausalLM, AutoTokenizer + + from soup_cli.commands.shrink import _fuse_adapter + + base_dir = tmp_path / "base" + tok = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct") + _tiny_llama(4, vocab_size=len(tok)).save_pretrained(str(base_dir)) + tok.save_pretrained(str(base_dir)) + + adapter_dir = tmp_path / "adapter" + base = AutoModelForCausalLM.from_pretrained(str(base_dir)) + peft_model = get_peft_model( + base, + LoraConfig(r=4, lora_alpha=8, target_modules=["q_proj", "v_proj"], + task_type=TaskType.CAUSAL_LM), + ) + peft_model.save_pretrained(str(adapter_dir)) + + _fuse_adapter(base_dir=str(base_dir), adapter_dir=str(adapter_dir)) + + # In-place overwrite yields a dense model — no adapter marker survives. + assert not (base_dir / "adapter_config.json").exists() + fused = AutoModelForCausalLM.from_pretrained(str(base_dir)) + assert fused.config.num_hidden_layers == 4 + + def test_dont_ship_exit_code_2(self, tmp_path, monkeypatch): + """A genuine perplexity regression past tolerance exits 2 (DON'T SHIP).""" + from typer.testing import CliRunner + + from soup_cli.cli import app + from soup_cli.commands import shrink as shrink_cmd + + # Force a regression: original 10.0 -> pruned 20.0 (ratio 2.0 >> tol). + seq = iter([10.0, 20.0]) + monkeypatch.setattr(shrink_cmd, "_perplexity", lambda *a, **k: next(seq)) + + monkeypatch.chdir(tmp_path) + model_dir = _write_tiny_model(tmp_path / "m_ds", layers=6) + calib = tmp_path / "calib.jsonl" + calib.write_text('{"text":"the quick brown fox"}\n', encoding="utf-8") + out_dir = tmp_path / "shrunk_ds" + r = CliRunner().invoke( + app, + ["shrink", "--model", model_dir, "--drop-layers", "2", + "--calib", "calib.jsonl", "--device", "cpu", + "--output-dir", str(out_dir), "--tolerance", "0.10"], + ) + assert r.exit_code == 2, (r.output, repr(r.exception))