# 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.