Skip to content

metiu1/prune_lab

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

21 Commits
 
 
 
 
 
 
 
 

Repository files navigation

prune_lab — compress LLMs on consumer hardware

⚠️ BETA — under active development. prune_lab is experimental: APIs, commands and output formats may change without notice. Some architectures are verified, others are still being tested. Use it for research/experimentation, not yet in production. Bug reports and contributions welcome in the issues.

Make large LLMs runnable on a normal PC — without retraining anything. prune_lab removes the redundant parts of a model (whole layers and, on Mixture-of-Experts, the experts the router never uses) and then quantizes to 4-bit. Everything runs with disk offload, so it works even with 16 GB of RAM and 4 GB of VRAM (tested on a GTX 1650).

No training, no datacenter GPU. Just surgery on the existing weights + quantization.


Table of contents


The idea in one formula

A model's weight is, in a nutshell:

total_weight  =  n_parameters  ×  bits_per_parameter
                 \____________/   \_________________/
                    pruning          quantization

They are two independent levers that multiply. If you remove 30% of the parameters and go from 16 to 4 bits, you don't add the savings: you multiply them.

0.70  (pruning)  ×  0.25  (16→4 bit)  ≈  0.18   →  ~5.5× smaller

The key point: not all parameters are needed. In a transformer many layers change the internal state very little (input ≈ output): they are redundant. In MoE models, the router always uses the same experts and ignores others. prune_lab measures which parts are superfluous and removes them for real (structured pruning: it deletes parameters, doesn't just zero them out — so memory actually drops).


The three levers

Lever What it removes How much On which models
Depth pruning whole redundant transformer layers ~25–30% all
Expert pruning the lowest-salience experts (gate prob mass) ~25% drop (see measured curve) MoE only
Q4 quantization weight precision (16 → 4 bit) ~4× on bits all

Depth pruning (ShortGPT + magnitude metric). For each layer we measure how much its contribution to the internal state matters. The classic ShortGPT method uses Block Influence, i.e. the angle between input and output:

BI_i = 1 − cos( hidden_in_i , hidden_out_i )

But in a transformer hidden_out = hidden_in + f(hidden_in): the residual dominates, so the cosine is ~1 almost everywhere and BI stays crushed toward 0 → it separates layers poorly. prune_lab adds (and uses as the primary signal) a magnitude metric, which doesn't saturate and also sees the layers that change the norm of the state (which the cosine ignores entirely):

importance_i = E[ ‖ hidden_out_i − hidden_in_i ‖ / ‖ hidden_in_i ‖ ]
             = E[ ‖ f(hidden_in_i) ‖ / ‖ hidden_in_i ‖ ]

Low importance = the layer adds little = redundant = removed first. Both values end up in layer_ranking.json; BI stays as a reference.

Expert pruning (MoE only): by salience, not by frequency. A MoE has dozens of experts per block but activates only a few per token (top-k). With hooks on the router, prune_lab accumulates, for each expert, the gate probability mass (the renormalized top-k softmax) — not the plain count of selections. The reason: an expert chosen rarely but with high probability when it is chosen weighs heavily on the output; counting it as "rare" would prune it by mistake. The probability mass ("salience") is a far more faithful proxy of importance than frequency. The lowest-salience experts are removed and the router is fixed up (along with the DeepSeek-style correction bias). It is the main lever on giant models (DeepSeek, GLM, Qwen3-MoE), because that's where most of the parameters live. (Router without explicit logits → salience automatically degrades to the count, without breaking.)

Multi-domain calibration. Both rankings above depend on the calibration data. Using a single one (wikitext) would prune layers/experts useful for code, math or other languages. prune_lab calibrates on a mix (prose + code + math + multilingual), so the ranking doesn't overfit to a single domain. The evaluation perplexity, on the other hand, stays on pure wikitext-2, to keep it comparable with the literature.

Quantization. The remaining weights go to 4-bit NF4 (double-quant).

⚠️ Breaking threshold (measured). Beyond ~30% of layers removed — or ~30% of experts on a fine-grained MoE — quality collapses without healing (light LoRA fine-tuning). Measured on OLMoE: 25% expert drop = +25.6% PPL (ok), 50% drop = +277% (broken). Below the threshold the loss is small. See Measured results and FAQ.


Installation

git clone https://github.com/metiu1/rete_llm.git
cd rete_llm
pip install -r prune_lab/requirements.txt   # torch, transformers, accelerate,
                                             # bitsandbytes, datasets, safetensors, gguf
  • Python 3.10+
  • 4-bit quantization (NF4) requires CUDA (bitsandbytes). Analysis, pruning and GGUF export (gguf, pure python) also run on CPU only: for Q4 without a GPU use pl.py export (below).
  • Everything uses automatic disk offload: the model does not need to fit in RAM.

What model can I feed it?

HuggingFace Transformers format — a causal LM (decoder) model. Two ways:

  • Hub ID: org/name, e.g. Qwen/Qwen2.5-3B. Downloaded automatically.
  • Local folder: a folder with config.json, the tokenizer files and the weights in *.safetensors (or pytorch_model.bin). You pass the folder path.

Supported architectures: Llama, Qwen/Qwen2/Qwen3(-MoE), Mistral, Mixtral, Gemma, Phi, DeepSeek, GLM and similar.

GGUF is not an input format. It is a deploy format: prune_lab produces it as output for llama.cpp/ollama, it does not read it as input.


Interactive mode (arrow-key TUI)

Don't want to remember the flags? Launch it with no arguments (or a command with no model): a navigable menu starts with the arrow keys ↑/↓, Enter to choose, and it asks you for the model and options via prompts.

python pl.py            # full menu: pick the action, then the parameters
python pl.py run        # goes straight to the pipeline and asks model + ratio…
python pl.py chat       # asks for the model and which variant to try

At the end it shows you the equivalent command (▶ Running: python pl.py run …) so you also learn the direct form. If the terminal doesn't support arrow keys, it automatically falls back to a numbered menu.


Quick start (a single command)

From any folder in the repo, a single entrypoint pl.py:

# full pipeline on a dense model: analyze → prune → quantize → evaluate
python pl.py run Qwen/Qwen2.5-0.5B --ratio 0.3

# pipeline on a MoE: also prune the least-used experts
python pl.py run Qwen/Qwen1.5-MoE-A2.7B --ratio 0.25 --keep-experts 0.5

# see what you've already produced
python pl.py list

# try a variant in interactive chat (streaming + token/s)
python pl.py chat Qwen/Qwen2.5-0.5B --variant quant4

pl.py figures out on its own whether the model is a MoE (it skips expert pruning if it isn't) and puts all output in prune_lab/work/ regardless of where you launch it. At the end it prints a weight/perplexity summary table:

====================================================
  SUMMARY  —  Qwen/Qwen2.5-0.5B
====================================================
  Layers:     24 -> 17 (ratio 0.3)
  Pruned weight (fp16): 0.71 GB
  Final Q4 weight:      0.28 GB
  ------------------------------------------------
  PPL base        : 13.201
  PPL pruned      : 18.940  (+43.5%)
  PPL final+Q4    : 19.702  (+49.2%)
====================================================

All commands

python pl.py <command> [arguments]
Command What it does Example
run end-to-end pipeline (analyze→prune→[experts]→quant→eval) python pl.py run Qwen/Qwen3-1.7B --ratio 0.4
list show worked models and ready variants python pl.py list
chat interactive try of a variant python pl.py chat Qwen/Qwen2.5-0.5B --variant quant4
analyze only rank layers + expert usage (step 1) python pl.py analyze <model>
prune only depth pruning (step 2) python pl.py prune <model> --ratio 0.3
experts only MoE expert pruning (step 2b) python pl.py experts <moe> --keep 0.5
surgery shard-level prune without loading the model (200B+, fused experts, multimodal; --stream for the giants) python pl.py surgery <model> --ratio 0 --keep-experts 0.5 --expert-metric weight-norm
quant only 4-bit NF4 quantization (step 3, needs CUDA) python pl.py quant <model> --source pruned
export quantize + export shard-level GGUF, ZERO-RAM (Llama/Mixtral, Qwen3-MoE, DeepSeek-V3); --keep-experts prunes and quantizes in one shot python pl.py export work/<slug>/surgery --quant q4_0
eval only perplexity (step 4) python pl.py eval work/<slug>/quant4

Main run flags

Flag Default Meaning
--ratio 0.3 0.3 fraction of layers to remove (0.3 = 30%)
--keep-experts 0.5 0.5 fraction of experts kept per block (MoE only)
--keep-edges 2 2 protects the first/last N layers (often critical)
--samples 32 32 calibration samples for the ranking
--skip-base-eval off skips the original model's PPL (faster)
--only analyze,prune all runs only a subset of steps

chat flags

Flag Default Meaning
--variant base base | pruned | pruned_moe | surgery | quant4 | quant4_base
--temp 0.7 0.0 0 = greedy, >0 = sampling
--max-new 200 200 max generated tokens
--prompt "..." one-shot non-interactive (for comparisons/scripts)
--raw off pure completion without chat template (base models)

Commands inside the chat: /reset /temp N /max N /exit.


How it works, step by step

pl.py run chains these 5 steps (also callable individually):

  1. analyze — loads the model in streaming/offload, does a few forward passes on the multi-domain calibration mix (prose + code + math + multilingual, with an offline fallback) and produces:

    • layer_ranking.json — layers sorted by importance (‖f(h)‖/‖h‖, with BI as a reference); includes calib_source and importance_metric
    • expert_ranking.json(MoE only) salience (gate prob mass) and counts for each expert, sorted from least salient
  2. prune — removes the first ratio × n_layer most redundant layers (respecting --keep-edges), stitches the ModuleList back together, updates the config and renumbers the internal indices. Saves to pruned/.

  3. experts (MoE only) — removes the lowest-salience experts per block (gate prob mass, not frequency), fixes up the router weights (and the DeepSeek-style correction bias). Saves to pruned_moe/.

  4. quant — quantizes the pruned model to 4-bit NF4 (double-quant). Saves to quant4/ and measures the on-disk weight.

  5. eval — perplexity with sliding window on wikitext-2 for base, pruned, pruned_moe, quant4 → comparison table in summary.json.


What it produces (the work/ folder)

Each model ends up in prune_lab/work/<slug>/ (where <slug> is the HF id with /__):

work/Qwen__Qwen2.5-0.5B/
├── layer_ranking.json     # layers sorted by redundancy
├── expert_ranking.json    # (MoE) usage of each expert
├── pruned/                # depth-pruned model (safetensors, fp16)
├── pruned_moe/            # (MoE) also expert-pruned
├── quant4/                # ⭐ final 4-bit model
├── quant4_base/           # (opt.) quant of the original, for comparison
└── summary.json           # weight + perplexity before/after

These outputs are large (whole models) and are not versioned: work/ and offload/ are in .gitignore. To free disk at the end keep only quant4/ and delete pruned/ and offload/.


Large models: shard-level surgery

A 200B+ model doesn't fit in RAM even just to be opened. surgery.py prunes by rewriting the safetensors shards directly, tensor by tensor, without ever loading the model: RAM peak ~2 GB regardless of size. The output shards are written as they go (not accumulated in RAM), so even the result of a pruned 235B doesn't have to stay in memory.

# 1) the analysis can also be done on an ALREADY quantized version of the model
#    (the layer ranking is robust to quantization)
python pl.py analyze <model_id> --samples 32 --seqlen 512

# 2) surgery on the BF16 shards: layers + experts in one shot
python pl.py surgery <model_id> --ratio 0.3 --keep-experts 0.5
# output: work/<slug>/surgery/  →  then GGUF Q4 export for deploy

Limited disk (--stream). The BF16 download of a 235B is ~470 GB: it doesn't fit on a consumer disk. With --stream, surgery downloads the input shards one at a time, processes them and deletes them right away (metadata aside): the disk peak is ~one shard + the output, not the whole model.

# process a 235B on ~60 GB free (requires salience from `analyze`)
python pl.py surgery Qwen/Qwen3-235B-A22B --stream --ratio 0 --keep-experts 0.75

Verified: surgery produces output bit-identical to the in-memory path (prune+experts) on both dense and MoE models. It handles separate and packed experts, routers with correction bias, layer/expert renumbering, first_k_dense_replace and multi-shard indexes.

Giant fused-expert MoE (gpt-oss, Qwen3.5-MoE) and multimodal models

Recent MoEs pack all the experts of a block into one 3D tensor [n_experts, …] (instead of a ModuleList of separate modules) and are often multimodal (…ForConditionalGeneration, weights under model.language_model.layers.*). prune_lab handles them:

  • Fused-expert detection: find_moe_blocks recognizes the block from the router (no ModuleList needed) and derives n_experts from the gate; surgery slices the 3D tensor on the kept experts.
  • Shared expert preserved: the always-active expert (mlp.shared_expert.*, Qwen3.5-style) is never pruned.
  • MTP head discarded: the Multi-Token-Prediction stacks (mtp.layers.*) have a layer numbering that collides with the decoders: they are excluded (--keep-mtp to keep them).
  • Nested config: updates num_hidden_layers/num_experts even when they are under text_config/language_model.
  • Hybrid stacks (Qwen3.5 = DeltaNet + attention): depth pruning would break the layer pattern → use --ratio 0 (expert pruning only), the real lever on these models.

Expert selection without a GPU (--expert-metric weight-norm). Salience analysis (step1) requires loading the model (a 122B needs datacenter GPUs). As an alternative, surgery can sort experts by the L2 norm of their weights, read directly from the shards: data-free, runs on CPU+disk, zero forward passes. A weaker proxy than salience (it doesn't look at real routing) but it makes expert pruning feasible on cheap hardware.

# Qwen3.5-122B-A10B: expert pruning only 256->128, data-free, without loading anything
python pl.py surgery Qwen/Qwen3.5-122B-A10B --ratio 0 --keep-experts 0.5 \
       --expert-metric weight-norm

On big hardware (vast.ai) for salience analysis raise the memory limits with the env vars PRUNE_MAX_GPU_GB / PRUNE_MAX_CPU_GB (defaults tuned for the local 4 GB VRAM / 16 GB RAM).


Honest hardware arithmetic

On 16 GB RAM / 4 GB VRAM a 220B model does not fit even pruned + Q4 (220B → ~110 GB in Q4; with 40% prune ~66 GB). The real routes:

  1. Exploit the MoE. Of a 220B only ~15–30B parameters are active per token. With expert pruning (this repo) + inference with experts on RAM/NVMe (ktransformers, llama.cpp --override-tensor) the model "runs" at ~1–5 s/token.
  2. More RAM. With 64–96 GB an expert-pruned + Q4 220B becomes usable.
  3. On current hardware, a realistic target is 20–40B (e.g. a MoE like Qwen3-30B-A3B) pruned + Q4.

Disk note. The BF16 download alone of a 220B is ~440 GB. Process by shards (surgery --stream: download/delete one shard at a time) or do the analysis on an already-quantized version.


Measured results

Numbers measured on this consumer machine (GTX 1650 4 GB VRAM, 16 GB RAM), not estimated. Raw data and reproducibility in prune_lab/evidence/.

End-to-end pipeline on a real MoEallenai/OLMoE-1B-7B (64 experts, top-8): expert-pruning by salience (keep 75%) + Q4.

before after Δ
Disk 13.7 GB ~3.0 GB −78 %
Perplexity (wikitext-2) 7.62 9.97 +30.8 %

Both within target (compression −75/80 %, PPL ≤ +25/35 %) on the same model, together.

The pruning curve is steep (measured on OLMoE): the operating point matters.

experts kept Δ PPL
100 % +0 %
75 % (drop 25 %) +25.6 % ✅ within target
50 % (drop 50 %) +277 % ❌ collapse

Fine-grained MoEs (many small experts) collapse beyond ~30% drop without healing: drop ~25%, not 50%. The lever that preserves the "brain" at low cost is quantization (Q4 alone: −72% disk, +11% PPL).

Consumer inference (ollama on a GTX 1650): a 9B Q4 of 5.6 GB that doesn't fit in VRAM (CPU+GPU offload) runs at 8.66 token/s — in/above the 4–7 tok/s target band.

Zero-RAM. bench_zero_ram.py: the private RAM (USS, the OOM metric) stays ~315 MB flat while the model doubles (3.3 → 6.6 GB) → peak independent of size. Verified on real Qwen3-235B weights (real shard): quantizing the embedding (151936, 4096) in row blocks keeps the peak at 674 MB (vs 3.3 GB materializing the whole fp32).

Honesty. The metrics are measured on real models manageable on this hardware (up to ~9–13 GB) and on real 235B weights; running a whole 100–220B+ (280 GB–1.3 TB + GGUF runtime) is not feasible on a consumer PC. The mechanism — and the code that holds up the giants' RAM and disk — is ready for anyone with the bandwidth/disk to launch it on a 235B (see commands above).


Deploy: GGUF export for llama.cpp / ollama

NF4 (quant) is handy and measurable but loads the whole model: on a 220B the offloaded modules stay fp16 and it needs CUDA — not a deploy route. That's why there's pl.py export: quantizer + GGUF export at the shard level, zero-RAM.

# already-pruned model (surgery output) -> GGUF Q4_0 (4-bit, runs without a GPU)
python pl.py export work/<slug>/surgery --quant q4_0

# or PRUNE (by salience) AND quantize in ONE pass, without ever writing the
# intermediate BF16 (essential on the giants: the pruned BF16 wouldn't fit on
# disk). Requires expert_ranking.json from `analyze`:
python pl.py export <local_folder> --quant q4_0 --keep-experts 0.75

ollama create my-model -f Modelfile   # Modelfile:  FROM ./model-q4_0.gguf

Giant from an HF ID on a small disk (--stream). With --stream the export downloads the shards one at a time, prunes+quantizes them and deletes them: the disk peak is ~one shard + the final GGUF. So a 235B (input ~470 GB, never all on disk) becomes a GGUF Q4 of a few tens of GB in one command, without a BF16 intermediate.

# 1) expert salience (once; on big hardware raise PRUNE_MAX_*)
python pl.py analyze  Qwen/Qwen3-235B-A22B
# 2) stream + prune + quant -> GGUF, in one pass, disk ~one shard + output
python pl.py export   Qwen/Qwen3-235B-A22B --stream --keep-experts 0.75 \
       --quant q4_0 --rank-file work/Qwen__Qwen3-235B-A22B/expert_ranking.json

Verified bit-identical to the non-streaming path (same quant, same pruning) on a multi-shard MoE. Only separate experts (Qwen3-MoE/DeepSeek/classic Mixtral); fused-expert MoEs already sit in a single shard → normal export.

One command for the whole validation. validate_giant.py chains everything and prints the 4-pillar table: it reads the base BF16 size from the HF metadata without downloading, runs analyze + stream-export while measuring the private RAM peak (USS) live, and — if ollama is present — measures the tok/s.

python validate_giant.py Qwen/Qwen3-235B-A22B --keep-experts 0.75 --quant q4_0
# -> #1 COMPRESSION: 470.2 GB -> ~xx GB (-yy%)
#    #4 ZERO-RAM:    USS peak <4000 MB (measured during the export)
#    #3 TOK/S:       n tok/s (ollama, on this HW)
#    #2 PPL:         compares pl.py eval <base> vs llama-perplexity <gguf>

Vocab included → runnable GGUF. By default export embeds the tokenizer in the GGUF (official gguf vocab classes: BPE/SentencePiece/HF), so the file is loadable and generatable directly by ollama/llama.cpp with no further steps. The token list is padded with UNUSED tokens up to the embedding rows (as the official convert does). The tokenizer.ggml.pre is set to default (GPT-2 fallback regex): the model loads and generates, but for byte-exact pretokenization matching the original model use the llama.cpp convert. Disable with --no-vocab.

How it stays under 4 GB of RAM. It reads the safetensors tensor by tensor; stacks the MoE experts into GGUF's 3D layout by writing the already-quantized bytes into a memmap on disk (256 experts don't fit in 4 GB → the peak stays one expert); the large 2D tensors (e.g. a giant's embedding) are read and quantized in row blocks directly from the file (get_slice), without ever materializing the whole float32 (bit-identical to quantizing everything at once). Measured: USS ~315 MB flat while the synthetic model doubles (bench_zero_ram.py), and 674 MB quantizing the real Qwen3-235B embedding (151936, 4096) from a real shard.

Verified architectures (tensor names and metadata taken from the official llama.cpp maps, not from assumptions): Llama / Mistral / Mixtral (llama), Qwen3-MoE (qwen3moe, with q_norm/k_norm and mixed dense+MoE layers), DeepSeek-V3 (deepseek2: MLA with kv_b_projk_b/v_b split, shared expert, exp_probs_b, leading_dense, sigmoid gating — ported 1:1 from conversion/deepseek.py). Other architectures → explicit guard, no unverified GGUF.

Q4_0 vs Q4_K_M. The gguf package (pure python) implements the "simple" quants (Q8_0/Q4_0), not the k-quants. --quant q4_0 is true 4-bit, zero-RAM, deployable right away. For the better Q4_K_M (same size, higher quality) export to f16/q8_0 and pass through llama-quantize:

python pl.py export work/<slug>/surgery --quant f16
./llama.cpp/llama-quantize model-f16.gguf model-q4km.gguf Q4_K_M

FAQ / common problems

can't open file '...\test_chat.py' You're using the old path. Now there's the single entrypoint: from any folder in the repo python pl.py chat <model> --variant quant4.

The chat spits random characters / nonsense text. Normal on small models pruned without healing. A 0.5B stripped of 30% of its layers breaks: quality collapses below the threshold (see The three levers). To understand what breaks, compare the variants:

python pl.py chat Qwen/Qwen2.5-0.5B --variant base        --prompt "Hi, who are you?"
python pl.py chat Qwen/Qwen2.5-0.5B --variant quant4_base --prompt "Hi, who are you?"  # quant only, no prune
python pl.py chat Qwen/Qwen2.5-0.5B --variant quant4      --prompt "Hi, who are you?"  # prune + quant

If quant4_base is good but quant4 isn't → it's the pruning's fault, by design. Small models (0.5–1B) only serve as a pipeline smoke-test, not as real targets. The real targets are 20–40B MoEs.

Output even worse on a base (non-instruct) model. Add --raw: base models have no chat template, applying one produces garbage.

fix_mistral_regex / tokenizer warning. Cosmetic, can be ignored (false positive on non-Mistral tokenizers).

NF4 requires CUDA. bitsandbytes quantization wants a GPU. Without one, prune on CPU and export to GGUF Q4 to quantize (above).

Recover quality after aggressive pruning (healing). Beyond ~30% of layers, a light LoRA fine-tuning (a few hours, little data) on the pruned model before quantization recovers most of the loss. It is not training from scratch.


File map

The whole pipeline lives in prune_lab/:

File What it does
pl.py single entrypoint with the subcommands (also ../pl.py from the root)
common.py loading with disk offload, MoE detection, utilities
calib.py multi-domain calibration (prose+code+math+multilingual) + wikitext-2 eval
step1_analyze.py layer importance (‖f(h)‖/‖h‖ + BI) + expert salience (gate prob mass)
step2_prune.py removes the chosen layers, saves the pruned model
step2b_prune_experts.py MoE: removes the least-used experts, fixes the router
surgery.py prune layers+experts on the safetensors shards, without loading the model (--stream for the giants)
step3_quantize.py quantizes to 4-bit NF4, measures the weight (needs CUDA)
export_gguf.py shard-level GGUF export, zero-RAM (prune+quant fused); Llama/Mixtral, Qwen3-MoE, DeepSeek-V3
step4_eval.py perplexity (sliding window)
run_pipeline.py end-to-end orchestrator + summary table
bench_zero_ram.py empirical test of the export's RAM peak (RSS/USS)
validate_giant.py one command that compresses a frontier MoE and prints the 4-pillar table (size, PPL, tok/s, zero-RAM)
test_chat.py interactive chat with streaming and speed stats
evidence/ measured results (JSON) of the 4 metrics

Additional technical detail (shard surgery, MoE edge cases) in prune_lab/README.md.


No model is retrained: prune_lab only does surgery on the existing weights + quantization. Structured pruning = memory that actually drops.

About

Compress large LLMs (MoE 200B+) on consumer hardware: structured depth + expert pruning and 4-bit quantization, with zero-RAM GGUF export for llama.cpp / ollama. Runs on 16GB RAM + 4GB VRAM. No retraining. [BETA]

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages