MoE-Watcher-Modifier
MoE-Watcher-Modifier is an open-source, model-agnostic toolkit that analyses which experts a Mixture-of-Experts model actually uses, ranks them by importance, and rewrites the checkpoint with fewer experts — shrinking the model without touching training or inference code.
- Model-agnostic Mixture-of-Experts analysis and pruning toolkit. Finds the experts your model actually uses and rewrites the checkpoint with fewer of them.
- Python
- Teams serving a large MoE model on hardware that can barely hold it
- Anyone running Qwen3-Next, Mixtral, DeepSeek, Phi-3.5-MoE or OLMoE locally
- Researchers studying expert utilisation and routing collapse
- Engineers who want real routing data from production traffic, not synthetic benchmarks
- torch >= 2.0.0
- safetensors >= 0.4.0
- rich >= 13.0.0 (optional — nicer terminal output)
- transformers >= 4.40.0 (optional — full-model mode only)
Overview
Most Mixture-of-Experts checkpoints ship far more experts than any single workload needs. MoE-Watcher-Modifier measures which experts your traffic actually routes to, ranks them, and writes out a smaller checkpoint containing only the ones that earn their place. It works on any safetensors MoE checkpoint and requires no changes to your training or serving code.
What it does
- Model-agnostic — works with any safetensors MoE checkpoint, schema auto-detected from config.json
- Router-only mode profiles a checkpoint in minutes on CPU, loading only the router tensors
- Full-model mode captures true routing decisions from your own prompts
- Daemon mode is a transparent proxy that accumulates routing stats from live traffic with no added response latency
- Rewrites checkpoints on CPU — memory use is bounded by the largest shard, not the whole model
- Emits a pruning manifest so every produced checkpoint is traceable back to its source
Quickstart
Install
pip install torch safetensors rich
pip install transformers # only for full-model mode1. Inspect the checkpoint
Auto-detects the schema and prints the layer and expert layout.
python3 moe_monitor.py inspect --model-dir /path/to/checkpoint2. Collect expert usage stats
Router-only mode needs no GPU and finishes in minutes — the recommended first pass.
python3 moe_monitor.py router-only \
--model-dir /path/to/checkpoint \
--output ./stats/report.json \
--keep-experts 128 \
--new-topk 4 \
--samples 1024 \
--strategy per-layer3. Dry-run the rewrite
Validates the plan against the checkpoint without writing anything.
python3 moe_prune.py dry-run \
--model-dir /path/to/checkpoint \
--plan ./stats/report-plan.json4. Write the pruned checkpoint
python3 moe_prune.py prune-checkpoint \
--model-dir /path/to/checkpoint \
--plan ./stats/report-plan.json \
--output-dir ./pruned-128experts \
--copy-support-filesSupported model families
The schema is detected automatically from the checkpoint's config.json. Pass --schema to override it when a model reports something unexpected.
| Schema | Models |
|---|---|
qwen3_next | Qwen3-Coder-Next, Qwen3-Next |
qwen_moe | Qwen1.5-MoE, Qwen2-MoE |
mixtral | Mixtral-8x7B, Mixtral-8x22B |
deepseek | DeepSeek-V2, DeepSeek-V3, DeepSeek-R1 |
phi3_moe | Phi-3.5-MoE |
olmoe | OLMoE |
Three ways to collect routing statistics
Router-only mode loads just the router weight tensors — a few megabytes — and probes each router with random unit-norm hidden states. No GPU, no full model load, minutes even on very large checkpoints. It captures structural routing preferences, which is enough for a first pruning pass.
Full-model mode loads the whole model and installs forward hooks on every MoE gate, then runs your real prompts through it. This captures true routing decisions on your actual workload, and needs enough RAM or VRAM to hold the model.
Daemon mode is the one worth reaching for in production. It runs as a transparent HTTP proxy in front of any OpenAI-compatible model server — ollama, vLLM, llama.cpp, LM Studio, text-generation-webui — and accumulates routing statistics from live user traffic in a background thread. Requests are forwarded untouched, so nothing is added to the response path.
How the daemon works
Your application points at the daemon instead of the backend. For each request the daemon extracts the prompt text, tokenises it with the checkpoint's own tokenizer, looks up token embeddings from embed_tokens.weight, forwards each token's hidden state through the router weights, and records the selected experts per layer. Every N requests it writes a report and prints a ready-to-run prune command.
Because it only ever loads the router weights and the embedding table, the daemon runs comfortably alongside the model server on the same box — Linux, macOS or Windows, AMD, NVIDIA or CPU-only.
What the rewriter actually does
The pruner drops expert tensors absent from the keep list, renames retained experts to compact IDs from 0 to target-1, slices each router weight matrix to match the retained expert IDs, updates config.json with the new expert count and top-k, and writes a pruning_manifest.json for traceability.
With --copy-support-files it also carries across the tokenizer, generation config and other auxiliary files, so the output directory is a complete, loadable checkpoint rather than a bag of tensors. No GPU is required and peak memory is bounded by the largest shard.
The iterative workflow
The strongest results come from iterating. Run router-only on the original checkpoint for an initial structural plan, prune, then load the now-smaller checkpoint — which may finally fit in memory — and run full-model mode on it for real routing statistics. Re-rank with a different keep count, prune again, and finish with a finetune or distillation pass on domain data to recover quality.
Each round replaces guesswork with a stronger routing signal than the round before it.
What is MoE-Watcher-Modifier?
It is an open-source Python toolkit for analysing and pruning Mixture-of-Experts models. It measures which experts a checkpoint actually routes to, ranks them by importance, and rewrites the checkpoint with only the experts worth keeping — reducing model size without any change to training or inference code.
Which MoE models does it support?
Any safetensors MoE checkpoint. Schemas ship for Qwen3-Next and Qwen3-Coder-Next, Qwen1.5-MoE and Qwen2-MoE, Mixtral-8x7B and 8x22B, DeepSeek-V2, V3 and R1, Phi-3.5-MoE, and OLMoE. The schema is auto-detected from config.json and can be overridden with --schema.
Do I need a GPU to prune a MoE model with it?
No. Router-only analysis and the checkpoint rewrite both run on CPU. Peak memory during the rewrite is bounded by the largest shard rather than the size of the whole model. A GPU is only needed for full-model mode, which loads the model to capture real routing decisions.
How do I get expert usage statistics from production traffic?
Run the daemon. It is a transparent HTTP proxy that sits in front of any OpenAI-compatible model server — ollama, vLLM, llama.cpp, LM Studio — and accumulates routing statistics from live requests in a background thread, adding no latency to responses. After N requests it prints a ranked expert table and the exact prune command to run.
Does pruning experts hurt model quality?
Removing experts is lossy, which is why the toolkit ranks experts on your own routing data rather than pruning blindly, and reports the expert coverage fraction so you can see how much routing mass the keep list captures. The recommended workflow finishes with a finetune or distillation pass on domain data to recover quality.
Is it safe to run against my only copy of a checkpoint?
The pruner never modifies the source checkpoint — it writes a new directory. There is also a dry-run subcommand that validates the plan against the checkpoint structure without writing any files.