bielquant's picture
Duplicate from pipenetwork/Inkling-Small-MLX-REAP25-4bit
f098b5e
|
Raw
History Blame
6.56 kB
---
license: apache-2.0
base_model: thinkingmachines/Inkling-Small
base_model_relation: quantized
pipeline_tag: image-text-to-text
library_name: mlx
tags:
- mlx
- moe
- multimodal
- inkling
- thinking-machines
- reap
- pruned
- audio-text-to-text
---
# Inkling-Small-MLX-REAP25-4bit
**Built with Inkling (Thinking Machines Lab).**
A **REAP-pruned**, 4-bit MLX build of
[thinkingmachines/Inkling-Small](https://huggingface.co/thinkingmachines/Inkling-Small):
each MoE layer keeps its **192 highest-saliency routed experts** (of 256), a
**25% expert prune**. No measurable loss — the build to take for a 128 gb mac.
**Code / loader:** [github.com/PipeNetwork/inkling-mlx](https://github.com/PipeNetwork/inkling-mlx)
## What is REAP pruning?
[REAP (Router-weighted Expert Activation Pruning, Cerebras, arXiv:2510.13999)](https://arxiv.org/abs/2510.13999)
ranks each routed expert by **saliency** = mean over the tokens that route to it of
`router_gate_weight × ‖expert_output‖₂` — its actual contribution to the residual
stream. The lowest-saliency experts are dropped; the router simply renormalizes over
the survivors (no weight surgery). The **2 shared "sink" experts, attention, and
embeddings are untouched.** Inkling routes **very uniformly** (routing entropy 0.908; only ~0.15 cold experts per layer once audio is included in the calibration), so it
is only *lightly* prunable — reflected below.
## Calibrated on text, images **and audio** (this matters)
Inkling-Small is multimodal, and expert saliency was profiled over a mixed corpus of
**text (code + 15 languages + reasoning), 200 real images, and 180 speech clips** run
through the full vision and audio paths. On this model that is not optional: **47.7
experts per layer are >50% audio-driven and 22.3 are >50% image-driven** — about 27% of
all experts serve primarily non-text input, and adding audio to the calibration dropped
the cold-expert count from 0.65 to 0.15 per layer. Experts that only ever fire on speech
look worthless to a text-only profiler and get pruned first, which is how the 975B model
lost speech transcription (word-overlap 0.88 → 0.57) while its text perplexity still
looked fine. Profiling all three modalities keeps them. On held-out tests this build
scores **vision 6/6** and **audio 0.896** word-overlap — both at the unpruned build's
level.
## Measured quality (4-bit)
| Build | Experts kept | Size | Text ppl | vs unpruned | Vision (image ID) | Audio (speech overlap) |
|---|---:|---:|---:|---:|---:|---:|
| [Inkling-Small-MLX-4bit](https://huggingface.co/pipenetwork/Inkling-Small-MLX-4bit) (unpruned) | 256 | ~148 GB | 5.452 | — | ✓ | ✓ |
| [Inkling-Small-MLX-REAP25-4bit](https://huggingface.co/pipenetwork/Inkling-Small-MLX-REAP25-4bit) | 192 | ~112 GB | 4.992 | −8.4% | 6/6 | 0.896 |
This build: **text perplexity 4.992 (−8.4% vs the unpruned 4-bit)**, **vision 6/6**
(held-out image ID), **audio 0.896** (held-out speech transcription word-overlap),
87.5% of router-weighted expert contribution retained.
**Read that −8.4% as "no measurable change", not as pruning improving the model.**
Perplexity was measured on two independent held-out sets: this one has REAP-25 at
**−8.4%** vs the unpruned 4-bit, a second one has it at **+0.55%**. Sets that small
disagree by a few percent in either direction, so the defensible claim is that a 25%
expert prune costs nothing measurable here — not that it helps.
**50% pruning is a different story and is deliberately not published.** It was built and
evaluated: text perplexity roughly **doubled** (+88% / +105% on the two sets) and speech
transcription fell from **0.874 to 0.702**, so it was dropped rather than shipped with a
warning. Inkling-Small is meaningfully less prunable than the 975B model, whose REAP-50
cost ~22%. Pruning is applied to the
already-quantized build; because expert subsetting is along the expert axis and
affine-quant groups run along the hidden axis, it is **bit-identical to pruning the bf16
source then requantizing**.
## Quantization scheme: affine int4 (not NVFP4 / MXFP4)
MLX supports FP4 modes and Thinking Machines ships an
[Inkling-NVFP4](https://huggingface.co/thinkingmachines/Inkling-NVFP4) checkpoint — so
for the record, we benchmarked round-trip reconstruction error (‖W − Ŵ‖ / ‖W‖ vs bf16)
on real Inkling expert weights:
| Scheme | bits/weight | reconstruction error |
|---|---:|---:|
| **affine int4** (group 64) | 4.50 | **~9.1%** |
| nvfp4 (group 16) | 4.50 | ~10.2% |
| mxfp4 (group 32) | 4.25 | ~12.3% |
Affine int4 is the most faithful: it is *asymmetric* (per-group scale **and** zero-point,
16 uniform levels), which centers on Inkling's near-Gaussian expert weights better than
symmetric FP4's fixed non-uniform levels (scale only, no zero-point). FP4's real payoff is
heavy-tailed *activations* and native Blackwell FP4 tensor cores — neither helps weight
fidelity on Apple Silicon, where MLX would dequantize FP4 anyway. So these builds use
affine int4; a Mac port of the NVFP4 checkpoint would be *lower* quality at best-equal size.
## ⚠️ Loading requires the bundled `inkling_mlx` loader
The `inkling_mm_model` architecture is **not** in stock `mlx-lm` / `mlx-vlm`, so this
repo bundles a minimal, numerically-validated MLX implementation under `inkling_mlx/`.
The reduced expert count is recorded in `config.json` (`n_routed_experts = 192`) and
the loader builds the model to match automatically.
```bash
pip install mlx mlx-lm transformers
```
```python
from inkling_mlx.load import load
from inkling_mlx.generate import greedy_generate
from transformers import AutoTokenizer
model, config = load("/path/to/this/repo") # eager wired load fits comfortably
tok = AutoTokenizer.from_pretrained("/path/to/this/repo", trust_remote_code=True)
ids = tok("The capital of France is")["input_ids"]
print(tok.decode(greedy_generate(model, config, ids, max_new_tokens=64)))
```
Needs an Apple-Silicon Mac with unified memory ≥ the size above. The smaller footprint is
the practical point: **~112 GB** fits a 128 GB Mac, where the unpruned 148 GB 4-bit build does not.
## Details
- Multimodal (HMLP vision + dMel audio towers + preprocessing) is included, same as the
base MLX build; the multi-token-prediction head is dropped.
- Quantized: attention / MLP / expert projections, embed+unembed, vision/audio matmuls.
Kept higher precision: MoE router, RMSNorms, the four short-convolutions per layer,
relative-position bias.
License: Apache-2.0 (inherits the base model).