feat(shrink): arch allowlist + prune_model_layers (v0.71.29)

src/soup_cli/utils/shrink.py

@@ -17,6 +17,7 @@ This module has two halves:
 from __future__ import annotations
 
 import math
+import re
 from dataclasses import asdict, dataclass
 
 from rich.panel import Panel
@@ -131,3 +132,77 @@ def render_shrink_panel(verdict: ShrinkVerdict) -> Panel:
         f"Healed: {'yes' if verdict.healed else 'no'}"
     )
     return Panel(body, title="soup shrink", border_style=color)
+
+
+# ---------------------------------------------------------------------------
+# Arch allowlist + prune (torch-lazy)
+# ---------------------------------------------------------------------------
+_ARCH_PATTERNS = {
+    "llama": re.compile(r"llama", re.I),
+    "qwen": re.compile(r"qwen", re.I),
+    "smollm": re.compile(r"smol", re.I),
+}
+SUPPORTED_SHRINK_ARCHS = tuple(_ARCH_PATTERNS)
+
+
+def shrink_arch_of(model: object) -> str:
+    """Return the supported family name for ``model`` 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.
+    """
+    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]])
+    for family, pattern in _ARCH_PATTERNS.items():
+        if pattern.search(haystack):
+            return family
+    raise ValueError(
+        f"soup shrink v1 supports {SUPPORTED_SHRINK_ARCHS}; got "
+        f"model_type={model_type!r} (unsupported). Open an issue to add it."
+    )
+
+
+def layer_list(model: object):
+    """Return ``model.model.layers`` (the decoder ``ModuleList``), arch-guarded."""
+    shrink_arch_of(model)  # raises on unsupported arch
+    try:
+        return model.model.layers  # type: ignore[attr-defined]
+    except AttributeError as exc:
+        raise ValueError("model has no .model.layers ModuleList") from exc
+
+
+def prune_model_layers(model: object, start: int, block_size: int) -> None:
+    """Drop decoder layers ``[start, start + block_size)`` in place.
+
+    Slices ``model.model.layers`` and patches ``config.num_hidden_layers``. The
+    first and last decoder layers are protected (they carry the most residual
+    transformation, per the paper), so the dropped block must stay within
+    ``[1, num_layers - 1)``. Callers MUST reload the model from the saved dir
+    before measuring/generating — slicing leaves each surviving layer's
+    ``self_attn.layer_idx`` stale, which ``from_pretrained`` reconstructs
+    correctly.
+    """
+    import torch.nn as nn
+
+    layers = layer_list(model)
+    n_total = len(layers)
+    if not isinstance(start, int) or isinstance(start, bool):
+        raise ValueError("start must be an int")
+    if not isinstance(block_size, int) or isinstance(block_size, bool):
+        raise ValueError("block_size must be an int")
+    if block_size < 1 or block_size >= n_total:
+        raise ValueError(f"block_size must be in [1, {n_total - 1}], got {block_size}")
+    end = start + block_size  # exclusive
+    if start < 1 or end > n_total - 1:
+        raise ValueError(
+            f"dropped block [{start}, {end}) must stay within [1, {n_total - 1}) "
+            "(the first and last layer are protected)"
+        )
+    kept = [layers[i] for i in range(n_total) if not (start <= i < end)]
+    model.model.layers = nn.ModuleList(kept)  # type: ignore[attr-defined]
+    model.config.num_hidden_layers = len(kept)  # type: ignore[attr-defined]

tests/test_v07129.py

@@ -116,3 +116,76 @@ class TestNoTopLevelTorch:
         assert not any(
             m.split(".")[0] in {"torch", "transformers", "peft"} for m in names
         ), names
+
+
+# ---------------------------------------------------------------------------
+# Task 2 — arch allowlist + prune_model_layers (torch, tiny CPU model)
+# ---------------------------------------------------------------------------
+def _tiny_llama(layers: int = 6):
+    from transformers import LlamaConfig, LlamaForCausalLM
+
+    cfg = LlamaConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=layers,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        vocab_size=128,
+        max_position_embeddings=64,
+    )
+    return LlamaForCausalLM(cfg)
+
+
+class TestPrune:
+    def test_arch_detected(self):
+        from soup_cli.utils.shrink import shrink_arch_of
+
+        assert shrink_arch_of(_tiny_llama()) == "llama"
+
+    def test_arch_rejects_unsupported(self):
+        from soup_cli.utils.shrink import shrink_arch_of
+
+        class _Cfg:
+            model_type = "gpt_neox"
+            architectures = ["GPTNeoXForCausalLM"]
+
+        class _M:
+            config = _Cfg()
+
+        with pytest.raises(ValueError, match="supports"):
+            shrink_arch_of(_M())
+
+    def test_prune_removes_block_and_patches_config(self):
+        from soup_cli.utils.shrink import prune_model_layers
+
+        m = _tiny_llama(6)
+        prune_model_layers(m, start=2, block_size=2)  # drop layers 2,3
+        assert len(m.model.layers) == 4
+        assert m.config.num_hidden_layers == 4
+
+    def test_prune_rejects_touching_last_layer(self):
+        from soup_cli.utils.shrink import prune_model_layers
+
+        m = _tiny_llama(6)
+        with pytest.raises(ValueError, match="protected"):
+            prune_model_layers(m, start=4, block_size=2)  # would include last (idx 5)
+
+    def test_prune_rejects_touching_first_layer(self):
+        from soup_cli.utils.shrink import prune_model_layers
+
+        m = _tiny_llama(6)
+        with pytest.raises(ValueError, match="protected"):
+            prune_model_layers(m, start=0, block_size=2)
+
+    def test_prune_rejects_block_too_large(self):
+        from soup_cli.utils.shrink import prune_model_layers
+
+        m = _tiny_llama(6)
+        with pytest.raises(ValueError, match="block_size"):
+            prune_model_layers(m, start=1, block_size=6)
+
+    def test_layer_list_arch_guarded(self):
+        from soup_cli.utils.shrink import layer_list
+
+        m = _tiny_llama(4)
+        assert len(layer_list(m)) == 4