File size: 4,043 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """Save / reload physically narrowed SparseMoE experts. Rebuild shapes before load."""
from __future__ import annotations
from typing import Any
import torch
from inference.model.dit.dit_module import SparseMoEMLP
from .surgery import materialize_expert_mid, resize_expert_mid
FORMAT = "davinci_tenp_v1"
def sparse_mlp_at(model, layer_idx: int) -> SparseMoEMLP:
mlp = model.block.layers[int(layer_idx)].mlp
if not isinstance(mlp, SparseMoEMLP):
raise TypeError(f"layer {layer_idx} mlp is {type(mlp).__name__}, expected SparseMoEMLP")
return mlp
def apply_keep_map(model, keep_by_layer_expert: dict[tuple[int, int], torch.Tensor]) -> dict:
"""Physically slice experts. keys are (layer_idx, expert_id)."""
layers: dict[str, dict] = {}
for (layer_idx, expert_id), keep in keep_by_layer_expert.items():
mlp = sparse_mlp_at(model, layer_idx)
expert = mlp.experts[int(expert_id)]
original = int(expert.down_proj.in_features)
materialize_expert_mid(expert, keep)
layers.setdefault(str(layer_idx), {"experts": []})
layers[str(layer_idx)]["experts"].append(
{
"expert_id": int(expert_id),
"original_width": original,
"kept_width": int(keep.numel()),
"keep_idx": keep.detach().cpu().long().tolist(),
}
)
return layers
def rebuild_mid_shapes(model, meta: dict) -> None:
layers = meta.get("layers") or {}
for layer_s, layer_meta in layers.items():
mlp = sparse_mlp_at(model, int(layer_s))
for spec in layer_meta["experts"]:
eid = int(spec["expert_id"])
kept = int(spec["kept_width"])
keep_idx = spec.get("keep_idx") or []
if keep_idx and len(keep_idx) != kept:
raise ValueError(
f"layer {layer_s} expert {eid}: keep_idx len {len(keep_idx)} != kept_width {kept}"
)
resize_expert_mid(mlp.experts[eid], kept)
def save_pruned_checkpoint(path: str, model, meta: dict) -> None:
payload = {
"format": FORMAT,
"model": model.state_dict(),
"tenp_meta": meta,
}
torch.save(payload, path)
def load_pruned_checkpoint(model, path: str, map_location="cpu") -> dict:
payload = torch.load(path, map_location=map_location, weights_only=False)
if payload.get("format") != FORMAT:
raise ValueError(f"unsupported checkpoint format {payload.get('format')!r}")
meta = payload.get("tenp_meta") or {}
rebuild_mid_shapes(model, meta)
model.load_state_dict(payload["model"], strict=True)
return meta
def layer_meta_from_prune(layer_idx: int, prune_meta: dict, extra: dict | None = None) -> dict:
return meta_from_prunes({int(layer_idx): prune_meta}, extra)
def meta_from_prunes(by_layer: dict[int, dict], extra: dict | None = None) -> dict:
layers: dict[str, Any] = {}
for layer_idx, prune_meta in sorted(by_layer.items()):
experts = []
for eid, spec in sorted(prune_meta.items()):
experts.append(
{
"expert_id": int(eid),
"original_width": int(spec["original_width"]),
"kept_width": int(spec["kept_width"]),
"keep_idx": list(spec["keep"]),
"n_tokens_mod": list(spec["n_tokens_mod"]),
"ols": False,
}
)
layers[str(int(layer_idx))] = {"experts": experts}
body: dict[str, Any] = {
"version": 1,
"score_type": (extra or {}).get("score_type", "projection"),
"modality_fuse": (extra or {}).get("modality_fuse", "fixed311"),
"p_mid": (extra or {}).get("p_mid"),
"ols": False,
"audio_primary_discount": False,
"layers": layers,
}
extra = extra or {}
for key in ("method", "variant", "allocation", "temporal", "redundancy"):
if key in extra:
body[key] = extra[key]
return body
|