Predictability-Aware Training for MoE Expert Prefetching

Can an MoE LLM be trained to be predictable — so that expert weights can be prefetched from disk before the router asks for them?

Inference engines like Colibri already run 744B-parameter MoE models on consumer hardware by keeping the dense backbone resident and streaming each expert's weights from disk when the router activates it. This is reactive: throughput is bottlenecked by fetch latency. This project trains a small activation-prediction model jointly with the LLM and folds its prediction accuracy into the LLM's training loss, so the large model is optimized not just for output quality but for being predictable.

Headline results (Tier A: 341M-param MoE, 16 experts, 3 seeds)

hit@k = fraction of tokens where the predictor's top-2 matches the router's realized top-2, measured h layers ahead (chance = 0.125). Between-seed std in parens; post-hoc control = fresh linear predictor on the frozen baseline backbone (the SOTA deployment pattern).

	h=1	h=2	h=4	val LM loss
baseline + post-hoc control	0.826 (0.003)	0.797 (0.002)	0.732 (0.004)	5.767 (0.006)
joint predictability training (lambda=0.3)	0.888 (0.002)	0.865 (0.002)	0.796 (0.004)	5.790 (0.006)

• +6.2-6.8 pts of expert-predictability at every horizon vs the linear post-hoc control (~15x seed noise); +2.9-5.2 pts vs the stronger ranking-aware control (MLP + margin loss, arXiv 2511.10676-style, 3 seeds, Exp 14) — the honest effect size. Quality cost: +0.033 +- 0.016 nats (3 seeds, init+data variance).
• A StickyMoE-style temporal-consistency loss (the closest training-time prior art) does NOT buy lookahead predictability: it monotonically reduces hit@k and raises router entropy in our formulation (Exp 7b).
• Predictability is a property of the backbone: fresh post-hoc predictors on the frozen joint backbone recover it fully (linear 0.890; ranking-MLP 0.930 vs 0.903 on baseline) — an inference engine can exploit it without any co-training.
• Trace-driven cache simulation (honest disk-queue economics): where disk slack exists, this accuracy converts to +3.2% tok/s and -31% misprefetch waste over the post-hoc control at Colibri-like geometry (Exp 5, corrected after red-team round 2).
• Boundary result (Tier C): the effect is pretraining-time only. A LoRA fine-tune on pretrained OLMoE-1B-7B at Exp-10 pressure left the backbone's intrinsic predictability unchanged (isolation test: 0.800 vs 0.799); at 10x pressure a small transfer appears (+1.3-2.6 pts linear probe) but remains 2.5-5x weaker than the pretraining effect under that probe — the method must be applied during pretraining, not retrofitted (Exps 10, 13).
• Full details in RESULTS.md (13 experiments, 3 red-team rounds); literature landscape in RESEARCH.md; method in PLAN.md.

Reproduce

uv venv .venv && uv pip install --python .venv/bin/python \
  --index-url https://rocm.nightlies.amd.com/v2/gfx1151/ --pre torch numpy
uv pip install --python .venv/bin/python datasets tiktoken

# joint training (treatment)
python train.py --tier B --horizons 1,2,4 --lambda-pred 0.3 --steps 6000 \
  --device cuda --save ckpt_joint.pt
# baseline + post-hoc predictor (control)
python train.py --tier B --horizons 1,2,4 --lambda-pred 0.0 --steps 6000 \
  --device cuda --save ckpt_base.pt
python train.py --tier B --horizons 1,2,4 --mode posthoc --ckpt ckpt_base.pt \
  --steps 2000 --device cuda

Developed on an AMD Ryzen AI Max+ 395 (Strix Halo, Radeon 8060S iGPU, gfx1151). For CPU-only: install the CPU torch wheel and use --device cpu.

Every experiment in RESULTS.md has a run chain in scripts/ (run under sg render on this box; || true after each leg absorbs the known ROCm teardown abort that fires after checkpoint save):

scripts/ chain	experiments
run_h124.sh, run_h124_rest.sh, run_lambda.sh	2, 3 (Tier B; Exp 1 predates scripts, logs only)
run_tierA.sh, run_tierA2.sh	4 (Tier A six arms + isolation test)
run_sim.sh	5, 8 (traces -> cache sim, TTFT, sensitivity grid)
run_seeds.sh	6 (3-seed replication)
run_remediation.sh	7 (ranking control s0, StickyMoE sweep, data-order seeds)
run_undertrain.sh, run_undertrain2.sh	9 (entropy ladder, 100M-token arms)
run_tierC.sh	10 (OLMoE LoRA A/B + isolation)
run_engine.sh	11 (export + O_DIRECT engine demo)
run_downstream.sh	12 (zero-shot downstream eval)
run_tierC_lam1.sh	13 (Tier C lambda=1.0 dose arm)
run_rank_seeds.sh	14 (3-seed ranking control)

Tracked small artifacts: results_cache_sim.csv, downstream_ckpt_*.json, paper_figure.png. Checkpoints and traces are gitignored (size); regenerate via the chains above (traces ~2 min via dump_traces.py).

Repository layout

model.py            GPT-style MoE + router + predictor + joint/ctrl losses
train.py            Tier A/B trainer (joint / posthoc modes, shuffled stream)
train_olmoe.py      Tier C trainer (OLMoE + LoRA + predictor, isolation modes)
cache_sim.py        trace-driven cache/prefetch simulator (honest disk queue)
dump_traces.py      routing-trace export for the simulator
toy_engine.py       O_DIRECT disk-resident inference engine (end-to-end demo)
export_engine_model.py  checkpoint -> engine format converter
eval_downstream.py  zero-shot downstream eval (HellaSwag/ARC/PIQA)
scripts/            run chains for every experiment in RESULTS.md
paper/              arXiv-ready LaTeX (main.tex + verified references.bib)
results_cache_sim.csv   simulator output tables
README.md / RESEARCH.md / PLAN.md / RESULTS.md / PAPER.md

Checkpoints, traces, logs, and engine blobs are local artifacts (gitignored); regenerate them via the scripts above.

Status

Complete: Tier A mechanism (3-seed), isolation test, sharpening control, StickyMoE baseline, 4x-training run, cache simulator + TTFT + sensitivity, real-engine demo, Tier C OLMoE boundary (two lambda values), downstream zero-shot quality eval. See RESULTS.md (13 experiments). Next: workshop writeup.