DistMoE-Qwen3.5-35B-A3B (router-free / AoE conversion, KL-repaired v0.2)
Qwen3.5-35B-A3B with the router removed. Every MoE layer's learned linear router is
replaced by a key-addressed AoE gate (PEER-style): a query projection + per-expert key
vectors, expert selection by dot-product similarity in a 128-d key space.
This revision (
v0.2, "KL-repaired") ships gates
plus KL-distillation-repaired weights:
the experts, attention projections, and norms have been finetuned to match the base router
model's output distribution under the new gating. The earlier gates-only conversion (base
weights byte-identical, ~45MB delta) is preserved at revision
v0-gates-only.
Why remove the router?
A router is a private lookup table: expert
i only means anything inside this one checkpoint.
A key-addressed gate makes every expert an
addressable artifact — its address is a vector
in key space, independent of the checkpoint it shipped in. That unlocks the
distmoe runtime:
- Stream & cache experts: run the backbone locally (~7% of weights), fetch int4-quantized
experts from a pool on demand, cache them on disk. The working set specializes toward your
task; network cost amortizes to ~zero as the cache warms.
- Grow the pool: new experts are trained and placed in key space (no router retrain) —
we've added domain experts at −32.6% domain perplexity with base experts untouched.
- Heterogeneous serving: the wire format is framework-free — a PyTorch pool serves an
MLX (Apple Silicon) backbone.
Usage (monolithic — runs like any HF model)
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4repo = "artivus-ai/DistMoE-Qwen3.5-35B-A3B"
5tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
6model = AutoModelForCausalLM.from_pretrained(
7 repo, dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
8
9msgs = [{"role": "user", "content": "Explain mixture-of-experts in two sentences."}]
10ids = tok(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True),
11 return_tensors="pt").input_ids.to(model.device)
12print(tok.decode(model.generate(ids, max_new_tokens=200)[0][ids.shape[1]:],
13 skip_special_tokens=True))
For distributed serving (backbone + streamed expert pool), see the
dist-moe repo.
How this checkpoint was made
- Imitation: per-layer key-gates trained by DAgger against the original router's
selections (top-8 agreement ~70–85% depending on layer).
- Gate repair: end-to-end LM-loss finetune of the gates only (experts frozen) on a
diverse mixture. This is the
v0-gates-only revision (generic held-out ppl 10.68).
- KL repair (this revision): token-level full-vocab KL distillation from the frozen
base router model, with gates + experts + attention/norms trainable (embeddings and
lm_head stay frozen, shared with the teacher). Three successive runs, with the final one
adding creative-prose coverage to the distillation corpus — corpus coverage, not
optimization, turned out to be the binding constraint.
Honest numbers (read before comparing)
Conversion is not free, but output-matching repair recovers most of it:
| model | generic held-out ppl |
|---|
| base Qwen (router) | 7.45 |
gates-only conversion (v0-gates-only) | 10.68 (+43% rel) |
| this revision (KL-repaired) | 8.46 (+13.5% rel) |
Repair progression: 10.68 → 9.97 (KL, experts) → 9.01 (wider corpus) → 8.94 (+attn/norms)
→ 8.46 (+creative-prose corpus coverage). The held-out criterion is creative prose —
deliberately far from the original repair mixture. Generation is coherent and factual in
our spot checks (greedy, chat template). Treat this as a research artifact demonstrating
router-free addressability, not a drop-in replacement for base Qwen.
Serving-k advisory (measured)
The config default is top_k=8, matching the base checkpoint contract. But the PPL-vs-k
curve on this conversion is U-shaped with its optimum at k=16-24 — raising k at serve
time is a free quality knob if you can afford the extra expert compute
(measured on the gates-only revision; the diffuse routing distribution that produces this
curve is unchanged by repair):
| serving k | held-out ppl (vs k=8) |
|---|
| 4 | +30-45% (catastrophic) |
| 8 (default) | baseline |
| 16-24 | -3 to -5% (optimum) |
| 32 | ~baseline |
| 64 | slightly worse |
1model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, ...)
2model.config.num_experts_per_tok = 16 # or pass top_k in config overrides
Two things NOT to do (measured, they hurt):
- Don't sharpen the gate softmax at serve time (temperature/logit scaling). The
routing distribution is diffuse by design and the expert mixture is co-adapted to it —
β=2 sharpening costs +10-19% ppl, β=8 costs ~10x.
- Don't prune "cold" experts. Traffic analysis shows zero dead experts; 99% coverage
requires 84-92% of the pool. The tail is load-bearing.
Companion repos: the expert pool
The pool is published as individually-addressable artifacts (per-expert int4 CBOR + key
vectors), split across two repos to stay under HF's per-repo file limits:
Note: the pool currently contains the original (pre-repair) base experts — it composes
with the v0-gates-only revision. A repaired-expert pool update will follow.
Files
model.safetensors-*.safetensors — base shards with KL-repaired experts and
attention/norm weights patched in (no longer byte-identical to base Qwen — see
v0-gates-only for that)
model-aoe-gates.safetensors — the 40 repaired AoE gates
modeling_dist_moe.py / configuration_dist_moe.py — trust_remote_code implementation
- The original router weights remain inside the shards (unreferenced by the index), but
since the experts are repaired, reconstructing the original base router model from this
revision alone is no longer possible — use
v0-gates-only or the base Qwen repo.
Lineage & license
Base model:
Qwen/Qwen3.5-35B-A3B (Apache-2.0).
Conversion, gates, repair, and modeling code: Nous Research / Artivus (Apache-2.0).