Upload 4 files

README.md

@@ -1,3 +1,189 @@
----
-license: h-research
----
+# LayerInteractions
+
+A small Python API for **order-2 interaction-aware transformer depth pruning** of Hugging Face causal language models.
+
+This repository packages the method used in the accompanying experiments. It is source code, not a dataset and not a web service.
+
+## Method
+
+Let `D0` be the calibration NLL of the full model, `Di` the NLL after deleting layer `i`, and `Dij` the NLL after deleting layers `i` and `j`. We form
+
+```text
+a_i  = D_i - D_0
+
+a_ij = D_ij - D_i - D_j + D_0
+```
+
+For a current deleted set `S`, the predicted marginal cost of deleting layer `i` is
+
+```text
+a_i + sum_{j in S} a_ij.
+```
+
+The greedy selector repeatedly deletes the layer with the smallest marginal predicted NLL increase. The implementation measures all single and pair deletions, checkpoints after every measurement, constructs the quadratic surrogate, and returns the retained layer set at any target depth.
+
+## Install directly from Hugging Face
+
+After uploading this folder to a Hugging Face repository:
+
+```bash
+pip install "git+https://huggingface.co/<HF_USERNAME>/LayerInteractions"
+```
+
+Or clone and install locally:
+
+```bash
+git clone https://huggingface.co/<HF_USERNAME>/LayerInteractions
+cd LayerInteractions
+pip install .
+```
+
+For the example scripts:
+
+```bash
+pip install ".[examples]"
+```
+
+## API
+
+```python
+from datasets import load_dataset
+from layer_interactions import Order2Pruner
+
+texts = [
+    x["text"]
+    for x in load_dataset(
+        "Salesforce/wikitext", "wikitext-2-raw-v1", split="train"
+    )
+    if x["text"].strip()
+]
+
+with Order2Pruner(
+    "allenai/OLMoE-1B-7B-0924",
+    dtype="bfloat16",
+) as pruner:
+    result = pruner.fit(
+        texts=texts,
+        n_sequences=32,
+        sequence_length=128,
+        seed=42,
+        checkpoint_path="order2_measurements.json",
+        max_delete=12,
+    )
+
+    print(result.select(8))
+    print(result.select(4))
+```
+
+The output of `select()` is
+
+```python
+{
+    "retained_layers": [...],
+    "deleted_layers": [...],
+}
+```
+
+## One-call pruning
+
+```python
+with Order2Pruner("allenai/OLMoE-1B-7B-0924") as pruner:
+    model, result = pruner.prune(
+        8,
+        texts=texts,
+        checkpoint_path="order2_measurements.json",
+    )
+```
+
+The returned `model` has the selected transformer blocks removed in place.
+
+## Save a pruned Hugging Face checkpoint
+
+```python
+pruner.save_pruned(
+    "OLMoE-Order2-8",
+    target_layers=8,
+    result=result,
+)
+```
+
+This writes the model, tokenizer and `layer_interaction_selection.json` using the normal Hugging Face `save_pretrained()` format.
+
+## Resuming long measurements
+
+Pair measurements are the expensive part. Pass a checkpoint path:
+
+```python
+result = pruner.fit(
+    texts=texts,
+    checkpoint_path="order2_measurements.json",
+    resume=True,
+)
+```
+
+The JSON file is updated after the baseline, every single-layer deletion, and every pair deletion. Re-running resumes from completed measurements.
+
+For a depth-`L` model, order-2 requires `L(L-1)/2` pair measurements. For example:
+
+```text
+L = 30 -> 435 pairs
+L = 16 -> 120 pairs
+```
+
+## Memory management
+
+By default, the API uses Hugging Face/Accelerate `device_map="auto"`. On CUDA, it reserves approximately 4 GiB of free GPU memory as activation headroom and makes the remaining memory available for weights. The balance can be controlled explicitly:
+
+```python
+pruner = Order2Pruner(
+    model_id,
+    gpu_memory_gib=16,
+    cpu_memory_gib=32,
+    activation_headroom_gib=4,
+)
+```
+
+The NLL is computed manually from the logits rather than by passing `labels=...`; this keeps the scoring path compatible with CPU/GPU-dispatched models.
+
+## Supported model layout
+
+Automatic layer discovery currently recognizes the common paths:
+
+```text
+model.layers
+transformer.h
+gpt_neox.layers
+```
+
+For another architecture, supply the path explicitly:
+
+```python
+Order2Pruner(model_id, layer_path="decoder.layers")
+```
+
+The DeepSeek-LLM-7B and OLMoE examples use the standard `model.layers` layout.
+
+## Examples
+
+- `examples/deepseek_7b.py`: 30 -> 15 and 30 -> 7.
+- `examples/olmoe_7b.py`: 16 -> 8 and 16 -> 4.
+
+## Public objects
+
+```python
+from layer_interactions import Order2Pruner, Order2Result
+```
+
+`Order2Result` exposes:
+
+```text
+baseline_nll
+single_nll
+pair_nll
+first_order
+second_order
+delete_order
+greedy_path
+```
+
+and can be saved or loaded as JSON.

UPLOAD_TO_HF.txt

@@ -0,0 +1,7 @@
+Create a Hugging Face repository named LayerInteractions and upload the CONTENTS of this folder.
+
+This is a code-only Python package. It does not require a Space, endpoint, server, database, or dataset repository.
+
+After upload, users can install it with:
+
+pip install "git+https://huggingface.co/<HF_USERNAME>/LayerInteractions"

VERSION.txt

@@ -0,0 +1 @@
+0.1.0

pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "layer-interactions"
+version = "0.1.0"
+description = "Second-order interaction pruning for Hugging Face causal language models"
+readme = "README.md"
+requires-python = ">=3.10"
+authors = [
+  {name = "Angshul Majumdar"}
+]
+dependencies = [
+  "torch>=2.2",
+  "transformers>=4.45",
+  "accelerate>=1.0",
+  "psutil>=5.9"
+]
+
+[project.optional-dependencies]
+examples = ["datasets>=3.0"]
+test = ["pytest>=8.0"]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["layer_interactions*"]