DistMoE-Qwen3.5-35B-A3B

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)

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

repo = "artivus-ai/DistMoE-Qwen3.5-35B-A3B"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    repo, dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)

msgs = [{"role": "user", "content": "Explain mixture-of-experts in two sentences."}]
ids = tok(tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True),
          return_tensors="pt").input_ids.to(model.device)
print(tok.decode(model.generate(ids, max_new_tokens=200)[0][ids.shape[1]:],
                 skip_special_tokens=True))

For distributed serving (backbone + streamed expert pool), see the dist-moe repo.

How this checkpoint was made

1. Imitation: per-layer key-gates trained by DAgger against the original router's selections (top-8 agreement ~70–85% depending on layer).
2. 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).
3. 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:

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):

model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, ...)
model.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:

• artivus-ai/DistMoE-Qwen3.5-35B-A3B-expert-pool — base set (256/layer) + spawned creative set (4/layer)
• artivus-ai/DistMoE-Qwen3.5-35B-A3B-expert-pool-donors — coder donor set (256/layer)

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).