inference_code/inference/pruning/apply.py

"""Slice SwiGLU mid-neurons from TENP / Wanda / SparseGPT. No OLS, no shared."""

from __future__ import annotations

import logging

import torch

from .original_structured import (
    sparsegpt_compensate_down,
    sparsegpt_mid_scores,
    wanda_mid_scores,
)
from .redundancy import contribution_redundancy, greedy_diversity_keep, n_keep_width
from .scoring import (
    DEFAULT_MODALITY_FUSE,
    allocate_global_keeps,
    fuse_global_modality_scores,
    fuse_modality_scores,
    fuse_norm,
    fuse_zscore,
    keep_indices,
    modality_lambdas,
    random_keep_indices,
)
from .surgery import materialize_expert_mid
from .temporal import (
    DEFAULT_N_MIN,
    DEFAULT_TEMPORAL_BETA,
    DEFAULT_TEMPORAL_TOPK,
    TEMPORAL_MODES,
    aggregate_temporal,
    fuse_bins_to_importance,
    mean_alpha,
    routing_sensitivity_alpha,
)

logger = logging.getLogger("tenp")

_MOD_NAMES = {0: "V", 1: "A", 2: "T"}
AUDIO = 1


def magnitude_score(expert) -> torch.Tensor:
    """Fallback when an expert saw no routed projection."""
    up = expert.up_gate_proj.weight.detach().float()
    d_int = up.shape[0] // 2
    paired = up.view(d_int, 2, up.shape[1]).norm(dim=(1, 2))
    down = expert.down_proj.weight.detach().float().norm(dim=0)
    return paired + down


def audio_primary_p_mid(n_tokens_mod: torch.Tensor, p_mid: float) -> tuple[float, bool]:
    total = int(n_tokens_mod.sum())
    if total <= 0:
        return p_mid, False
    if float(n_tokens_mod[AUDIO]) / total > 0.5:
        return float(min(0.50, max(0.15, p_mid - 0.10))), True
    return p_mid, False


def score_table(stats, expert, score: str, percdamp: float = 0.01) -> torch.Tensor:
    if score == "projection":
        return stats.projection.detach()
    if score == "energy":
        return stats.energy.detach()
    if score == "wanda":
        return wanda_mid_scores(stats.energy, expert.down_proj.weight)
    if score == "sparsegpt":
        if stats.gram is None:
            raise ValueError("SparseGPT structured pruning requires collected Gram statistics")
        return sparsegpt_mid_scores(
            stats.gram, stats.energy, expert.down_proj.weight, percdamp=percdamp
        )
    raise ValueError(f"score_table does not handle {score!r}")


def _bin_score_table(stats, expert, score: str, percdamp: float = 0.01) -> torch.Tensor | None:
    """Per-bin [K, 3, d_int] table, or None when the score has no bin stats."""
    if stats.projection_bin is None or int(stats.num_bins) <= 1:
        return None
    if score == "projection":
        return stats.projection_bin.detach()
    if score == "energy":
        if stats.energy_bin is None:
            return None
        return stats.energy_bin.detach()
    if score == "wanda":
        if stats.energy_bin is None:
            return None
        tables = [
            wanda_mid_scores(stats.energy_bin[k], expert.down_proj.weight)
            for k in range(int(stats.num_bins))
        ]
        return torch.stack(tables, dim=0)
    return None


def layer_sensitivity(collector) -> torch.Tensor | None:
    """Token-weighted mean of expert D_{e,k} at a layer."""
    sens = []
    weights = []
    for stats in collector.stats.values():
        if stats.sensitivity_sum is None or stats.n_tokens_bin is None:
            continue
        d = stats.sensitivity()
        n = stats.n_tokens_bin.float()
        sens.append(d)
        weights.append(n)
    if not sens:
        return None
    stacked = torch.stack(sens, dim=0)
    w = torch.stack(weights, dim=0)
    return (stacked * w).sum(dim=0) / w.sum(dim=0).clamp_min(1.0)


