AIMER, short for Absolute mean over root mean square IMportance for Expert Ranking, is a simple, weight-only criterion for post-training pruning of mixture-of-experts (MoE) language models. Instead of collecting router statistics or activation traces on a calibration set, AIMER ranks experts directly from pretrained weights, prunes them uniformly layer by layer, saves a standard Hugging Face checkpoint, and supports downstream evaluation with the same benchmark stack used in the paper.

🔍 Overview

This repository gives you a minimal end-to-end pipeline:

pretrained MoE checkpoint
        |
        v
  AIMER expert scoring
        |
        v
uniform layer-wise pruning
        |
        v
  pruned HF checkpoint
        |
        v
evaluation with lm-eval / EvalPlus / LiveCodeBench / EvalScope / WildBench

✨ Highlights

• Calibration-free pruning. No calibration corpus, router statistics, or activation collection is required for AIMER scoring.
• Weight-only ranking. Expert scores are computed directly from pretrained expert parameters.
• Uniform expert pruning. The released CLI prunes the same number of experts from each detected MoE layer.
• Standard checkpoint export. Pruned models are saved with save_pretrained(...) and can be evaluated directly.
• Paper-aligned evaluation wrappers. The repo includes scripts for multiple-choice QA, coding, math, and WildBench-style evaluation.

🗂️ Repository At A Glance

🧩 MoE families in this repo

🛠️ Installation

bash scripts/build.sh

⚡ Quick Start

✂️ Prune a model with AIMER

Use the pruning CLI directly when you only want the pruned checkpoint and score tables:

python src/reap/calib_free_prune.py \
  --model-name Qwen/Qwen3-30B-A3B-Instruct-2507 \
  --metric aimer \
  --sparsity-ratio 0.25 \
  --output-dir artifacts/calib-free/Qwen3-30B-A3B-Instruct-2507/aimer_0.250_uniform \
  --device-map auto \
  --torch-dtype bfloat16 \
  --metric-device auto \
  --local-files-only false \
  --trust-remote-code true

Behavior worth knowing:

• Higher AIMER scores are treated as more removable, so those experts are pruned.
• Pruning is always uniform across detected MoE layers.
• --local-files-only defaults to true, so set it to false if the model is not already cached locally.
• --metric-device cpu is useful if you want scoring to run from CPU copies after model load.

▶️ Prune and evaluate in one command

For the paper-style workflow, use the wrapper:

bash experiments/calib-free-cli.sh \
  0,1 \
  Qwen/Qwen3-30B-A3B-Instruct-2507 \
  aimer \
  0.25 \
  42

This wrapper:

1. sets CUDA_VISIBLE_DEVICES
2. prunes the checkpoint with AIMER
3. writes the pruned model under artifacts/calib-free/...
4. runs the evaluation pipeline

The first positional argument is the CUDA device mask. The wrapper also derives the serving port from the first visible GPU as 8300 + first_device.

Minimal example:

bash scripts/command_calibfree.sh

Important default:

• experiments/calib-free-cli.sh enables all benchmark groups by default, including WildBench.

📊 Evaluation

You can evaluate either an original checkpoint or a pruned checkpoint directly:

python src/reap/eval.py \
  --model-name artifacts/calib-free/Qwen3-30B-A3B-Instruct-2507/aimer_0.250_uniform \
  --seed 42 \
  --use-server true \
  --vllm_port 8300 \
  --run-lm-eval true \
  --run-evalplus true \
  --run-livecodebench true \
  --run-math true \
  --run-wildbench false \
  --results_dir artifacts/calib-free/Qwen3-30B-A3B-Instruct-2507/aimer_0.250_uniform/eval_vllm

📦 Output Layout

Pruned checkpoints are saved under:

artifacts/calib-free/<model-name>/<metric>_<ratio>_uniform/

Typical contents:

🏗️ Project Structure

AIMER/
|-- config/
|-- experiments/
|-- scripts/
|-- src/reap/
|   |-- args.py
|   |-- calib_free_prune.py
|   |-- eval.py
|   `-- model_util.py
`-- third-party/

🙏 Acknowledgement

This release builds on infrastructure adapted from CerebrasResearch/reap, the official codebase for REAP: Router-weighted Expert Activation Pruning for SMoE compression. We thank the REAP authors for open-sourcing their codebase and evaluation pipeline, which helped make this release possible.

📄 License

This repository is released under the Apache 2.0 License. See LICENSE.