| """Apply REAP's expert selection to an original checkpoint (any expert dtype: NVFP4, FP8, BF16). |
| |
| REAP's prune.py picks, per layer, the n experts with the lowest saliency (torch.topk(saliency, n, |
| largest=False) on the observer data) and keeps the rest in index order. This recomputes that exact |
| selection from the saved observations, checks it against the router rows of REAP's own pruned |
| checkpoint (--reference), and writes the source checkpoint with only the kept experts: per-expert |
| tensors (`...layers.<l>.mlp.experts.<e>.*`; DeepSeek-V4.1 `layers.<l>.ffn.experts.<e>.*`) renumbered |
| 0..K-1 in kept order, the router (`...mlp.gate.weight`; DeepSeek `ffn.gate.weight/bias/bias_vl`) sliced |
| to the kept rows, the MTP draft layers (`mtp.*`) dropped, |
| every other tensor and file unchanged, and num_experts (DeepSeek: n_routed_experts) = K in the |
| (text) config. The kept indices |
| go to <out>/reap_selection.json. --pad-to N appends N - K experts no token can reach (zero weights, |
| zero router rows, router bias -1e4; DeepSeek selects by score + bias with biases 9.6-14), for kernels |
| that take only some expert counts; the model computes exactly what the K-expert model does. |
| |
| external/reap/.venv/bin/python external/reap/experiments/apex/apply_selection.py \ |
| --observations <artifacts>/all/observations_1024_cosine-seed_42.pt --ratio 0.5 \ |
| --source <HF snapshot> --reference <REAP pruned dir> --out <dir> |
| """ |
|
|
| import argparse |
| import json |
| import re |
| import shutil |
| from pathlib import Path |
|
|
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import save_file |
|
|
| |
| |
| EXPERT = re.compile(r"^((?:.*\.)?layers\.(\d+)\.(?:mlp|ffn)\.experts\.)(\d+)(\..*)$") |
| ROUTER = re.compile(r"^(?:.*\.)?layers\.(\d+)\.(?:mlp|ffn)\.gate\.(?:weight|bias|bias_vl)$") |
| SHARD_BYTES = 5 * 2**30 |
|
|
|
|
| def selection(observations: Path, ratio: float, metric: str) -> dict[int, list[int]]: |
| """REAP prune.py's choice: per layer, prune the n lowest-saliency experts, keep the rest.""" |
| data = torch.load(observations, weights_only=False) |
| kept = {} |
| for layer, state in data.items(): |
| num_experts = state["expert_frequency"].shape[0] |
| n = int(num_experts * ratio) |
| _, pruned = torch.topk(state[metric], n, largest=False) |
| pruned = set(pruned.tolist()) |
| kept[int(layer)] = [i for i in range(num_experts) if i not in pruned] |
| return kept |
|
|
|
|
| def reference_routers(reference: Path) -> dict[int, torch.Tensor]: |
| out = {} |
| for f in sorted(reference.glob("*.safetensors")): |
| with safe_open(f, "pt") as fh: |
| for k in fh.keys(): |
| if (m := ROUTER.match(k)) and k.endswith(".weight"): |
| out[int(m.group(1))] = fh.get_tensor(k) |
| return out |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| parser.add_argument("--observations", required=True) |
| parser.add_argument("--ratio", type=float, required=True, help="share of experts pruned per layer") |
| parser.add_argument("--metric", default="reap") |
| parser.add_argument("--source", required=True, help="original checkpoint directory") |
| parser.add_argument("--reference", default=None, help="REAP's pruned checkpoint (router check)") |
| parser.add_argument("--out", required=True) |
| parser.add_argument( |
| "--pad-to", type=int, default=None, |
| help="append never-routed experts up to this count (DeepSeek: vLLM's sqrtsoftplus top-k kernel " |
| "takes 64/128/192/256/384 experts): zero weights, zero router rows, router bias -1e4", |
| ) |
| args = parser.parse_args() |
|
|
| source, out = Path(args.source), Path(args.out) |
| kept = selection(Path(args.observations), args.ratio, args.metric) |
| k = {len(v) for v in kept.values()} |
| assert len(k) == 1, k |
| k = k.pop() |
| index = json.loads((source / "model.safetensors.index.json").read_text())["weight_map"] |
|
|
| routers = {} |
| for name, f in index.items(): |
| if (m := ROUTER.match(name)) and not name.startswith("mtp.") and name.endswith(".weight"): |
| with safe_open(source / f, "pt") as fh: |
| routers[int(m.group(1))] = fh.get_tensor(name) |
| assert sorted(routers) == sorted(kept), (sorted(routers)[:5], sorted(kept)[:5]) |
| if args.reference: |
| ref = reference_routers(Path(args.reference)) |
| assert sorted(ref) == sorted(kept), "reference layers differ" |
| for layer, rows in kept.items(): |
| same = torch.equal(routers[layer][rows].to(ref[layer].dtype), ref[layer]) |
| assert same, f"layer {layer}: kept set differs" |
| print(f"selection matches REAP's pruned checkpoint in all {len(kept)} layers") |
|
|
| out.mkdir(parents=True, exist_ok=False) |
| new_index, shard, size, n_shard = {}, {}, 0, 0 |
|
|
| def flush() -> None: |
| nonlocal shard, size, n_shard |
| if shard: |
| n_shard += 1 |
| name = f"model-{n_shard:05d}.safetensors" |
| save_file(shard, out / name, metadata={"format": "pt"}) |
| new_index.update({key: name for key in shard}) |
| shard, size = {}, 0 |
|
|
| position = {layer: {e: j for j, e in enumerate(rows)} for layer, rows in kept.items()} |
| pad = 0 if args.pad_to is None else args.pad_to - k |
| assert pad >= 0, (args.pad_to, k) |
| templates: dict[int, dict[str, tuple[str, torch.Tensor]]] = {} |
| by_file: dict[str, list[str]] = {} |
| for name, f in index.items(): |
| by_file.setdefault(f, []).append(name) |
| dropped = 0 |
| for f, names in sorted(by_file.items()): |
| with safe_open(source / f, "pt") as fh: |
| for name in sorted(names): |
| if name.startswith("mtp."): |
| dropped += 1 |
| continue |
| if m := EXPERT.match(name): |
| layer, expert = int(m.group(2)), int(m.group(3)) |
| if expert not in position[layer]: |
| dropped += 1 |
| continue |
| key, tensor = f"{m.group(1)}{position[layer][expert]}{m.group(4)}", fh.get_tensor(name) |
| if pad and m.group(4) not in templates.setdefault(layer, {}): |
| templates[layer][m.group(4)] = (m.group(1), torch.zeros_like(tensor)) |
| elif (m := ROUTER.match(name)) and int(m.group(1)) in kept: |
| key, tensor = name, fh.get_tensor(name)[kept[int(m.group(1))]].contiguous() |
| if pad: |
| fill = -1e4 if not name.endswith(".weight") else 0.0 |
| extra = torch.full((pad, *tensor.shape[1:]), fill, dtype=tensor.dtype) |
| tensor = torch.cat([tensor, extra]) |
| else: |
| key, tensor = name, fh.get_tensor(name) |
| shard[key] = tensor |
| size += tensor.numel() * tensor.element_size() |
| if size >= SHARD_BYTES: |
| flush() |
| for layer, suffixes in sorted(templates.items()): |
| for suffix, (prefix, zero) in suffixes.items(): |
| for j in range(k, k + pad): |
| shard[f"{prefix}{j}{suffix}"] = zero.clone() |
| size += zero.numel() * zero.element_size() |
| if size >= SHARD_BYTES: |
| flush() |
| flush() |
| total = sum((out / f).stat().st_size for f in set(new_index.values())) |
| (out / "model.safetensors.index.json").write_text( |
| json.dumps({"metadata": {"total_size": total}, "weight_map": dict(sorted(new_index.items()))}, indent=2) |
| ) |
| config = json.loads((source / "config.json").read_text()) |
| text = config.get("text_config", config) |
| text["n_routed_experts" if "n_routed_experts" in text else "num_experts"] = k + pad |
| (out / "config.json").write_text(json.dumps(config, indent=2) + "\n") |
| for f in source.iterdir(): |
| if f.is_file() and not f.name.endswith(".safetensors") and f.name not in ( |
| "model.safetensors.index.json", |
| "config.json", |
| ): |
| shutil.copy(f, out / f.name) |
| (out / "reap_selection.json").write_text( |
| json.dumps( |
| { |
| "observations": str(Path(args.observations).resolve()), |
| "ratio_pruned": args.ratio, |
| "metric": args.metric, |
| "source": str(source.resolve()), |
| "kept_per_layer": k, |
| "padding_experts": pad, |
| "kept": {str(layer): rows for layer, rows in sorted(kept.items())}, |
| } |
| ) |
| + "\n" |
| ) |
| print(f"{out}: {k} experts per layer kept, {dropped} tensors dropped, {total / 1e9:.1f} GB") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
