Model Pruning and Sparsity

"Perfection is achieved, not when there is nothing more to add, but when there is nothing left to take away."

Pip, Minimalist AI Agent

1. Why Pruning Matters for LLM Deployment

Large language models are expensive to serve because they are large. A 70B parameter model at FP16 precision occupies roughly 140 GB of GPU memory, requiring multi-GPU setups just to load the weights. Every forward pass multiplies input activations through these weight matrices, so the number of nonzero weights directly determines both memory footprint and compute cost. Pruning attacks this problem at the source: it sets a fraction of weights to zero, reducing the effective model size without changing the architecture.

The practical appeal of pruning is that it operates on a fully trained model. Unlike distillation, which requires training a new smaller model from scratch, pruning preserves the original model's architecture and most of its learned representations. This makes it particularly attractive when you have a fine-tuned model (Section 14.1) that performs well on your task and you need to compress it for deployment without retraining.

The sparsity opportunity. Research consistently shows that 50% to 80% of weights in large transformers can be removed with minimal quality degradation. The reason is that pretraining creates heavily over-parameterized models, where many weights encode redundant or rarely used knowledge. For a specific deployment task (say, customer service chat), the fraction of truly essential weights is even smaller. This means that task-specific pruning can often achieve higher sparsity ratios than general-purpose pruning while maintaining quality on the target domain.

Pruning also complements the API cost optimization strategies from Section 10.3. If you self-host a pruned model behind your own API endpoint using frameworks like vLLM or TGI (Section 9.4), you reduce per-request compute costs directly, potentially making self-hosting competitive with commercial API pricing for high-volume workloads.

2. Unstructured Pruning

Unstructured pruning removes individual weights anywhere in the model's weight matrices. The resulting sparse matrices have zeros scattered throughout, with no requirement that the zeros form regular patterns. This flexibility allows unstructured pruning to achieve very high sparsity (often 70% to 90%) because the algorithm can selectively remove whichever weights matter least, regardless of their position.

2.1 Magnitude Pruning

The simplest and oldest pruning strategy is magnitude pruning: sort all weights by their absolute value and set the smallest ones to zero. The intuition is straightforward: weights near zero contribute little to the output of any matrix multiplication. Despite its simplicity, magnitude pruning remains a competitive baseline for moderate sparsity levels (up to 50%).

# Magnitude pruning with PyTorch: zero out the smallest 50% of weights
# in every Linear layer, then make the pruning permanent.
import torch
import torch.nn.utils.prune as prune
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")

# Apply magnitude pruning to all Linear layers
for name, module in model.named_modules():
 if isinstance(module, torch.nn.Linear):
 prune.l1_unstructured(module, name="weight", amount=0.5) # 50% sparsity

# Check sparsity of a specific layer
layer = model.model.layers[0].self_attn.q_proj
zeros = (layer.weight == 0).sum().item()
total = layer.weight.numel()
print(f"Sparsity: {zeros / total:.1%}") # Should be ~50.0%

# Make pruning permanent (remove the reparameterization)
for name, module in model.named_modules():
 if isinstance(module, torch.nn.Linear):
 prune.remove(module, "weight")

The limitation of magnitude pruning is that it ignores context. A weight with small absolute value might still be critical if it sits on a high-activation pathway. This is why more sophisticated methods have emerged.

2.2 Movement Pruning

Movement pruning (Sanh et al., 2020) takes a different approach: instead of looking at weight magnitudes, it tracks which weights are moving toward zero during fine-tuning. Weights that the optimization process is "trying to eliminate" are pruned, while weights that are growing in magnitude are retained. This produces better results than magnitude pruning when combined with task-specific fine-tuning, because the pruning decision aligns with the model's adaptation to the target task.

2.3 SparseGPT

SparseGPT (Frantar & Alistarh, 2023) was a breakthrough for LLM pruning because it can prune massive models in a single pass without any retraining. The algorithm works layer by layer, solving an optimization problem that minimizes the reconstruction error introduced by pruning. For each layer, SparseGPT computes the optimal set of weights to remove and adjusts the remaining weights to compensate for the removed ones. The key insight is that this layer-wise approach scales linearly with model size, making it feasible to prune models with hundreds of billions of parameters in hours rather than weeks.

