| """Original SparseGPT one-shot unstructured pruning for nn.Linear.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| import time |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| class SparseGPT: |
| """Hessian-aware unstructured pruning (Frantar & Alistarh, arXiv:2301.00774).""" |
|
|
| def __init__(self, layer: nn.Linear): |
| self.layer = layer |
| self.dev = layer.weight.device |
| w = layer.weight.data |
| self.rows = w.shape[0] |
| self.columns = w.shape[1] |
| self.H = torch.zeros((self.columns, 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 len(inp.shape) == 3: |
| inp = inp.reshape((-1, inp.shape[-1])) |
| inp = inp.t() |
| if self.nsamples > 0: |
| self.H *= self.nsamples / (self.nsamples + tmp) |
| self.nsamples += tmp |
| inp = math.sqrt(2 / self.nsamples) * inp.float() |
| self.H += inp.matmul(inp.t()) |
|
|
| def fasterprune( |
| self, |
| sparsity: float, |
| prunen: int = 0, |
| prunem: int = 0, |
| blocksize: int = 128, |
| percdamp: float = 0.01, |
| verbose: bool = True, |
| ): |
| w = self.layer.weight.data.clone().float() |
| h = self.H |
| self.H = None |
| dead = torch.diag(h) == 0 |
| h[dead, dead] = 1 |
| w[:, dead] = 0 |
|
|
| losses = torch.zeros(self.rows, device=self.dev) |
| mean_diag = torch.mean(torch.diag(h)).clamp_min(1e-8) |
| damp = float(percdamp) * float(mean_diag) |
| diag = torch.arange(self.columns, device=self.dev) |
| hinv = None |
| for _ in range(16): |
| trial = h.clone() |
| trial[diag, diag] += damp |
| chol, info = torch.linalg.cholesky_ex(trial) |
| info_i = int(info.item()) if info.ndim == 0 else int(info.max().item()) |
| if info_i == 0: |
| inv = torch.cholesky_inverse(chol) |
| inv = 0.5 * (inv + inv.transpose(-1, -2)) |
| inv[diag, diag] += damp |
| inv_chol, info2 = torch.linalg.cholesky_ex(inv, upper=True) |
| info2_i = int(info2.item()) if info2.ndim == 0 else int(info2.max().item()) |
| if info2_i == 0: |
| hinv = inv_chol |
| break |
| damp *= 10.0 |
| if hinv is None: |
| raise torch.linalg.LinAlgError( |
| f"SparseGPT Hessian not SPD after damping; columns={self.columns}" |
| ) |
|
|
| mask = None |
| tick = time.time() |
|
|
| for i1 in range(0, self.columns, blocksize): |
| i2 = min(i1 + blocksize, self.columns) |
| count = i2 - i1 |
| w1 = w[:, i1:i2].clone() |
| q1 = torch.zeros_like(w1) |
| err1 = torch.zeros_like(w1) |
| losses1 = torch.zeros_like(w1) |
| hinv1 = hinv[i1:i2, i1:i2] |
|
|
| if prunen == 0: |
| if mask is not None: |
| mask1 = mask[:, i1:i2] |
| else: |
| tmp = w1 ** 2 / (torch.diag(hinv1).reshape((1, -1))) ** 2 |
| thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)] |
| mask1 = tmp <= thresh |
| else: |
| mask1 = torch.zeros_like(w1) == 1 |
|
|
| for i in range(count): |
| wi = w1[:, i] |
| d = hinv1[i, i] |
| if prunen != 0 and i % prunem == 0: |
| tmp = w1[:, i:(i + prunem)] ** 2 / ( |
| torch.diag(hinv1)[i:(i + prunem)].reshape((1, -1)) |
| ) ** 2 |
| mask1.scatter_(1, i + torch.topk(tmp, prunen, dim=1, largest=False)[1], True) |
| q = wi.clone() |
| q[mask1[:, i]] = 0 |
| q1[:, i] = q |
| losses1[:, i] = (wi - q) ** 2 / d ** 2 |
| err = (wi - q) / d |
| w1[:, i:] -= err.unsqueeze(1).matmul(hinv1[i, i:].unsqueeze(0)) |
| err1[:, i] = err |
|
|
| w[:, i1:i2] = q1 |
| losses += torch.sum(losses1, 1) / 2 |
| w[:, i2:] -= err1.matmul(hinv[i1:i2, i2:]) |
|
|
| sparsity_val = 1.0 - float((w != 0).sum()) / float(w.numel()) |
| if verbose: |
| if self.dev.type == "cuda": |
| torch.cuda.synchronize() |
| print( |
| f"sparsegpt {self.layer}: time {time.time() - tick:.2f}s " |
| f"error {torch.sum(losses).item():.4f} " |
| f"sparsity {sparsity_val:.3f}" |
| ) |
| self.layer.weight.data = w.reshape(self.layer.weight.shape).to(self.layer.weight.dtype) |
| return sparsity_val |
|
|
| def free(self): |
| self.H = None |
| torch.cuda.empty_cache() |
|
|
