Upload 6 files

examples/deepseek_7b.py

@@ -0,0 +1,25 @@
+from datasets import load_dataset
+from layer_interactions import Order2Pruner
+
+MODEL = "deepseek-ai/deepseek-llm-7b-base"
+
+texts = [
+    row["text"]
+    for row in load_dataset(
+        "Salesforce/wikitext", "wikitext-2-raw-v1", split="train"
+    )
+    if row["text"].strip()
+]
+
+with Order2Pruner(MODEL, dtype="bfloat16") as pruner:
+    result = pruner.fit(
+        texts=texts,
+        n_sequences=32,
+        sequence_length=128,
+        seed=42,
+        checkpoint_path="deepseek_order2.json",
+        max_delete=23,  # enough for 30 -> 7
+    )
+
+    print("30 -> 15:", result.select(15))
+    print("30 -> 7:", result.select(7))

examples/olmoe_7b.py

@@ -0,0 +1,28 @@
+from datasets import load_dataset
+from layer_interactions import Order2Pruner
+
+MODEL = "allenai/OLMoE-1B-7B-0924"
+
+texts = [
+    row["text"]
+    for row in load_dataset(
+        "Salesforce/wikitext", "wikitext-2-raw-v1", split="train"
+    )
+    if row["text"].strip()
+]
+
+with Order2Pruner(MODEL, dtype="bfloat16") as pruner:
+    result = pruner.fit(
+        texts=texts,
+        n_sequences=32,
+        sequence_length=128,
+        seed=42,
+        checkpoint_path="olmoe_order2.json",
+        max_delete=12,  # enough for 16 -> 4
+    )
+
+    print("16 -> 8:", result.select(8))
+    print("16 -> 4:", result.select(4))
+
+    # Optional: materialize and save one compressed checkpoint.
+    # pruner.save_pruned("olmoe_order2_8", 8, result=result)

layer_interactions/__init__.py

@@ -0,0 +1,5 @@
+from .pruner import Order2Pruner
+from .result import Order2Result
+
+__all__ = ["Order2Pruner", "Order2Result"]
+__version__ = "0.1.0"

layer_interactions/pruner.py

