🤖 LLM Pruning: Movement Pruning, SparseGPT, and Wanda

import torch
import torch.nn as nn


class MovementPrunedLinear(nn.Module):
    """
    Linear layer with movement-based pruning via learned soft gates.

    The effective weight is: w_eff = sigmoid(alpha) * w
    Alpha is jointly trained with w. After fine-tuning, weights with
    alpha < 0 (sigmoid(alpha) < 0.5) are pruned.

    Regularizer: L1 on sigmoid(alpha) drives gates toward 0 (sparse).
    """

    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features) * 0.02)
        self.alpha = nn.Parameter(torch.zeros(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.bias = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate = torch.sigmoid(self.alpha)
        return nn.functional.linear(x, gate * self.weight, self.bias)

    def movement_regularizer(self) -> torch.Tensor:
        """ℓ₁ on gates — drives some alpha → -∞ (gate → 0), pruning those weights."""
        return torch.sigmoid(self.alpha).sum()

    def apply_mask(self, threshold: float = 0.5) -> None:
        """Hard-prune: zero out weights where sigmoid(alpha) < threshold."""
        with torch.no_grad():
            gate = torch.sigmoid(self.alpha)
            self.weight.data.mul_(gate.ge(threshold).float())
            self.alpha.data.fill_(float("inf"))  # freeze surviving gates at 1


def movement_prune_training_loss(
    model: nn.Module,
    task_loss: torch.Tensor,
    lam: float,
) -> torch.Tensor:
    """Add movement regularizer to the task loss during fine-tuning."""
    reg = sum(
        m.movement_regularizer()
        for m in model.modules()
        if isinstance(m, MovementPrunedLinear)
    )
    return task_loss + lam * reg

import torch
import torch.nn as nn
from typing import Optional


class SparseGPTPruner:
    """
    One-shot layer-wise pruning following SparseGPT (Frantar & Alistarh, 2023).

    Usage:
        pruner = SparseGPTPruner(layer)
        # Register hook and run calibration forward passes
        for batch in calibration_loader:
            _ = model(batch)  # triggers hook
        pruner.prune(sparsity=0.5)
        pruner.remove_hook()
    """

    def __init__(self, layer: nn.Linear, damping_factor: float = 1e-2):
        self.layer = layer
        self.damping = damping_factor
        self.H: Optional[torch.Tensor] = None
        self.n_samples = 0
        self._hook = layer.register_forward_hook(self._accumulate_H)

    def _accumulate_H(
        self,
        module: nn.Module,
        inp: tuple,
        out: torch.Tensor,
    ) -> None:
        x = inp[0].detach().float()
        if x.dim() == 3:
            x = x.reshape(-1, x.size(-1))  # (batch * seq, d_in)
        n = x.size(0)
        if self.H is None:
            d = x.size(1)
            self.H = torch.zeros(d, d, device=x.device, dtype=torch.float32)
        self.H.addmm_(x.T, x)
        self.n_samples += n

    def remove_hook(self) -> None:
        self._hook.remove()

    def prune(self, sparsity: float, block_size: int = 128) -> None:
        """
        Apply SparseGPT pruning: sequential OBS column-wise within blocks.

        Args:
            sparsity: fraction of weights to zero out
            block_size: number of columns to process jointly (trade accuracy for speed)
        """
        assert self.H is not None, "Run calibration forward passes first"
        self.remove_hook()

        d_in = self.H.size(0)
        W = self.layer.weight.data.float().clone()  # (d_out, d_in)

        # Normalize and damp
        H = self.H / self.n_samples
        H += self.damping * torch.eye(d_in, device=H.device)

        # Cholesky decomposition for stable inversion
        try:
            L = torch.linalg.cholesky(H)
            H_inv = torch.cholesky_inverse(L)
        except torch.linalg.LinAlgError:
            H_inv = torch.linalg.pinv(H)

        n_prune_per_row = round(sparsity * d_in)
        mask = torch.ones(W.shape, dtype=torch.bool, device=W.device)

        # Process each output row independently
        for row in range(W.size(0)):
            w = W[row].clone()
            H_inv_local = H_inv.clone()
            active = torch.ones(d_in, dtype=torch.bool, device=w.device)

            for _ in range(n_prune_per_row):
                # OBS saliency for active weights
                diag = H_inv_local.diag().clamp(min=1e-8)
                saliency = w.pow(2) / (2.0 * diag)
                saliency[~active] = float("inf")
                j = int(saliency.argmin())

                # OBS weight correction
                scale = w[j] / H_inv_local[j, j].clamp(min=1e-8)
                w -= scale * H_inv_local[:, j]
                w[j] = 0.0
                active[j] = False

                # Rank-1 Woodbury update of H_inv (Schur complement formula)
                hjj = H_inv_local[j, j].clamp(min=1e-8)
                H_inv_local -= torch.outer(
                    H_inv_local[:, j], H_inv_local[j, :]
                ) / hjj

            W[row] = w
            mask[row] = active

        self.layer.weight.data = W.to(self.layer.weight.dtype)
        # Store mask for potential downstream use
        self.mask = mask

    def prune_2_4(self) -> None:
        """
        Apply NVIDIA 2:4 structured sparsity: exactly 2 zeros per 4 consecutive weights.
        Uses OBS saliency within each group of 4 to select which 2 to prune.
        """
        assert self.H is not None
        self.remove_hook()

        d_in = self.H.size(0)
        W = self.layer.weight.data.float().clone()
        H = self.H / self.n_samples
        H += self.damping * torch.eye(d_in, device=H.device)
        H_inv = torch.cholesky_inverse(torch.linalg.cholesky(H))

        for row in range(W.size(0)):
            w = W[row]
            # Process in groups of 4
            for g in range(0, d_in, 4):
                grp = slice(g, min(g + 4, d_in))
                w_grp = w[grp]
                H_inv_grp = H_inv[grp, :][:, grp]
                diag_grp = H_inv_grp.diag().clamp(min=1e-8)
                saliency = w_grp.pow(2) / (2.0 * diag_grp)
                # Prune 2 lowest-saliency within group
                n_grp = len(w_grp)
                n_prune = max(0, n_grp - 2)
                if n_prune > 0:
                    prune_local = saliency.topk(n_prune, largest=False).indices
                    w_grp[prune_local] = 0.0
            W[row] = w

        self.layer.weight.data = W.to(self.layer.weight.dtype)

import torch
import torch.nn as nn
from typing import Optional


class WandaPruner:
    """
    Wanda pruning: saliency = |W_ij| * ||X_{:,j}||_2.
    Requires one forward pass on calibration data; no weight updates.
    """

    def __init__(self, layer: nn.Linear):
        self.layer = layer
        self._activation_norms: Optional[torch.Tensor] = None
        self._n_samples = 0
        self._hook = layer.register_forward_hook(self._accumulate_activations)

    def _accumulate_activations(
        self,
        module: nn.Module,
        inp: tuple,
        out: torch.Tensor,
    ) -> None:
        x = inp[0].detach().float()
        if x.dim() == 3:
            x = x.reshape(-1, x.size(-1))  # (n_tokens, d_in)
        if self._activation_norms is None:
            self._activation_norms = torch.zeros(x.size(1), device=x.device)
        # Accumulate sum of squares: ||X_{:,j}||_2^2 = sum_t X_{tj}^2
        self._activation_norms.add_(x.pow(2).sum(dim=0))
        self._n_samples += x.size(0)

    def remove_hook(self) -> None:
        self._hook.remove()

    def prune(self, sparsity: float) -> None:
        """
        Prune `sparsity` fraction of weights per output row using Wanda saliency.
        Pruning is row-local: each output neuron has its own threshold.
        """
        assert self._activation_norms is not None, "Run calibration data first"
        self.remove_hook()

        # ||X_{:,j}||_2 = sqrt(sum of squares)
        act_norms = self._activation_norms.sqrt()  # (d_in,)

        W = self.layer.weight.data.float()  # (d_out, d_in)
        n_prune = round(sparsity * W.size(1))

        # Saliency: |W_ij| * ||X_{:,j}||_2
        # Broadcast act_norms (d_in,) across output rows
        saliency = W.abs() * act_norms.unsqueeze(0)  # (d_out, d_in)

        # Per-row threshold (row-local pruning)
        thresholds = saliency.kthvalue(n_prune, dim=1).values  # (d_out,)
        mask = saliency.gt(thresholds.unsqueeze(1))  # (d_out, d_in)

        self.layer.weight.data = (W * mask.float()).to(self.layer.weight.dtype)


# ──────────────────────────────────────────────────────────────────────────────
# Convenience: prune an entire model's linear layers with Wanda
# ──────────────────────────────────────────────────────────────────────────────

def wanda_prune_model(
    model: nn.Module,
    calibration_loader: torch.utils.data.DataLoader,
    sparsity: float,
    device: str = "cuda",
    n_calibration_batches: int = 4,
) -> None:
    """
    Prune all Linear layers in `model` using Wanda at the given sparsity.
    Runs calibration forward passes, then prunes each layer in one shot.
    """
    model.eval()
    model.to(device)

    # Register Wanda hooks on all Linear layers
    pruners = {
        name: WandaPruner(module)
        for name, module in model.named_modules()
        if isinstance(module, nn.Linear)
    }

    # Calibration forward passes (no grad needed)
    with torch.no_grad():
        for i, batch in enumerate(calibration_loader):
            if i >= n_calibration_batches:
                break
            if isinstance(batch, (list, tuple)):
                batch = batch[0]
            model(batch.to(device))

    # Prune each layer
    for name, pruner in pruners.items():
        pruner.prune(sparsity=sparsity)

    print(f"Pruned {len(pruners)} linear layers to {sparsity:.0%} sparsity")