File size: 2,372 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 | """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
|