@@ -0,0 +1,456 @@
+from __future__ import annotations
+
+import gc
+import itertools
+import json
+import os
+from pathlib import Path
+import random
+import shutil
+import tempfile
+from typing import Iterable, Sequence
+
+import psutil
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase
+
+from .result import Order2Result
+
+
+_LAYER_PATHS = (
+    "model.layers",       # Llama/Mistral/DeepSeek/OLMoE and compatible models
+    "transformer.h",      # GPT-2 style
+    "gpt_neox.layers",    # GPT-NeoX style
+)
+
+
+def _resolve_dtype(dtype: str | torch.dtype) -> torch.dtype:
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    table = {
+        "bfloat16": torch.bfloat16,
+        "bf16": torch.bfloat16,
+        "float16": torch.float16,
+        "fp16": torch.float16,
+        "float32": torch.float32,
+        "fp32": torch.float32,
+    }
+    try:
+        return table[dtype.lower()]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported dtype: {dtype}") from exc
+
+
+def _get_attr_path(obj, path: str):
+    for part in path.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def _set_attr_path(obj, path: str, value) -> None:
+    parts = path.split(".")
+    parent = obj
+    for part in parts[:-1]:
+        parent = getattr(parent, part)
+    setattr(parent, parts[-1], value)
+
+
+class Order2Pruner:
+    """Second-order interaction layer pruning for Hugging Face causal LMs.
+
+    The method measures the NLL after deleting each layer and each pair of
+    layers, forms a quadratic deletion-loss surrogate, then greedily deletes
+    the layer with minimum current marginal predicted NLL increase.
+    """
+
+    def __init__(
+        self,
+        model: str | Path | PreTrainedModel,
+        tokenizer: str | Path | PreTrainedTokenizerBase | None = None,
+        *,
+        dtype: str | torch.dtype = "bfloat16",
+        layer_path: str | None = None,
+        device_map: str | dict | None = "auto",
+        gpu_memory_gib: int | None = None,
+        cpu_memory_gib: int | None = None,
+        activation_headroom_gib: float = 4.0,
+        offload_dir: str | Path | None = None,
+        local_files_only: bool = False,
+        trust_remote_code: bool = False,
+    ) -> None:
+        self.model_source = model
+        self.tokenizer_source = tokenizer
+        self.dtype = _resolve_dtype(dtype)
+        self.layer_path = layer_path
+        self.device_map = device_map
+        self.gpu_memory_gib = gpu_memory_gib
+        self.cpu_memory_gib = cpu_memory_gib
+        self.activation_headroom_gib = float(activation_headroom_gib)
+        self.local_files_only = local_files_only
+        self.trust_remote_code = trust_remote_code
+        self.model: PreTrainedModel | None = model if isinstance(model, PreTrainedModel) else None
+        self.tokenizer: PreTrainedTokenizerBase | None = (
+            tokenizer if isinstance(tokenizer, PreTrainedTokenizerBase) else None
+        )
+        self._owns_offload_dir = offload_dir is None
+        self.offload_dir = Path(offload_dir) if offload_dir else Path(
+            tempfile.mkdtemp(prefix="layer_interactions_offload_")
+        )
+        self.offload_dir.mkdir(parents=True, exist_ok=True)
+        self._resolved_layer_path: str | None = None
+        self.result: Order2Result | None = None
+
+    @property
+    def depth(self) -> int:
+        self._ensure_loaded()
+        return len(self._layers())
+
+    def _auto_max_memory(self) -> dict | None:
+        if not torch.cuda.is_available():
+            return None
+
+        free_gib = torch.cuda.mem_get_info()[0] / 2**30
+        gpu = self.gpu_memory_gib
+        if gpu is None:
+            gpu = max(1, int(free_gib - self.activation_headroom_gib))
+
+        cpu_free_gib = psutil.virtual_memory().available / 2**30
+        cpu = self.cpu_memory_gib
+        if cpu is None:
+            cpu = max(2, int(cpu_free_gib - 4.0))
+
+        return {0: f"{gpu}GiB", "cpu": f"{cpu}GiB"}
+
+    def _ensure_loaded(self) -> None:
+        if self.tokenizer is None:
+            source = self.tokenizer_source or self.model_source
+            if isinstance(source, PreTrainedModel):
+                raise ValueError("Pass a tokenizer when model is an already-instantiated model.")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                source,
+                local_files_only=self.local_files_only,
+                trust_remote_code=self.trust_remote_code,
+                use_fast=True,
+            )
+            if self.tokenizer.pad_token_id is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        if self.model is None:
+            kwargs = dict(
+                local_files_only=self.local_files_only,
+                trust_remote_code=self.trust_remote_code,
+                dtype=self.dtype,
+                low_cpu_mem_usage=True,
+            )
+            if self.device_map is not None:
+                kwargs["device_map"] = self.device_map
+            max_memory = self._auto_max_memory()
+            if max_memory is not None and self.device_map == "auto":
+                kwargs.update(
+                    max_memory=max_memory,
+                    offload_folder=str(self.offload_dir),
+                    offload_state_dict=True,
+                    offload_buffers=True,
+                )
+            self.model = AutoModelForCausalLM.from_pretrained(self.model_source, **kwargs)
+
+        self.model.config.use_cache = False
+        if hasattr(self.model, "generation_config"):
+            self.model.generation_config.use_cache = False
+        self.model.eval()
+        self._resolve_layer_path()
+
+    def _resolve_layer_path(self) -> str:
+        if self._resolved_layer_path is not None:
+            return self._resolved_layer_path
+        self._ensure_model_exists_for_resolution()
+        candidates = (self.layer_path,) if self.layer_path else _LAYER_PATHS
+        for path in candidates:
+            if path is None:
+                continue
+            try:
+                value = _get_attr_path(self.model, path)
+            except AttributeError:
+                continue
+            if isinstance(value, (torch.nn.ModuleList, list, tuple)):
+                self._resolved_layer_path = path
+                return path
+        raise ValueError(
+            "Could not find transformer layers automatically. "
+            "Pass layer_path, e.g. layer_path='model.layers'."
+        )
+
+    def _ensure_model_exists_for_resolution(self) -> None:
+        if self.model is None:
+            raise RuntimeError("Model has not been loaded.")
+
+    def _layers(self):
+        self._ensure_model_exists_for_resolution()
+        path = self._resolved_layer_path or self._resolve_layer_path()
+        return _get_attr_path(self.model, path)
+
+    def _set_layers(self, layers: Sequence[torch.nn.Module]) -> None:
+        self._ensure_model_exists_for_resolution()
+        path = self._resolved_layer_path or self._resolve_layer_path()
+        _set_attr_path(self.model, path, torch.nn.ModuleList(list(layers)))
+
+        for k, block in enumerate(self._layers()):
+            if hasattr(block, "layer_idx"):
+                block.layer_idx = k
+            if hasattr(block, "self_attn") and hasattr(block.self_attn, "layer_idx"):
+                block.self_attn.layer_idx = k
+
+        if hasattr(self.model.config, "num_hidden_layers"):
+            self.model.config.num_hidden_layers = len(layers)
+
+    def calibration_batches(
+        self,
+        texts: Sequence[str],
+        *,
+        n_sequences: int = 32,
+        sequence_length: int = 128,
+        seed: int = 42,
+        tokenizer_chunk_size: int = 4096,
+    ) -> list[torch.Tensor]:
+        self._ensure_loaded()
+        joined = "\n\n".join(str(x) for x in texts if str(x).strip())
+        enc = self.tokenizer(
+            joined,
+            add_special_tokens=False,
+            truncation=True,
+            max_length=tokenizer_chunk_size,
+            return_overflowing_tokens=True,
+            return_attention_mask=False,
+        )
+        ids = torch.tensor(
+            list(itertools.chain.from_iterable(enc["input_ids"])), dtype=torch.long
+        )
+        max_start = len(ids) - sequence_length - 1
+        if max_start < 0:
+            raise ValueError(
+                f"Calibration corpus has {len(ids)} tokens; need at least {sequence_length + 1}."
+            )
+        if max_start + 1 < n_sequences:
+            raise ValueError(
+                f"Not enough distinct start positions for {n_sequences} calibration sequences."
+            )
+        starts = random.Random(seed).sample(range(max_start + 1), n_sequences)
+        return [ids[s : s + sequence_length].unsqueeze(0) for s in starts]
+
+    def _input_device(self) -> torch.device:
+        emb = self.model.get_input_embeddings()
+        for p in emb.parameters():
+            if p.device.type != "meta":
+                return p.device
+        return torch.device("cpu")
+
+    def score_nll(self, batches: Iterable[torch.Tensor]) -> float:
+        self._ensure_loaded()
+        dev = self._input_device()
+        total = 0.0
+        ntok = 0
+        with torch.inference_mode():
+            for cpu_x in batches:
+                x = cpu_x.to(dev)
+                output = self.model(input_ids=x, use_cache=False)
+                logits = output.logits[:, :-1, :]
+                target = x[:, 1:].to(logits.device)
+                loss = F.cross_entropy(
+                    logits.reshape(-1, logits.shape[-1]).float(),
+                    target.reshape(-1),
+                    reduction="sum",
+                )
+                total += float(loss.detach().cpu())
+                ntok += int(target.numel())
+                del x, output, logits, target, loss
+        if ntok == 0:
+            raise ValueError("No calibration tokens were scored.")
+        return total / ntok
+
+    @staticmethod
+    def _atomic_write(path: Path, obj: dict) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = path.with_suffix(path.suffix + ".tmp")
+        tmp.write_text(json.dumps(obj, indent=2))
+        os.replace(tmp, path)
+
+    def fit(
+        self,
+        *,
+        texts: Sequence[str] | None = None,
+        batches: Sequence[torch.Tensor] | None = None,
+        n_sequences: int = 32,
+        sequence_length: int = 128,
+        seed: int = 42,
+        checkpoint_path: str | Path | None = None,
+        resume: bool = True,
+        max_delete: int | None = None,
+    ) -> Order2Result:
+        """Measure baseline/single/pair NLLs and construct the greedy order-2 path."""
+        self._ensure_loaded()
+        if batches is None:
+            if texts is None:
+                raise ValueError("Pass either texts=... or batches=...")
+            batches = self.calibration_batches(
+                texts,
+                n_sequences=n_sequences,
+                sequence_length=sequence_length,
+                seed=seed,
+            )
+
+        depth = self.depth
+        ckpt = Path(checkpoint_path) if checkpoint_path else None
+        state = {
+            "method": "order-2 interaction greedy",
+            "depth": depth,
+            "baseline_nll": None,
+            "single_nll": {},
+            "pair_nll": {},
+            "complete": False,
+        }
+        if ckpt and resume and ckpt.exists():
+            loaded = json.loads(ckpt.read_text())
+            if int(loaded.get("depth", depth)) != depth:
+                raise ValueError("Checkpoint depth does not match the loaded model.")
+            state.update(loaded)
+            state.setdefault("single_nll", {})
+            state.setdefault("pair_nll", {})
+
+        original = list(self._layers())
+        try:
+            if state["baseline_nll"] is None:
+                self._set_layers(original)
+                state["baseline_nll"] = self.score_nll(batches)
+                if ckpt:
+                    self._atomic_write(ckpt, state)
+
+            for i in range(depth):
+                key = str(i)
+                if key in state["single_nll"]:
+                    continue
+                kept = [b for j, b in enumerate(original) if j != i]
+                self._set_layers(kept)
+                state["single_nll"][key] = self.score_nll(batches)
+                if ckpt:
+                    self._atomic_write(ckpt, state)
+                self._set_layers(original)
+
+            total_pairs = depth * (depth - 1) // 2
+            for i in range(depth):
+                for j in range(i + 1, depth):
+                    key = f"{i},{j}"
+                    if key in state["pair_nll"]:
+                        continue
+                    kept = [b for q, b in enumerate(original) if q not in (i, j)]
+                    self._set_layers(kept)
+                    state["pair_nll"][key] = self.score_nll(batches)
+                    if ckpt:
+                        self._atomic_write(ckpt, state)
+                    self._set_layers(original)
+                    print(f"pair {key:>7s} | {len(state['pair_nll'])}/{total_pairs}", flush=True)
+
+            result = Order2Result(
+                depth=depth,
+                baseline_nll=float(state["baseline_nll"]),
+                single_nll={int(k): float(v) for k, v in state["single_nll"].items()},
+                pair_nll={
+                    tuple(int(x) for x in k.split(",")): float(v)
+                    for k, v in state["pair_nll"].items()
+                },
+            )
+            result.build_interactions().build_greedy_path(max_delete=max_delete)
+            self.result = result
+            if ckpt:
+                result.save_json(ckpt)
+            return result
+        finally:
+            self._set_layers(original)
+
+    def select(self, target_layers: int, result: Order2Result | None = None) -> dict[str, list[int]]:
+        result = result or self.result
+        if result is None:
+            raise RuntimeError("Call fit() first or pass result=...")
+        return result.select(target_layers)
+
+    def apply(
+        self,
+        target_layers: int,
+        *,
+        result: Order2Result | None = None,
+    ) -> PreTrainedModel:
+        """Apply a selection to the currently loaded model in place."""
+        self._ensure_loaded()
+        result = result or self.result
+        if result is None:
+            raise RuntimeError("Call fit() first or pass result=...")
+        selection = result.select(target_layers)
+        original = list(self._layers())
+        retained = selection["retained_layers"]
+        self._set_layers([original[i] for i in retained])
+        return self.model
+
+    def prune(
+        self,
+        target_layers: int,
+        *,
+        texts: Sequence[str] | None = None,
+        batches: Sequence[torch.Tensor] | None = None,
+        n_sequences: int = 32,
+        sequence_length: int = 128,
+        seed: int = 42,
+        checkpoint_path: str | Path | None = None,
+        resume: bool = True,
+    ) -> tuple[PreTrainedModel, Order2Result]:
+        result = self.fit(
+            texts=texts,
+            batches=batches,
+            n_sequences=n_sequences,
+            sequence_length=sequence_length,
+            seed=seed,
+            checkpoint_path=checkpoint_path,
+            resume=resume,
+            max_delete=self.depth - target_layers,
+        )
+        model = self.apply(target_layers, result=result)
+        return model, result
+
+    def save_pruned(
+        self,
+        output_dir: str | Path,
+        target_layers: int,
+        *,
+        result: Order2Result | None = None,
+        safe_serialization: bool = True,
+    ) -> Path:
+        """Apply a selection and save the pruned model/tokenizer with HF save_pretrained()."""
+        model = self.apply(target_layers, result=result)
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        model.save_pretrained(output_dir, safe_serialization=safe_serialization)
+        if self.tokenizer is not None:
+            self.tokenizer.save_pretrained(output_dir)
+        (output_dir / "layer_interaction_selection.json").write_text(
+            json.dumps((result or self.result).select(target_layers), indent=2)
+        )
+        return output_dir
+
+    def close(self, *, drop_model: bool = False) -> None:
+        if drop_model:
+            self.model = None
+            self.tokenizer = None
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            try:
+                torch.cuda.ipc_collect()
+            except Exception:
+                pass
+        if self._owns_offload_dir:
+            shutil.rmtree(self.offload_dir, ignore_errors=True)
+
+    def __enter__(self) -> "Order2Pruner":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close(drop_model=True)