def neuron_importance(
    stats,
    expert,
    score: str,
    modality_fuse: str = DEFAULT_MODALITY_FUSE,
    temporal: str = "none",
    n_min: int = DEFAULT_N_MIN,
    temporal_beta: float = DEFAULT_TEMPORAL_BETA,
    temporal_topk: int = DEFAULT_TEMPORAL_TOPK,
    layer_sens: torch.Tensor | None = None,
    percdamp: float = 0.01,
) -> torch.Tensor:
    """Fused neuron score. `temporal=none` is the original summed-table TENP."""
    if temporal not in TEMPORAL_MODES:
        raise ValueError(f"unknown temporal {temporal!r}")
    table = score_table(stats, expert, score, percdamp=percdamp)
    lambdas = modality_lambdas(stats.n_tokens_mod, modality_fuse)
    use_bins = (
        temporal != "none"
        and score in ("projection", "energy", "wanda")
        and stats.projection_bin is not None
        and int(stats.num_bins) > 1
    )
    zscore = fuse_zscore(modality_fuse)
    norm = fuse_norm(modality_fuse)
    if not use_bins:
        return fuse_modality_scores(table, lambdas=lambdas, zscore=zscore, norm=norm)
    bin_table = _bin_score_table(stats, expert, score, percdamp=percdamp)
    if bin_table is None:
        return fuse_modality_scores(table, lambdas=lambdas, zscore=zscore, norm=norm)
    n_bin = stats.n_tokens_bin
    importance = fuse_bins_to_importance(bin_table, n_bin, lambdas=lambdas, n_min=n_min)
    if temporal == "mean":
        alpha = mean_alpha(n_bin, n_min=n_min)
        return aggregate_temporal(importance, alpha=alpha, mode="mean")
    if temporal == "ta":
        return aggregate_temporal(
            importance, mode="ta", beta=temporal_beta, temporal_topk=temporal_topk
        )
    d_e = stats.sensitivity()
    alpha = routing_sensitivity_alpha(
        n_bin,
        d_e,
        n_min=n_min,
        layer_sensitivity=layer_sens,
    )
    return aggregate_temporal(importance, alpha=alpha, mode="st")


def select_keep(
    score: torch.Tensor,
    p_mid: float,
    expert,
    stats,
    redundancy: bool,
) -> torch.Tensor:
    if not redundancy:
        return keep_indices(score, p_mid)
    n_keep = n_keep_width(int(score.numel()), p_mid)
    r_plus = contribution_redundancy(stats.gram, expert.down_proj.weight)
    return greedy_diversity_keep(score, r_plus, n_keep)


def tc_keep_indices(
    score: torch.Tensor,
    raw_signed_score: torch.Tensor,
    criticality: torch.Tensor,
    p_mid: float,
    multiple: int = 4,
) -> torch.Tensor:
    """TC-TENP keep set: raw-negative safe region, then temporal risk.

    When the raw-negative region can fill the pruning budget, the original
    TENP selection is returned unchanged. Otherwise all raw-negative neurons
    are removed and remaining deletions are ordered lexicographically by
    ``(criticality, TENP score)``.
    """
    s = score.detach().float().reshape(-1)
    raw = raw_signed_score.detach().float().reshape(-1)
    c = criticality.detach().float().reshape(-1)
    if raw.shape != s.shape or c.shape != s.shape:
        raise ValueError(
            f"TC vectors must share shape, got score={tuple(s.shape)} "
            f"raw={tuple(raw.shape)} criticality={tuple(c.shape)}"
        )
    n_keep = n_keep_width(int(s.numel()), p_mid, multiple=multiple)
    n_delete = int(s.numel()) - n_keep
    safe = torch.nonzero(raw <= 0, as_tuple=False).reshape(-1)
    if int(safe.numel()) >= n_delete:
        return keep_indices(s, p_mid, multiple=multiple)

    positive = torch.nonzero(raw > 0, as_tuple=False).reshape(-1)
    n_extra = n_delete - int(safe.numel())
    # Stable two-pass sort implements lexicographic (criticality, TENP score).
    by_score = positive[torch.argsort(s[positive], stable=True)]
    ordered = by_score[torch.argsort(c[by_score], stable=True)]
    remove = torch.cat([safe, ordered[:n_extra]])
    keep_mask = torch.ones(int(s.numel()), dtype=torch.bool, device=s.device)
    keep_mask[remove] = False
    keep = torch.nonzero(keep_mask, as_tuple=False).reshape(-1)
    keep, _ = torch.sort(keep)
    return keep


