diff --git a/src/soup_cli/utils/shrink.py b/src/soup_cli/utils/shrink.py index f3f77c9..5d9f044 100644 --- a/src/soup_cli/utils/shrink.py +++ b/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] diff --git a/tests/test_v07129.py b/tests/test_v07129.py index 9f9610a..7cedc72 100644 --- a/tests/test_v07129.py +++ b/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