| """Accumulate TENP projection (and energy control) on routed tokens. No shared expert.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
|
|
| import torch |
|
|
| from .score import PROJ_EPS, gated_energy, gated_signed_projection |
| from .temporal import DEFAULT_NUM_BINS, timestep_to_bin |
|
|
| VIDEO, AUDIO, TEXT = 0, 1, 2 |
| VALID_MODALITIES = (VIDEO, AUDIO, TEXT) |
| DEFAULT_PERTURB_RATIO = 0.05 |
| DEFAULT_PERTURB_REPEATS = 3 |
| DEFAULT_TOKEN_CACHE_PER_BIN = 0 |
|
|
|
|
| def _as_2d_table(t: torch.Tensor, d_int: int) -> torch.Tensor: |
| if t.ndim == 2 and t.shape[0] == 3: |
| return t |
| raise ValueError(f"expected [3, d_int={d_int}] table, got {tuple(t.shape)}") |
|
|
|
|
| @dataclass |
| class ExpertStats: |
| projection: torch.Tensor |
| energy: torch.Tensor |
| gram: torch.Tensor | None = None |
| n_tokens: int = 0 |
| n_tokens_mod: torch.Tensor = None |
| projection_bin: torch.Tensor | None = None |
| energy_bin: torch.Tensor | None = None |
| n_tokens_bin: torch.Tensor | None = None |
| n_tokens_mod_bin: torch.Tensor | None = None |
| sensitivity_sum: torch.Tensor | None = None |
| sensitivity_count: torch.Tensor | None = None |
| rel_abs_z: torch.Tensor | None = None |
| token_cache_per_bin: int = 0 |
| cache_z: list | None = None |
| cache_y: list | None = None |
| cache_gate: list | None = None |
| cache_mod: list | None = None |
| _cache_seen: list = field(default_factory=list) |
|
|
| @property |
| def num_bins(self) -> int: |
| if self.projection_bin is None: |
| return 1 |
| return int(self.projection_bin.shape[0]) |
|
|
| @classmethod |
| def zeros( |
| cls, |
| d_int: int, |
| device, |
| num_bins: int = DEFAULT_NUM_BINS, |
| collect_gram: bool = True, |
| token_cache_per_bin: int = DEFAULT_TOKEN_CACHE_PER_BIN, |
| ) -> "ExpertStats": |
| k = max(int(num_bins), 1) |
| cap = max(int(token_cache_per_bin), 0) |
| gram = ( |
| torch.zeros(d_int, d_int, device=device, dtype=torch.float32) |
| if collect_gram |
| else None |
| ) |
| return cls( |
| projection=torch.zeros(3, d_int, device=device, dtype=torch.float32), |
| energy=torch.zeros(3, d_int, device=device, dtype=torch.float32), |
| gram=gram, |
| n_tokens_mod=torch.zeros(3, device=device, dtype=torch.long), |
| projection_bin=torch.zeros(k, 3, d_int, device=device, dtype=torch.float32), |
| energy_bin=torch.zeros(k, 3, d_int, device=device, dtype=torch.float32), |
| n_tokens_bin=torch.zeros(k, device=device, dtype=torch.long), |
| n_tokens_mod_bin=torch.zeros(k, 3, device=device, dtype=torch.long), |
| sensitivity_sum=torch.zeros(k, device=device, dtype=torch.float32), |
| sensitivity_count=torch.zeros(k, device=device, dtype=torch.float32), |
| rel_abs_z=torch.zeros(d_int, device=device, dtype=torch.float32), |
| token_cache_per_bin=cap, |
| cache_z=[None] * k if cap else None, |
| cache_y=[None] * k if cap else None, |
| cache_gate=[None] * k if cap else None, |
| _cache_seen=[0] * k if cap else [], |
| ) |
|
|
| def _ensure_bins(self, d_int: int) -> None: |
| k = self.num_bins |
| device = self.projection.device |
| if self.projection_bin is None: |
| self.projection_bin = torch.zeros(k, 3, d_int, device=device, dtype=torch.float32) |
| self.projection_bin[0] = self.projection |
| if self.energy_bin is None: |
| self.energy_bin = torch.zeros(k, 3, d_int, device=device, dtype=torch.float32) |
| self.energy_bin[0] = self.energy |
| if self.n_tokens_bin is None: |
| self.n_tokens_bin = torch.zeros(k, device=device, dtype=torch.long) |
| self.n_tokens_bin[0] = int(self.n_tokens) |
| if self.n_tokens_mod_bin is None: |
| self.n_tokens_mod_bin = torch.zeros(k, 3, device=device, dtype=torch.long) |
| self.n_tokens_mod_bin[0] = self.n_tokens_mod |
| if self.sensitivity_sum is None: |
| self.sensitivity_sum = torch.zeros(k, device=device, dtype=torch.float32) |
| if self.sensitivity_count is None: |
| self.sensitivity_count = torch.zeros(k, device=device, dtype=torch.float32) |
| if self.rel_abs_z is None: |
| self.rel_abs_z = torch.zeros(d_int, device=device, dtype=torch.float32) |
|
|
| def sensitivity(self, eps: float = 1e-8) -> torch.Tensor: |
| self._ensure_bins(int(self.projection.shape[-1])) |
| return self.sensitivity_sum / self.sensitivity_count.clamp_min(eps) |
|
|
| def update( |
| self, |
| z: torch.Tensor, |
| y: torch.Tensor, |
| gate: torch.Tensor, |
| modality: torch.Tensor, |
| W_down: torch.Tensor, |
| bin_id: int = 0, |
| collect_sensitivity: bool = False, |
| perturb_ratio: float = DEFAULT_PERTURB_RATIO, |
| perturb_repeats: int = DEFAULT_PERTURB_REPEATS, |
| generator: torch.Generator | None = None, |
| ) -> None: |
| valid = (modality >= 0) & (modality <= TEXT) |
| if int(valid.sum()) == 0: |
| return |
| z = z[valid] |
| y = y[valid] |
| gate = gate[valid] |
| modality = modality[valid] |
| q = gated_signed_projection(z, y, gate, W_down) |
| e = gated_energy(z, gate) |
| d_int = int(z.shape[-1]) |
| self._ensure_bins(d_int) |
| k = self.num_bins |
| b = min(max(int(bin_id), 0), k - 1) |
| if self.gram is not None: |
| z_g = z.float() * gate.float().reshape(-1, 1) |
| self.gram.add_(z_g.T @ z_g) |
| n = int(q.shape[0]) |
| self.n_tokens += n |
| self.n_tokens_bin[b] += n |
| y_norm = y.float().norm(dim=-1, keepdim=True).clamp_min(PROJ_EPS) |
| self.rel_abs_z.add_((z.float().abs() / y_norm).sum(dim=0)) |
| for m in VALID_MODALITIES: |
| mask = modality == m |
| n_m = int(mask.sum()) |
| if n_m == 0: |
| continue |
| self.n_tokens_mod[m] += n_m |
| self.n_tokens_mod_bin[b, m] += n_m |
| q_sum = q[mask].sum(dim=0) |
| e_sum = e[mask].sum(dim=0) |
| self.projection[m] += q_sum |
| self.energy[m] += e_sum |
| self.projection_bin[b, m] += q_sum |
| self.energy_bin[b, m] += e_sum |
| if collect_sensitivity: |
| d = perturbation_sensitivity( |
| z, y, W_down, perturb_ratio, perturb_repeats, generator |
| ) |
| self.sensitivity_sum[b] += float(d) * n |
| self.sensitivity_count[b] += n |
| self._maybe_cache(z, y, gate, modality, b) |
|
|
| def _maybe_cache( |
| self, |
| z: torch.Tensor, |
| y: torch.Tensor, |
| gate: torch.Tensor, |
| modality: torch.Tensor, |
| bin_id: int, |
| ) -> None: |
| """Reservoir of `token_cache_per_bin` tokens per (bin, modality). |
| |
| Per-modality reservoirs keep enough audio/text tokens (they are ~5% of routed |
| tokens) to measure per-modality output error after pruning. `cache_mod[bin]` |
| holds the modality id aligned with `cache_z[bin]`. Tensors are stored fp16. |
| """ |
| cap = int(self.token_cache_per_bin) |
| if cap <= 0: |
| return |
| k = self.num_bins |
| if self.cache_z is None: |
| self.cache_z = [None] * k |
| self.cache_y = [None] * k |
| self.cache_gate = [None] * k |
| self.cache_mod = [None] * k |
| self._cache_seen = [[0, 0, 0] for _ in range(k)] |
| if self.cache_mod is None: |
| self.cache_mod = [None] * k |
| if not isinstance(self._cache_seen[bin_id], list): |
| self._cache_seen[bin_id] = [int(self._cache_seen[bin_id]), 0, 0] |
| z_all = z.detach().to(torch.float16).cpu() |
| y_all = y.detach().to(torch.float16).cpu() |
| g_all = gate.detach().float().reshape(-1).cpu() |
| m_all = modality.detach().to(torch.int8).reshape(-1).cpu() |
| cur_z = self.cache_z[bin_id] |
| cur_mod = self.cache_mod[bin_id] |
| if cur_z is not None and cur_mod is None: |
| cur_mod = torch.full((int(cur_z.shape[0]),), -1, dtype=torch.int8) |
| parts_z, parts_y, parts_g, parts_m = [], [], [], [] |
| for m in VALID_MODALITIES: |
| sel = (m_all == m).nonzero(as_tuple=False).reshape(-1) |
| n = int(sel.numel()) |
| if cur_z is not None: |
| keep_m = (cur_mod == m).nonzero(as_tuple=False).reshape(-1) |
| z_m, y_m, g_m = cur_z[keep_m], self.cache_y[bin_id][keep_m], self.cache_gate[bin_id][keep_m] |
| else: |
| z_m = torch.empty(0, int(z_all.shape[-1]), dtype=torch.float16) |
| y_m = torch.empty(0, int(y_all.shape[-1]), dtype=torch.float16) |
| g_m = torch.empty(0, dtype=torch.float32) |
| seen = int(self._cache_seen[bin_id][m]) |
| if n > 0: |
| filled = int(z_m.shape[0]) |
| if filled < cap: |
| take = min(n, cap - filled) |
| z_m = torch.cat([z_m, z_all[sel[:take]]], dim=0) |
| y_m = torch.cat([y_m, y_all[sel[:take]]], dim=0) |
| g_m = torch.cat([g_m, g_all[sel[:take]]], dim=0) |
| sel = sel[take:] |
| seen += take |
| rem = int(sel.numel()) |
| if rem > 0: |
| |
| |
| counts = torch.arange(1, rem + 1, dtype=torch.float64) + float(seen) |
| r = (torch.rand(rem, dtype=torch.float64) * counts).long() |
| hit = (r < cap).nonzero(as_tuple=False).reshape(-1) |
| for i in hit.tolist(): |
| j = int(sel[i]) |
| slot = int(r[i]) |
| z_m[slot] = z_all[j] |
| y_m[slot] = y_all[j] |
| g_m[slot] = g_all[j] |
| seen += rem |
| self._cache_seen[bin_id][m] = seen |
| parts_z.append(z_m) |
| parts_y.append(y_m) |
| parts_g.append(g_m) |
| parts_m.append(torch.full((int(z_m.shape[0]),), int(m), dtype=torch.int8)) |
| self.cache_z[bin_id] = torch.cat(parts_z, dim=0) |
| self.cache_y[bin_id] = torch.cat(parts_y, dim=0) |
| self.cache_gate[bin_id] = torch.cat(parts_g, dim=0) |
| self.cache_mod[bin_id] = torch.cat(parts_m, dim=0) |
|
|
| def state_dict(self) -> dict: |
| def _cpu_list(items): |
| if items is None: |
| return None |
| out = [] |
| for t in items: |
| out.append(None if t is None else t.detach().cpu()) |
| return out |
|
|
| return { |
| "projection": self.projection.detach().cpu(), |
| "energy": self.energy.detach().cpu(), |
| "gram": None if self.gram is None else self.gram.detach().cpu(), |
| "n_tokens": int(self.n_tokens), |
| "n_tokens_mod": self.n_tokens_mod.detach().cpu(), |
| "projection_bin": None if self.projection_bin is None else self.projection_bin.detach().cpu(), |
| "energy_bin": None if self.energy_bin is None else self.energy_bin.detach().cpu(), |
| "n_tokens_bin": None if self.n_tokens_bin is None else self.n_tokens_bin.detach().cpu(), |
| "n_tokens_mod_bin": None if self.n_tokens_mod_bin is None else self.n_tokens_mod_bin.detach().cpu(), |
| "sensitivity_sum": None if self.sensitivity_sum is None else self.sensitivity_sum.detach().cpu(), |
| "sensitivity_count": None if self.sensitivity_count is None else self.sensitivity_count.detach().cpu(), |
| "rel_abs_z": None if self.rel_abs_z is None else self.rel_abs_z.detach().cpu(), |
| "token_cache_per_bin": int(self.token_cache_per_bin), |
| "cache_z": _cpu_list(self.cache_z), |
| "cache_y": _cpu_list(self.cache_y), |
| "cache_gate": _cpu_list(self.cache_gate), |
| "cache_mod": _cpu_list(self.cache_mod), |
| } |
|
|
| @classmethod |
| def from_state(cls, payload: dict, device="cpu") -> "ExpertStats": |
| projection = payload["projection"].to(device=device, dtype=torch.float32) |
| energy = payload["energy"].to(device=device, dtype=torch.float32) |
| d_int = int(projection.shape[-1]) |
| stats = cls( |
| projection=_as_2d_table(projection, d_int), |
| energy=_as_2d_table(energy, d_int), |
| gram=( |
| None |
| if payload.get("gram") is None |
| else payload["gram"].to(device=device, dtype=torch.float32) |
| ), |
| n_tokens=int(payload["n_tokens"]), |
| n_tokens_mod=payload["n_tokens_mod"].to(device=device, dtype=torch.long), |
| ) |
| proj_bin = payload.get("projection_bin") |
| if proj_bin is None: |
| stats.projection_bin = stats.projection.unsqueeze(0).clone() |
| stats.energy_bin = stats.energy.unsqueeze(0).clone() |
| stats.n_tokens_bin = torch.tensor( |
| [stats.n_tokens], device=device, dtype=torch.long |
| ) |
| stats.n_tokens_mod_bin = stats.n_tokens_mod.unsqueeze(0).clone() |
| stats.sensitivity_sum = torch.zeros(1, device=device, dtype=torch.float32) |
| stats.sensitivity_count = torch.zeros(1, device=device, dtype=torch.float32) |
| else: |
| stats.projection_bin = proj_bin.to(device=device, dtype=torch.float32) |
| energy_bin = payload.get("energy_bin") |
| stats.energy_bin = ( |
| stats.energy.unsqueeze(0).clone() |
| if energy_bin is None |
| else energy_bin.to(device=device, dtype=torch.float32) |
| ) |
| n_bin = payload.get("n_tokens_bin") |
| k = int(stats.projection_bin.shape[0]) |
| stats.n_tokens_bin = ( |
| torch.zeros(k, device=device, dtype=torch.long) |
| if n_bin is None |
| else n_bin.to(device=device, dtype=torch.long) |
| ) |
| n_mod_bin = payload.get("n_tokens_mod_bin") |
| stats.n_tokens_mod_bin = ( |
| torch.zeros(k, 3, device=device, dtype=torch.long) |
| if n_mod_bin is None |
| else n_mod_bin.to(device=device, dtype=torch.long) |
| ) |
| sens_sum = payload.get("sensitivity_sum") |
| stats.sensitivity_sum = ( |
| torch.zeros(k, device=device, dtype=torch.float32) |
| if sens_sum is None |
| else sens_sum.to(device=device, dtype=torch.float32) |
| ) |
| sens_n = payload.get("sensitivity_count") |
| stats.sensitivity_count = ( |
| torch.zeros(k, device=device, dtype=torch.float32) |
| if sens_n is None |
| else sens_n.to(device=device, dtype=torch.float32) |
| ) |
| rel = payload.get("rel_abs_z") |
| stats.rel_abs_z = ( |
| torch.zeros(d_int, device=device, dtype=torch.float32) |
| if rel is None |
| else rel.to(device=device, dtype=torch.float32) |
| ) |
| stats.token_cache_per_bin = int(payload.get("token_cache_per_bin") or 0) |
| cache_z = payload.get("cache_z") |
| if cache_z is None: |
| stats.cache_z = None |
| stats.cache_y = None |
| stats.cache_gate = None |
| stats.cache_mod = None |
| else: |
| stats.cache_z = [None if t is None else t.float().cpu() for t in cache_z] |
| cache_y = payload.get("cache_y") or [None] * len(cache_z) |
| cache_g = payload.get("cache_gate") or [None] * len(cache_z) |
| cache_m = payload.get("cache_mod") or [None] * len(cache_z) |
| stats.cache_y = [None if t is None else t.float().cpu() for t in cache_y] |
| stats.cache_gate = [None if t is None else t.float().cpu() for t in cache_g] |
| stats.cache_mod = [None if t is None else t.to(torch.int8).cpu() for t in cache_m] |
| return stats |
|
|
|
|
| def stacked_token_cache( |
| stats: ExpertStats, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: |
| """Concatenate per-bin token reservoirs. Returns (z, y, gate) or None.""" |
| if not stats.cache_z: |
| return None |
| zs, ys, gs = [], [], [] |
| for z, y, g in zip(stats.cache_z, stats.cache_y or [], stats.cache_gate or []): |
| if z is None or int(z.numel()) == 0: |
| continue |
| zs.append(z.float()) |
| ys.append(y.float()) |
| gs.append(g.float().reshape(-1)) |
| if not zs: |
| return None |
| return torch.cat(zs, dim=0), torch.cat(ys, dim=0), torch.cat(gs, dim=0) |
|
|
|
|
| def stacked_token_cache_mod( |
| stats: ExpertStats, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None: |
| """Like stacked_token_cache but also returns the per-token modality id (-1 if unknown).""" |
| packed = stacked_token_cache(stats) |
| if packed is None: |
| return None |
| z, y, g = packed |
| mods = [] |
| for i, zb in enumerate(stats.cache_z): |
| if zb is None or int(zb.numel()) == 0: |
| continue |
| m = None if stats.cache_mod is None or i >= len(stats.cache_mod) else stats.cache_mod[i] |
| if m is None or int(m.numel()) != int(zb.shape[0]): |
| m = torch.full((int(zb.shape[0]),), -1, dtype=torch.int8) |
| mods.append(m.to(torch.int8)) |
| return z, y, g, torch.cat(mods, dim=0) |
|
|
|
|
| def perturbation_sensitivity( |
| z: torch.Tensor, |
| y: torch.Tensor, |
| W_down: torch.Tensor, |
| ratio: float, |
| repeats: int, |
| generator: torch.Generator | None, |
| ) -> float: |
| """Relative L2 of randomly masking ρ mid-neurons, averaged over repeats.""" |
| n, d_int = z.shape |
| del n |
| n_mask = max(1, int(round(float(ratio) * d_int))) |
| n_mask = min(n_mask, d_int) |
| w = W_down.detach().to(device=z.device, dtype=torch.float32) |
| z_f = z.float() |
| y_norm = y.float().norm(dim=-1).clamp_min(PROJ_EPS) |
| acc = 0.0 |
| r = max(int(repeats), 1) |
| for _ in range(r): |
| idx = torch.randperm(d_int, generator=generator)[:n_mask].to(device=z.device) |
| delta = z_f.index_select(1, idx) @ w.index_select(1, idx).T |
| acc += float((delta.norm(dim=-1) / y_norm).mean()) |
| return acc / r |
|
|
|
|
| _perturbation_sensitivity = perturbation_sensitivity |
|
|
|
|
| @dataclass |
| class LayerCollector: |
| d_int: int |
| num_experts: int |
| device: torch.device |
| mlp: object |
| stats: dict[int, ExpertStats] = field(default_factory=dict) |
| num_bins: int = DEFAULT_NUM_BINS |
| collect_sensitivity: bool = True |
| perturb_ratio: float = DEFAULT_PERTURB_RATIO |
| perturb_repeats: int = DEFAULT_PERTURB_REPEATS |
| collect_gram: bool = True |
| token_cache_per_bin: int = DEFAULT_TOKEN_CACHE_PER_BIN |
| current_timestep: float | None = None |
| current_step: int | None = None |
| num_steps: int | None = None |
|
|
| def __post_init__(self): |
| if self.mlp is None: |
| raise ValueError("TENP collector needs mlp to read down_proj.weight") |
| self.num_bins = max(int(self.num_bins), 1) |
| self.token_cache_per_bin = max(int(self.token_cache_per_bin), 0) |
| self.stats = { |
| eid: ExpertStats.zeros( |
| self.d_int, |
| self.device, |
| num_bins=self.num_bins, |
| collect_gram=bool(self.collect_gram), |
| token_cache_per_bin=self.token_cache_per_bin, |
| ) |
| for eid in range(self.num_experts) |
| } |
| self._sens_gen = torch.Generator(device="cpu") |
| self._sens_gen.manual_seed(0) |
|
|
| def _down_weight(self, expert_id: int) -> torch.Tensor: |
| return self.mlp.experts[int(expert_id)].down_proj.weight |
|
|
| def current_bin(self) -> int: |
| mlp = self.mlp |
| t = getattr(mlp, "diffusion_timestep", None) |
| if t is None: |
| t = self.current_timestep |
| step_index = getattr(mlp, "diffusion_step_index", None) |
| if step_index is None: |
| step_index = self.current_step |
| num_steps = getattr(mlp, "diffusion_num_steps", None) |
| if num_steps is None: |
| num_steps = self.num_steps |
| return timestep_to_bin( |
| t, |
| self.num_bins, |
| step_index=step_index, |
| num_steps=num_steps, |
| ) |
|
|
| def __call__( |
| self, |
| expert_id: int, |
| modality: torch.Tensor, |
| gate: torch.Tensor, |
| z: torch.Tensor, |
| y: torch.Tensor, |
| ) -> None: |
| self.stats[int(expert_id)].update( |
| z, |
| y, |
| gate, |
| modality, |
| W_down=self._down_weight(expert_id), |
| bin_id=self.current_bin(), |
| collect_sensitivity=self.collect_sensitivity, |
| perturb_ratio=self.perturb_ratio, |
| perturb_repeats=self.perturb_repeats, |
| generator=self._sens_gen, |
| ) |
|
|
| def state_dict(self) -> dict: |
| return { |
| "d_int": int(self.d_int), |
| "num_experts": int(self.num_experts), |
| "num_bins": int(self.num_bins), |
| "stats": {str(eid): st.state_dict() for eid, st in self.stats.items()}, |
| } |
|
|
| @classmethod |
| def from_state(cls, payload: dict, mlp, device="cpu") -> "LayerCollector": |
| col = cls( |
| d_int=int(payload["d_int"]), |
| num_experts=int(payload["num_experts"]), |
| device=torch.device(device), |
| mlp=mlp, |
| num_bins=int(payload.get("num_bins") or DEFAULT_NUM_BINS), |
| collect_sensitivity=False, |
| collect_gram=False, |
| token_cache_per_bin=0, |
| ) |
| for eid_s, st in payload["stats"].items(): |
| col.stats[int(eid_s)] = ExpertStats.from_state(st, device=device) |
| if col.stats: |
| col.num_bins = next(iter(col.stats.values())).num_bins |
| return col |
|
|
