diff --git a/scripts/derive-protected-experts.py b/scripts/derive-protected-experts.py new file mode 100755 index 0000000..43d9bcb --- /dev/null +++ b/scripts/derive-protected-experts.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +"""derive-protected-experts.py — rewrite the recipe's protected-expert section +from a calibration recording. RUNS ON DANKSERVER (venv: numpy). + +The law (docs/MOE-W4.md "Protected experts", restated in the recipe): +an expert whose share of its layer's down-input energy is at least 3% AND +at least three times its share of the routing is kept bf16; the count must +be EVEN on each rank's half of the layer (the four-bit GEMM takes the rest +as two equal tables by parity) — an odd half gains its next most energetic +non-qualifying expert. + +Data: the calibration dir's per-expert down Grams, +gram.r.blk.L.ffn_down_exps..bin — 32-byte RADGRAM2 header +(magic, k, rank, rows, domain) then k*k f32, named with the checkpoint's +GLOBAL expert numbering (arch/common/rad_block_moe_fp8.h: ge = +expert_base + e). rows (header) is the expert's routing count; the trace +of the Gram is its down-input energy. Rank halves: 0-255, 256-511 +(contiguous expert_base split at TP=2). + +Output: the cast-rule lines (recipe format) on stdout, a report on stderr. + +Usage: derive-protected-experts.py [--min-share 0.03] + [--min-route-mult 3.0] +""" +import argparse +import glob +import os +import re +import struct +import sys + +import numpy as np + +NAME_RE = re.compile(r"^gram\.r(\d)\.blk\.(\d+)\.ffn_down_exps\.(\d+)\.bin$") + + +def read_gram(path): + """(k, rank, rows, trace) from a RADGRAM2 file (full matrix on disk).""" + with open(path, "rb") as f: + hdr = f.read(32) + magic, k, rank, rows, domain = struct.unpack("<8sIIIQI", hdr[:32]) + if magic != b"RADGRAM2": + raise SystemExit(f"not a RADGRAM2 file: {path}") + n = k * k + a = np.frombuffer(f.read(n * 4), dtype=np.float32) + if a.size != n: + raise SystemExit(f"short gram file: {path}") + tr = float(a.reshape(k, k).diagonal().sum()) + return k, rank, int(rows), tr + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("calib_dir") + ap.add_argument("--min-share", type=float, default=0.03) + ap.add_argument("--min-route-mult", type=float, default=3.0) + a = ap.parse_args() + + # layer -> global expert -> (rank, rows, energy) + data = {} + for path in sorted(glob.glob(os.path.join( + a.calib_dir, "gram.r*.ffn_down_exps.*.bin"))): + m = NAME_RE.match(os.path.basename(path)) + if not m: + continue + rank, layer, expert = int(m.group(1)), int(m.group(2)), int(m.group(3)) + k, _, rows, tr = read_gram(path) + data.setdefault(layer, {})[expert] = (rank, rows, tr) + + if not data: + raise SystemExit("no per-expert down grams found") + + lines, report = [], [] + total = 0 + for layer in sorted(data): + experts = data[layer] + tot_e = sum(v[2] for v in experts.values()) + tot_r = sum(v[1] for v in experts.values()) + if tot_e <= 0 or tot_r <= 0: + report.append(f"blk.{layer}: empty recording, skipped") + continue + qual, pool = [], [] + for e, (rank, rows, tr) in experts.items(): + es, rs = tr / tot_e, rows / max(tot_r, 1) + rec = (e, rank, es, rs) + if es >= a.min_share and es >= a.min_route_mult * rs: + qual.append(rec) + elif es > 0: + pool.append(rec) + chosen = {r[0]: (r[2], r[3]) for r in qual} + for half in (0, 1): # rank halves by the contiguous expert_base split + n_half = sum(1 for e in chosen if e // 256 == half) + if n_half % 2: + pool.sort(key=lambda r: -r[2]) + for cand in pool: + if cand[1] == half and cand[0] not in chosen: + chosen[cand[0]] = (cand[2], cand[3]) + report.append( + f"blk.{layer}: half {half} odd -> added e{cand[0]} " + f"(energy {cand[2]*100:.2f}%)") + break + for e in sorted(chosen): + es, rs = chosen[e] + lines.append(f"blk.{layer}.ffn_*_exps.{e}.weight" + f" cast dtype=bf16") + report.append(f"blk.{layer} e{e:3d} (rank {e // 256}): " + f"energy {es*100:5.2f}%, routing {rs*100:5.3f}%, " + f"ratio {es/max(rs,1e-12):6.1f}") + total += len(chosen) + + print("\n".join(lines)) + print("\n".join(report), file=sys.stderr) + print(f"total protected: {total}", file=sys.stderr) + + +if __name__ == "__main__": + main()