moe-compress

A unified Mixture-of-Experts (MoE) expert-compression library implementing

• REAP — Router-weighted Expert Activation Pruning (Lasby et al., 2025, arXiv:2510.13999)
• REAM — Router-weighted Expert Activation Merging (Jha et al., 2026; arXiv:2604.04356)

behind a single API, with a registry/adapter architecture designed for adding further methods (HC-SMoE, M-SMoE, ...) and model architectures.

Both methods share one calibration pipeline, one saliency core, and one model-surgery layer:

REAP:  score experts with S_j = mean_{x in X_j} g_j(x) * ||f_j(x)||_2,
       drop the lowest-scoring experts + their router rows
REAM:  keep the top-S_j experts as protected centroids, greedily group the
       rest onto them by (gate-logit + gated-output) similarity, align
       neurons with the Hungarian algorithm, merge with saliency weights

The numerical core is ported from the official REAM reference implementation (Samsung, Copyright (c) 2026 — see attribution headers in src/moe_compress/core/) and verified against it weight-for-weight (tests/test_golden_parity.py).

Installation

conda create -n moe_comp_env python=3.12 -y
conda activate moe_comp_env
pip install torch --index-url https://download.pytorch.org/whl/cu128
pip install -e .

Requires Python >= 3.10, PyTorch >= 2.4 and transformers >= 4.57 (transformers >= 5.x for ZAYA1 / LFM2.5 support).

Quickstart (Python API)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from moe_compress import compress, CalibrationConfig

name = "LiquidAI/LFM2.5-8B-A1B"
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(
    name, dtype=torch.bfloat16, device_map="cuda").eval()

# REAP: prune 25% of experts per layer
model = compress(
    model,
    method="reap",
    compression_ratio=0.25,
    calibration_data=CalibrationConfig(
        datasets=["c4", "math", "code"],
        mix_ratio=[0.0, 0.3, 0.7],   # C4 : Math : Code
        num_samples=3072,
        seq_len=512,
    ),
    tokenizer=tokenizer,
)
model.save_pretrained("out/lfm25-reap25")

# REAM: merge instead of prune (reuses the same saliency machinery)
model = compress(model, method="ream", compression_ratio=0.25,
                 calibration_data=batch,   # or a CalibrationConfig
                 group_size=16)            # REAM's C hyperparameter

calibration_data accepts either a pre-tokenized batch {"input_ids": (N, L), "attention_mask": (N, L)} or a CalibrationConfig (datasets are then streamed and tokenized on the fly; .pt files produced by the original REAM calibration_data.py are supported via pretokenized_dir/tokenizer_tag).

Quickstart (CLI)

# REAP (pruning, remove 25% of experts per layer)
moe-compress compress \
    --model LiquidAI/LFM2.5-8B-A1B \
    --method reap --ratio 0.25 \
    --save-path output/lfm25-reap25

# REAM (merging, remove 25% of experts per layer)
moe-compress compress \
    --model LiquidAI/LFM2.5-8B-A1B \
    --method ream --ratio 0.25 --group-size 16 \
    --save-path output/lfm25-ream25

Calibration defaults to the paper's C4:Math:Code = 0.0:0.3:0.7 mixture (3072 samples x 512 tokens). To customize the mixture and size, add:

# supplementary: custom calibration mixture and size
    --datasets c4,math,code --mix-ratio 0.2,0.3,0.5 --num-samples 256

# other models: just change --model
moe-compress compress --model Zyphra/ZAYA1-8B --method reap --ratio 0.25 \
    --save-path output/zaya1-reap25

# YAML config with flag overrides (defaults < YAML < flags)
moe-compress compress --config configs/ream_lfm2.yaml --ratio 0.5

# pre-tokenize calibration data into reusable .pt files
moe-compress prepare-data --model Qwen/Qwen3-30B-A3B-Instruct-2507 \
    --tag qwen3 --output-dir data

Example configs live in configs/.

Supported methods

method="ream", merging="none" reproduces REAP; grouping="hcsmoe" gives the HC-SMoE baseline; group_size=0 gives MC-SMoE-style grouping; saliency="freq" gives frequency-based variants.

Supported model architectures

Unknown-but-standard MoE models fall back to a generic adapter; custom architectures can pass compress(model, adapter=MyAdapter(model), ...). See docs/DESIGN.md ("Extending the library") for how to add methods, adapters, and calibration datasets.

DeepSeek-V4

pip install -e ".[fp8]"        # transformers >= 5.14 + the FP8/FP4 kernels
CC=$(which gcc) moe-compress compress --config configs/reap_deepseek_v4.yaml

accelerate is a hard requirement (the FP8 quantizer refuses to load a checkpoint without it) and is a base dependency. The kernels package is optional: the adapter falls back to a dequantizing reference forward whenever the kernel cannot serve a module — not installed, weights off CUDA, or the kernel raising at call time (DeepGEMM rejects fp32 activations). The fallback is correct but much slower, and it skips the kernel's activation quantization, so calibration statistics come out marginally cleaner than inference would produce.

CC matters: Triton compiles a host launcher, and if it picks up NVIDIA's nvc from an HPC SDK on PATH the build fails on -Wno-psabi, which pushes every module onto the slow fallback. Pointing CC at gcc avoids that. Nothing breaks without it — you just lose the fast path, and a RuntimeWarning per module says so.

Checkpoints larger than memory (--streaming)

Compression is already a layer-at-a-time pass; only loading and saving assume the model is resident. --streaming removes that assumption — each decoder layer is read from the checkpoint when its turn comes and written back out as soon as it is done, so peak memory is the skeleton plus one layer rather than the whole model:

moe-compress compress --config configs/reap_deepseek_v4.yaml --streaming

DeepSeek-V4-Flash is 155 GiB; streaming compresses it inside a 100 GiB job. Results are identical to the resident path — tests/test_streaming.py asserts bit-equality of every expert tensor, router row and hash table for both REAP and REAM. The cost is roughly one extra pass over the weights.

Offloading is not an alternative. device_map="auto" with an offload_folder loads fine, but accelerate hands offloaded modules back as meta tensors: compression's in-place edits (tensor[keep_idx]) then succeed, change nothing and raise nothing, and the saved model is empty where it was edited. compress() rejects a meta-device model rather than let that happen.

Without --streaming, --device-map auto spreads the model over GPU and host RAM using --max-memory-gpu / --max-memory-cpu. Leave GPU headroom for activations. There is no distributed compression: extra nodes do not give one run more memory.

Streaming requires the adapter to declare uniform_layer_inputs — that every decoder layer receives the same masks and position embeddings, which is what lets the probe stop at layer 0. It is set for DeepSeek-V4 (verified against DeepseekV4Model.forward) and off elsewhere, so other architectures get a clear error instead of silently wrong statistics.

Three things are specific to this architecture:

• The checkpoint stays quantized. Routed experts are FP4 (e2m1 packed two per byte, one UE8M0 scale per [1, 32] row group), everything else FP8 e4m3 on [128, 128] blocks; expanding a ~280B model to bf16 is not an option, so experts are dequantized one at a time. REAP is bit-exact — dropping an expert slices both the weight and its scale grid. REAM is not: a merged expert is snapped back onto the FP4 grid, and with only eight e2m1 magnitudes that snap is comparable to the original quantization error, which eats into what merging buys over pruning. Compare the two before assuming REAM wins here; singleton groups are left untouched so only genuinely merged experts pay it.
• The first three layers are hash-routed. tid2eid maps token id to expert ids directly. They are pruned to the same width as the rest (the config has one global n_routed_experts, so leaving them wide yields an unloadable checkpoint) and the table is rewritten onto the survivors: to the merge centroid under REAM, and under REAP to the surviving expert whose mean output on the dropped expert's own routed tokens is closest. Rows may end up naming one expert twice, which every transformers expert dispatch accumulates per (token, slot) — that expert simply gets a proportionally larger share.
• Hash-layer saliency is entirely corpus-driven. The table is near-uniform over experts, so table frequency carries no signal; an expert owning token ids absent from the calibration corpus is scored on nothing and gets dropped. That damage barely moves perplexity on general text and lands on inputs with rare tokens (other languages, code punctuation, jargon). The reconstruction report's hash_coverage block makes the bias visible — unseen_share_dropped above unseen_share_kept means the ranking preferred experts the corpus happened to exercise.

Measuring where quality was lost

pytest proves the surgery is correct; it says nothing about quality. Setting recon_probe_tokens caches a bounded slice of calibration data at each MoE block, then re-runs the compressed block on that same input and reports the relative error. It costs one extra block forward per layer and needs no full-model inference.

