Add opt-in per-expert causal-ablation harness (ABLATE_SCORE) for moe() - #621
Conversation
Off-by-default research instrument: ABLATE_SCORE=<manifest> runs a teacher-forced prefill per item with a chosen (layer,expert) cell ablated in one of three modes (contribution-zero / route-around / module-swap) and reads out the final logits (ABLATE_OUT, coli-ablate/1). Inert unless enabled (guarded by g_abl.mode); standalone unit gate tests/test_ablate proves inertness + per-mode semantics (18/18). Causal complement to the Expert Atlas (JustVugg#175).
|
Nice harness — the design is right: the three moe() hooks are `if(g_abl.mode==...)` guards that short-circuit when `mode==0`, so an un-ablated build runs the exact original path (byte-identical by construction), the ablation is per-item with an explicit reset (no leak), and it's fully opt-in via `ABLATE_SCORE`. The standalone `test_ablate.c` faithful-miniature is a good call, and all 9 CI checks that ran are green. One thing before I merge, because this touches the hottest loop in the engine and I don't merge core-path changes on inspection alone: the token-exact oracle didn't run on this PR (it's absent from the checks — a CI-trigger quirk on your branch, not a failure), and that's exactly the check that proves moe() output is unchanged with ablation off. Could you confirm byte-identity empirically? Either:
A one-line note that decode tok/s is unchanged with the harness off (the added branches are predicted-not-taken, so I expect zero regression) would close it. Do that and this merges — the code itself already looks correct. |
|
|
…l.h in the colibri rule Two conflicting hunks, resolved by a maintainer rather than asking for a rebase. TEST_BINS: dev no longer has a hand-written list, so this branch's entry is dropped and its test is picked up by its own rule. The colibri$(EXE) rule: both sides were right. This branch adds abl.h to the prerequisites; dev added $(VK_OBJ)/$(VK_SPV) and the Vulkan backend rules. Merged rather than picked -- dev's line with abl.h reinstated -- so the ablation harness still rebuilds when its header changes and the Vulkan build keeps working. make test-c passes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Resolved by a maintainer instead of asking for another rebase. The only conflicting hunk was the hand-written TEST_BINS line, which dev no longer has: gates are derived from the build rules, so this branch's tests are picked up by their own rules and the manual list entry is dropped. No other file conflicted and no commit on this branch was rewritten. Verified before pushing: every test this branch adds a rule for is in the gate set, and make test-c passes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
An off-by-default research instrument that measures the causal effect of individual experts on the model's real output — a complement to the routing-affinity Expert Atlas (#175): the Atlas shows what routes where; this shows what an expert actually does.
What it adds
c/abl.h(new, ~100 lines, self-contained, zero deps): a per-item ablation config with three modes + a fail-closed per-item reset. Holds no cross-item state.c/colibri.c— four small guarded insertions:#include "abl.h"+static Abl g_abl = {0};norm_topk): mode 2 route-around (drop the ablated cell so survivors renormalise torouted_scale) and mode 3 module-swap (remap the slot's expert id at the same routed weight) — both act before counters/norm_topk/routed_scale/ROUTE_TRACE, so downstream accounting reflects the post-ablation routing.if(!nr) continue;): mode 1 contribution — skip this expert's matmul+accumulate so its weighted output is zero (routing/counters/trace unchanged).run_ablate_score()— a teacher-forced path reached viaABLATE_SCORE=<manifest>that, per item, resetsg_abl, configures the ablated(layer,expert)cells, runs one prefill, and reads out the final logits at every target position toABLATE_OUT(JSONL,coli-ablate/1: exact per-token NLL, gold-vs-best logit margin, argmax correctness, top-32 for an approximate next-token KL). Dispatched just beforeSCORE. IfROUTE_TRACEis set, the existing tracer dumps the post-ablation routing for free.c/Makefile:abl.hadded to thecolibriprerequisite list; a standalonetests/test_ablatetarget inTEST_BINS.c/tests/test_ablate.c(new): a no-model unit gate that drives a faithful miniature ofmoe()'s routing+accumulate against the sameabl.hthe engine links.Motivation
colibri already reasons about which experts matter —
EXPERT_BUDGET(#254), the Expert Atlas (#175, routing affinity only),CACHE_ROUTE. What's missing is a causal instrument: "if this specific expert is dropped / zeroed / substituted for this token, how do the final logits move?" — the evidence needed to tell a safe-to-drop expert from a load-bearing one, i.e. to give expert-pruning/budgeting an empirical basis rather than a blind gate-weight cap. The harness produces the evidence; it doesn't act on it (no pruning claim bundled).Testing (against HEAD
81f08a0)make colibricompiles clean under the project's-O3 -march=native -Wall -Wextra(0 warnings from the patch);ABLATE_SCORE/ABLATE_OUT/coli-ablate/1present in the binary.make tests/test_ablate→ ALL PASS (18/18): P1 mode-0 byte-identical to the un-hooked reference (inert when off); P2 ablating an unused cell is a no-op in all 3 modes with the(layer,expert)key respected; P3abl_reset()restores OFF (no spec leaks to the next item); S1/S2/S3 each mode has exactly its intended arithmetic. Added toTEST_BINSso the inertness guarantee can be a permanent CI gate.Caveats / scope
g_abl.mode, only ever set on theABLATE_SCOREpath; with it unset the forward is unchanged.COLI_CUDA/COLI_METALan expert may compute on the GPU before that skip (matches the deterministic-CPU-forward discipline). Modes 2/3 edit the selected set in FASE A and are backend-agnostic. If you'd prefer mode 1 backend-agnostic too, an alternative is zeroingws[]for the ablated cell right afterrouted_scale— one line, all backends honour it — at the cost of the route trace showing weight 0 for that cell. Flagged for your call.Prepared for review — happy to adjust scope, in particular (a) whether mode-1 should be made backend-agnostic (see caveat), and (b) whether you'd like
tests/test_ablatein the defaultmake testgate (it is, viaTEST_BINS). Heads-up: there may be a delay before I can respond to review comments.