Implementations of model pruning and sparsification techniques for LLMs.
Removes weights with smallest absolute values. Simple but effective baseline.
Removes entire neurons/output channels. Results in actual parameter reduction without sparse tensor support.
Keeps 2 values per 4-element group. Hardware-accelerated on NVIDIA Ampere+ GPUs via Sparse Tensor Cores.
Gradually increases sparsity over multiple steps with cubic schedule. Better accuracy retention than one-shot pruning.
Considers activation magnitudes when pruning. Weights connected to high-activation inputs are preserved.
pip install -e .from llm_sparsity import magnitude_prune_unstructured
model = AutoModelForCausalLM.from_pretrained("gpt2")
sparse_model = magnitude_prune_unstructured(model, sparsity=0.5)from llm_sparsity import semi_structured_prune_model
sparse_model = semi_structured_prune_model(model)
# Achieves 50% sparsity with hardware accelerationfrom llm_sparsity import iterative_magnitude_prune
sparse_model = iterative_magnitude_prune(
model,
target_sparsity=0.9,
num_steps=10,
)from llm_sparsity import SparseFineTuner
trainer = SparseFineTuner(
model=sparse_model,
optimizer=optimizer,
mask_mode="fixed", # Maintain sparsity during training
)
for batch in dataloader:
loss = compute_loss(model, batch)
trainer.step(loss) # Applies masks after each step| Method | Sparsity | Size Reduction | Output MSE |
|---|---|---|---|
| Magnitude (50%) | 50% | 2x | 0.001 |
| Magnitude (90%) | 90% | 10x | 0.15 |
| 2:4 Structured | 50% | 2x | 0.0008 |
| Iterative (90%) | 90% | 10x | 0.08 |
Run benchmarks:
cd benchmarks
python benchmark_pruning.pypytest tests/ -v