moe-compress compress --config configs/reap_deepseek_v4.yaml \
    --recon-probe-tokens 4096 --recon-report output/recon.json

# REAP vs REAM on the same calibration data, with the summary tables
python scripts/recon_deepseek_v4.py --ratio 0.25 --out output/recon_r25.json

Per layer the report gives rel_l2 (block output — what the rest of the network sees, damped by the untouched shared expert), rel_l2_routed (routed experts only — the undamped effect of the surgery), requant_rel_error (how far the FP4 snap moved a merged expert, REAM only) and, for hash layers, hash_coverage. Because each layer is scored on its own cached input rather than on drifted hidden states, the numbers attribute error per layer and do not sum to an end-to-end error.

Calibration data

Registered sources: c4 (allenai/c4, streamed), math (AI-MO/NuminaMath-1.5, cn_k12/olympiads subsets), code (bigcode/the-stack-smol) — the REAM paper's mixture. Mixing ratios are free parameters (mix_ratio must align with datasets by position and drives the MC-vs-generation trade-off studied in the paper).

Results: DeepSeek-V4-Flash

Full benchmark runs on DeepSeek-V4-Flash (256 routed experts, 156 GiB, MXFP4 experts) compressed to 178 and 132 experts per layer with both methods. Served with SGLang on GH200s (4 nodes for the base model, 2 for the compressed ones).

Change from the base model, in points:

REAP is close enough to the base model that the difference is mostly not measurable. REAP 178 returns identical pass@1 on both code benchmarks — 0.8720 and 0.7407, not one problem different across 542 — and gives up 0.8 points of GSM8K. REAP 132 halves the checkpoint, 156 GiB to 79, and still lands within 2 points on GSM8K and inside the noise floor elsewhere. REAM pays a much larger price, and one that grows sharply with how much is removed: −8.6 points of GSM8K at 178 experts, −25.6 at 132.

Read the ordering off GSM8K, not the mean. It is the only one of the four that reproduces to ±0.25 points (see the caveats below), and on it the ranking is exactly what the sizes predict:

base 0.9484 > REAP 178 0.9401 > REAP 132 0.9295 > REAM 178 0.8620 > REAM 132 0.6922

The places where a compressed model appears to beat the base model — REAP 132's +2.44 on HumanEval+ is 4 problems out of 164 — are inside the reproducibility of those benchmarks. "Not measurably worse" is the claim they support; "better" is not.

All five rows come from one serving configuration. An earlier revision of this table had the base model's EvalPlus figures from a different one, and re-running them to match moved HumanEval+ by a single problem — which was worth 0.6 points of apparent advantage to every compressed model. Matching the configuration is not a formality at this resolution.

Metrics: GSM8K exact_match,strict-match, MATH-500 math_verify,none, HumanEval+/MBPP+ pass@1_plus, all greedy (n=1). MATH-500 also carries an exact_match,none figure that runs 4-7 points lower — do not mix the two.

Two caveats worth carrying over to any comparison of your own:

• MATH-500 is not reproducible to better than ~5 points. Re-running the same model on the same 500 problems moved the score by −2.8, +3.4 and −4.2 points across three models. That is generation non-determinism, not sampling error — long chains have more places to diverge. GSM8K over 1319 problems stayed within ±0.25 points in the same experiment. Read MATH-500 differences below ~5 points as noise (the +0.80 above included).
• Greedy decoding does not make scores independent of serving settings. Batch composition and prefix caching change floating-point reduction order, which moves logits in the low bits, which occasionally flips an argmax and changes the whole continuation. Usually sub-point noise — but in an earlier measurement round a run of REAM 178 under a saturated KV cache scored GSM8K 0.6755 against 0.8476 on re-measurement, a 17-point error. The tell was the gap between exact_match,strict-match and exact_match,flexible-extract: 1.3 points in the bad run, 0.08 after. Check that gap on every run; more than a point means the generations, not the model, are the problem.

Repository layout

src/moe_compress/       the library (see docs/DESIGN.md)
  core/                 saliency, grouping, alignment, observer, pipeline
  models/               ModelAdapter base + per-architecture adapters
  compressors/          BaseCompressor + REAP/REAM (COMPRESSORS registry)
  data/                 calibration dataset registry + batch building
  eval/                 benchmark catalog + lm-eval/EvalPlus backends
