Model Pruning & Sparsity

"A sculptor does not add marble to create a statue. She removes everything that is not the statue."

Quant, Elegantly Subtractive AI Agent

1. Why Pruning Matters for LLMs

Modern LLMs contain billions of parameters, yet research consistently shows that a large fraction of these weights contribute minimally to model output. The Lottery Ticket Hypothesis (Frankle & Carlin, 2019) demonstrated that dense networks contain sparse subnetworks that, when trained in isolation, match the full model's performance. For LLMs, this insight translates into a practical opportunity: we can remove 50% to 90% of weights with surprisingly little degradation in quality, yielding faster inference and lower memory consumption.

Pruning complements quantization, the technique covered in Section 9.1. While quantization reduces the number of bits per weight, pruning reduces the number of weights altogether. The two techniques are orthogonal and composable: you can prune a model and then quantize the remaining weights, achieving compression ratios that neither technique could reach alone.

Why most weights do not matter. Consider a 70B parameter model trained on trillions of tokens. During training, the model must handle every possible topic, language, and reasoning pattern in its training data. But at inference time, any given query activates only a small subset of the model's knowledge. The remaining weights are effectively dead code for that query. This is why Mixture-of-Experts architectures (discussed in Section 07.2) are so effective: they formalize this observation by routing each token to only a fraction of the model's parameters. Pruning achieves a similar effect after training, by permanently removing the weights that contribute least across a representative set of inputs. The connection to LoRA adapters (Section 15.1) is also worth noting: both pruning and LoRA exploit the fact that the effective dimensionality of the model's computation is much lower than the number of parameters suggests.

1.1 Pruning vs. Quantization vs. Distillation

2. Unstructured Pruning

Unstructured pruning sets individual weight values to zero without any constraint on their position within the weight matrix. This is the simplest and most flexible form of pruning: any weight in any position can be removed. The resulting weight matrix has the same shape as the original but contains many zeros (a sparse matrix). The challenge is that standard GPU hardware cannot efficiently skip over scattered zeros, so unstructured pruning often requires specialized sparse matrix libraries to translate sparsity into actual speedups.

2.1 Magnitude Pruning

The oldest and most intuitive pruning criterion is weight magnitude: remove the weights closest to zero. The logic is straightforward. A weight near zero multiplied by any activation produces a near-zero output, so removing it should have minimal effect. Magnitude pruning typically proceeds iteratively: prune a small percentage of weights, fine-tune the model briefly to recover quality, then prune again. Code Fragment 9.5.1 below puts this into practice.

# Unstructured magnitude pruning: zero out the smallest 30% of weights
# in each Linear layer using PyTorch's built-in prune module.
import torch
import torch.nn.utils.prune as prune

# Load a pretrained model (using a small example for illustration)
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")

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

# Check sparsity of a specific layer
layer = model.model.decoder.layers[0].self_attn.q_proj
total = layer.weight.nelement()
zeros = (layer.weight == 0).sum().item()
print(f"Sparsity: {zeros / total:.1%}")
# Sparsity: 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")

2.2 SparseGPT: One-Shot Pruning for Large Models

Iterative magnitude pruning requires multiple rounds of pruning and fine-tuning, which is prohibitively expensive for models with billions of parameters. SparseGPT (Frantar & Alistarh, 2023) solves this with a one-shot approach: it prunes the model in a single pass without any fine-tuning. The key idea is to compensate for each pruned weight by adjusting the remaining weights in the same row. SparseGPT frames this as a layer-wise sparse regression problem, solving for the optimal remaining weights given the pruning mask.

The algorithm processes each layer independently. For a weight matrix W and a calibration set of inputs X, SparseGPT minimizes the reconstruction error ||WX - W'X||, where W' is the pruned and adjusted weight matrix. Using an efficient solver based on the Hessian inverse (computed from the calibration data), SparseGPT can prune a 175-billion-parameter model in under 4 hours on a single GPU. At 50% unstructured sparsity, SparseGPT applied to LLaMA-65B increases perplexity by less than 0.5 points, a negligible degradation for most applications.

2.3 Wanda: Pruning Without Weight Updates

Wanda (Sun et al., 2024) takes an even simpler approach than SparseGPT. Instead of adjusting remaining weights after pruning, Wanda uses a pruning criterion that considers both the weight magnitude and the norm of the corresponding input activations. The importance score for weight w_ij is simply |w_ij| * ||x_j||, where ||x_j|| is the L2 norm of the j-th input feature computed over a small calibration set. Weights with low importance scores are pruned.

