File size: 4,338 Bytes
f8cf4de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | # 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.
|