configs/                example YAML configs
docs/                   ANALYSIS.md (paper<->code), DESIGN.md (architecture)
tests/                  pytest suite (tiny in-memory models, no downloads)
scripts/                GPU smoke tests
paper/                  the REAP and REAM papers (PDF)

Benchmarking (moe-compress eval)

vLLM-backed benchmarks for base and compressed models, selected individually by name:

Names are case- and separator-insensitive (arc-c = ARC_c). A group name (mc, math, code) expands to its members, and any other name is passed through as an lm-eval task id, so tasks outside the table (piqa, ...) work too. The + variants report pass@1 on EvalPlus' extended tests; requesting HumanEval and HumanEval+ together costs one generation pass.

# install the eval dependencies (vllm 0.25.1 matches torch 2.11)
pip install -e ".[eval]"
# prebuilt flashinfer kernels for vLLM (avoids fragile on-the-fly JIT builds;
# match the installed flashinfer-python version)
pip install flashinfer-jit-cache==0.6.13 --extra-index-url https://flashinfer.ai/whl/cu130

# individual benchmarks
moe-compress eval --model output/lfm25-reap25 --tasks HumanEval+,MBPP+,GSM8K,BoolQ

# whole groups, reduced scale (per-task sample limit)
moe-compress eval --model output/lfm25-reap25 --tasks mc,math,code --limit 50

# full run from a YAML config
moe-compress eval --config configs/eval.yaml

# compare two result directories (e.g. base vs compressed)
moe-compress eval-diff output/eval/lfm25-base output/eval/lfm25-reap25

Models too large for one node (--engine server)

--engine vllm builds the engine inside each backend's process, which caps the model at what one host's GPUs hold and loads the weights once per backend. For anything larger, stand up one OpenAI-compatible server and point the benchmarks at it. --engine server talks plain HTTP, so any engine will do:

pip install -e ".[eval]"

# one rank per node; each finds the others through --dist-init-addr
python -m sglang.launch_server --model-path output/deepseek-v4-flash-reap25 \
    --served-model-name dsv4 --tp-size 4 \
    --nnodes 4 --node-rank $RANK --dist-init-addr <head>:5000 \
    --host 0.0.0.0 --port 30000 &

moe-compress eval --model output/deepseek-v4-flash-reap25 --model-tag dsv4 \
    --engine server --base-url http://<head>:30000/v1 --served-model-name dsv4 \
    --tasks GSM8K,MATH-500,HumanEval+,MBPP+

The results table above was produced this way. SGLang rather than vllm serve for DeepSeek-V4 specifically: it has the architecture in-tree and its multi-node mode is --nnodes/--node-rank over plain torch.distributed, with no Ray and no pipeline-parallel P2P path. vLLM's multi-node pipeline parallelism wedged mid-decode on this model. Neither is a requirement of this package — it only needs an endpoint.

This is not merely a convenience. EvalPlus' vLLM provider hard-codes its engine keyword arguments — tensor_parallel_size, dtype, trust_remote_code, enable_prefix_caching and nothing else — so it cannot be given pipeline parallelism or a Ray cluster, and vLLM offers no environment variable to force the executor backend. Reaching the model over HTTP is the only way to benchmark HumanEval+/MBPP+ on a model that spans nodes. lm-eval goes through local-completions, EvalPlus through its openai backend, and both share the one loaded copy.

Tokenization stays local (the --model path is used as the tokenizer), so few-shot prompt construction does not depend on the server.

Caveat: only generative benchmarks are verified over this path. Loglikelihood scoring (the mc group) needs the endpoint's echo + logprobs support and has not been tested; use the in-process engine for those.

Ready-made Miyabi PBS scripts for both the compression and the multi-node benchmark live in the private companion repository under miyabi_script/.

Resuming an interrupted run (--use-cache)

A full lm-eval pass on a large model takes longer than a short scheduler queue allows, and a server that dies mid-benchmark otherwise costs the whole run. --use-cache keeps every generation in SQLite as it arrives, so re-running the same command picks up where the last one stopped:

moe-compress eval --model output/deepseek-v4-flash-reap25 --model-tag dsv4 \
    --engine server --base-url http://<head>:30000/v1 --served-model-name dsv4 \
    --tasks GSM8K,MATH-500 --use-cache

Cached requests: 1204, Requests remaining: 615