layer_interactions/result.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import json
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class Order2Result:
+    """Measurements and selections produced by the order-2 interaction method."""
+
+    depth: int
+    baseline_nll: float
+    single_nll: dict[int, float]
+    pair_nll: dict[tuple[int, int], float]
+    first_order: dict[int, float] = field(default_factory=dict)
+    second_order: dict[tuple[int, int], float] = field(default_factory=dict)
+    delete_order: list[int] = field(default_factory=list)
+    greedy_path: list[dict[str, Any]] = field(default_factory=list)
+
+    def build_interactions(self) -> "Order2Result":
+        d0 = float(self.baseline_nll)
+        self.first_order = {
+            i: float(self.single_nll[i]) - d0 for i in range(self.depth)
+        }
+        self.second_order = {}
+        for i in range(self.depth):
+            for j in range(i + 1, self.depth):
+                self.second_order[(i, j)] = (
+                    float(self.pair_nll[(i, j)])
+                    - float(self.single_nll[i])
+                    - float(self.single_nll[j])
+                    + d0
+                )
+        return self
+
+    def build_greedy_path(self, max_delete: int | None = None) -> "Order2Result":
+        if not self.first_order or not self.second_order:
+            self.build_interactions()
+
+        if max_delete is None:
+            max_delete = self.depth - 1
+        if not 0 <= max_delete < self.depth:
+            raise ValueError(f"max_delete must be in [0, {self.depth - 1}]")
+
+        deleted: list[int] = []
+        deleted_set: set[int] = set()
+        path: list[dict[str, Any]] = []
+        cumulative = 0.0
+
+        for step in range(max_delete):
+            candidates: list[tuple[float, int]] = []
+            for i in range(self.depth):
+                if i in deleted_set:
+                    continue
+                interaction = sum(
+                    self.second_order[tuple(sorted((i, j)))] for j in deleted
+                )
+                marginal = self.first_order[i] + interaction
+                candidates.append((marginal, i))
+
+            marginal, chosen = min(candidates, key=lambda z: (z[0], z[1]))
+            deleted.append(chosen)
+            deleted_set.add(chosen)
+            cumulative += marginal
+            path.append(
+                {
+                    "step": step + 1,
+                    "deleted_layer": chosen,
+                    "marginal_predicted_nll_change": float(marginal),
+                    "cumulative_predicted_nll_change": float(cumulative),
+                }
+            )
+
+        self.delete_order = deleted
+        self.greedy_path = path
+        return self
+
+    def select(self, target_layers: int) -> dict[str, list[int]]:
+        if not 1 <= target_layers <= self.depth:
+            raise ValueError(f"target_layers must be in [1, {self.depth}]")
+        n_delete = self.depth - target_layers
+        if len(self.delete_order) < n_delete:
+            self.build_greedy_path(max_delete=n_delete)
+        deleted = list(self.delete_order[:n_delete])
+        deleted_set = set(deleted)
+        retained = [i for i in range(self.depth) if i not in deleted_set]
+        return {"retained_layers": retained, "deleted_layers": deleted}
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "method": "order-2 interaction greedy",
+            "depth": self.depth,
+            "baseline_nll": float(self.baseline_nll),
+            "single_nll": {str(k): float(v) for k, v in self.single_nll.items()},
+            "pair_nll": {f"{i},{j}": float(v) for (i, j), v in self.pair_nll.items()},
+            "first_order_delta": {str(k): float(v) for k, v in self.first_order.items()},
+            "second_order_interaction": {
+                f"{i},{j}": float(v) for (i, j), v in self.second_order.items()
+            },
+            "delete_order": list(self.delete_order),
+            "greedy_path": list(self.greedy_path),
+            "complete": True,
+        }
+
+    def save_json(self, path: str | Path) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = path.with_suffix(path.suffix + ".tmp")
+        tmp.write_text(json.dumps(self.to_dict(), indent=2))
+        tmp.replace(path)
+
+    @classmethod
+    def from_dict(cls, obj: dict[str, Any]) -> "Order2Result":
+        pair = {}
+        for key, value in obj.get("pair_nll", {}).items():
+            i, j = (int(x) for x in key.split(","))
+            pair[(i, j)] = float(value)
+        second = {}
+        for key, value in obj.get("second_order_interaction", {}).items():
+            i, j = (int(x) for x in key.split(","))
+            second[(i, j)] = float(value)
+        result = cls(
+            depth=int(obj["depth"]),
+            baseline_nll=float(obj["baseline_nll"]),
+            single_nll={int(k): float(v) for k, v in obj.get("single_nll", {}).items()},
+            pair_nll=pair,
+            first_order={
+                int(k): float(v) for k, v in obj.get("first_order_delta", {}).items()
+            },
+            second_order=second,
+            delete_order=[int(x) for x in obj.get("delete_order", [])],
+            greedy_path=list(obj.get("greedy_path", [])),
+        )
+        return result
+
+    @classmethod
+    def load_json(cls, path: str | Path) -> "Order2Result":
+        return cls.from_dict(json.loads(Path(path).read_text()))

