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

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

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:

pip install "git+https://huggingface.co/<HF_USERNAME>/LayerInteractions"

Or clone and install locally:

git clone https://huggingface.co/<HF_USERNAME>/LayerInteractions
cd LayerInteractions
pip install .

For the example scripts:

pip install ".[examples]"

API

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

{
    "retained_layers": [...],
    "deleted_layers": [...],
}

One-call pruning

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

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:

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:

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:

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:

model.layers
transformer.h
gpt_neox.layers

For another architecture, supply the path explicitly:

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

from layer_interactions import Order2Pruner, Order2Result

Order2Result exposes:

baseline_nll
single_nll
pair_nll
first_order
second_order
delete_order
greedy_path

and can be saved or loaded as JSON.