| """Original Wanda unstructured pruning for nn.Linear (Sun et al., 2023).""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| class WrappedGPT: |
| """Accumulate per-input-channel squared L2, matching others/wanda WrappedGPT.""" |
|
|
| def __init__(self, layer: nn.Linear): |
| self.layer = layer |
| self.dev = layer.weight.device |
| self.columns = layer.weight.data.shape[1] |
| self.scaler_row = torch.zeros((self.columns), device=self.dev) |
| self.nsamples = 0 |
|
|
| def add_batch(self, inp: torch.Tensor, out: torch.Tensor): |
| if len(inp.shape) == 2: |
| inp = inp.unsqueeze(0) |
| tmp = inp.shape[0] |
| if isinstance(self.layer, nn.Linear): |
| if len(inp.shape) == 3: |
| inp = inp.reshape((-1, inp.shape[-1])) |
| inp = inp.t() |
| self.scaler_row *= self.nsamples / (self.nsamples + tmp) |
| self.nsamples += tmp |
| inp = inp.type(torch.float32) |
| self.scaler_row += torch.norm(inp, p=2, dim=1) ** 2 / self.nsamples |
|
|
|
|
| def prune_linear_wanda(linear: nn.Linear, scaler_row: torch.Tensor, sparsity: float) -> float: |
| """Zero the lowest |W| * sqrt(scaler) entries in each output row. Return realized sparsity.""" |
| if not 0.0 <= float(sparsity) < 1.0: |
| raise ValueError(f"sparsity must be in [0, 1), got {sparsity}") |
| w = linear.weight.data |
| metric = torch.abs(w) * torch.sqrt( |
| scaler_row.reshape((1, -1)).to(device=w.device, dtype=torch.float32) |
| ) |
| n_prune = int(metric.shape[1] * float(sparsity)) |
| mask = torch.zeros_like(metric, dtype=torch.bool) |
| if n_prune > 0: |
| indices = torch.sort(metric, dim=-1, stable=True)[1][:, :n_prune] |
| mask.scatter_(1, indices, True) |
| w[mask] = 0 |
| return float(mask.sum()) / float(mask.numel()) |
|
|
