Angshul/higher-order-pruner

Higher-Order Pruner

A reusable Python pruning method hosted in a Hugging Face model repository. This repository is not a pruned checkpoint and does not contain benchmark results. It supplies a direct code API for pruning compatible unquantized PyTorch models.

The method combines calibration-derived KFAC-diagonal second-order information with a clipped diagonal higher-order finite-difference estimate, exact global sparsity allocation, and monotone held-out refinement. It does not initialize from Wanda, SparseGPT, magnitude pruning, or a previously pruned checkpoint.

Install from this model repository

pip install "git+https://huggingface.co/OWNER/REPOSITORY"

Then load the method configuration from the Hub and use it in code:

from higher_order_pruner import HigherOrderPruner

pruner = HigherOrderPruner.from_pretrained("OWNER/REPOSITORY")

result = pruner.prune_hf_model(
    "Qwen/Qwen3-4B-Base",
    calibration_texts=my_calibration_texts,
    validation_texts=my_heldout_texts,
    sparsity=0.50,
    output_dir="./pruned-model",
    sequence_length=512,
    device="cuda",
)
print(result["report"]["achieved_sparsity"])

pruned-model/checkpoint is a normal Hugging Face checkpoint saved with save_pretrained and safetensors.

Generic PyTorch API

The core is not restricted to language models. Supply any differentiable scalar loss and batches already placed on the desired device:

from higher_order_pruner import HigherOrderPruner

pruner = HigherOrderPruner.from_pretrained("OWNER/REPOSITORY")

def loss_fn(model, batch):
    x, target = batch
    return (model(x) - target).square().mean()

report = pruner.prune(
    model,
    calibration_batches,
    loss_fn,
    sparsity=0.50,
    validation_batches=heldout_batches,
    progress_dir="./pruning-progress",
)

For nonstandard architectures, pass explicit module_names=[...]. The method prunes unique floating-point torch.nn.Linear.weight tensors. It refuses quantized weights and does not promise that arbitrary custom kernels, tied non-Linear tensors, or remote-code-only architectures are compatible.

Important behavior

• Calibration statistics and higher-order estimates are computed from the caller's model and data at runtime; no Qwen/Falcon statistics are embedded.
• The requested zero count is enforced exactly across selected weights.
• Existing exact zeros are included in that budget.
• The HF adapter auto-tunes its measured curvature kernel for throughput under a 95% VRAM ceiling. It does not require a specific GPU model.
• Unstructured zeros require an appropriate sparse inference kernel to produce inference speedups.

Local development

pip install -e ".[test]"
pytest -q

upload_to_hub.py creates a new Hugging Face model repository and uploads this package. It refuses to overwrite an existing repository.