tests/test_result.py

@@ -0,0 +1,34 @@
+from layer_interactions.result import Order2Result
+
+
+def test_interactions_and_selection():
+    # D0 = 1.0. First-order deletion costs: [0.1, 0.2, 0.05]
+    # Pair (0,2) has a positive interaction of +0.3, so after deleting 2,
+    # deleting 1 is cheaper than deleting 0.
+    r = Order2Result(
+        depth=3,
+        baseline_nll=1.0,
+        single_nll={0: 1.1, 1: 1.2, 2: 1.05},
+        pair_nll={
+            (0, 1): 1.3,
+            (0, 2): 1.45,
+            (1, 2): 1.25,
+        },
+    )
+    r.build_interactions().build_greedy_path(max_delete=2)
+    assert r.delete_order == [2, 1]
+    assert r.select(1)["retained_layers"] == [0]
+
+
+def test_roundtrip_dict():
+    r = Order2Result(
+        depth=2,
+        baseline_nll=1.0,
+        single_nll={0: 1.1, 1: 1.2},
+        pair_nll={(0, 1): 1.35},
+    ).build_interactions().build_greedy_path(max_delete=1)
+    rr = Order2Result.from_dict(r.to_dict())
+    assert rr.depth == 2
+    assert rr.single_nll == r.single_nll
+    assert rr.pair_nll == r.pair_nll
+    assert rr.delete_order == r.delete_order