BitSmith

INT4 LLM Quantization + Structured Pruning — Research + C++ Inference Engine

Three-phase project: Phase 1 researches post-training quantization methods and measures their accuracy cost. Phase 2 ships a real inference engine — INT4 packed weights, ARM NEON SIMD kernels, no PyTorch at runtime. Phase 3 adds graph-theory-guided structured pruning — physically removing attention heads and MLP neurons to give real FLOP reduction.

---

Key Results (Phase 1)

Method	Bits	Size	Perplexity	vs FP16
FP16 baseline	16	248.9 MB	35.05	—
Uniform PTQ	8	166.6 MB	35.36	+0.9%
AWQ	8	166.6 MB	35.13	+0.3%
Uniform PTQ	4	124.1 MB	1,978	+5,544%
AWQ	4	124.1 MB	188	+437%
Uniform PTQ	2	102.9 MB	~5×10¹⁷	collapse
AWQ	2	102.9 MB	~98 M	still bad

AWQ is ~10× more accurate than Uniform PTQ at INT4 — activation-aware scaling protects the weights that matter most.

Model: GPT-2 Small (124 M params) · Dataset: WikiText-2 · Group size: 128

---

Key Results (Phase 3 — Structured Pruning)

Head pruning (global importance ranking, PageRank bottleneck guard):

Config	Heads Pruned	Size	Perplexity	vs FP16
Head 10%	2/144	250.7 MB	73.7	+126%
Head 20%	9/144	247.9 MB	81.3	+149%
Head 30%	20/144	243.5 MB	105.5	+223%
Head 50%	44/144	233.9 MB	276.6	+746%

MLP neuron pruning (per-layer weight-product importance: ‖c_fc_col‖ × ‖c_proj_row‖):

Config	Neurons Pruned	Size	Perplexity	vs FP16
Neuron 10%	3,684/36,864	240.0 MB	56.2	+72%
Neuron 20%	7,368/36,864	228.5 MB	92.9	+184%
Neuron 30%	11,064/36,864	217.0 MB	173.6	+431%
Neuron 50%	18,432/36,864	194.0 MB	755.3	+2,212%

Note on combined pruning: pruning BOTH heads and neurons simultaneously at >10% per layer compounds errors within each layer too aggressively for GPT-2 Small (which has limited redundancy). For meaningful size reduction at acceptable quality, combine pruning + INT4: .venv/bin/python3 scripts/run_pruning_benchmark.py --head-ratios 0.3 --neuron-ratios 0 --quantize

The unique twist: unlike unstructured pruning (zeros scattered in the matrix), structured pruning physically removes rows and columns — the matrix is smaller, the FLOP count is lower, and every clock cycle saved is real.

---

What's Inside

BitSmith/
├── src/
│   ├── quantization/
│   │   ├── base.py          # Per-group asymmetric quant math + Conv1D support
│   │   │                    # quantize_to_packed() → INT4 nibble bytes for C++ engine
│   │   ├── uniform.py       # Standard PTQ — round-to-nearest on every linear layer
│   │   └── awq.py           # AWQ — calibration hooks, activation stats, α grid search
│   ├── pruning/             # ── Phase 3: Structured Pruning ──
│   │   ├── importance.py    # Calibration hooks → head/neuron importance scores
│   │   ├── graph.py         # NetworkX DiGraph + PageRank-based pruning decisions
│   │   └── structured.py    # Physical weight removal (resize rows/columns in-place)
│   └── benchmark.py         # Orchestrator: loads model/data, runs all combos, saves JSON
├── csrc/                    # ── Phase 2: C++ inference engine ──
│   ├── ops.hpp              # Op declarations
│   ├── ops_neon.cpp         # ARM NEON INT4 kernel + NEON layernorm/f32_matvec
│   ├── ops_scalar.cpp       # Scalar fallback (non-ARM platforms)
│   ├── weights_io.hpp/cpp   # Binary .weights format (BTSM magic, 52-byte header)
│   ├── gpt2.hpp/cpp         # Full GPT-2 prefill forward pass in C++
│   └── bindings.cpp         # pybind11 — exposes BitSmithEngine to Python
├── python/
│   ├── export_weights.py    # AWQ-quantize → pack INT4 nibbles → write .weights file
│   └── bitsmith_engine.py   # Python wrapper: load(), forward(), generate()
├── dashboard/
│   └── app.py               # Streamlit + Plotly dashboard (Phases 1–3)
├── scripts/
│   ├── run_benchmark.py          # Phase 1 CLI entry point
│   ├── run_engine_demo.py        # Phase 2: export weights → load C++ engine → generate text
│   └── run_pruning_benchmark.py  # Phase 3: sweep pruning ratios, save pruning_results.json
├── CMakeLists.txt           # Builds bitsmith_cpp.so (pybind11, NEON detection)
├── results/                 # benchmark_results.json + pruning_results.json
├── Makefile
└── requirements.txt