def global_keep_maps(
    targets: list[tuple[int, object]],
    collectors: dict[int, object],
    p_mid: float,
    score: str,
    modality_fuse: str = DEFAULT_MODALITY_FUSE,
    generator: torch.Generator | None = None,
    percdamp: float = 0.01,
    temporal: str = "none",
    n_min: int = DEFAULT_N_MIN,
    temporal_beta: float = DEFAULT_TEMPORAL_BETA,
    temporal_topk: int = DEFAULT_TEMPORAL_TOPK,
) -> dict[int, dict[int, torch.Tensor]]:
    """One keep budget over all routed experts, matching miniSeedance global-layers."""
    keys: list[tuple[int, int]] = []
    for layer_idx, mlp in targets:
        for eid, _expert in enumerate(mlp.experts):
            keys.append((int(layer_idx), int(eid)))
    if not keys:
        raise ValueError("no MoE experts to allocate")
    d_int = int(targets[0][1].experts[0].down_proj.in_features)
    if score == "random":
        noise = torch.rand(len(keys), d_int, generator=generator)
        keeps = allocate_global_keeps(noise, p_mid)
    else:
        tables = []
        dead = []
        mags = []
        lambdas = modality_lambdas(None, modality_fuse)
        fused_rows = []
        for layer_idx, eid in keys:
            mlp = next(m for li, m in targets if int(li) == int(layer_idx))
            expert = mlp.experts[eid]
            stats = collectors[int(layer_idx)].stats[eid]
            table = score_table(stats, expert, score, percdamp=percdamp)
            tables.append(table.detach().float().cpu())
            dead.append(float(table.abs().sum()) == 0)
            mags.append(magnitude_score(expert).cpu())
            logger.info(
                "layer %d expert %s n_mod(V,A,T)=%s",
                layer_idx, eid, [int(x) for x in stats.n_tokens_mod.tolist()],
            )
            if temporal != "none" and float(table.abs().sum()) != 0:
                fused_rows.append(
                    neuron_importance(
                        stats,
                        expert,
                        score,
                        modality_fuse=modality_fuse,
                        temporal=temporal,
                        n_min=n_min,
                        temporal_beta=temporal_beta,
                        temporal_topk=temporal_topk,
                        layer_sens=layer_sensitivity(collectors[int(layer_idx)]),
                        percdamp=percdamp,
                    ).detach().float().cpu()
                )
            elif temporal != "none":
                fused_rows.append(torch.zeros(d_int))
        if temporal == "none":
            fused = fuse_global_modality_scores(
                torch.stack(tables, dim=0),
                lambdas=lambdas,
                dead_mask=torch.tensor(dead, dtype=torch.bool),
                magnitude=torch.stack(mags, dim=0),
            )
        else:
            fused = torch.stack(fused_rows, dim=0)
            dead_t = torch.tensor(dead, dtype=torch.bool)
            if bool(dead_t.any()):
                fused[dead_t] = -1e6
        keeps = allocate_global_keeps(fused, p_mid)
    maps: dict[int, dict[int, torch.Tensor]] = {}
    for (layer_idx, eid), keep in zip(keys, keeps):
        maps.setdefault(int(layer_idx), {})[int(eid)] = keep
        logger.info(
            "global keep layer %d expert %s n=%d",
            layer_idx, eid, int(keep.numel()),
        )
    return maps


