File size: 8,936 Bytes
3d1bd28 | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """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
# Qwen: ...layers.<l>.mlp.experts.<e>.* and ...mlp.gate.weight; DeepSeek-V4.1: layers.<l>.ffn.experts.<e>.*
# and layers.<l>.ffn.gate.{weight, bias, bias_vl} (all indexed by expert)
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]]] = {} # layer -> suffix -> (prefix, 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() # safetensors refuses shared storage
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, # ids k..k+pad-1: zero weights, router bias -1e4, never routed
"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()
|