---

Quick Start — Phase 1 (Python Benchmark)

# 1. Clone and enter the repo
git clone <repo-url>
cd BitSmith

# 2. Create and populate the virtual environment
python3 -m venv .venv
make install          # pip install -r requirements.txt into .venv

# 3. Run the full benchmark (~10 min on Apple MPS / ~30 min on CPU)
make benchmark

# Optional: smoke-test with just INT8 uniform first (~2 min)
make benchmark-fast

# 4. Open the dashboard
make dashboard        # → http://localhost:8501

Benchmark CLI options

python scripts/run_benchmark.py --help

  --bits    N [N ...]     bit-widths to test  (default: 8 4 2)
  --methods {uniform,awq} methods to run      (default: uniform awq)
  --output  DIR           output directory    (default: results/)
  --device  {auto,cpu,cuda,mps}               (default: auto)
  --quiet                 suppress progress bars

---

Quick Start — Phase 2 (C++ Inference Engine)

Requirements: Apple Silicon Mac (or any ARM64 Linux). CMake is installed automatically via pip.

# Build the C++ extension (one-time, ~30 s)
make engine

# Export AWQ INT4 weights + run greedy generation demo
# (~5 min first run: downloads WikiText-2, runs AWQ calibration)
make engine-demo

Or step by step:

# Export weights only (saves model.weights, ~124 MB)
.venv/bin/python3 python/export_weights.py

# Run demo with a custom prompt
.venv/bin/python3 scripts/run_engine_demo.py \
    --prompt "The future of AI is" \
    --max-new-tokens 80

Expected output:

[demo] Device       : NEON
[demo] Architecture : 12 layers  12 heads  d_model=768
[demo] Quantization : INT4  group_size=128
──────────────────────────────────────────────────────────
The future of AI is not to replace humans but to augment
them. The goal is to make humans more capable...
──────────────────────────────────────────────────────────
[demo] 76 new tokens in 13.2s  (5.8 tok/s)  device=NEON

---

Quick Start — Phase 3 (Structured Pruning)

# Quick smoke-test: 30% head pruning, 32 calib seqs (~5 min on CPU)
make prune-fast

# Full sweep: head-only (10/20/30/50%) + neuron-only (10/20/30/50%) + combined 30%+INT4
# Runs three passes with --append, writes results/pruning_results.json (~30 min on CPU)
make prune-benchmark

# Custom: 30% heads + INT4 AWQ after pruning
.venv/bin/python3 scripts/run_pruning_benchmark.py \
  --head-ratios 0.3 --neuron-ratios 0 --quantize

# View results in the dashboard (Phase 3 panel appears automatically)
make dashboard

Pruning CLI options

python scripts/run_pruning_benchmark.py --help

  --head-ratios   R [R ...]  head pruning fractions    (default: 0.3 0.5 0.7)
  --neuron-ratios R [R ...]  neuron pruning fractions  (default: 0.3 0.5 0.7)
  --quantize                 apply AWQ INT4 after pruning
  --n-calib       N          calibration sequences     (default: 100)
  --n-eval        N          evaluation sequences      (default: 256)
  --output        DIR        output directory          (default: results/)
  --device        {auto,cpu,cuda,mps}
  --quiet                    suppress progress bars
  --append                   append to existing pruning_results.json

---

How It Works

Phase 1 — Quantization Methods

Uniform PTQ

Standard post-training quantization. For every linear layer weight matrix:

