diff --git a/src/soup_cli/utils/shrink.py b/src/soup_cli/utils/shrink.py index 5d9f044..a07b2a5 100644 --- a/src/soup_cli/utils/shrink.py +++ b/src/soup_cli/utils/shrink.py @@ -206,3 +206,125 @@ def prune_model_layers(model: object, start: int, block_size: int) -> None: 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] + + +# --------------------------------------------------------------------------- +# Importance scan + block selection + drop-count resolution +# --------------------------------------------------------------------------- +_IMPORTANCE_MAX_LENGTH = 512 +_DEFAULT_MAX_PROMPTS = 256 + + +def resolve_drop_count(num_layers: int, *, drop_ratio, drop_layers) -> int: + """Resolve the block **count** from exactly one of ratio / explicit count. + + ``drop_layers = round(drop_ratio * num_layers)`` when a ratio is given. + Validated against the position bound ``1 <= count <= num_layers - 2`` (the + first and last layer are always protected). + """ + if (drop_ratio is None) == (drop_layers is None): + raise ValueError("set exactly one of --drop-ratio / --drop-layers") + if drop_layers is not None: + if isinstance(drop_layers, bool) or not isinstance(drop_layers, int): + raise ValueError("drop_layers must be an int") + count = drop_layers + else: + if isinstance(drop_ratio, bool) or not isinstance(drop_ratio, (int, float)): + raise ValueError("drop_ratio must be a number") + if not (0.0 < float(drop_ratio) < 1.0): + raise ValueError("drop_ratio must be in (0, 1)") + count = round(float(drop_ratio) * num_layers) + max_count = num_layers - 2 # protect first + last + if not (1 <= count <= max_count): + raise ValueError( + f"drop count {count} out of range [1, {max_count}] for {num_layers} " + "layers (first and last layer protected)" + ) + return count + + +def compute_layer_importance( + model: object, + tokenizer: object, + prompts, + *, + block_size: int, + device: str, + max_prompts: int = _DEFAULT_MAX_PROMPTS, +) -> list[LayerImportance]: + """Rank every position-valid contiguous block by residual angular distance. + + One ``output_hidden_states=True`` forward per calibration prompt captures + every layer boundary at once (``hidden_states`` has ``num_layers + 1`` + entries; index 0 is the embedding output, index ``k`` the output of decoder + layer ``k - 1``). For a block ``[L, L + block_size)`` the residual entering + the block is ``hidden_states[L]`` and the residual leaving it is + ``hidden_states[L + block_size]``. The block importance is the mean, over + every non-pad token across the whole calib set, of the per-token angular + distance ``arccos(cos) / pi`` between those two residuals (Gromov et al.). + + Only position-valid starts ``L in [1, num_layers - block_size - 1]`` are + scored (first and last layer protected). Returns the blocks sorted by + ascending distance (safest to drop first). + """ + import torch + + n_layers = int(model.config.num_hidden_layers) # type: ignore[attr-defined] + valid_starts = list(range(1, n_layers - block_size)) # L in [1, n-bs-1] + if not valid_starts: + raise ValueError( + f"block_size {block_size} leaves no position-valid block for " + f"{n_layers} layers (first and last layer protected)" + ) + prompt_list = [p for p in prompts if isinstance(p, str) and p.strip()][:max_prompts] + if not prompt_list: + raise ValueError("calib prompts must contain at least one non-empty string") + + sums = {s: 0.0 for s in valid_starts} + counts = {s: 0 for s in valid_starts} + model.eval() # type: ignore[attr-defined] + with torch.no_grad(): + for text in prompt_list: + raw = tokenizer( + text, + return_tensors="pt", + truncation=True, + max_length=_IMPORTANCE_MAX_LENGTH, + ) + inputs = {k: v.to(device) for k, v in dict(raw).items()} + out = model(**inputs, output_hidden_states=True) # type: ignore[operator] + hidden = out.hidden_states + if len(hidden) != n_layers + 1: + raise ValueError( + f"expected {n_layers + 1} hidden states, got {len(hidden)}" + ) + mask = inputs.get("attention_mask") + for start in valid_starts: + h_in = hidden[start][0].to(torch.float32) # [seq, D] + h_out = hidden[start + block_size][0].to(torch.float32) + cos = torch.nn.functional.cosine_similarity(h_in, h_out, dim=-1) + cos = cos.clamp(-1.0, 1.0) + dist = torch.arccos(cos) / torch.pi # [seq] + if mask is not None: + keep = mask[0].to(torch.bool) + dist = dist[keep] + sums[start] += float(dist.sum().item()) + counts[start] += int(dist.numel()) + + imps = [ + LayerImportance( + start=s, + block_size=block_size, + angular_distance=(sums[s] / counts[s]) if counts[s] else float("inf"), + ) + for s in valid_starts + ] + imps.sort(key=lambda x: x.angular_distance) + return imps + + +def select_drop_block(importances) -> LayerImportance: + """Return the least-important (min angular-distance) candidate block.""" + if not importances: + raise ValueError("no importance scores to select from") + return min(importances, key=lambda x: x.angular_distance) diff --git a/tests/test_v07129.py b/tests/test_v07129.py index 7cedc72..50b61bd 100644 --- a/tests/test_v07129.py +++ b/tests/test_v07129.py @@ -189,3 +189,169 @@ class TestPrune: m = _tiny_llama(4) assert len(layer_list(m)) == 4 + + +# --------------------------------------------------------------------------- +# Task 3 — importance scan (off-by-one pinned) + selection + drop count +# --------------------------------------------------------------------------- +class TestImportance: + def test_off_by_one_boundary_indices(self): + """Pin: block [L, L+n) uses hidden_states[L] and hidden_states[L+n]; + hidden_states has num_layers+1 entries. A fake model returns per-layer + constant hidden states so the boundary maths is asserted, not just that + the pass runs.""" + import torch + + from soup_cli.utils import shrink + + num_layers = 4 + # len must be num_layers+1; index 0 = embeddings, index k = layer k-1 out. + hs = tuple(torch.ones(1, 3, 8) * (k + 1) for k in range(num_layers + 1)) + + class _Cfg: + model_type = "llama" + architectures = ["LlamaForCausalLM"] + num_hidden_layers = num_layers + + class _Out: + hidden_states = hs + + class _Model: + config = _Cfg() + + def eval(self): + return self + + def __call__(self, **kw): + assert kw.get("output_hidden_states") is True + return _Out() + + class _Tok: + def __call__(self, text, **kw): + return { + "input_ids": torch.ones(1, 3, dtype=torch.long), + "attention_mask": torch.ones(1, 3, dtype=torch.long), + } + + imps = shrink.compute_layer_importance( + _Model(), _Tok(), ["hi"], block_size=1, device="cpu" + ) + # valid starts for n=1, num_layers=4: L in [1, 4-1-1=2] -> {1, 2} + starts = sorted(i.start for i in imps) + assert starts == [1, 2] + # constant (colinear) vectors -> cos == 1 -> angular distance 0. + assert all(abs(i.angular_distance) < 1e-6 for i in imps) + assert all(i.block_size == 1 for i in imps) + + def test_hidden_states_length_mismatch_raises(self): + import torch + + from soup_cli.utils import shrink + + class _Cfg: + model_type = "llama" + architectures = ["LlamaForCausalLM"] + num_hidden_layers = 4 + + class _Out: + hidden_states = tuple(torch.ones(1, 3, 8) for _ in range(3)) # wrong len + + class _Model: + config = _Cfg() + + def eval(self): + return self + + def __call__(self, **kw): + return _Out() + + class _Tok: + def __call__(self, text, **kw): + return { + "input_ids": torch.ones(1, 3, dtype=torch.long), + "attention_mask": torch.ones(1, 3, dtype=torch.long), + } + + with pytest.raises(ValueError, match="hidden states"): + shrink.compute_layer_importance( + _Model(), _Tok(), ["hi"], block_size=1, device="cpu" + ) + + def test_select_drop_block_min(self): + from soup_cli.utils.shrink import LayerImportance, select_drop_block + + imps = [ + LayerImportance(1, 2, 0.9), + LayerImportance(3, 2, 0.1), + LayerImportance(5, 2, 0.5), + ] + chosen = select_drop_block(imps) + assert chosen.start == 3 and chosen.angular_distance == 0.1 + + def test_select_drop_block_empty_raises(self): + from soup_cli.utils.shrink import select_drop_block + + with pytest.raises(ValueError): + select_drop_block([]) + + def test_resolve_drop_count_ratio(self): + from soup_cli.utils.shrink import resolve_drop_count + + assert resolve_drop_count(30, drop_ratio=0.25, drop_layers=None) == 8 # round(7.5) + assert resolve_drop_count(30, drop_ratio=None, drop_layers=6) == 6 + + def test_resolve_drop_count_rejects_both_or_neither(self): + from soup_cli.utils.shrink import resolve_drop_count + + with pytest.raises(ValueError, match="exactly one"): + resolve_drop_count(30, drop_ratio=0.25, drop_layers=6) + with pytest.raises(ValueError, match="exactly one"): + resolve_drop_count(30, drop_ratio=None, drop_layers=None) + + def test_resolve_drop_count_position_bound(self): + from soup_cli.utils.shrink import resolve_drop_count + + with pytest.raises(ValueError, match="range"): + resolve_drop_count(4, drop_ratio=None, drop_layers=3) # > num_layers-2 + + def test_resolve_drop_count_ratio_bounds(self): + from soup_cli.utils.shrink import resolve_drop_count + + with pytest.raises(ValueError): + resolve_drop_count(30, drop_ratio=1.5, drop_layers=None) + with pytest.raises(ValueError): + resolve_drop_count(30, drop_ratio=0.0, drop_layers=None) + + def test_compute_importance_no_valid_starts_raises(self): + import torch + + from soup_cli.utils import shrink + + class _Cfg: + model_type = "llama" + architectures = ["LlamaForCausalLM"] + num_hidden_layers = 4 + + class _Out: + hidden_states = tuple(torch.ones(1, 3, 8) for _ in range(5)) + + class _Model: + config = _Cfg() + + def eval(self): + return self + + def __call__(self, **kw): + return _Out() + + class _Tok: + def __call__(self, text, **kw): + return { + "input_ids": torch.ones(1, 3, dtype=torch.long), + "attention_mask": torch.ones(1, 3, dtype=torch.long), + } + + with pytest.raises(ValueError, match="position-valid"): + shrink.compute_layer_importance( + _Model(), _Tok(), ["hi"], block_size=3, device="cpu" + )