inference_code/inference/pruning/surgery.py

"""SwiGLU pair indexing and physical Linear surgery. No residual-width slice."""

from __future__ import annotations

import torch
from torch import nn


def round_to_multiple(n: int, multiple: int) -> int:
    if multiple <= 0:
        raise ValueError("multiple must be positive")
    return int(round(n / multiple) * multiple)


def interleave_swiglu_indices(keep_idx: torch.Tensor) -> torch.Tensor:
    """Map kept SwiGLU neurons to interleaved `up_gate_proj` rows: (2i, 2i+1)."""
    keep_idx = keep_idx.long().reshape(-1)
    return torch.stack((2 * keep_idx, 2 * keep_idx + 1), dim=1).reshape(-1)


pair_rows = interleave_swiglu_indices


@torch.no_grad()
def materialize_expert_mid(expert, keep_idx: torch.Tensor) -> None:
    """Replace expert Linears with a physically narrower SwiGLU mid width."""
    keep_idx = keep_idx.to(device=expert.down_proj.weight.device, dtype=torch.long)
    keep_idx, _ = torch.sort(keep_idx)
    rows = interleave_swiglu_indices(keep_idx)
    old_up = expert.up_gate_proj
    old_down = expert.down_proj
    new_up = nn.Linear(
        old_up.in_features,
        int(rows.numel()),
        bias=False,
        device=old_up.weight.device,
        dtype=old_up.weight.dtype,
    )
    new_down = nn.Linear(
        int(keep_idx.numel()),
        old_down.out_features,
        bias=False,
        device=old_down.weight.device,
        dtype=old_down.weight.dtype,
    )
    new_up.weight.copy_(old_up.weight.index_select(0, rows))
    new_down.weight.copy_(old_down.weight.index_select(1, keep_idx))
    expert.up_gate_proj = new_up
    expert.down_proj = new_down
    if hasattr(expert, "intermediate_size"):
        expert.intermediate_size = int(keep_idx.numel())


def resize_expert_mid(expert, kept_width: int) -> None:
    """Allocate empty narrow Linears (load weights afterwards)."""
    kept = int(kept_width)
    old_up = expert.up_gate_proj
    old_down = expert.down_proj
    expert.up_gate_proj = nn.Linear(
        old_up.in_features,
        2 * kept,
        bias=False,
        device=old_up.weight.device,
        dtype=old_up.weight.dtype,
    )
    expert.down_proj = nn.Linear(
        kept,
        old_down.out_features,
        bias=False,
        device=old_down.weight.device,
        dtype=old_down.weight.dtype,
    )
    if hasattr(expert, "intermediate_size"):
        expert.intermediate_size = kept