Frank
c13a2568d4
The swift15 rule (top 7% by imatrix energy, bar-only) left the plane size to the distribution. The w5nl64 arm trades a *pinned host pool* against KL, so the operator needs to say how many GiB of bf16 experts the pool can afford and have the set land there: --target-count walks the ranked candidates until the byte bill matches, keeps the bar-only pass for experts that are disproportionate for their size, and reports the pool's ceiling instead of padding when the count cannot be reached. TP=2 cuts the expert plane into two 256-expert rank halves, and an odd count in either half unbalances the placement, so the parity audit is a refusal (exit 2, 'PARITY LAW BROKEN'), not a silent fix-up: the driver stops rather than converting a payload whose expert placement no one priced.
245 lines
10 KiB
Python
Executable File
245 lines
10 KiB
Python
Executable File
#!/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<rank>.blk.L.ffn_down_exps.<e>.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 <calib dir> [--min-share 0.03]
|
|
[--min-route-mult 3.0] [--target-count 0]
|
|
--target-count N fills the byte bill: the bars pick the base set exactly as
|
|
the bar-only pass does, then the energy ranking admits experts one at a time
|
|
(each with its parity partner) until the count lands. The energy bar NEVER
|
|
relaxes, so asking for 48 bf16 experts never quietly buys 80; the ratio bar
|
|
walks down only when the ranking cannot produce that many candidates.
|
|
0 takes the bars as given.
|
|
"""
|
|
import argparse
|
|
import collections
|
|
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$")
|
|
CAST_RE = re.compile(r"^blk\.(\d+)\.ffn_\*_exps\.(\d+)\.weight")
|
|
|
|
|
|
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("<8sIIQI", hdr[:28])
|
|
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 select(data, min_share, min_route_mult):
|
|
"""(cast lines, report lines, count) for one pair of bars."""
|
|
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 >= min_share and es >= 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)
|
|
return lines, report, total
|
|
|
|
|
|
def parity_repair(chosen, by_half, report, only=None):
|
|
"""Add the next-best candidate to every half — or to `only` — holding an odd count.
|
|
|
|
The four-bit GEMM takes each half's remainder as two equal tables, so an odd count
|
|
of bf16 experts on a half is not a style problem: the two tables would not match.
|
|
`by_half` is energy-descending, so the first unchosen candidate is the next most
|
|
energetic, which is the expert worth keeping. Returns how many were added.
|
|
"""
|
|
added = 0
|
|
for key in ([only] if only is not None else sorted(by_half)):
|
|
layer, half = key
|
|
n = sum(1 for (l, x) in chosen if l == layer and x // 256 == half)
|
|
if n % 2 == 0:
|
|
continue
|
|
for e, es, rs in by_half.get(key, []):
|
|
if (layer, e) not in chosen:
|
|
chosen[(layer, e)] = (es, rs)
|
|
added += 1
|
|
report.append(f"blk.{layer}: half {half} odd -> added e{e} "
|
|
f"(energy {es*100:.2f}%)")
|
|
break
|
|
else:
|
|
report.append(f"blk.{layer}: half {half} holds {n} and no candidate is "
|
|
f"left to make it even")
|
|
return added
|
|
|
|
|
|
def select_to_count(data, min_share, min_route_mult, target):
|
|
"""(lines, report, total, cutoff) for the set that LANDS the count.
|
|
|
|
The bars decide the base set; the count is then filled from the energy ranking,
|
|
one expert at a time, each addition taking its half's next-best candidate with it
|
|
so every rank's half stays even. Only the RATIO bar walks down to find candidates
|
|
— it is the law that says an expert is disproportionate and not merely busy, and
|
|
the caller's byte bill has to be filled — while the ENERGY bar holds, because that
|
|
bar is what the extra GiB of bf16 is supposed to buy. `cutoff` is the energy share
|
|
the set bottoms out at: the bar this campaign actually converted with.
|
|
"""
|
|
report = []
|
|
shares, ranked = {}, []
|
|
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
|
|
for e, (_rank, rows, tr) in experts.items():
|
|
es, rs = tr / tot_e, rows / max(tot_r, 1)
|
|
if es <= 0:
|
|
continue
|
|
shares[(layer, e)] = (es, rs)
|
|
ranked.append((es, layer, e, rs))
|
|
ranked.sort(key=lambda r: (-r[0], r[1], r[2]))
|
|
by_half = {}
|
|
for es, layer, e, rs in ranked:
|
|
by_half.setdefault((layer, e // 256), []).append((e, es, rs))
|
|
|
|
chosen = {k: v for k, v in shares.items()
|
|
if v[0] >= min_share and v[0] >= min_route_mult * v[1]}
|
|
total = len(chosen)
|
|
total += parity_repair(chosen, by_half, report)
|
|
|
|
mult = min_route_mult
|
|
while True:
|
|
pool = [(es, l, e, rs) for es, l, e, rs in ranked
|
|
if es >= mult * rs and (l, e) not in chosen]
|
|
if len(pool) + total >= target or mult <= 0.0:
|
|
break
|
|
mult -= 0.5
|
|
report.append(f"ratio bar {mult + 0.5:g} leaves {len(pool)} fill candidates "
|
|
f"for {target - total} slots; relaxing toward 0")
|
|
for es, layer, e, rs in pool:
|
|
if total >= target:
|
|
break
|
|
if (layer, e) in chosen:
|
|
continue
|
|
chosen[(layer, e)] = (es, rs)
|
|
total += 1
|
|
report.append(f"blk.{layer}: count fill -> added e{e} (energy {es*100:.2f}%)")
|
|
total += parity_repair(chosen, by_half, report, only=(layer, e // 256))
|
|
lines, cutoff = [], None
|
|
for (layer, e) in sorted(chosen):
|
|
es, rs = chosen[(layer, e)]
|
|
cutoff = es if cutoff is None else min(cutoff, es)
|
|
lines.append(f"blk.{layer}.ffn_*_exps.{e}.weight cast dtype=bf16")
|
|
return lines, report, total, cutoff
|
|
|
|
|
|
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)
|
|
ap.add_argument("--target-count", type=int, default=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")
|
|
|
|
if a.target_count > 0:
|
|
lines, report, total, cutoff = select_to_count(
|
|
data, a.min_share, a.min_route_mult, a.target_count)
|
|
bar = f"energy cutoff {cutoff:.3g}" if cutoff is not None else "no candidates"
|
|
else:
|
|
lines, report, total = select(data, a.min_share, a.min_route_mult)
|
|
bar = f"min-share {a.min_share:.3g}"
|
|
|
|
# The emitted lines ARE the contract with the four-bit GEMM: every rank's half of
|
|
# every layer must hold an even number of bf16 experts. Judged on the output and
|
|
# not on the internal sets, so a slip in either pass is caught here rather than at
|
|
# engine load time.
|
|
halves = collections.Counter()
|
|
for rule in lines:
|
|
m = CAST_RE.match(rule)
|
|
if m:
|
|
halves[(int(m.group(1)), int(m.group(2)) // 256)] += 1
|
|
odd = {key: n for key, n in halves.items() if n % 2}
|
|
if odd:
|
|
print("PARITY LAW BROKEN: " + ", ".join(
|
|
f"blk.{lay} half {half} holds {n} protected experts"
|
|
for (lay, half), n in sorted(odd.items())), file=sys.stderr)
|
|
raise SystemExit(2)
|
|
|
|
print("\n".join(lines))
|
|
print("\n".join(report), file=sys.stderr)
|
|
print(f"total protected: {total} (target {a.target_count or 'none'}, {bar}, "
|
|
f"min-route-mult {a.min_route_mult:g})", file=sys.stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|