| """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) |
|
|