1. Split the weight rows into groups of 128 elements (per-group quantization)
2. For each group compute: scale = (max − min) / (2ᵇ − 1), zero_point = round(−min / scale)
3. W_q = clamp(round(W / scale + zp), 0, 2ᵇ−1) → dequantize back to float for inference

No calibration data needed. Fast but accuracy collapses below INT8.

AWQ (Activation-aware Weight Quantization)

Lin et al., 2023 — arxiv.org/abs/2306.00978

Not all weights are equally important. Weights that multiply large activations contribute disproportionately to the output — protecting them is worth the extra effort.

Algorithm:

1. Calibration — run 100 samples from WikiText-2 through the model with forward hooks to record x_max[c] = mean |activation[:, c]| per input channel.

2. Scale search — for each linear layer, grid-search α ∈ {0.1, 0.3, 0.5, 0.7, 0.9}:

   scales  = x_max ^ α / mean(x_max ^ α)   # per input-channel scaling vector
   W_sc    = W × scales                      # scale columns up before quant
   W_q_sc  = dequant(quant(W_sc))
   W_q     = W_q_sc / scales                 # un-scale after quant
   error   = ‖W_q − W‖_F
   

   Pick the α that minimises the Frobenius error.

3. Apply — replace each layer's weight with W_q for the best α.

---

Phase 2 — C++ Inference Engine

Binary .weights Format

Header (52 bytes, little-endian):
  magic[4]       "BTSM"
  version[4]     uint32 = 1
  arch[16]       "gpt2\0..." (null-padded)
  n_layers[4]    uint32
  n_heads[4]     uint32
  d_model[4]     uint32
  vocab_size[4]  uint32
  max_seq_len[4] uint32
  bits[4]        uint32
  group_size[4]  uint32

Tensors (repeated):
  name_len[4]    uint32
  name[name_len] char
  dtype[1]       uint8  (0 = float32, 1 = packed_int4)
  ndim[4]        uint32
  shape[ndim×4]  uint32 each
  nbytes[8]      uint64
  data[nbytes]   raw bytes

INT4 Nibble Packing

packed_byte[k] = w[2k] | (w[2k+1] << 4)
lo nibble (bits 3:0) = even-indexed weight
hi nibble (bits 7:4) = odd-indexed weight

ARM NEON INT4 Kernel (the centrepiece)

Processes 16 INT4 values (8 bytes) per NEON iteration with 4 parallel accumulator registers to hide FP multiply-accumulate latency (~4 cycles on Apple M-series):

vld1_u8 (8 bytes)         → packed uint8x8_t
vand_u8(packed, 0x0F)     → lo nibbles [w0,w2,w4,...,w14]
vshr_n_u8(packed, 4)      → hi nibbles [w1,w3,w5,...,w15]
vmovl_u8 → vmovl_u16 → vcvtq_f32_u32   (widen to float)
(wf - vzero) × vscale     (dequantize)
vuzp1q_f32 / vuzp2q_f32   (deinterleave x into even/odd)
vmlaq_f32 × 4 accumulators (fused multiply-accumulate)
vaddvq_f32                 (horizontal sum)

#ifdef BITSMITH_NEON guards all intrinsics; ops_scalar.cpp provides a portable fallback for non-ARM platforms.

GPT-2 Forward Pass

Full prefill mode: token sequence → logits at last position.

Per transformer block:

1. LayerNorm → c_attn (QKV) via int4_matvec
2. Split Q, K, V → scaled dot-product attention with causal mask
3. c_proj via int4_matvec + residual
4. LayerNorm → mlp_fc via int4_matvec → GELU
5. mlp_proj via int4_matvec + residual

Final: LayerNorm + weight-tied LM head (wte transpose multiply) → logits[vocab_size].

---

Phase 3 — Structured Pruning + Graph Theory

Importance Scoring

• Head importance — forward hook on c_attn output: extract the V segment [seq, d_model], reshape to [seq, n_heads, head_dim], compute mean |V[:, h, :]| per head across all calibration samples. A head with near-zero V output contributes nothing to the residual stream regardless of its Q/K patterns.

• Neuron importance — calibration-free weight-product metric:

  importance[k] = ‖c_fc_weight[:, k]‖₂ × ‖c_proj_weight[k, :]‖₂
  

  This captures the actual output reconstruction cost of neuron k — a neuron's impact is proportional to both how strongly it is activated (large c_fc column norm) and how much its output affects the residual stream (large c_proj row norm). Mean GELU activation magnitude was evaluated but performs worse than random pruning on GPT-2 Small: context-specialized neurons that fire rarely on calibration data carry large c_proj weights and are disproportionately important.

