| |
| """Write an LLAMA_EXPERT_LOCK_LIST file: the most used kept experts of the streamed layers. |
| |
| Slots are positions in the pruned model's expert dimension = index in the ascending list of kept expert ids |
| (the order prune_experts.py writes). Usage counts come from mimo_expert_importance.json (count_<source>). |
| |
| usage: python hot_lock_list.py --importance mimo_expert_importance.json --prune prune_..._keep144.json \ |
| --layers 14-30 --frac 0.6 [--source onpol] --out hot60.txt |
| """ |
| import argparse |
| import json |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--importance", required=True) |
| ap.add_argument("--prune", required=True, help="prune list of the model ({'pruned': {layer: [ids]}})") |
| ap.add_argument("--layers", required=True, help="streamed layers, e.g. 14-30") |
| ap.add_argument("--frac", type=float, default=0.6, help="share of kept experts to lock per layer") |
| ap.add_argument("--source", default="onpol", help="usage source: onpol, de, baekpica, ...") |
| ap.add_argument("--n-expert", type=int, default=256) |
| ap.add_argument("--out", required=True) |
| a = ap.parse_args() |
| imp = json.load(open(a.importance, encoding="utf-8"))["layers"] |
| pruned = json.load(open(a.prune, encoding="utf-8"))["pruned"] |
| lo, hi = (int(x) for x in a.layers.split("-")) |
| cover = [] |
| with open(a.out, "w", encoding="ascii") as f: |
| for L in range(lo, hi + 1): |
| kept = sorted(set(range(a.n_expert)) - set(pruned[str(L)])) |
| use = {r["expert"]: r.get(f"count_{a.source}", 0.0) for r in imp[str(L)]} |
| order = sorted(kept, key=lambda e: -use.get(e, 0.0)) |
| hot = order[: round(a.frac * len(kept))] |
| slots = sorted(kept.index(e) for e in hot) |
| f.write(f"{L} " + " ".join(map(str, slots)) + "\n") |
| tot = sum(use.get(e, 0.0) for e in kept) or 1.0 |
| cover.append(sum(use.get(e, 0.0) for e in hot) / tot) |
| print(f"{a.out}: layers {lo}-{hi}, {a.frac:.0%} of kept experts per layer, " |
| f"covers {sum(cover) / len(cover):.1%} of {a.source} routings (min {min(cover):.1%})") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