# Using SparseGPT via the SparseML library
# Install: pip install sparseml
from sparseml.transformers import SparseAutoModelForCausalLM
from sparseml.transformers import oneshot

# One-shot pruning with SparseGPT
oneshot(
 model="meta-llama/Llama-3.2-1B",
 dataset="open_platypus",
 recipe="recipe.yaml", # Specifies 50% unstructured sparsity
 output_dir="./pruned-model",
 num_calibration_samples=512,
)

# The recipe.yaml would contain:
# sparsity_modifiers:
# SparseGPTModifier:
# sparsity: 0.5
# sequential_update: true
# targets: ["re:model.layers.\\d+.self_attn", "re:model.layers.\\d+.mlp"]

3. Structured Pruning

Structured pruning removes entire structural components of the model: full rows or columns from weight matrices, entire attention heads, or even whole transformer layers. Unlike unstructured pruning, which produces irregular sparsity patterns that require specialized sparse matrix libraries, structured pruning produces a genuinely smaller dense model that runs efficiently on standard hardware with no special kernel support.

3.1 Attention Head Pruning

Multi-head attention (Section 4.2) distributes computation across multiple parallel attention heads. Research has shown that many of these heads are redundant: removing 20% to 40% of attention heads often has minimal impact on downstream performance. Michel et al. (2019) demonstrated that in some BERT models, a single attention head per layer suffices for most tasks.

For LLMs, head pruning is typically guided by an importance score computed over a calibration dataset. The score measures how much each head's output contributes to the model's final predictions. Heads with the lowest importance scores are removed, and the model architecture is updated accordingly.

# Gradient-based attention head importance scoring for structured pruning.
# Accumulates gradient magnitudes across Q, K, V projections per head
# to produce a per-layer importance matrix guiding removal decisions.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")

# Compute head importance via gradient-based scoring
def compute_head_importance(model, dataloader, num_heads=32, num_layers=16):
 """Score each attention head by its gradient-weighted contribution."""
 head_importance = torch.zeros(num_layers, num_heads)

 model.eval()
 for batch in dataloader:
 outputs = model(**batch, output_attentions=True)
 loss = outputs.loss
 loss.backward()

 for layer_idx in range(num_layers):
 attn_module = model.model.layers[layer_idx].self_attn
 # Accumulate gradient magnitudes for Q, K, V projections
 for proj in [attn_module.q_proj, attn_module.k_proj, attn_module.v_proj]:
 grad = proj.weight.grad
 if grad is not None:
 # Reshape to (num_heads, head_dim, input_dim) and sum
 head_dim = grad.shape[0] // num_heads
 grad_per_head = grad.view(num_heads, head_dim, -1)
 head_importance[layer_idx] += grad_per_head.abs().sum(dim=(1, 2))

 model.zero_grad()

 return head_importance / len(dataloader)

# After computing importance, prune the lowest-scoring heads
# by zeroing their Q, K, V, and O projection rows/columns

3.2 Layer Pruning

Layer pruning removes entire transformer layers from the model stack. This is the most aggressive form of structured pruning: removing a layer eliminates all its parameters (self-attention, feed-forward network, layer norms) in one operation. A 32-layer model pruned to 24 layers is 25% smaller and 25% faster, with no sparse kernel overhead.

The challenge is selecting which layers to remove. Shallow layers (close to the input) tend to encode syntactic features, while deep layers (close to the output) encode more abstract, task-specific representations. Middle layers are often the most redundant. Studies on LLaMA models have shown that removing 8 out of 32 layers from the middle of the stack preserves over 90% of benchmark performance, while removing layers from the first or last quarter causes severe degradation.

3.3 Width Pruning

Width pruning reduces the hidden dimension of the model by removing entire neurons (columns) from feed-forward layers or reducing the embedding dimension. This produces a uniformly narrower model. Width pruning is less common than head or layer pruning for LLMs because it requires careful handling of residual connections, but it can be effective when combined with knowledge distillation to recover lost quality.

4. The Wanda Method

Wanda (Sun et al., 2024), which stands for Pruning by Weights and Activations, introduced an elegantly simple idea: instead of using only weight magnitudes to decide what to prune, multiply each weight by the corresponding input activation norm. A weight might be small in absolute terms, but if it sits on a pathway with large activations, removing it would significantly change the layer's output. Conversely, a large weight on a pathway with near-zero activations contributes nothing and can be safely pruned.

