Inkling-MLX-REAP25-4bit

Built with Inkling (Thinking Machines Lab).

A REAP-pruned, 4-bit MLX build of thinkingmachines/Inkling: each MoE layer keeps its 192 highest-saliency routed experts (of 256), a 25% expert prune. Sweet spot — clears the 512 gb memory cliff.

Code / loader: github.com/PipeNetwork/inkling-mlx

What is REAP pruning?

REAP (Router-weighted Expert Activation Pruning, Cerebras, arXiv: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.87 overlap (vs 0.57 text+image), at no extra text cost.

Measured quality (4-bit)

This build: text perplexity 3.946 (+3.0% vs the unpruned 4-bit), vision 6/6 (held-out image ID), audio 0.87 (held-out speech transcription word-overlap), 90.3% 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 checkpoint — so for the record, we benchmarked round-trip reconstruction error (‖W − Ŵ‖ / ‖W‖ vs bf16) on real Inkling expert weights:

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.

pip install mlx mlx-lm transformers

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: ~402 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).