The intuition is elegant: a large weight connected to an input feature that is always near zero is less important than a small weight connected to a highly active feature. By incorporating activation statistics, Wanda captures the actual contribution of each weight during inference rather than relying on magnitude alone. Despite its simplicity (no weight updates, no iterative optimization), Wanda matches or approaches SparseGPT quality at 50% sparsity while running orders of magnitude faster. Code Fragment 9.5.2 below puts this into practice.

# Activation-aware pruning: register a forward hook to measure per-neuron
# activation magnitude, then prune channels with the lowest average activation.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "facebook/opt-1.3b"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Collect activation norms from calibration data
calibration_texts = [
 "The quick brown fox jumps over the lazy dog.",
 "Machine learning models require careful optimization.",
 # ... more calibration examples (typically 128 samples)
]

activation_norms = {}

def hook_fn(name):
 def hook(module, input, output):
 # input[0] shape: (batch, seq_len, hidden_dim)
 x = input[0].detach().float()
 norm = x.norm(dim=(0, 1)) # norm across batch and sequence
 if name in activation_norms:
 activation_norms[name] += norm
 else:
 activation_norms[name] = norm
 return hook

# Register hooks on all linear layers
hooks = []
for name, module in model.named_modules():
 if isinstance(module, torch.nn.Linear):
 hooks.append(module.register_forward_hook(hook_fn(name)))

# Run calibration data through the model
for text in calibration_texts:
 inputs = tokenizer(text, return_tensors="pt")
 with torch.no_grad():
 model(**inputs)

# Remove hooks
for h in hooks:
 h.remove()

# Compute Wanda importance scores and prune
sparsity = 0.5
for name, module in model.named_modules():
 if isinstance(module, torch.nn.Linear) and name in activation_norms:
 W = module.weight.data.float()
 # Importance = |weight| * activation_norm
 importance = W.abs() * activation_norms[name].unsqueeze(0)

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

 # Create and apply pruning mask
 mask = importance > threshold
 module.weight.data *= mask.half()
 print(f"{name}: pruned {(~mask).sum().item()} / {mask.numel()} weights")

3. Structured Pruning

Structured pruning removes entire structural units from the model: rows or columns of weight matrices, entire attention heads, or even full transformer layers. Unlike unstructured pruning, structured pruning produces a genuinely smaller model that runs faster on standard hardware without any special sparse matrix support. The trade-off is that structured pruning is coarser-grained, so it typically achieves less compression at a given quality level compared to unstructured approaches.

3.1 What Can Be Pruned

• Attention heads: Research by Michel et al. (2019) showed that many attention heads can be removed after training with minimal quality loss. In some models, over 40% of heads are redundant.
• Feed-forward neurons: Entire rows of the intermediate feed-forward layer can be removed. Since the FFN is typically the largest component of each transformer block (4x the hidden dimension), this offers significant compression.
• Transformer layers: For very deep models, removing entire layers (sometimes called layer pruning or depth pruning) can reduce both memory and latency substantially. Models like LLaMA-2-70B can lose several layers with minimal degradation on many benchmarks.
• Embedding dimensions: Reducing the hidden dimension uniformly across the model, sometimes called width pruning, produces a model that is architecturally identical to a smaller variant.

3.2 Structured Pruning with Importance Scores

The key question in structured pruning is how to score the importance of each structural unit. Common approaches include:

• Taylor expansion: Approximate the change in loss from removing a unit using a first-order Taylor expansion. The importance score for a unit with parameters θ is |θ · ∇θL|, which can be computed from a single backward pass over calibration data.
• Activation norm: Score each neuron or head by the average magnitude of its output activations. Units that consistently produce near-zero outputs are candidates for removal.
• Cosine similarity: Identify and merge redundant heads by measuring the cosine similarity between their output representations. Highly similar heads are likely redundant.

4. NVIDIA 2:4 Structured Sparsity

NVIDIA's Ampere architecture (A100 and newer GPUs) introduced hardware support for a specific sparsity pattern called 2:4 structured sparsity. The rule is simple: in every group of 4 consecutive weights, exactly 2 must be zero. This 50% sparsity pattern is restrictive compared to arbitrary unstructured pruning, but it has a crucial advantage: the GPU's Sparse Tensor Cores can process 2:4 sparse matrices at nearly 2x the throughput of dense matrices with no software overhead.

4.1 How 2:4 Sparsity Works at the Hardware Level

The GPU stores a 2:4 sparse matrix using two components: a compressed data array containing only the non-zero values (half the original size) and a small metadata array encoding which 2 of the 4 positions in each group are non-zero (2 bits per group). During matrix multiplication, the Sparse Tensor Core uses the metadata to select the correct input activations to multiply with the compressed weights, skipping the zero positions entirely. This selection happens in hardware, adding no software overhead and requiring no changes to the surrounding computation.

