usterquant's picture
Duplicate from pipenetwork/Inkling-MLX-REAP12-4bit
e1e5908
|
Raw
History Blame
5.94 kB
---
license: apache-2.0
base_model: thinkingmachines/Inkling
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-MLX-REAP12-4bit
**Built with Inkling (Thinking Machines Lab).**
A **REAP-pruned**, 4-bit MLX build of
[thinkingmachines/Inkling](https://huggingface.co/thinkingmachines/Inkling):
each MoE layer keeps its **225 highest-saliency routed experts** (of 256), a
**12% expert prune**. Free lunch β€” text, vision and audio intact.
**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.922;
only ~1 cold expert per layer under multimodal calibration), so it is only *lightly*
prunable β€” reflected below.
## Calibrated on text, images **and audio** (this matters)
Inkling 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. This is deliberate: a **text-only** calibration prunes
experts that ground *visual* features (a Pallas's cat β†’ *"brown bear"*, a golf ball β†’
*"butterfly"*); adding only text+image then leaves *audio*-grounding experts unprotected
(speech transcription word-overlap fell from 0.88 to 0.57 at 25% pruning) β€” all while
text perplexity looked fine the whole time. Profiling over all three modalities keeps
every expert that matters to any of them. On held-out tests this build scores **vision
6/6** (vs 2/6 text-only) and **audio 0.88** overlap (vs 0.57 text+image), at no extra
text cost.
## Measured quality (4-bit)
| Build | Experts kept | Size | Text ppl | vs unpruned | Vision (image ID) | Audio (speech overlap) |
|---|---:|---:|---:|---:|---:|---:|
| [Inkling-MLX-4bit](https://huggingface.co/pipenetwork/Inkling-MLX-4bit) (unpruned) | 256 | ~490 GB | 3.830 | β€” | βœ“ | βœ“ |
| [Inkling-MLX-REAP12-4bit](https://huggingface.co/pipenetwork/Inkling-MLX-REAP12-4bit) | 225 | ~470 GB | 3.806 | -0.6% | 6/6 | 0.88 |
| [Inkling-MLX-REAP25-4bit](https://huggingface.co/pipenetwork/Inkling-MLX-REAP25-4bit) | 192 | ~402 GB | 3.946 | +3.0% | 6/6 | 0.87 |
| [Inkling-MLX-REAP50-4bit](https://huggingface.co/pipenetwork/Inkling-MLX-REAP50-4bit) | 128 | ~272 GB | 4.682 | +22.2% | 5/6 | 0.87 |
This build: **text perplexity 3.806 (-0.6% vs the unpruned 4-bit)**, **vision 6/6**
(held-out image ID), **audio 0.88** (held-out speech transcription word-overlap),
96.2% of router-weighted expert contribution retained. 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 = 225`) 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
(vs the 496 GB unpruned 4-bit) is the practical point: **~470 GB** loads eager/wired-resident
on a 512 GB machine without the memory-ceiling thrash.
## 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).