File size: 1,805 Bytes
fc3971e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """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())
|