REAP MLX

REAP MLX applies Router-weighted Expert Activation Pruning (REAP) to MLX-LM mixture-of-experts models on Apple Silicon. It loads an MLX-LM model, collects selected-expert activation statistics on calibration data, prunes low-saliency experts in place, saves the pruned model, reloads it, and writes validation telemetry for the run.

Use it when you want local MoE pruning experiments on Apple Silicon without a CUDA, PyTorch, vLLM, plotting, or paper-reproduction stack.

• MLX-only runtime: the execution path is built around mlx and mlx-lm.
• Adapter-based architecture support: model-family differences are isolated in small adapters under src/reap/model_adapters.py.
• Calibration-driven pruning: observation records selected routes, router scores, output norms, weighted activation metrics, and max activations.
• Save/reload validation: runs through mlx_lm.save, reloads the artifact, validates expert counts, and can run a generation smoke test.
• Structured telemetry: each run writes validation-metrics.json with model metadata, timings, memory samples, throughput, pruning decisions, and artifact sizes.
• Import-light package: importing reap must not import MLX, MLX-LM, datasets, Torch, or vLLM.

Prerequisites:

• Apple Silicon macOS environment
• Python 3.12 or 3.13
• uv
• An MLX-LM compatible MoE model, local path, or Hugging Face repo

Install the development environment:

git clone git@github.com:egesabanci/reap-mlx.git
cd reap-mlx
uv sync --group dev

Check the CLI:

uv run python -m reap.entrypoint --help
# after install / uv sync:
uv run reap-mlx --help

Run the focused test suite:

uv run python -m pytest -q tests/test_mlx_*.py

The experiment wrapper uses LiquidAI/LFM2.5-8B-A1B-MLX-4bit by default and writes artifacts under artifacts/mlx/.

MAX_SAMPLES=8 MAX_SEQ_LENGTH=1024 uv run \
  bash experiments/mlx-pruning.sh \
  LiquidAI/LFM2.5-8B-A1B-MLX-4bit \
  theblackcat102/evol-codealpaca-v1 \
  reap \
  0.25 \
  42

Equivalent direct CLI call:

uv run python -m reap.entrypoint \
  --model-name LiquidAI/LFM2.5-8B-A1B-MLX-4bit \
  --dataset-name theblackcat102/evol-codealpaca-v1 \
  --prune-method reap \
  --compression-ratio 0.25 \
  --max-samples 8 \
  --max-seq-length 1024 \
  --seed 42 \
  --output-dir artifacts/mlx/lfm2-smoke \
  --verbose

The output directory contains the saved MLX-LM artifact and a metrics file:

artifacts/mlx/lfm2-smoke/
  config.json
  *.safetensors or *.npz
  tokenizer files
  validation-metrics.json

The pruning pipeline is intentionally linear and inspectable.

Only these model types are accepted. Architectures that reuse a supported module layout but route differently (for example DeepSeek, GLM4-MoE, or MiniMax) are rejected up front instead of being pruned with the wrong routing.

mlx-lm 0.32 changed the LFM2-MoE and PhiMoE routing. The routers replay whichever routing the installed mlx-lm implements, and the replay parity check stops any run whose routing they do not reproduce.

Adding a model family should usually mean adding or extending an adapter in src/reap/model_adapters.py, router logic in src/reap/router.py, and pruning coverage in tests/test_mlx_model_adapters.py, tests/test_mlx_router.py, and tests/test_mlx_prune.py.

All MoE layers must end with the same retained expert count because mlx-lm stores a single global num_experts. Selective flags are validated before calibration. Layers skipped from the primary --prune-method are still width-matched via expert_frequency so save/reload stays valid. Use --target-experts N to retain a fixed expert count on heterogeneous native layouts.

--prune-method selects the per-expert saliency score used to rank experts. Higher scores are kept.

For LFM2 models, the replay follows the installed mlx-lm. Under mlx-lm 0.31 with use_expert_bias=True, residual routing uses bias-adjusted scores, so REAP saliency accumulation uses pure softmax probabilities at the selected indices instead, which keeps the additive bias out of the ranking. mlx-lm 0.32 uses the Hugging Face routing (sigmoid gates, with expert_bias only choosing the experts), whose weights carry no bias and feed saliency directly; see observation and metrics.

--compression-ratio must be in [0, 1). For each MoE layer, REAP MLX prunes int(num_experts * compression_ratio) experts and always keeps at least one expert.

Run:

uv run python -m reap.entrypoint --help

Core options:

Selective pruning (--skip-layer-indices, --prune-layer-indices, --per-layer-ratios), calibration minimums (--min-calibration-samples, --allow-partial-calibration), resume strictness (--strict-resume / --no-strict-resume), and the observer replay check (--skip-replay-check) are described in the full CLI reference.

Every run writes structured validation telemetry. The JSON includes:

• run configuration and runtime versions;
• model architecture, adapter, expert count, top-k, and quantization metadata;
• calibration sample counts, token counts, and token throughput;
• per-phase timings for load, calibration, observe, prune, save, reload, and smoke validation;
• MLX memory samples when available;
• observer summaries and saliency inputs;
• per-layer retained and removed experts;
• output artifact sizes;
• reload validation and generation smoke results.

Use this file to compare pruning settings, calibration size, runtime cost, and artifact size across runs.

Maintainer-focused reference documentation is available in docs/index.md. It covers the pipeline, architecture, model adapters, calibration, observation metrics, pruning semantics, save/reload validation, CLI behavior, telemetry, and development workflow.

REAP MLX is an independent MLX-focused implementation. Some implementation decisions are inspired by the original CerebrasResearch/reap repository.

The pruning method is based on the paper REAP the Experts: Why Pruning Prevails for One-Shot MoE compression.

reap-mlx/
  README.md
  LICENSE
  pyproject.toml
  uv.lock
  experiments/
    mlx-pruning.sh
  docs/
    index.md
    *.md
  src/
    reap/
      checkpoint.py
      data.py
      entrypoint.py
      eval_nll.py
      metrics.py
      model_adapters.py
      observer.py
      prune.py
      router.py
      save.py
      validation_metrics.py
  tests/
    test_mlx_*.py        # test_mlx_e2e.py runs tiny real mlx-lm models

Install and sync dependencies:

uv sync --group dev

Run tests:

uv run python -m pytest -q tests/test_mlx_*.py

Run lint, import-safety, and CLI smoke checks:

uv run ruff check src tests
uv run python -c "import reap; import reap.entrypoint"
uv run python -m reap.entrypoint --help
uv lock --check

Check formatting hygiene before committing:

git diff --check

• Optional held-out NLL comparison of the unpruned and pruned model: --eval-calibration-nll
• Task benchmarks are out of scope for this repository; see docs/eval.md for how to evaluate a saved artifact with external harnesses.

REAP MLX is released under the MIT License. See LICENSE.