| """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 | |