Graph Construction (NetworkX)

The transformer is modelled as a compact nx.DiGraph (158 nodes, ~300 edges):

• Nodes: embed (source), head_Li_Hj (one per attention head), mlp_Li (one per MLP block, importance = mean neuron score), output (sink)
• Edges: embed → heads[0], heads[i] → mlp[i], mlp[i] → heads[i+1], last layer → output
• Edge weights = source node importance score

PageRank (pure-Python power iteration, no scipy) identifies structural bottlenecks — nodes many information paths flow through, even if their raw activation magnitude is low. These are protected from pruning.

Heads are ranked globally by importance score (with a per-layer cap so no single layer loses more than head_ratio of its heads). Neurons are ranked per-layer — early-layer neurons have systematically lower absolute weight norms than later layers regardless of their functional importance, so global neuron ranking would concentrate all pruning in the first few layers.

Physical Pruning

Pruning physically resizes weight tensors (no zeros left in place):

MLP neuron pruning (clean — no forward-pass changes):

c_fc  Conv1D [d_model, 4*d_model] → [d_model, n_keep]
c_proj Conv1D [4*d_model, d_model] → [n_keep, d_model]

Attention head pruning (patches num_heads + split_size):

c_attn Conv1D [d_model, 3*d_model] → [d_model, 3*n_keep*head_dim]
c_proj Conv1D [d_model, d_model]   → [n_keep*head_dim, d_model]
block.attn.num_heads  = n_keep
block.attn.split_size = n_keep * head_dim

GPT2Attention.forward uses split_size and num_heads directly — no HuggingFace internals hacked.

---

Dashboard (Phases 1–3)

Five panels at http://localhost:8501:

Panel	What it shows
Pareto scatter (log scale)	Model Size (MB) on x-axis, Perplexity on y-axis. Bubble size = bit-width. Dotted line = Pareto frontier.
Bar charts	Size reduction % vs FP16; PPL degradation % vs FP16 (log scale).
Results table	All metrics, sortable. Best method per bit-width highlighted in green.
AWQ vs Uniform gap	Bar chart of PPL improvement from AWQ over Uniform at each bit-width.
Phase 3 — Pruning	Combined Pareto chart (pruned vs quantized), perplexity-vs-ratio sweep, full pruning results table. Appears automatically when results/pruning_results.json exists.

---

Implementation Notes

• GPT-2 uses Conv1D (from transformers.pytorch_utils), not nn.Linear. Its weight is stored transposed [in, out]. The get_weight_2d / set_weight_2d helpers in src/quantization/base.py normalise this transparently everywhere.
• Phase 1 uses simulated quantization — weights are stored as dequantized float32 so the standard PyTorch forward pass works. This measures information loss, not inference throughput.
• Phase 2 is real INT4 — the C++ engine packs two 4-bit values per byte and dequantizes on the fly during matrix-vector multiply. No PyTorch dependency at inference time.
• Phase 3 pruning is truly structured — weight tensors are resized (not zeroed). A pruned model is a strictly smaller PyTorch model that the standard forward pass runs on without modification.
• Head pruning patches num_heads, split_size, and c_attn.nf on each GPT2Attention block after resizing the weights. HuggingFace's Conv1D uses .nf (output features) inside its forward() view call — failing to update it causes a shape mismatch at runtime. No HuggingFace source modifications required beyond these attribute patches.
• MPS support (Phases 1 & 3) — device auto-detected; runs on Apple Silicon GPU out of the box.
• Build system — pybind11 is installed via pip into the venv; make engine uses .venv/bin/cmake (universal arm64/x86 binary) with -DCMAKE_OSX_ARCHITECTURES=arm64 to produce a native arm64 extension.

---

Requirements

• Python 3.11+
• PyTorch 2.1+
• NetworkX 3.0+ (Phase 3 graph analysis)
• C++17 compiler (Xcode Command Line Tools on macOS, Phase 2 only)
• See requirements.txt for the full list (pybind11, networkx included)