Reduce LLM size and accelerate inference using pruning techniques like Wanda and SparseGPT. Use when compressing models without retraining, achieving 50% sparsity with minimal accuracy loss, or enabling faster inference on hardware accelerators. Covers unstructured pruning, structured pruning, N:M sparsity, magnitude pruning, and one-shot methods.
Instalar com Codex ou Claude Copie este prompt, cole no Codex, Claude ou outro assistente e deixe que ele revise a página da skill e instale para você.
Um comando direto ignora o prompt de revisão. Verifique a origem antes de executá-lo.
O comando permanece em uma só linha. Role horizontalmente para revisá-lo antes de copiar.
Prefere uma cópia local? Baixe os arquivos disponíveis atualmente no SkillsMP.
Explorador de arquivos
2 arquivos
Exibindo SKILL.md
SKILL.md
Instruções da origem · Visualização somente leitura
name
model-pruning
description
Reduce LLM size and accelerate inference using pruning techniques like Wanda and SparseGPT. Use when compressing models without retraining, achieving 50% sparsity with minimal accuracy loss, or enabling faster inference on hardware accelerators. Covers unstructured pruning, structured pruning, N:M sparsity, magnitude pruning, and one-shot methods.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# Load model
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-7b-hf",
torch_dtype=torch.float16,
device_map="cuda"
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
# Calibration data (small dataset for activation statistics)
calib_data = [
"The quick brown fox jumps over the lazy dog.",
"Machine learning is transforming the world.",
"Artificial intelligence powers modern applications.",
]
# Wanda pruning functiondefwanda_prune(model, calib_data, sparsity=):
activations = {}
():
():
activations[name] = [].detach().().mean(dim=)
hook
hooks = []
name, module model.named_modules():
(module, torch.nn.Linear):
hooks.append(module.register_forward_hook(hook_fn(name)))
model.()
torch.no_grad():
text calib_data:
inputs = tokenizer(text, return_tensors=).to(model.device)
model(**inputs)
hook hooks:
hook.remove()
name, module model.named_modules():
(module, torch.nn.Linear) name activations:
W = module.weight.data
act = activations[name]
importance = W.() * act.unsqueeze()
threshold = torch.quantile(importance.flatten(), sparsity)
mask = importance >= threshold
W *= mask.()
model
pruned_model = wanda_prune(model, calib_data, sparsity=)
pruned_model.save_pretrained()
0.5
"""
Wanda: Prune by weight magnitude × input activation.
Args:
sparsity: Fraction of weights to prune (0.5 = 50%)
"""
# 1. Collect activation statistics
def
hook_fn
name
def
hook
module, input, output
# Store input activation norms
input
0
abs
0
return
# Register hooks for all linear layers
for
in
if
isinstance
# Run calibration data
eval
with
for
in
"pt"
# Remove hooks
for
in
# 2. Prune weights based on |weight| × activation
for
in
if
isinstance
and
in
# Compute importance: |weight| × activation
abs
0
# Flatten and find threshold
# Create mask
# Apply mask (prune)
float
return
# Apply Wanda pruning (50% sparsity, one-shot, no retraining)
0.5
# Save
"./llama-2-7b-wanda-50"
SparseGPT (Second-Order Pruning)
Source: arXiv 2301.00774
from sparsegpt import SparseGPT
# Load model
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
# Initialize SparseGPT
pruner = SparseGPT(model)
# Calibration data
calib_data = load_calibration_data() # ~128 samples# Prune (one-shot, layer-wise reconstruction)
pruned_model = pruner.prune(
calib_data=calib_data,
sparsity=0.5, # 50% sparsity
prunen=0, # Unstructured (0) or N:M structured
prunem=0,
percdamp=0.01, # Damping for Hessian inverse
)
# Results: Near-lossless pruning at 50% sparsity
N:M Structured Pruning (Hardware Accelerator)
defnm_prune(weight, n=2, m=4):
"""
N:M pruning: Keep N weights per M consecutive weights.
Example: 2:4 = keep 2 out of every 4 weights.
Compatible with NVIDIA sparse tensor cores (2:4, 4:8).
"""# Reshape weight into groups of M
shape = weight.shape
weight_flat = weight.flatten()
# Pad to multiple of M
pad_size = (m - weight_flat.numel() % m) % m
weight_padded = F.pad(weight_flat, (0, pad_size))
# Reshape into (num_groups, m)
weight_grouped = weight_padded.reshape(-1, m)
# Find top-N in each group
_, indices = torch.topk(weight_grouped.abs(), n, dim=-1)
# Create mask
mask = torch.zeros_like(weight_grouped)
mask.scatter_(1, indices, 1.0)
# Apply mask
weight_pruned = weight_grouped * mask
# Reshape back
weight_pruned = weight_pruned.flatten()[:weight_flat.numel()]
return weight_pruned.reshape(shape)
# Apply 2:4 sparsity (NVIDIA hardware)for name, module in model.named_modules():
ifisinstance(module, torch.nn.Linear):
module.weight.data = nm_prune(module.weight.data, n=2, m=4)
# 50% sparsity, 2× speedup on A100 with sparse tensor cores
defgradual_prune(model, initial_sparsity=0.0, final_sparsity=0.5, num_steps=100):
"""Gradually increase sparsity during training."""for step inrange(num_steps):
# Current sparsity
current_sparsity = initial_sparsity + (final_sparsity - initial_sparsity) * (step / num_steps)
# Prune at current sparsityfor module in model.modules():
ifisinstance(module, torch.nn.Linear):
weight = module.weight.data
threshold = torch.quantile(weight.abs().flatten(), current_sparsity)
mask = weight.abs() >= threshold
weight *= mask.float()
# Train one step
train_step(model)
return model
Strategy 2: Layer-wise Pruning
deflayer_wise_prune(model, sparsity_per_layer):
"""Different sparsity for different layers."""# Early layers: Less pruning (more important)# Late layers: More pruning (less critical)
sparsity_schedule = {
"layer.0": 0.3, # 30% sparsity"layer.1": 0.4,
"layer.2": 0.5,
"layer.3": 0.6, # 60% sparsity
}
for name, module in model.named_modules():
ifisinstance(module, torch.nn.Linear):
# Find layer indexfor layer_name, sparsity in sparsity_schedule.items():
if layer_name in name:
# Prune at layer-specific sparsity
prune_layer(module, sparsity)
breakreturn model
Strategy 3: Iterative Pruning + Fine-tuning
defiterative_prune_finetune(model, target_sparsity=0.5, iterations=5):
"""Prune gradually with fine-tuning between iterations."""
current_sparsity = 0.0
sparsity_increment = target_sparsity / iterations
for i inrange(iterations):
# Increase sparsity
current_sparsity += sparsity_increment
# Prune
prune_model(model, sparsity=current_sparsity)
# Fine-tune (recover accuracy)
fine_tune(model, epochs=2, lr=1e-5)
return model
# Results: Better accuracy than one-shot at high sparsity
# One-shot, no retraining → Wanda or SparseGPTif no_retraining_budget:
use_method = "wanda"# Faster# Best quality → SparseGPTif need_best_quality:
use_method = "sparsegpt"# More accurate# Hardware speedup → N:M structuredif need_speedup:
use_method = "nm_prune"# 2:4 or 4:8
3. Avoid Common Pitfalls
# ❌ Bad: Pruning without calibration data
prune_random(model) # No activation statistics# ✅ Good: Use calibration data
prune_wanda(model, calib_data)
# ❌ Bad: Too high sparsity in one shot
prune(model, sparsity=0.9) # Massive accuracy loss# ✅ Good: Gradual or iterative
iterative_prune(model, target=0.9, steps=10)