Writes are committed per response, so a hard kill loses nothing already generated. Scores are still computed from the complete request set, so the last run in a series is the one that writes lm_eval.json. EvalPlus resumes on its own and needs no flag.

The flag takes no path: the cache lives under <output-dir>/<model-tag>/lm_eval_cache/ and is per model by construction. lm-eval keys cached responses on the prompt alone — no model name, no checkpoint path — so a shared file would let one model score with another's answers and report it as a normal result. A cache_meta.json alongside the database catches the case the path cannot, one tag re-pointed at a different checkpoint, and refuses to run.

Delete the directory when re-measuring a model whose serving setup changed (concurrency, KV budget, prefix caching). Those settings do not change the cache key but can change the output — see the second caveat under the results table above.

Prompt protocol

lm-eval's defaults (raw completion prompts, no BOS token, short generation budget) understate instruction- and reasoning-tuned models badly — enough to swamp the effect of compression you are trying to measure. Two protocols, differing by task type:

• Multiple choice — --add-bos-token only, for models trained with a BOS token; without it they score near chance. The chat template hurts loglikelihood scoring, so leave it off.
• Generative reasoning (GSM8K, MATH-500) — add --apply-chat-template --num-fewshot 0 --max-gen-toks 8192 --max-model-len 16384 so the model can think in the format it was trained on.

# multiple choice
moe-compress eval --model output/model-reap25 --tasks MMLU,ARC-c,HellaSwag,WinoGrande \
    --add-bos-token --output-dir output/eval/mc

# reasoning (separate --output-dir: both use the lm_eval backend, so a second
# run in the same directory would overwrite the first one's lm_eval.json)
moe-compress eval --model output/model-reap25 --tasks GSM8K,MATH-500 \
    --add-bos-token --apply-chat-template --num-fewshot 0 \
    --max-gen-toks 8192 --max-model-len 16384 --output-dir output/eval/math

Whichever protocol you pick, use the same one for every model you intend to compare.

The two protocols also change which answer filter works. GSM8K's strict-match filter demands the few-shot #### 42 format and scores 0 on chat-formatted answers, so under the chat protocol flexible-extract is the real number. A headline metric of exactly 0 with a non-zero sibling filter is detected and reported instead, with a warning; the summary line names the metric each score came from.

Both backends run on vLLM by default. --engine hf switches them to transformers (lm-eval's hf model, EvalPlus --backend hf), which is considerably slower but supports architectures vLLM has no model class for:

moe-compress eval --model output/model-reap25 --engine hf \
    --tasks mc,math,code --limit 50

Requested benchmarks are grouped by backend and written to output/eval/<model-tag>/<backend>.json (lm_eval.json, evalplus.json). When more than one backend is involved each runs in a fresh process (vLLM engines do not reliably release GPU memory in-process). Note --limit applies to each of MMLU's 57 subtasks. eval-diff merges all result files in a directory, so runs that covered different benchmarks still compare on what they share. Further backends (LiveCodeBench, SWE-Bench, BFCL, ...) can be added by registering an Evaluator subclass in EVALUATORS and listing their benchmarks in src/moe_compress/eval/tasks.py.

--suites mc,math,code remains accepted as a legacy spelling of --tasks with group names.

Tests and smoke runs

pytest tests/                       # fast, CPU, no downloads (~5 s)
ls scripts/                         # per-model GPU smoke tests

The GPU smoke scripts compress with both methods at 25%, then report parameter reduction, C4 perplexity before/after, a greedy generation sample, runtime and peak VRAM, and verify save/reload.

Parity tests compare the ported math against the original REAM reference implementation; they run when an untracked working copy is present at output/ream (or in the git history of this repository) and skip otherwise.

Documentation

• docs/ANALYSIS.md — precise mapping from the papers' equations to code, shared-component analysis, and known implementation deviations.
• docs/DESIGN.md — architecture, registry/adapter design, the capture-and-replay pipeline, and design patterns adopted from ms-swift/verl/slime.

License and attribution

• Ported computation from the REAM reference implementation, Copyright (c) 2026 Samsung Electronics Co., Ltd. — see the LICENSE file and the attribution headers on the ported files.
• src/moe_compress/core/hc_smoe.py contains HC-SMoE clustering code, Copyright (c) 2023 UNITES Lab.
• The Cerebras REAP repository (https://github.com/CerebrasResearch/reap) served as a design reference for observers/tests.