def prune_moe_layer(
    mlp,
    collector,
    p_mid: float,
    modality_fuse: str = DEFAULT_MODALITY_FUSE,
    score: str = "projection",
    apply_audio_discount: bool = False,
    generator: torch.Generator | None = None,
    keep_map: dict | None = None,
    compensate: str = "none",
    percdamp: float = 0.01,
    temporal: str = "none",
    redundancy: bool = False,
    n_min: int = DEFAULT_N_MIN,
    temporal_beta: float = DEFAULT_TEMPORAL_BETA,
    temporal_topk: int = DEFAULT_TEMPORAL_TOPK,
) -> dict:
    """Prune routed experts in `mlp`. Never drops a whole expert. Never runs OLS.

    `keep_map` is an optional {expert_id: LongTensor} from a global budget.
    `compensate` is "none" or "sparsegpt" (OBS on down_proj before slice).
    `apply_audio_discount` stays off for MoE_4000/16: no expert is audio-primary.
    Default `temporal=none` and `redundancy=False` is the original TENP protocol.
    """
    if score not in ("projection", "energy", "wanda", "sparsegpt", "random"):
        raise ValueError(f"unknown score {score!r}")
    if compensate not in ("none", "sparsegpt"):
        raise ValueError(f"unknown compensate {compensate!r}")
    if temporal not in TEMPORAL_MODES:
        raise ValueError(f"unknown temporal {temporal!r}")
    meta = {}
    layer_sens = layer_sensitivity(collector)
    for eid, expert in enumerate(mlp.experts):
        stats = collector.stats[eid]
        n_mod = [int(x) for x in stats.n_tokens_mod.tolist()]
        logger.info(
            "expert %s n_tokens=%d n_mod(V,A,T)=%s",
            eid, stats.n_tokens, n_mod,
        )
        table = None if score == "random" else score_table(
            stats, expert, score, percdamp=percdamp
        )
        if table is not None:
            for m, name in _MOD_NAMES.items():
                if float(table[m].abs().sum()) == 0:
                    logger.info("expert %s skip empty modality %s", eid, name)
        p_use = p_mid
        audio_pri = False
        if apply_audio_discount:
            p_use, audio_pri = audio_primary_p_mid(stats.n_tokens_mod, p_mid)
            if audio_pri:
                logger.info("audio-primary expert %s p_mid %.2f -> %.2f", eid, p_mid, p_use)
        if keep_map is not None:
            keep = keep_map[eid].to(expert.up_gate_proj.weight.device)
        elif score == "random":
            d_int = int(expert.down_proj.in_features)
            keep = random_keep_indices(d_int, p_use, generator=generator).to(
                expert.up_gate_proj.weight.device
            )
        elif table is None or float(table.abs().sum()) == 0:
            logger.info("expert %s dead %s; magnitude fallback", eid, score)
            sc = magnitude_score(expert)
            keep = select_keep(sc, p_use, expert, stats, redundancy=False).to(
                expert.up_gate_proj.weight.device
            )
        else:
            lambdas = modality_lambdas(stats.n_tokens_mod, modality_fuse)
            logger.info(
                "expert %s modality_fuse=%s lambdas=%s temporal=%s redundancy=%s",
                eid, modality_fuse, lambdas, temporal, redundancy,
            )
            sc = neuron_importance(
                stats,
                expert,
                score,
                modality_fuse=modality_fuse,
                temporal=temporal,
                n_min=n_min,
                temporal_beta=temporal_beta,
                temporal_topk=temporal_topk,
                layer_sens=layer_sens,
                percdamp=percdamp,
            )
            keep = select_keep(sc, p_use, expert, stats, redundancy=redundancy).to(
                expert.up_gate_proj.weight.device
            )
        used_sparsegpt = False
        if (
            compensate == "sparsegpt"
            and stats.gram is not None
            and float(stats.energy.abs().sum()) > 0
        ):
            compensated = sparsegpt_compensate_down(
                expert.down_proj.weight, stats.gram, keep, percdamp=percdamp,
            )
            expert.down_proj.weight.data.copy_(compensated.to(expert.down_proj.weight.dtype))
            used_sparsegpt = True
        original_width = int(expert.down_proj.in_features)
        materialize_expert_mid(expert, keep)
        meta[eid] = {
            "keep": keep.cpu().tolist(),
            "n_tokens": stats.n_tokens,
            "n_tokens_mod": n_mod,
            "ols": False,
            "sparsegpt_compensate": used_sparsegpt,
            "score": score,
            "audio_primary": audio_pri,
            "p_mid": float(p_use),
            "original_width": original_width,
            "kept_width": int(keep.numel()),
            "temporal": temporal,
            "redundancy": bool(redundancy),
        }
    return meta