REAP expert pruning

REAP (Router-weighted Expert Activation Pruning) is a one-shot method — no gradient training, no fine-tuning. It runs a calibration set through the model, scores every expert by the mean of router_weight × activation_norm, and drops the least-salient ones. Because only a few experts fire per token, many are redundant.

What it buys you: MoE models keep the same number of experts-per-token, so pruning reduces total size / memory footprint, not decode speed. That makes it the right tool for fitting a model that's otherwise too big for a Spark's 128 GB — not for making an already-fitting model faster.

The calibration saliency scores are sums over tokens, so they're additive: you can split the calibration set across machines and sum the partial stats.

Setup

git clone https://github.com/CerebrasResearch/reap && cd reap
# Use a venv whose torch has CUDA built in — check first:
python -c "import torch; print(torch.cuda.is_available())"   # must be True
pip install accelerate datasets scikit-learn umap-learn safetensors

REAP's default imports pull heavy eval-only dependencies (vllm, lm_eval, evalplus). They're only needed for the optional --do_eval step; the fork used here makes them lazy so a pure prune run doesn't need them installed.

Recipe 1 — Nemotron-3-Nano-30B (fits in RAM)

A 31.6 B MoE with 128 routed experts. At ~59 GB (BF16) it loads comfortably into the Spark's 128 GB, so this is the straightforward path.

PYTHONPATH=src python -m reap.layerwise_prune \
  --model_name /path/to/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 \
  --dataset_name "theblackcat102/evol-codealpaca-v1:256" \
  --prune_method reap \
  --compression_ratio 0.5 \
  --batch_size 1 \
  --model_max_length 512 \
  --run_observer_only false \
  --results_dir out/nemotron

• --compression_ratio 0.5 is the fraction to prune → 64 of 128 experts are removed, leaving 64.
• dataset:256 streams 256 calibration samples (the :N suffix sets the count).
• The layer-wise observer processes one block at a time, so peak GPU memory is one block plus activations.

Result: 128 → 64 experts, 59 GB → 32 GB, coherent output preserved. Published at sapidlabs/Nemotron-3-Nano-30B-A3B-REAP-64e.

Spark gotchas (Nemotron-H)

• mamba-ssm not required. Nemotron-H is a hybrid Mamba/Attention/MoE. The repo's own modeling code hard-requires the mamba-ssm CUDA kernels, but the native transformers modeling has a pure-torch fallback. Load with trust_remote_code=False to use it (you'll see "falling back to the naive implementation" — that's expected, just slower).
• Nemotron-H is native in recent transformers, so no custom code is needed at all.

Recipe 2 — Hy3 340B (exceeds RAM → disk streaming)

Hunyuan Hy3 is HYV3ForCausalLM: 80 layers, 192 experts, ~295 B params, 557 GB on disk across 99 shards. That doesn't fit in one Spark's 128 GB — or even two Sparks combined (242 GB). The model must stream from disk.

The fork adds a --disk_stream mode: the model is built on the meta device (zero weights resident), and each decoder layer's weights are materialized from the on-disk safetensors shards only while that layer is being processed, then released back to meta. Peak RAM is ~one layer (~7 GB) regardless of the 557 GB total — verified running the full 295 B model in ~32 GB RAM.

Because even the pruned model is bigger than RAM, this is a two-step flow: stream-calibrate to save the saliency stats, then apply the prune as a checkpoint-to-checkpoint transform.

Step 1 — calibrate (streaming). --run_observer_only stops after calibration and writes the per-expert saliency to observations_*.pt, so the expensive streaming pass is done once.

PYTHONPATH=src PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
  python -m reap.layerwise_prune \
  --model_name /path/to/Hy3 \
  --dataset_name "theblackcat102/evol-codealpaca-v1:64" \
  --prune_method reap \
  --batch_size 1 \
  --model_max_length 512 \
  --disk_stream true \
  --run_observer_only true \
  --results_dir out/hy3

--disk_stream true builds the model on meta and streams each layer from the shards (needs a local checkpoint with model.safetensors.index.json).

Step 2 — apply the prune (streaming). REAP's normal prune step reloads the whole model into RAM — the same wall again. reap.stream_prune instead reads the shards, keeps the top experts per layer by saliency, re-indexes the survivors, and writes new shards holding only ~one shard in RAM:

PYTHONPATH=src python -m reap.stream_prune \
  --model_name /path/to/Hy3 \
  --observer_state out/hy3/Hy3/.../observations_1024_cosine.pt \
  --compression_ratio 0.75 \
  --out_dir /path/to/Hy3-REAP-48e

• --compression_ratio 0.75 prunes 75% → 192 experts become 48 (a "REAP75").

Result: 192 → 48 experts, 557 GB → 157 GB, loads cleanly with a normal from_pretrained. (Still BF16 — a 4-bit quant pass takes it to ~40 GB to actually run on one Spark.)

Why streaming beats the alternatives on a Spark

• Don't load the whole model to cuda. The GB10's memory is unified — CPU and GPU share one 128 GB pool. Loading a 557 GB model to cuda OOMs, and even loading it to cpu and moving blocks to GPU starves the CUDA allocator. Streaming keeps only one ~7 GB layer resident, so it sidesteps the wall.
• Weight-conversion matters. Hy3 stores experts unfused per-expert in the checkpoint (experts.N.gate_proj/up_proj/down_proj) but the model fuses them into batched gate_up_proj/down_proj at load time, plus renames router.gate → gate, shared_mlp → shared_experts, expert_bias → e_score_correction_bias. The streaming loader reproduces that conversion exactly (validated bit-for-bit against a normal from_pretrained load).

Running a pruned model that still exceeds one Spark

The BF16 Hy3-REAP is 157 GB — smaller than 557 GB, but still over one Spark's 128 GB. Two ways to actually run it:

• Quantize to 4-bit (~40 GB) and serve on a single Spark. This is the path to on-device inference, but tooling for a brand-new arch like hy_v3 is the catch (the base model's 4-bit was MLX/Apple; llama.cpp GGUF support may not exist yet).
• Pipeline-parallel across two Sparks. 157 GB doesn't fit in one 128 GB Spark but does fit across two (242 GB combined) — vllm serve --pipeline-parallel-size 2 puts ~40 layers (~78 GB) on each node, activations hopping over the fabric. vLLM 0.24 supports HYV3ForCausalLM natively. The multi-node setup has real environment friction (NCCL-over-fabric, Ray login-env PATH for the flashinfer JIT's ninja, Ray's memory monitor false-OOMing on unified memory, leaked placement groups) — all documented in the DGX Spark field notes under Clustering.

General gotchas

• compression_ratio is the pruned fraction, not the kept fraction. 0.5 removes half (keeps half); 0.75 removes three-quarters (keeps a quarter).
• Use the CUDA venv. A venv with CPU-only torch will run — silently, on the CPU — an order of magnitude slower. Always check torch.cuda.is_available().
• Clear __pycache__ after editing REAP. Stale .pyc files cost hours here: a running process kept executing old bytecode and ignoring source fixes. find src -name __pycache__ -type d -exec rm -rf {} + when in doubt.
• MoE calibration is compute-heavy — every expert's activation is computed for every calibration token to score it, so runtime scales with experts × layers × tokens. Start with a small calibration set to validate the pipeline, then scale.