The Wanda importance score for weight w_ij in a linear layer is:

score($w_{ij}$) = |$w_{ij}$| · ||$X_{j}$||₂

where X_j is the j-th column of the input activation matrix collected over the calibration set. This is computed once per layer using a small calibration dataset (128 examples typically suffice), making Wanda nearly as fast as simple magnitude pruning while achieving quality comparable to the much more expensive SparseGPT.

A small numeric example shows how activation context changes pruning decisions:

# Wanda score: numeric walkthrough for 4 weights
import numpy as np

weights = np.array([0.01, -0.8, 0.3, -0.05])
activation_norms = np.array([50.0, 0.1, 2.0, 40.0])

scores = np.abs(weights) * activation_norms
print("weights: ", weights)
print("activation norms:", activation_norms)
print("Wanda scores: ", scores)
# Wanda scores: [0.5, 0.08, 0.6, 2.0]
# Magnitude pruning would drop w0 (0.01) and w3 (-0.05) first.
# Wanda drops w1 (-0.8) instead because its activation norm is tiny.

Code Fragment 10.5.9 implements Wanda pruning from scratch, showing the importance computation and threshold-based masking for a single linear layer.

# Wanda pruning implementation sketch
# Reference: https://github.com/locuslab/wanda

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def wanda_prune_layer(weight, activation_norms, sparsity_ratio=0.5):
 """
 Prune a linear layer using the Wanda criterion.

 Args:
 weight: (out_features, in_features) weight matrix
 activation_norms: (in_features,) L2 norm of input activations
 sparsity_ratio: fraction of weights to prune
 """
 # Compute importance: |weight| * activation_norm per input feature
 importance = weight.abs() * activation_norms.unsqueeze(0)

 # Find the threshold for the desired sparsity
 num_prune = int(sparsity_ratio * weight.numel())
 threshold = torch.kthvalue(importance.flatten(), num_prune).values

 # Create mask: keep weights above threshold
 mask = importance > threshold
 return weight * mask.float(), mask

# Collect activation norms over calibration data
def collect_activation_norms(model, calibration_loader, layer_idx):
 """Hook into a specific layer to collect input activation norms."""
 norms = []

 def hook_fn(module, input, output):
 # input[0] shape: (batch, seq_len, hidden_dim)
 norms.append(input[0].float().norm(dim=(0, 1)))

 layer = model.model.layers[layer_idx].mlp.gate_proj
 handle = layer.register_forward_hook(hook_fn)

 with torch.no_grad():
 for batch in calibration_loader:
 model(**batch)

 handle.remove()
 return torch.stack(norms).mean(dim=0)

For production use, the llmcompressor library wraps Wanda (and SparseGPT) into a single high-level call:

# Library shortcut: Wanda pruning via llmcompressor (pip install llmcompressor)
from llmcompressor.modifiers.pruning import WandaPruningModifier
from llmcompressor import oneshot

recipe = WandaPruningModifier(sparsity=0.5, block_size=128)
oneshot(
 model="meta-llama/Llama-3.2-1B",
 recipe=recipe,
 dataset="ultrachat-200k",
 output_dir="./llama-1b-wanda-50",
)

5. Sparse Model Inference

Pruning a model is only half the battle. To realize actual speedups, the inference engine must exploit the sparsity pattern. A sparse weight matrix stored in dense format takes exactly the same amount of memory and compute as the original. Three approaches exist for translating sparsity into real performance gains.

5.1 Sparse Matrix Formats and Kernels

Unstructured sparse matrices are stored in compressed formats like CSR (Compressed Sparse Row) or CSC (Compressed Sparse Column) that skip zero values entirely. Libraries like torch.sparse and specialized CUDA kernels can perform matrix multiplications using these formats. However, on GPUs, the overhead of indirect indexing in sparse formats often negates the benefit of skipping zero multiplications, unless sparsity exceeds 90% to 95%. This is why unstructured sparsity at moderate levels (50% to 80%) frequently fails to deliver wall-clock speedups on current hardware, despite reducing the theoretical compute by half or more.

