AverageMetaheuristicsEnjoyer/moe-routing-drift-checkpoints

MoE routing drift — checkpoints

Weights behind a 2x2 experiment: adaptation (none / GEPA / prompt-tuning / prefix-tuning) crossed with router retraining (frozen gate / gate retrained), on two MoE bases. Measurements, per-example scores and the write-ups live in the companion dataset repo moe-routing-drift-results.

Two kinds of artifact, and neither is a full model:

	file	what it holds
PEFT arm	adapter/adapter_model.safetensors + adapter_config.json	one tensor prompt_embeddings; prompt-tuning [100, hidden], prefix-tuning [100, layers*2*kv_dim]
retrained router	checkpoints/epoch_*/router.safetensors	only the MoE gate matrices: model.layers.N.mlp.gate.weight (+ expert_bias on Ling)

The experts, attention and embeddings are never touched — the router checkpoint is 0.061% of Ling and 0.041% of Qwen, so any quality change comes from which expert sees which token, not from the model learning something new.

Every run saves epochs 1, 2, 3, 4, 5, 7, 9, 11. Checkpoints are selected by quality on val, never by val loss — the two disagreed three times in this project. The selected epoch per cell is in manifest.json.

Loading

A PEFT arm (a subfolder, so pass subfolder=):

from peft import PeftModel
model = PeftModel.from_pretrained(
    base, "AverageMetaheuristicsEnjoyer/moe-routing-drift-checkpoints",
    subfolder="ling-mini-2.0/arms/prompt-tuning_ml-r1000_s42/checkpoints/epoch_005/adapter")

A router checkpoint is a partial state dict, not a model. Match gate modules by layer suffix and load non-strictly — this also works on a PEFT-wrapped model:

from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

p = hf_hub_download("AverageMetaheuristicsEnjoyer/moe-routing-drift-checkpoints",
                    "ling-mini-2.0/routers/gepa_s42/checkpoints/epoch_001/router.safetensors")
state = load_file(p)
gates = {n: m for n, m in model.named_modules() if n.endswith("mlp.gate")}
for name, mod in gates.items():
    suffix = name.split("model.layers.")[-1]           # "N.mlp.gate"
    mod.weight.data.copy_(state[f"model.layers.{suffix}.weight"].to(mod.weight.dtype))
    key = f"model.layers.{suffix}.expert_bias"
    if key in state and hasattr(mod, "expert_bias"):   # Ling only; Qwen has no such buffer
        mod.expert_bias.data.copy_(state[key].to(mod.expert_bias.dtype))

Base revisions these were trained against, pinned: inclusionAI/Ling-mini-2.0 at 920c3fd9916e3d5e543fc4f609e827cad8a32983, Qwen/Qwen3-30B-A3B-Instruct-2507 at 0d7cf23991f47feeb3a57ecb4c9cee8ea4a17bfe. manifest.json carries a sha256 for every file here — a corrupted base shard cost this project a day of chasing NaNs, so verify before you debug.

Layout

ling-mini-2.0/arms/{prompt,prefix}-tuning_ml-r1000_s42[_on-router]/
ling-mini-2.0/routers/{base,gepa,prompt,prefix}_s42/ base_gamma0_s42/
                      {base,gepa}_s4{3,4}/ prompt_weak-ep{2,3}_s42/
qwen3-30b-a3b-instruct-2507/arms/... routers/...
ling-mini-2.0-aqua/arms/...        # AQuA-RAT reasoning leg
manifest.json                      # path, size, sha256, run config, which 2x2 cell

Task: civil_comments multi-label toxicity, 1000 training examples, seed 42, 11 epochs, batch 4 x accum 8. Arms keep their own papers recipes (Adafactor 0.3 for prompt-tuning, AdamW 5e-5 for prefix-tuning), 100 virtual tokens each. Router: AdamW 3e-5, cross-entropy only — no auxiliary balancing loss anywhere.

_on-router arms are the reverse training order: the gate is calibrated on the base first, then the arm is trained on top of it with the gate frozen.

Caveats

• Arm strength is not matched. Each arm follows its own recipe, so comparing cells across arms mixes mechanism with adaptation strength. All eight epochs of every arm are published so a strength-matched comparison can be built without retraining.
• One seed per PEFT cell. Base and GEPA cells are replicated at seeds 42/43/44 (spread 0.008 against an effect of 0.13); the prompt and prefix cells are not.
• Router LR 3e-5 was not swept. Early checkpoints are the hedge.
• Ling expert_bias accumulation. The DeepSeek-style balancing step was added directly into a bfloat16 buffer, where a 1e-4 step rounds away for most experts. Runs marked in manifest.json predate the float32 fix; base_fixedbias_s42 is the re-run.
• GEPA is a text prompt, not weights — it lives in the results repo.