Model Pruning

import torchfrom transformers import AutoModelForCausalLM, AutoTokenizer
# Load modelmodel = 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 functiondef wanda_prune(model, calib_data, sparsity=0.5):    """    Wanda: Prune by weight magnitude × input activation.
    Args:        sparsity: Fraction of weights to prune (0.5 = 50%)    """    # 1. Collect activation statistics    activations = {}
    def hook_fn(name):        def hook(module, input, output):            # Store input activation norms            activations[name] = input[0].detach().abs().mean(dim=0)        return hook
    # Register hooks for 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    model.eval()    with torch.no_grad():        for text in calib_data:            inputs = tokenizer(text, return_tensors="pt").to(model.device)            model(**inputs)
    # Remove hooks    for hook in hooks:        hook.remove()
    # 2. Prune weights based on |weight| × activation    for name, module in model.named_modules():        if isinstance(module, torch.nn.Linear) and name in activations:            W = module.weight.data            act = activations[name]
            # Compute importance: |weight| × activation            importance = W.abs() * act.unsqueeze(0)
            # Flatten and find threshold            threshold = torch.quantile(importance.flatten(), sparsity)
            # Create mask            mask = importance >= threshold
            # Apply mask (prune)            W *= mask.float()
    return model
# Apply Wanda pruning (50% sparsity, one-shot, no retraining)pruned_model = wanda_prune(model, calib_data, sparsity=0.5)
# Savepruned_model.save_pretrained("./llama-2-7b-wanda-50")

from sparsegpt import SparseGPT
# Load modelmodel = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
# Initialize SparseGPTpruner = SparseGPT(model)
# Calibration datacalib_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

def nm_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():    if isinstance(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

# Prune weights with smallest absolute valuesimportance = weight.abs()threshold = torch.quantile(importance, sparsity)mask = importance >= threshold

# Importance = |weight| × input_activationimportance = weight.abs() * activation# Better than magnitude alone (considers usage)

# Uses Hessian (second derivative) for importance# More accurate but computationally expensiveimportance = weight^2 / diag(Hessian)

# Unstructured (random)# [1, 0, 1, 0, 1, 1, 0, 0]# Pros: Flexible, high quality# Cons: No speedup
# Structured (block)# [1, 1, 0, 0, 1, 1, 0, 0]# Pros: Hardware friendly# Cons: More accuracy loss
# N:M (semi-structured)# [1, 0, 1, 0] [1, 1, 0, 0]  (2:4 pattern)# Pros: Hardware speedup + good quality# Cons: Requires specific hardware (NVIDIA)

def gradual_prune(model, initial_sparsity=0.0, final_sparsity=0.5, num_steps=100):    """Gradually increase sparsity during training."""    for step in range(num_steps):        # Current sparsity        current_sparsity = initial_sparsity + (final_sparsity - initial_sparsity) * (step / num_steps)
        # Prune at current sparsity        for module in model.modules():            if isinstance(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

def layer_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():        if isinstance(module, torch.nn.Linear):            # Find layer index            for layer_name, sparsity in sparsity_schedule.items():                if layer_name in name:                    # Prune at layer-specific sparsity                    prune_layer(module, sparsity)                    break
    return model

def iterative_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 in range(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

from transformers import Trainer, TrainingArguments
def production_pruning_pipeline(    model_name="meta-llama/Llama-2-7b-hf",    target_sparsity=0.5,    method="wanda",  # or "sparsegpt"):    # 1. Load model    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)    tokenizer = AutoTokenizer.from_pretrained(model_name)
    # 2. Load calibration data    calib_dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1000]")
    # 3. Apply pruning    if method == "wanda":        pruned_model = wanda_prune(model, calib_dataset, sparsity=target_sparsity)    elif method == "sparsegpt":        pruner = SparseGPT(model)        pruned_model = pruner.prune(calib_dataset, sparsity=target_sparsity)
    # 4. (Optional) Fine-tune to recover accuracy    training_args = TrainingArguments(        output_dir="./pruned-model",        num_train_epochs=1,        per_device_train_batch_size=4,        learning_rate=1e-5,        bf16=True,    )
    trainer = Trainer(        model=pruned_model,        args=training_args,        train_dataset=finetune_dataset,    )
    trainer.train()
    # 5. Save    pruned_model.save_pretrained("./pruned-llama-7b-50")    tokenizer.save_pretrained("./pruned-llama-7b-50")
    return pruned_model
# Usagepruned_model = production_pruning_pipeline(    model_name="meta-llama/Llama-2-7b-hf",    target_sparsity=0.5,    method="wanda")

from lm_eval import evaluator
# Evaluate pruned vs original modeloriginal_results = evaluator.simple_evaluate(    model="hf",    model_args="pretrained=meta-llama/Llama-2-7b-hf",    tasks=["arc_easy", "hellaswag", "winogrande"],)
pruned_results = evaluator.simple_evaluate(    model="hf",    model_args="pretrained=./pruned-llama-7b-50",    tasks=["arc_easy", "hellaswag", "winogrande"],)
# Compareprint(f"Original: {original_results['results']['arc_easy']['acc']:.3f}")print(f"Pruned:   {pruned_results['results']['arc_easy']['acc']:.3f}")print(f"Degradation: {(original_results - pruned_results):.3f}")
# Typical results at 50% sparsity:# - Wanda: <1% accuracy loss# - SparseGPT: <0.5% accuracy loss# - Magnitude: 2-3% accuracy loss

# Conservative (safe)sparsity = 0.3  # 30%, <0.5% loss
# Balanced (recommended)sparsity = 0.5  # 50%, ~1% loss
# Aggressive (risky)sparsity = 0.7  # 70%, 2-5% loss
# Extreme (model-dependent)sparsity = 0.9  # 90%, significant degradation

# 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

# ❌ Bad: Pruning without calibration dataprune_random(model)  # No activation statistics
# ✅ Good: Use calibration dataprune_wanda(model, calib_data)
# ❌ Bad: Too high sparsity in one shotprune(model, sparsity=0.9)  # Massive accuracy loss
# ✅ Good: Gradual or iterativeiterative_prune(model, target=0.9, steps=10)