5.2 NVIDIA Sparse Tensor Cores and 2:4 Sparsity

NVIDIA's Ampere (A100), Ada Lovelace (L40, RTX 4090), and Hopper (H100) GPU architectures include sparse tensor cores that natively accelerate a specific sparsity pattern: 2:4 structured sparsity. In this pattern, exactly 2 out of every 4 consecutive weights must be zero. The hardware stores only the 2 nonzero values plus a small index, achieving a genuine 2x speedup for matrix multiplications with no software overhead.

# Convert all Linear layers to NVIDIA 2:4 semi-structured sparsity.
# In each group of 4 weights, the 2 smallest are zeroed for hardware acceleration.
import torch
from torch.sparse import to_sparse_semi_structured, SparseSemiStructuredTensor

# Enable the fast path for semi-structured sparsity
SparseSemiStructuredTensor._FORCE_CUTLASS = True

model = AutoModelForCausalLM.from_pretrained(
 "meta-llama/Llama-3.2-1B",
 torch_dtype=torch.float16,
 device_map="cuda"
)

# Apply 2:4 sparsity pattern to all linear layers
for name, module in model.named_modules():
 if isinstance(module, torch.nn.Linear):
 # Prune to 2:4 pattern (keep 2 largest per group of 4)
 weight = module.weight.data
 # Reshape into groups of 4 along the input dimension
 w_groups = weight.view(-1, weight.shape[1] // 4, 4)
 # Find the 2 smallest values in each group
 _, indices = w_groups.abs().topk(2, dim=-1, largest=False)
 mask = torch.ones_like(w_groups, dtype=torch.bool)
 mask.scatter_(-1, indices, False)
 weight.view_as(w_groups).masked_fill_(~mask, 0)

 # Convert to semi-structured sparse format for hardware acceleration
 module.weight = torch.nn.Parameter(
 to_sparse_semi_structured(module.weight.data)
 )

The 2:4 pattern is particularly powerful because it is the only sparsity pattern that delivers guaranteed hardware speedups across all batch sizes and sequence lengths on supported GPUs. For teams deploying on NVIDIA hardware, 2:4 sparsity is often the most practical pruning strategy.

5.3 DeepSparse and Specialized Runtimes

For CPU deployment, Neural Magic's DeepSparse runtime is specifically designed to exploit unstructured sparsity. It converts sparse weight matrices into an optimized execution plan that skips zero computations on CPUs, where the memory access patterns make sparse computation more favorable than on GPUs. DeepSparse can deliver near-GPU throughput for highly sparse models (90%+) running on commodity CPUs, making it an attractive option for teams without GPU infrastructure.

6. Combining Pruning with Quantization and Distillation

The three compression techniques, pruning, quantization, and distillation, are largely orthogonal and can be stacked for aggressive compression. The order in which you apply them matters significantly.

6.1 The Recommended Pipeline

Step 1: Prune first. Apply SparseGPT or Wanda to achieve the target sparsity. Pruning works best on full-precision weights where the importance signal is cleanest.

Step 2: Quantize the sparse model. After pruning, quantize the remaining nonzero weights to INT4 or INT8 using GPTQ or AWQ (Section 9.1). The calibration data for quantization should reflect the pruned model's behavior, not the original dense model's.

Step 3: Optionally, distill. If quality after pruning and quantization is insufficient, use the original dense model as a teacher to fine-tune the compressed student. Even a few hundred steps of distillation can recover 2 to 5 percentage points on benchmarks.

6.2 Sparse Quantized Formats

The combination of 2:4 sparsity and INT8 quantization is particularly well supported on NVIDIA hardware, achieving a combined 4x speedup (2x from sparsity times 2x from quantization) with a single fused kernel. For even more aggressive compression, some teams apply 2:4 sparsity at INT4 precision, though this requires careful quality validation.

7. Practical Guidelines

Choosing the right pruning approach depends on your deployment hardware, latency requirements, and acceptable quality loss. The following decision framework covers the most common scenarios.

7.1 Which Method When

7.2 Sparsity Targets and Quality Tradeoffs

The 50% sweet spot. Across most methods and models, 50% sparsity is the point where you get meaningful compression with minimal quality loss. Below 50%, the compression savings are often not worth the engineering effort. Above 70%, quality degradation becomes noticeable on challenging benchmarks, though task-specific performance may still be acceptable.

Layer-wise sparsity allocation. Not all layers are equally pruneable. Embedding layers and the final language modeling head should generally be left dense (or pruned very conservatively), as they have outsized impact on model quality. Middle transformer layers tolerate higher sparsity than the first and last few layers. Tools like SparseGPT and Wanda support per-layer sparsity targets to exploit this observation.

Calibration data matters. The choice of calibration data for one-shot methods (SparseGPT, Wanda) significantly affects which weights are identified as important. Use calibration data that resembles your deployment distribution. If your application is code generation, calibrate on code samples, not Wikipedia articles. Even 128 well-chosen calibration examples can outperform 1024 poorly chosen ones.

7.3 Validation Strategy

After pruning, always validate on a held-out evaluation set that represents your production workload. Standard benchmarks (MMLU, HellaSwag) provide a general quality signal, but task-specific evaluation (Section 25.1) is essential because pruning can affect different capabilities unevenly. A model that retains 95% of its general benchmark score might have lost 20% on a specific subtask that matters for your application.

8. Exercises

pip install datasets transformers datasketch

from datasets import load_dataset

# Load a subset of a public corpus
ds = load_dataset("wikitext", "wikitext-103-raw-v1", split="train")
print(f"Original size: {len(ds):,} examples")
print(f"Sample: {ds[100]['text'][:200]}")

# Quality filters: remove empty lines and very short documents
def quality_filter(example):
 text = example["text"].strip()
 if len(text) < 50: # too short
 return False
 if text.startswith("="): # section headers only
 return False
 return True

filtered = ds.filter(quality_filter, num_proc=4)
print(f"After quality filter: {len(filtered):,} examples")
print(f"Removed: {len(ds) - len(filtered):,} ({(len(ds)-len(filtered))/len(ds)*100:.1f}%)")

import hashlib

seen_hashes = set()
duplicates = 0

def dedup_exact(example):
 global duplicates
 text_hash = hashlib.md5(example["text"].encode()).hexdigest()
 if text_hash in seen_hashes:
 duplicates += 1
 return False
 seen_hashes.add(text_hash)
 return True

deduped_exact = filtered.filter(dedup_exact, num_proc=1) # single-proc for hash set
print(f"After exact dedup: {len(deduped_exact):,} examples")
print(f"Exact duplicates removed: {duplicates:,}")

from datasketch import MinHash, MinHashLSH

# Build LSH index
lsh = MinHashLSH(threshold=0.8, num_perm=128)
minhashes = {}

def compute_minhash(text, num_perm=128):
 m = MinHash(num_perm=num_perm)
 for word in text.lower().split():
 m.update(word.encode("utf-8"))
 return m

# Index a subset (full corpus would take longer)
subset = deduped_exact.select(range(min(5000, len(deduped_exact))))
near_dupes = set()

for i, example in enumerate(subset):
 mh = compute_minhash(example["text"])
 # Check for near-duplicates already in the index
 candidates = lsh.query(mh)
 if candidates:
 near_dupes.add(i)
 else:
 lsh.insert(str(i), mh)

print(f"Near-duplicates found: {len(near_dupes)} out of {len(subset)}")
print(f"Near-duplicate rate: {len(near_dupes)/len(subset)*100:.2f}%")

# Remove near-duplicates
keep_indices = [i for i in range(len(subset)) if i not in near_dupes]
clean_data = subset.select(keep_indices)
print(f"Clean dataset size: {len(clean_data):,}")

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
block_size = 512

def tokenize_and_chunk(examples):
 tokens = tokenizer(examples["text"], truncation=False)["input_ids"]
 # Concatenate all tokens, then split into fixed-size blocks
 all_tokens = [t for doc in tokens for t in doc]
 chunks = [all_tokens[i:i+block_size]
 for i in range(0, len(all_tokens) - block_size + 1, block_size)]
 return {"input_ids": chunks}

tokenized = clean_data.map(
 tokenize_and_chunk, batched=True, remove_columns=clean_data.column_names
)
print(f"Training chunks: {len(tokenized):,} (each {block_size} tokens)")
print(f"Total tokens: {len(tokenized) * block_size:,}")