File size: 4,021 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Original Wanda/SparseGPT importance adapted to SwiGLU mid-neuron groups.
The algorithms remain unstructured at the score level, but the selected
coordinates are removed as complete SwiGLU neurons so checkpoint shapes shrink.
"""
from __future__ import annotations
import torch
def wanda_mid_scores(energy: torch.Tensor, down_weight: torch.Tensor) -> torch.Tensor:
"""Return classic Wanda scores for V/A/T and each SwiGLU mid coordinate."""
if energy.ndim != 2 or energy.shape[0] != 3:
raise ValueError(f"energy must be [3, d_int], got {tuple(energy.shape)}")
if down_weight.ndim != 2 or down_weight.shape[1] != energy.shape[1]:
raise ValueError(
f"down_weight must be [hidden, {energy.shape[1]}], got {tuple(down_weight.shape)}"
)
col_norm = down_weight.detach().float().norm(dim=0)
return col_norm.unsqueeze(0) * energy.detach().float().clamp_min(0).sqrt()
def _damped_hinv(gram: torch.Tensor, percdamp: float) -> torch.Tensor:
"""Upper-triangular Cholesky factor of damped H^{-1} (miniSeedance SparseGPT)."""
if gram.ndim != 2 or gram.shape[0] != gram.shape[1]:
raise ValueError(f"gram must be square, got {tuple(gram.shape)}")
h = gram.detach().float().clone()
h = (h + h.T) * 0.5
diag = torch.diag(h)
dead = diag == 0
if bool(dead.any()):
ids = torch.arange(h.shape[0], device=h.device)
h[ids[dead], ids[dead]] = 1.0
diag = torch.diag(h)
damping = float(percdamp) * diag.clamp_min(1e-8).mean()
ids = torch.arange(h.shape[0], device=h.device)
inv = None
for _ in range(8):
trial = h.clone()
trial[ids, ids] += damping
chol, info = torch.linalg.cholesky_ex(trial)
if int(info.max()) == 0:
inv = torch.cholesky_inverse(chol)
break
damping *= 10.0
if inv is None:
raise torch.linalg.LinAlgError("unable to stabilize Gram matrix for SparseGPT")
return torch.linalg.cholesky(inv, upper=True)
def sparsegpt_mid_scores(
gram: torch.Tensor,
energy: torch.Tensor,
down_weight: torch.Tensor,
percdamp: float = 0.01,
) -> torch.Tensor:
"""Return OBS/SparseGPT salience for grouped SwiGLU mid coordinates."""
if energy.ndim != 2 or energy.shape[0] != 3:
raise ValueError(f"energy must be [3, d_int], got {tuple(energy.shape)}")
if down_weight.ndim != 2 or down_weight.shape[1] != energy.shape[1]:
raise ValueError(
f"down_weight must be [hidden, {energy.shape[1]}], got {tuple(down_weight.shape)}"
)
hinv = _damped_hinv(gram, percdamp)
d = torch.diag(hinv).clamp_min(1e-8)
base = down_weight.detach().float().pow(2).sum(dim=0) / d.pow(2)
return energy.detach().float().clamp_min(0) * base.unsqueeze(0)
def sparsegpt_compensate_down(
down_weight: torch.Tensor,
gram: torch.Tensor,
keep_idx: torch.Tensor,
percdamp: float = 0.01,
eps: float = 1e-8,
) -> torch.Tensor:
"""OBS-style compensation on down_proj before physical slice (miniSeedance SparseGPT)."""
keep_idx = keep_idx.long()
w = down_weight.detach().clone().float()
d_int = w.shape[1]
keep_set = set(keep_idx.tolist())
drop = [i for i in range(d_int) if i not in keep_set]
if not drop:
return w.to(down_weight.dtype)
work = gram.device
if work.type == "cpu" and torch.cuda.is_available():
work = torch.device("cuda")
w = w.to(work)
gram = gram.to(work)
hinv = _damped_hinv(gram, percdamp)
d_diag = torch.diag(hinv).clamp_min(eps)
importance = w.pow(2).sum(dim=0) / d_diag.pow(2)
drop_t = torch.tensor(drop, device=w.device, dtype=torch.long)
drop_sorted = drop_t[torch.argsort(importance[drop_t])]
for i in drop_sorted.tolist():
col = w[:, i]
d = hinv[i, i]
err = col / d
w[:, i:] -= err.unsqueeze(1).matmul(hinv[i, i:].unsqueeze(0))
w[:, i] = 0
return w.to(device=down_weight.device, dtype=down_weight.dtype)
|