Figure 9.5.1: The 2:4 structured sparsity pattern. Every group of 4 weights retains exactly 2 non-zero values. The compressed representation stores only the non-zero values plus lightweight index metadata.

4.2 Applying 2:4 Sparsity in Practice

NVIDIA's ASP (Automatic Sparsity) library makes it straightforward to apply 2:4 sparsity during fine-tuning. The typical workflow is: (1) start from a dense pretrained model, (2) compute importance scores (magnitude or Wanda-style), (3) apply the 2:4 pruning mask, and (4) fine-tune for a few hundred steps to recover quality. The fine-tuning step is important because the 2:4 constraint forces pruning decisions at a coarser granularity than optimal unstructured pruning, so the model benefits from adapting to the specific sparsity pattern. Code Fragment 9.5.3 below puts this into practice.

# NVIDIA ASP (Automatic Sparsity): apply 2:4 structured sparsity to all
# Linear layers, enabling hardware-accelerated sparse matrix multiplication.
import torch
from transformers import AutoModelForCausalLM
# NVIDIA's Automatic Sparsity library
# pip install nvidia-pyindex asp
from apex.contrib.sparsity import ASP

model = AutoModelForCausalLM.from_pretrained(
 "meta-llama/Llama-2-7b-hf",
 torch_dtype=torch.float16,
 device_map="auto"
)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)

# Initialize ASP: analyzes model layers for 2:4 compatibility
ASP.init_model_for_pruning(
 model,
 mask_calculator="m4n2_1d", # 2:4 pattern (2 non-zero per 4)
 verbosity=2, # show which layers are pruned
 whitelist=[torch.nn.Linear], # only prune linear layers
 allow_recompute_mask=False # lock mask after initial computation
)

# Initialize optimizer with ASP awareness
ASP.init_optimizer_for_pruning(optimizer)

# Compute and apply the 2:4 pruning mask (one-shot)
ASP.compute_sparse_masks()

# Fine-tune for a few hundred steps to recover quality
# The mask is applied automatically during forward/backward passes
for step, batch in enumerate(calibration_dataloader):
 outputs = model(**batch)
 loss = outputs.loss
 loss.backward()
 optimizer.step()
 optimizer.zero_grad()
 if step >= 500:
 break

print("2:4 sparse model ready for Sparse Tensor Core inference")

5. Pruning + Quantization: The Compression Stack

Pruning and quantization target different axes of model compression and combine naturally. A practical compression pipeline might look like this: (1) apply SparseGPT or Wanda at 50% sparsity, (2) quantize the remaining weights to 4-bit using GPTQ or AWQ, and (3) deploy with a runtime that supports both sparse and quantized operations. The combined compression can reach 8x to 16x reduction in memory footprint. This makes otherwise impractical models viable for the hybrid ML/LLM architectures discussed in Chapter 12, where a compressed LLM runs alongside specialized smaller models.

5.1 Combining SparseGPT with GPTQ

The SparseGPT and GPTQ algorithms share the same mathematical framework (layer-wise optimal weight adjustment using the Hessian inverse), so they can be combined into a single pass. The joint algorithm first determines the pruning mask using the Wanda or magnitude criterion, then jointly optimizes the quantization grid and the remaining weight values to minimize reconstruction error. Knowledge distillation provides yet another compression axis that can be layered on top of pruning and quantization. This joint optimization produces better results than applying pruning and quantization sequentially because the quantization step can compensate for pruning errors and vice versa.

6. When to Prune vs. Quantize

Choosing between pruning and quantization (or using both) depends on your deployment constraints, hardware capabilities, and quality requirements. Here is a practical decision framework:

• Memory-constrained, standard GPU: Start with quantization (4-bit GPTQ or AWQ). This gives the best compression-to-quality ratio on hardware without sparse kernel support.
• Latency-constrained, NVIDIA Ampere+: Apply 2:4 structured sparsity for a nearly free 2x throughput gain, then add INT8 quantization. This maximizes throughput while maintaining quality.
• CPU or edge deployment: Use structured pruning to create a genuinely smaller model (fewer attention heads, narrower FFN layers), then quantize to INT8. The Neural Magic DeepSparse engine can also accelerate unstructured sparsity on CPUs.
• Maximum compression needed: Apply SparseGPT at 50% sparsity followed by GPTQ 4-bit quantization. This stacks both techniques for 8x+ compression.
• Quality is paramount: Use 2:4 structured sparsity alone (the most conservative pruning option) or skip pruning entirely and rely solely on quantization at higher bit-widths (8-bit or 6-bit).

Exercises