tools: expert_atlas — confound-controlled probe harness for the GLM-5.2 expert atlas (#175) (#218)

Probe sweep + affinity analysis + leave-one-prompt-out validation, so anyone can build and
cross-validate the atlas on their own box rather than trusting one machine.

The four traps this harness exists to control (each silently corrupts the atlas):

  --topp     prunes experts by cumulative probability. Measured, same prompt:
             topp=0   -> 21,000 selections across 7,587 distinct experts
             topp=0.7 -> 11,944 selections across 4,687 distinct experts
             It hides 38% of the experts, and it is the recommended speed setting.
  MTP/DRAFT  eusage is incremented inside moe(), BEFORE verification, so rejected
             speculative drafts count experts routed for text never emitted.
  .coli_usage is loaded at startup and accumulates, so a naive STATS dump contains all
             prior history rather than this run.
  autocorrelation: routing within one run is highly correlated, so an expert firing 38
             times during one prompt is ONE observation, not 38. Entropy/chi-square on raw
             selections certifies single-prompt flukes as perfect specialists — analyze.py
             therefore requires affinity to replicate across a category independent prompts.

Result on GLM-5.2 744B int4 (Zen5, CPU routing path), 10 topics x 3 prompts x 64 tokens:

  leave-one-prompt-out accuracy   29/30 = 96.7%   (chance 10%)
  strong specialists (spec>=0.5)  1,041 / 13,260  (7.9%)
  specialisation vs depth         layer 3 ~0.07 -> layers 18-58 ~0.19-0.27
  replication gate rejected       587 single-prompt flukes

The one miss is the interesting part: a Chinese-language poetry prompt classifies as poetry,
not Chinese — routing follows the task over the language.
This commit is contained in:
Dennis Paul
2026-07-15 07:41:24 +02:00
committed by GitHub
parent 416570339f
commit a28f31fa3b
5 changed files with 403 additions and 0 deletions
+62
View File
@@ -0,0 +1,62 @@
# Expert Atlas — what does each of the 19,456 experts actually do?
Probe harness for #175. Runs a set of topic-tagged prompts, dumps each run's expert-routing
histogram, and turns them into a per-expert topic-affinity vector.
```bash
cd c
export COLI_MODEL=/path/to/glm52_i4
./tools/expert_atlas/sweep.sh # 30 probes (10 topics x 3 prompts)
python3 tools/expert_atlas/analyze.py --stats atlas_out/stats --out atlas_out/experts.json
python3 tools/expert_atlas/validate.py atlas_out/stats 200 # leave-one-prompt-out check
```
## Read this before you trust any atlas
Four things silently corrupt this measurement. The sweep script controls all of them; if you
roll your own, don't skip them.
| trap | effect | control |
|---|---|---|
| **`--topp`** | prunes experts by cumulative probability — measured: it hides **38% of the distinct experts** (7,587 → 4,687). It is also the *recommended speed setting*. | `TOPP=0` |
| **speculative drafts** | `eusage` is incremented inside `moe()`, *before* verification, so **rejected** drafts count. Those are experts routed for text the model never emitted. | `MTP=0 DRAFT=0` |
| **`.coli_usage`** | is loaded at startup and accumulates, so a naive `STATS` dump contains **all prior history**, not this run. | remove per run (script backs it up and restores) |
| **autocorrelation** | routing inside one run is highly correlated — the same context routes to the same experts token after token. An expert firing 38 times during one prompt is **one** observation, not 38. Chi-square/entropy on raw selections will certify single-prompt flukes as perfect specialists. | `analyze.py` requires the affinity to **replicate across a category's independent prompts** |
The CUDA expert tier is also not run-to-run deterministic, so the sweep uses `--gpu none`. Tier
config only decides *where weights live*, not what the router picks, so this costs nothing.
## Method
`analyze.py`:
1. `n[e][c]` — selections of expert *e* while running category *c*
2. `f[e][c] = n[e][c] / N[c]` — normalise by **category size** (prefill routes the prompt too, so a
verbose category would otherwise look busier)
3. `p(c|e)` — renormalise into a topic distribution per expert, i.e. base-rate corrected. Ranking
by raw count instead just rediscovers which experts are popular in general.
4. `spec(e) = 1 H(p(c|e)) / log C` — 0 = generalist, 1 = fires for exactly one topic
5. **replication gate** — an expert is only a candidate specialist for *c* if it fires in ≥2 of *c*'s
independent prompts
`validate.py` — leave-one-prompt-out. Build each category's top-K specialist set from its *other*
prompts, then check which set the held-out prompt's routing actually lands in. If specialisation
were an artifact of prompt wording, the held-out prompt would not prefer its own category.
## Result on GLM-5.2 744B int4 (Zen5, CPU routing path)
- **Leave-one-prompt-out: 29/30 = 96.7%** (chance 10%). Specialisation is a property of the topic,
not of prompt wording.
- The single miss is instructive: `写一首关于秋天的短诗` ("write a short poem about autumn") is
classified **poetry**, not Chinese — routing follows the **task** over the **language**.
- **Only 7.9% of experts are strong specialists** (spec ≥ 0.5). The "one expert = one topic"
picture is wrong for ~92% of them.
- Specialisation **rises with depth**: layer 3 ≈ 0.07 (generalist, token/syntax level) → layers
1858 ≈ 0.190.27.
- The replication gate removed **587** experts that looked like flawless specialists on one prompt.
## Extending it
`probes.json` is the whole probe set — add categories and prompts. Use **3+ prompts per category
with varied phrasing**, or the replication gate has nothing to check and you are back to measuring
one prompt. Keep prompt lengths in a similar band across categories.
+126
View File
@@ -0,0 +1,126 @@
#!/usr/bin/env python3
"""GLM-5.2 Expert Atlas — affinity with CROSS-PROMPT REPLICATION (#175).
Fixes a trap in the naive version. Routing selections inside one run are heavily correlated:
the same context routes to the same experts token after token. So an expert that fires 38
times during a single 'code' prompt is ONE effective observation, not 38 independent draws.
Treat them as independent and a chi-square will happily certify a single-prompt fluke as a
perfect specialist (spec=1.000, lift=10.0) on 38 selections. It is measuring one prompt.
So specialisation is only claimed when it REPLICATES across the independent prompts of a
category:
- each category has R prompts (here 3), each run is one replicate
- share[e][run] = selections of e in that run / total selections in that run
- an expert is a candidate specialist for category c only if it fires in >= MIN_RUNS of
c's runs (default 2/3) -> it is a property of the TOPIC, not of one prompt's wording
- affinity uses the MEAN share across a category's runs, so one hot run cannot carry it
- reliability = (runs in top category) / (runs in that category)
Also reports the generalist/specialist split by layer depth, which is an average over
thousands of experts and is robust to the above.
"""
import argparse, glob, json, math, os
from collections import defaultdict
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--stats", default="stats")
ap.add_argument("--min-count", type=int, default=30)
ap.add_argument("--min-runs", type=int, default=2, help="must fire in >= this many of the top category's runs")
ap.add_argument("--out", default="experts.json")
a = ap.parse_args()
# run[(cat,idx)][(layer,expert)] = count ; run_tot[(cat,idx)] = total
run_counts, run_tot = defaultdict(dict), defaultdict(int)
for path in sorted(glob.glob(os.path.join(a.stats, "*.txt"))):
base = os.path.basename(path)[:-4]
cat, idx = base.rsplit("_", 1)
for line in open(path):
p = line.split()
if len(p) != 3:
continue
l, e, n = int(p[0]), int(p[1]), int(p[2])
run_counts[(cat, idx)][(l, e)] = n
run_tot[(cat, idx)] += n
cats = sorted({c for c, _ in run_counts})
runs_of = {c: sorted(i for cc, i in run_counts if cc == c) for c in cats}
C = len(cats)
print(f"categories ({C}): " + ", ".join(f"{c}[{len(runs_of[c])}]" for c in cats))
experts = {k for d in run_counts.values() for k in d}
print(f"experts seen: {len(experts):,}\n")
atlas, dropped_sparse, dropped_unrepl = [], 0, 0
for key in experts:
total = sum(run_counts[r].get(key, 0) for r in run_counts)
if total < a.min_count:
dropped_sparse += 1
continue
# mean share per category across its runs (a single hot run cannot carry the category)
mean_share, fired_runs = {}, {}
for c in cats:
shares, fired = [], 0
for i in runs_of[c]:
n = run_counts[(c, i)].get(key, 0)
shares.append(n / max(1, run_tot[(c, i)]))
if n > 0:
fired += 1
mean_share[c] = sum(shares) / len(shares)
fired_runs[c] = fired
s = sum(mean_share.values())
if s <= 0:
continue
p = {c: mean_share[c] / s for c in cats}
top = max(cats, key=lambda c: p[c])
# REPLICATION GATE: the affinity must show up in >= min-runs of that category's prompts
if fired_runs[top] < a.min_runs:
dropped_unrepl += 1
continue
H = -sum(v * math.log(v) for v in p.values() if v > 0)
atlas.append({
"layer": key[0], "expert": key[1], "total": total,
"spec": round(1.0 - H / math.log(C), 4),
"top_topic": top,
"top_lift": round(p[top] * C, 2),
"reliability": f"{fired_runs[top]}/{len(runs_of[top])}",
"p": {c: round(p[c], 4) for c in cats},
})
atlas.sort(key=lambda r: (-r["spec"], -r["total"]))
print(f"dropped {dropped_sparse:,} sparse (<{a.min_count} sel)")
print(f"dropped {dropped_unrepl:,} UNREPLICATED (fired in <{a.min_runs} runs of their top topic)")
print(f"kept {len(atlas):,} experts\n")
print("=== most specialised, replicated across prompts ===")
print(f"{'layer':>5} {'exp':>4} {'sel':>6} {'spec':>6} {'lift':>6} {'repl':>5} topic")
for r in atlas[:20]:
print(f"{r['layer']:>5} {r['expert']:>4} {r['total']:>6} {r['spec']:>6.3f} "
f"{r['top_lift']:>6.2f} {r['reliability']:>5} {r['top_topic']}")
print("\n=== specialisation vs layer depth (mean over experts; robust to the above) ===")
by_layer = defaultdict(list)
for r in atlas:
by_layer[r["layer"]].append(r["spec"])
ls = sorted(by_layer)
for L in ls[::max(1, len(ls)//13)]:
v = by_layer[L]
print(f" layer {L:>3} n={len(v):>4} spec {sum(v)/len(v):.3f} {'#'*int(60*sum(v)/len(v))}")
print("\n=== experts owned per topic (replicated only) ===")
own = defaultdict(int)
for r in atlas:
own[r["top_topic"]] += 1
for c in sorted(own, key=lambda x: -own[x]):
print(f" {c:<14} {own[c]:>5}")
strong = [r for r in atlas if r["spec"] >= 0.5]
print(f"\nstrong specialists (spec >= 0.5, replicated): {len(strong):,} / {len(atlas):,} "
f"({100*len(strong)/max(1,len(atlas)):.1f}%)")
json.dump({"categories": cats, "experts": atlas}, open(a.out, "w"), indent=1)
print(f"wrote {a.out}")
if __name__ == "__main__":
main()
+54
View File
@@ -0,0 +1,54 @@
{
"_comment": "Probe set for the GLM-5.2 Expert Atlas (#175). 10 categories x 3 prompts. Prompts are deliberately varied in phrasing within a category so the affinity vector reflects the TOPIC, not one prompt's surface form. Lengths are kept within a narrow band across categories: prefill routes the prompt tokens too, so a long prompt in one category would inflate its counts.",
"code_python": [
"Write a Python function that merges two sorted lists into one sorted list.",
"In Python, explain the difference between a generator and a list comprehension.",
"Refactor this Python snippet to avoid a nested loop: for a in xs: for b in ys: if a==b: out.append(a)"
],
"code_sql": [
"Write a SQL query that returns the top 5 customers by total order value.",
"Explain what a LEFT JOIN does differently from an INNER JOIN in SQL.",
"Write SQL to add an index on the email column of a users table and explain when it helps."
],
"math_proof": [
"Prove that the square root of 2 is irrational.",
"Show that the sum of the first n odd numbers equals n squared.",
"Explain why the derivative of e^x is e^x."
],
"chinese": [
"请用三句话解释什么是机器学习。",
"写一首关于秋天的短诗。",
"中国的四大发明是什么?请简要说明。"
],
"german": [
"Erkläre in drei Sätzen, wie ein Verbrennungsmotor funktioniert.",
"Schreibe eine kurze formelle E-Mail, in der du einen Termin absagst.",
"Was ist der Unterschied zwischen Dativ und Akkusativ im Deutschen?"
],
"poetry": [
"Write a short poem about the sea at night.",
"Compose four lines of verse about an empty train station.",
"Write a haiku about the first snow of winter."
],
"law": [
"Explain in plain terms what consideration means in contract law.",
"What is the difference between a misdemeanor and a felony?",
"Summarize what the term 'force majeure' covers in a commercial contract."
],
"medicine": [
"Explain the difference between type 1 and type 2 diabetes.",
"What are the common symptoms of iron deficiency anemia?",
"Describe how a vaccine produces immunity."
],
"json_format": [
"Return a JSON object with keys name, age, and email for a fictional user. Output only JSON.",
"Convert this to JSON: name Alice, roles admin and editor, active true. Output only JSON.",
"Write a JSON schema for an object with a required string field 'id' and an optional integer 'count'."
],
"casual_chat": [
"Hey, what's a good way to spend a rainy Sunday afternoon?",
"I'm feeling pretty tired today. Any tips to get through the afternoon?",
"What's your favourite kind of weather, and why?"
]
}
+82
View File
@@ -0,0 +1,82 @@
#!/usr/bin/env bash
# GLM-5.2 Expert Atlas — probe sweep (#175).
#
# cd c && ./tools/expert_atlas/sweep.sh [probes.json] [outdir]
#
# Env: COLI_MODEL, NGEN (default 64), COLI (default ./coli)
#
# ---------------------------------------------------------------------------------------------
# THE CONFOUNDS THIS SCRIPT EXISTS TO CONTROL. Each one silently corrupts the atlas.
#
# TOPP=0 --topp prunes experts by cumulative probability. Measured on GLM-5.2 int4,
# same prompt, only top-p changed:
# topp=0 -> 21,000 selections across 7,587 distinct experts
# topp=0.7 -> 11,944 selections across 4,687 distinct experts
# It hides 38% of the experts. It is also the RECOMMENDED SPEED SETTING, so this
# is very easy to walk into: with top-p on you profile the pruner, not the model.
#
# MTP=0 Speculative drafts route experts for tokens that are later REJECTED and never
# DRAFT=0 emitted (eusage is incremented inside moe(), before verification). Those counts
# would describe text the model never produced.
#
# --gpu none The CUDA expert tier is not run-to-run deterministic (VRAM placement shifts
# between runs). Routing on the CPU path is reproducible. The tier only decides
# where weights live, not what the router picks, so CPU costs nothing here.
#
# --temp 0 Greedy: deterministic continuation, reproducible atlas.
#
# rm .coli_usage before EVERY run
# eusage is LOADED from <model>/.coli_usage at startup and written back at exit, so
# a naive STATS dump contains ALL PRIOR HISTORY, not this run. Removing it per run
# makes each dump exactly one probe. The user's learned cache is backed up and
# restored on exit.
#
# Prompt lengths are kept in a narrow band across categories: prefill routes the prompt tokens
# too, so a verbose category would otherwise simply look "busier".
# ---------------------------------------------------------------------------------------------
set -uo pipefail
HERE="$(cd "$(dirname "$0")" && pwd)"
PROBES="${1:-$HERE/probes.json}"
OUT="${2:-./atlas_out}"
COLI="${COLI:-./coli}"
NGEN="${NGEN:-64}"
MODEL="${COLI_MODEL:?set COLI_MODEL to the snapshot directory}"
[ -f "$PROBES" ] || { echo "no probe file: $PROBES" >&2; exit 1; }
[ -x "$COLI" ] || { echo "no coli at $COLI (run from c/, or set COLI=)" >&2; exit 1; }
mkdir -p "$OUT/stats"
export COLI_MODEL="$MODEL" MTP=0 DRAFT=0 TOPP=0
USAGE="$MODEL/.coli_usage"
BACKUP="$OUT/.coli_usage.backup"
[ -f "$USAGE" ] && cp "$USAGE" "$BACKUP" && echo "backed up $USAGE"
restore(){ [ -f "$BACKUP" ] && cp "$BACKUP" "$USAGE" && echo "restored $USAGE"; }
trap restore EXIT
python3 - "$PROBES" > "$OUT/runlist.tsv" <<'PY'
import json, sys
for cat, prompts in json.load(open(sys.argv[1])).items():
if cat.startswith('_'):
continue
for i, p in enumerate(prompts):
print(f"{cat}\t{i}\t{p}")
PY
n=$(wc -l < "$OUT/runlist.tsv"); i=0
echo "$n probes -> $OUT/stats"
while IFS=$'\t' read -r cat idx prompt; do
i=$((i+1))
dst="$OUT/stats/${cat}_${idx}.txt"
[ -s "$dst" ] && { echo " [$i/$n] $cat/$idx (cached)"; continue; }
rm -f "$USAGE" # start from an EMPTY routing history
STATS="$dst" "$COLI" run "$prompt" --ngen "$NGEN" --ctx 4096 --gpu none --temp 0 \
> "$OUT/stats/${cat}_${idx}.log" 2>&1
echo " [$i/$n] $cat/$idx $(grep -aoE '[0-9]+ selections across [0-9]+ distinct experts' \
"$OUT/stats/${cat}_${idx}.log" | head -1)"
done < "$OUT/runlist.tsv"
echo
echo "next:"
echo " python3 $HERE/analyze.py --stats $OUT/stats --out $OUT/experts.json"
echo " python3 $HERE/validate.py $OUT/stats 200"
+79
View File
@@ -0,0 +1,79 @@
#!/usr/bin/env python3
"""Leave-one-prompt-out validation of the Expert Atlas (#175).
"Replicates across the 3 prompts I picked" is not the same as "generalises". The atlas is
only real if a specialist set learned from SOME prompts predicts routing on a prompt it has
never seen.
Protocol, for every category c and every held-out prompt h of c:
1. build c's top-K specialist set from c's OTHER prompts only (h excluded entirely)
2. do the same for all other categories (using ALL their prompts — they never saw h either)
3. on the held-out run h, measure what share of routing selections land in each set
4. the atlas works if c's own set wins on h
If specialisation were an artifact of prompt wording, the held-out prompt would not prefer
its own category's set. Chance is 1/C.
"""
import glob, json, os, sys
from collections import defaultdict
STATS = sys.argv[1] if len(sys.argv) > 1 else "stats"
K = int(sys.argv[2]) if len(sys.argv) > 2 else 200 # specialists per category
runs, tot = {}, {}
for path in sorted(glob.glob(os.path.join(STATS, "*.txt"))):
cat, idx = os.path.basename(path)[:-4].rsplit("_", 1)
d = {}
t = 0
for line in open(path):
p = line.split()
if len(p) == 3:
d[(int(p[0]), int(p[1]))] = int(p[2])
t += int(p[2])
runs[(cat, idx)] = d
tot[(cat, idx)] = t
cats = sorted({c for c, _ in runs})
idxs = {c: sorted(i for cc, i in runs if cc == c) for c in cats}
C = len(cats)
def specialists(cat, exclude):
"""Top-K experts by lift for `cat`, computed WITHOUT the excluded run."""
share = defaultdict(lambda: defaultdict(float))
for c in cats:
used = [i for i in idxs[c] if not (c == cat and i == exclude)]
for i in used:
for k, n in runs[(c, i)].items():
share[k][c] += n / max(1, tot[(c, i)]) / len(used)
scored = []
for k, per in share.items():
s = sum(per.values())
if s <= 0:
continue
p = per[cat] / s
if per[cat] > 0:
scored.append((p, k))
scored.sort(reverse=True)
return {k for _, k in scored[:K]}
print(f"leave-one-prompt-out, {C} categories, top-{K} specialists per category")
print(f"chance = {100.0/C:.1f}%\n")
hits = 0
trials = 0
for c in cats:
for h in idxs[c]:
sets = {cc: specialists(cc, h if cc == c else None) for cc in cats}
held = runs[(c, h)]
htot = max(1, tot[(c, h)])
scores = {cc: sum(held.get(k, 0) for k in sets[cc]) / htot for cc in cats}
win = max(scores, key=scores.get)
ok = win == c
hits += ok
trials += 1
own = 100 * scores[c]
best_other = 100 * max(v for cc, v in scores.items() if cc != c)
print(f" {c:<12} prompt {h} own-set {own:5.2f}% best-other {best_other:5.2f}% "
f"-> {'HIT ' if ok else 'MISS'} (predicted {win})")
print(f"\naccuracy: {hits}/{trials} = {100*hits/trials:.1f}% (chance {100.0/C:.1f}%)")