Skip to content

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

19 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Inkling · MLX

Run Thinking Machines Lab's Inkling — a 975B-total / 41B-active sparse-MoE, natively multimodal model — and its smaller sibling Inkling-Small (276B-total / 12B-active) on Apple Silicon with MLX.

Inkling collection Inkling-Small collection MLX Model Arch License

This repo is a from-scratch MLX port of Inkling's inkling_mm_model architecture (not in stock mlx-lm/mlx-vlm) plus a streaming HF→MLX converter/quantizer that never loads the ~1.9 TB bf16 checkpoint into RAM.

Both family members share one architecture and one code path — the port reads every dimension from config.json, so Inkling-Small needed no code changes at all:

Inkling Inkling-Small
total / active params 975B / 41B 276B / 12B
hidden · layers 6144 · 66 4096 · 42
attn heads (global · SWA-kv) 64 · 16 32 · 8
dense · expert intermediate 24576 · 3072 16384 · 2048
muP logit divisor 24.0 16.0
experts (routed / top-k / shared) 256 / 6 / 2 256 / 6 / 2
> What is the capital of France?
The capital of France is **Paris**.

> Write one sentence about the ocean.
The ocean covers more than seventy percent of Earth's surface and remains
one of the least explored frontiers on our planet.

— generated by the 4-bit build on an M3 Ultra.

✨ Features

🍎 Native Apple Silicon pure MLX; no CUDA, no PyTorch for inference
🧩 Full architecture hybrid local/global attention, per-head QK-norm, relative-position bias, 4 short-convs/layer, sigmoid MoE router (256 experts, top-6 + 2 shared "sink"), muP logits
🖼️ Multimodal towers HMLP vision + dMel audio encoders ported & numerically validated
📦 3 / 4 / 6 / 8-bit + bf16 standard MLX affine group quant; 4-bit runs the full 975B model in ~496 GB, or Inkling-Small in ~148 GB (REAP-25: ~112 GB)
👥 Both family members one code path for Inkling (975B-A41B) and Inkling-Small (276B-A12B) — all dims read from config.json
🌊 Streaming convert quantizes tensor-by-tensor — never holds the 1.9 TB model in memory
Validated fp32 parity vs reference + coherent real generation (see Validation)

🗂️ Which build do I download?

Browse either family as a collection — Inkling · Inkling-Small — or pick from the tables below.

Inkling (975B-A41B):

Build Size Runs on
4-bit ~496 GB 512 GB Mac (M3 Ultra)
6-bit ~717 GB ≥ ~768 GB
8-bit ~937 GB ≥ ~1 TB

The quantized quality is effectively lossless here — even 4-bit reproduces the model's structure with ~100% confidence on the correct next token.

Inkling-Small (276B-A12B) — perplexity measured on one fixed held-out set, so the columns are directly comparable:

Build Size Text ppl vs 8-bit Runs on
8-bit ~280 GB 5.569 ≥ 384 GB
6-bit ~214 GB 5.569 0.0% ≥ 256 GB
4-bit ~148 GB 5.452 −2.1% 192 GB Mac
REAP-25 4-bit ~112 GB 4.992 −8.4% † 128 GB Mac
3-bit ⚠️ ~116 GB 6.706 +20.4% 128 GB Mac (tight)

† No measurable change, not an improvement — see REAP builds below.

No bf16 build is published for Inkling-Small, deliberately. The MLX bf16 conversion is bit-identical to the upstream checkpoint (name-mapping and layout only; the dtype cast is a no-op — verified with scripts/verify_bf16_passthrough.py), so it carries nothing thinkingmachines/Inkling-Small does not already have, and at ~527 GB it does not fit a 512 GB Mac. Regenerate it in ~3 minutes with scripts/convert_all.sh if you want a local requant source.

4-bit is the one to take. Its perplexity sits at the 6/8-bit level — the −2.1% is eval-set noise, not 4-bit genuinely beating higher precision; the honest reading is "no measurable loss down to 4-bit." Going up to 6- or 8-bit buys nothing here.

For a 128 GB Mac, take REAP-25, not 3-bit. They are the same size (112 vs 116 GB) and the gap is not close: REAP-25 costs nothing measurable, while 3-bit costs +20.4% perplexity, loops on repetition and emits stray glyphs after its answers. 3-bit remains published as the fallback for anyone who wants an unpruned expert set at that footprint, but it is marked experimental for good reason.

🧬 REAP-pruned builds (smaller, expert-pruned)

REAP (Cerebras, arXiv:2510.13999) drops the lowest-saliency routed experts per MoE layer, where saliency = mean over active tokens of router_gate × ‖expert_output‖₂. The router renormalizes over survivors; the 2 shared experts, attention, and embeddings are untouched.

Calibrate on text, images and audio. 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 matters: a text-only calibration prunes experts that ground visual features (a Pallas's cat → "brown bear", a golf ball → "butterfly"), and text+image alone leaves audio-grounding experts unprotected (speech transcription overlap fell 0.88 → 0.57 at 25% pruning) — all while text perplexity looked fine. Profiling over every modality keeps each expert that matters to any of them. Inkling routes very uniformly (entropy 0.922; ~0 cold experts/layer under full multimodal calibration), so it is only lightly prunable:

Build Experts kept Size Text ppl vs unpruned Vision (image ID) Audio (speech)
4-bit (unpruned) 256 ~496 GB 3.830
REAP-12 225 ~470 GB 3.806 −0.6% (free) 6/6 0.88
REAP-25 192 ~402 GB 3.946 +3.0% 6/6 0.87
REAP-50 128 ~272 GB 4.682 +22.2% ⚠️ 5/6 0.87

(Text-only calibration scored 2/6 vision; text+image left audio at 0.57 — hence the full multimodal recalibration, which recovers both at no text cost.) REAP-25 is the sweet spot: a real size cut that clears the 512 GB memory cliff (comfortable eager/wired load) for ~3% text perplexity, with vision and audio intact. REAP-50 keeps audio but text and fine-grained vision degrade — experimental only.

Inkling-Small: REAP-25 is free, REAP-50 is not

Same procedure, its own calibration and its own numbers — prunability does not transfer between family members:

Build Experts kept Size Text ppl vs unpruned Vision Audio
4-bit (unpruned) 256 ~148 GB 5.452 6/6 0.874
REAP-25 192 ~112 GB 4.992 −8.4% / +0.55% 6/6 0.896
REAP-50 (not published) 128 ~76 GB 10.238 +88% / +105% 6/6 0.702

† Measured on two independent held-out sets, which disagree in sign. The claim is "no measurable cost at 25% pruning", not that pruning improves the model.

Inkling-Small is less prunable than the 975B model, and the profile says why: with audio in the calibration it has only 0.15 cold experts per layer (vs ~1), because 47.7 experts/layer are >50% audio-driven and 22.3 are >50% image-driven — ~27% of experts serve primarily non-text input. There is very little dead weight to remove. At 25% the saliency ranking still finds it; at 50% it is cutting into live audio-grounding experts, and speech transcription drops 0.874 → 0.702 while perplexity doubles. So REAP-50 was built, measured, and dropped rather than published — unlike on the big model, where +22% was defensible as an experimental build.

Worth noting: saliency retention mispredicted this. Small retains less saliency than the big model at every ratio (87.5% vs 90.3% at 25%), which suggested it would pay more — it paid less at 25% and far more at 50%. Retention is a routing statistic, not a quality metric; the eval is what decides.

Reproduce with scripts/fetch_calib_assets.shscripts/build_calib.pyscripts/profile_experts_mm.pyscripts/prune_build.py <usage.npz>scripts/eval_build.py.

📐 Architecture

tokens ─▶ embed ─▶ embed-norm ─┐
                                ▼
   66 × ┌───────────────────────────────────────────────┐
        │ pre-norm ▶ attention(hybrid SWA/global, QK-norm,│
        │            rel-pos bias, +short-conv) ▶ +resid  │
        │ pre-norm ▶ MoE(256 experts, top-6 + 2 shared)   │
        │            or dense SwiGLU (layers 0–1) ▶ +resid │
        └───────────────────────────────────────────────┘
                                ▼
                 final-norm ▶ ÷muP ▶ unembed
   (image/audio features scatter into the token stream before the stack)

🧬 What makes Inkling architecturally unusual

inkling_mm_model is not a Llama variant with extra towers. Eight mechanisms needed implementing from scratch, and several are easy to get subtly wrong.

Fused gate+up weights are stored ROW-INTERLEAVED

The single most important detail in this repo. Every fused gate+up MLP weight — mlp.w13_dn (dense), experts.w13_weight, shared_experts.shared_w13_weight — is stored interleaved row-wise: [g0, u0, g1, u1, …]. De-interleave as gate = rows 0::2, up = rows 1::2.

The obvious contiguous split, [:half] / [half:], silently scrambles gate↔up in every layer and produces confident, fluent, wrong output — identical in bf16 and fp32, so precision experiments don't reveal it. The transformers PR skeleton has the same omission, so single-layer parity against it passes while both are wrong. The authoritative loader is SGLang's deinterleave_w13.

Hybrid attention — most layers never see the whole sequence

sliding_window_size: 512, local_layer_ids: [0,1,2,3,4, 6,7,8,9,10, 12, …]

Layers alternate in blocks of six: five sliding-window layers over 512 tokens, then one global. 55 of 66 layers are local in Inkling (35 of 42 in Small). The two paths have independent head counts — Inkling runs 64 global heads over 8 KV heads but 64 SWA heads over 16, so the GQA ratio differs by path.

Global layers additionally apply log-scaling for long context: queries and the position bias are multiplied by 1 + 0.1·log(max(n/128000, 1)), a no-op below 128k tokens and a gentle temperature correction above it. Sliding-window layers skip it — their span is fixed at 512, so there is nothing to correct for.

Four short convolutions per layer

sconv_kernel_size: 4

Depthwise causal convolutions of width 4 sit on k, v, the attention input, and the MLP input — four per layer, 264 in the full model. They give each position a small learned mixing window before attention or the MLP sees it. Checkpoint layout is [C, 1, K]; MLX's conv1d wants [C, K, 1].

They are also stateful during decode: each needs a rolling 3-token history, so an incremental cache has to carry conv state alongside the KV cache.

Per-head QK normalization changes the attention scale

Queries and keys are RMS-normalized per head before the dot product. Because both operands then have unit RMS, the usual 1/√d softmax scale is wrong — Inkling uses 1/d. Using 1/√d gives an over-sharp distribution that still produces plausible text, which is why this is worth stating explicitly.

Relative-position bias conditioned on the hidden state

d_rel: 16, rel_extent: 1024

Not a fixed ALiBi-style slope. Each layer holds a bank of bias-vs-distance profiles [d_rel, rel_extent], and every query projects its own d_rel-dim vector that mixes those profiles into a bias value per backward distance. The bias is zero outside 0 ≤ distance < rel_extent, and sliding-window layers use rel_extent = 512 to match their span.

Sigmoid router with two shared "sink" experts

n_routed_experts: 256, num_experts_per_tok: 6, n_shared_experts: 2, shared_expert_sink: true, route_scale: 8.0, norm_after_topk: true

The router scores with a sigmoid, not a softmax, plus a per-expert correction bias and a global scale. Top-6 of 256 are selected, weights are renormalized after the top-k (norm_after_topk), then scaled by 8. Two shared experts run for every token — and because shared_expert_sink is set, their gate rows live in the same router matrix as the routed experts, appended after the 256. Any code that subsets the router (pruning) must keep those trailing rows.

Layers 0–1 are dense SwiGLU (dense_mlp_idx: 2); everything above is MoE.

muP logits and a padded vocabulary

The final hidden state is divided by logits_mup_width_multiplier (24.0 for Inkling, 16.0 for Small) before the unembed — a muP-style width correction that keeps logit scale stable across model sizes. The embedding matrix is padded to 201,024 rows while only 200,058 are real (unpadded_vocab_size); logits are truncated to the real vocab, and sampling the padding would otherwise be possible.

Multimodal towers scatter into the token stream

Vision is HMLP — no ViT. Images are cut into 40px patches (temporal patch size 2 for video), then run through four linear+norm layers straight into the text hidden size. Audio is dMel: a log-mel spectrogram is discretized into 16 bins per mel channel across 80 channels (range −7 to 2), embedded, and normalized.

Both produce soft tokens that are scattered into the token sequence at placeholder positions before the decoder stack runs — so the decoder sees one flat stream and never knows which entries came from pixels or audio. One 40px image patch costs one token, so image resolution translates directly into prompt length.

🛠️ Prerequisites

  • macOS on Apple Silicon with enough unified memory for the build above
  • Python 3.10+
git clone https://github.com/PipeNetwork/inkling-mlx.git
cd inkling-mlx
pip install -r requirements.txt        # mlx, mlx-lm, transformers, numpy, scipy

For near-capacity builds, raise the GPU wired-memory limit once per boot:

sudo sysctl iogpu.wired_limit_mb=500000   # ~524 GB; adjust to your RAM

🚀 Quick start

from inkling_mlx.load import load
from inkling_mlx.generate import greedy_generate
from transformers import AutoTokenizer

model, config = load("/path/to/Inkling-MLX-4bit")
tok = AutoTokenizer.from_pretrained("/path/to/Inkling-MLX-4bit", trust_remote_code=True)
tok.chat_template = open("/path/to/Inkling-MLX-4bit/chat_template.jinja").read()

ids = tok.apply_chat_template(
    [{"role": "user", "content": "What is the capital of France?"}],
    add_generation_prompt=True, reasoning_effort="none", tokenize=True)["input_ids"]
print(tok.decode(greedy_generate(model, config, ids, max_new_tokens=64)))

Inkling is a reasoning model: its chat template injects a Thinking effort level system message. Use reasoning_effort{none, minimal, low, medium, high, max}.

Images & audio

InklingProcessor handles the full preprocessing (image patchify + CLIP-normalize; audio log-mel → dMel bins) and inserts the placeholder soft-tokens:

from inkling_mlx.processing import InklingProcessor
from PIL import Image

proc = InklingProcessor(tok, open("/path/to/model/chat_template.jinja").read())
inputs = proc.apply([{"role": "user", "content": [
    {"type": "image", "image": Image.open("cat.jpg")},
    {"type": "text",  "text": "What's in this image?"},
    # or {"type": "audio", "audio": <16kHz mono np.ndarray>, "sampling_rate": 16000}
]}])
out = greedy_generate(model, config, inputs["input_ids"], max_new_tokens=128,
                      pixel_values=inputs.get("pixel_values"),
                      audio_input_ids=inputs.get("audio_input_ids"))
print(tok.decode(out[len(inputs["input_ids"]):]))

Speech transcription is verbatim — two held-out LibriSpeech clips through the 4-bit REAP-25 build ("Transcribe the speech in this audio."), each returned in ~7 s:

Reference Model output
…WHEN DEATH LIKE SOME REMORSELESS CREDITOR SEIZES ON ALL WE FONDLY THOUGHT OUR OWN… "…when death, like some remorseless creditor, seizes on all we fondly thought our own…"
…ACCORDING TO SALVIAN AND HIS CONTEMPORARIES THE VANDAL CONQUERORS WORKED IN NORTH AFRICA… "…according to Salvian and his contemporaries, the Vandal conquerors worked in North Africa…"

Both scored 1.00 content-word overlap — including obscure proper nouns ("Salvian") and archaic phrasing — with punctuation and casing added. (This fidelity is because the pruning was calibrated on audio; the text+image-only build paraphrased and dropped proper nouns — see below.)

⚡ Performance

  • Load mode. load() eagerly materializes weights wired-resident by default, so forwards don't re-read the mmap. For a model near your RAM ceiling (e.g. the 496 GB 4-bit build on a 512 GB Mac), the eager copy may not fit — pass --lazy (CLI) / load(path, lazy=True) to mmap instead (lower peak RAM, but the first forward pays the ~weight-read and decode can thrash near capacity).

  • Attention uses the fused scaled_dot_product_attention kernel.

  • Multimodal prefill scales with image patches (one 40 px patch = one soft-token). Big images become long prompts; cap resolution with max_long_edge to trade a little detail for a much shorter prefill:

    proc.apply(messages, max_long_edge=512)   # ~130 patches instead of ~450 for a 960px image

🔄 Converting / quantizing yourself

# one build
python -m inkling_mlx.convert_cli --src /path/Inkling-src --dst out-4bit --bits 4
# standard sweep (4/6/8-bit + bf16), smallest first
scripts/convert_all.sh /path/Inkling-src       /path/out  Inkling
scripts/convert_all.sh /path/Inkling-Small-src /path/out  Inkling-Small

# REAP-pruned build: profile expert saliency, then convert with --prune
scripts/fetch_calib_assets.sh /path/out                   # imagenette + LibriSpeech
python scripts/build_calib.py                             # -> calib_wide.json
python scripts/profile_experts_mm.py /path/out-4bit       # -> expert_usage_mm.npz
python scripts/prune_experts.py 0.25                      # -> keep_indices.npz (25% prune)
python -m inkling_mlx.convert_cli --src /path/Inkling-src --dst out-reap25 --bits 4 \
       --prune /path/keep_indices.npz
# (or prune an already-quantized build in one pass, bit-identically: scripts/prune_build.py)

The script suite is shared between both family members. Point it at a sweep with:

export INKLING_OUT=/path/inkling-small-out       # build root + calibration assets
export INKLING_SRC=/path/Inkling-Small-src       # HF source checkpoint
export INKLING_PREFIX=Inkling-Small              # build-dir name prefix

Defaults reproduce the original 975B paths, so existing invocations are unchanged.

Quantized: attention/MLP/expert projections, embeddings, vision/audio matmuls. Kept high-precision: the MoE router, RMSNorms, the four short-convs, and the relative-position bias. Conversion is streaming, so it runs in bounded memory regardless of model size.

Why affine int4 and not MXFP4 / NVFP4?

MLX also supports floating-point 4-bit modes (mx.quantize(mode="mxfp4"|"nvfp4")), and Thinking Machines ships an Inkling-NVFP4 checkpoint — so it's a fair question. 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's asymmetric (per-group scale and zero-point, 16 uniform levels), so it 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 (MLX would dequantize FP4 anyway, as there's no FP4 compute). So these builds use affine int4; a Mac port of the NVFP4 checkpoint would be lower quality at best-equal size.

✅ Validation

Check Result
fp32 decoder-layer parity vs reference max |Δ| ~4e-3 on real weights, both attention types + full MoE
fp32 parity at Inkling-Small's config shape rel max Δ 2.3e-06, argmax agreement 100%
Vision + audio towers vs reference ~1e-5
InklingProcessor preprocessing ~1e-7 (patchify/normalize, log-mel→dMel)
bf16 build vs source checkpoint bit-identical, 87/87 sampled tensors (verify_bf16_passthrough.py)
Real generation coherent and correct at 3/4/6/8-bit (above)
Held-out image ID 6/6 unpruned and REAP-12/25 (multimodal calibration)
Held-out speech transcription 0.87–0.90 word overlap
python tests/parity.py                # decoder math vs a standalone torch reference
python scripts/smoke_test.py <build>  # real weights, generation + speed + peak memory
python scripts/eval_build.py <build>  # audio + image ID + text ppl in one load

Two bugs worth knowing about

Both were in the tooling, both silent, and both would have shipped wrong numbers:

  • eval_build.py scored absent image probes as failures. It skipped probes whose file was missing but still divided by the full probe count — a build scoring 5/5 would be published as "5/6". Now scores against the number actually evaluated and names what it skipped.
  • build_calib.py silently shrank. It listed specific project paths for its code corpus; when those projects moved, the corpus quietly got smaller rather than erroring, weakening every calibration built from it. Now scans roots via INKLING_CODE_ROOTS.

📁 Project structure

inkling-mlx/
├── inkling_mlx/            # the MLX model + converter/loader/generation
│   ├── model.py text.py layers.py attention.py moe.py
│   ├── vision.py audio.py common.py cache.py config.py
│   └── convert.py load.py generate.py convert_cli.py
├── scripts/               # convert_all.sh, verify_shards.py, smoke_test.py
├── tests/                 # parity / cache / driver / single-layer reference checks
└── requirements.txt

🐛 Troubleshooting

Problem Fix
Incoherent / multilingual output ensure you loaded the chat_template.jinja and used apply_chat_template
OOM while loading raise iogpu.wired_limit_mb; use a smaller build; close other apps
No module named inkling_mlx.models... this arch isn't in stock mlx-lm — load via this repo's load()
Very slow first token MLX pages weights from disk on first use; subsequent tokens are faster

📊 What we learned measuring these builds

Quant tolerance does not transfer between models. On Inkling-Small, 4-bit shows no measurable loss against 8-bit. On DeepSeek-V4 — same pipeline, same affine int4 — 4-bit costs a significant 4.4%. Measure per model; do not carry a bit-width recommendation across architectures.

Saliency retention is a poor predictor of pruning damage. It is monotonic with damage but wildly non-linear, and it has now mispredicted in both directions:

Model Prune Saliency retained Actual cost
Inkling-Small 25% 87.5% none measurable
Inkling-Small 50% 71.3% +88% ppl, audio 0.874 → 0.702 (dropped)
Inkling (975B) 25% 90.3% +3.0%
DeepSeek-V4 25% 96.2% +4.5%

Inkling-Small retains less saliency than DeepSeek-V4 at every ratio yet tolerates 25% pruning for free, while V4 does not. Treat the curve as a screen for whether pruning is viable, then measure.

Text perplexity cannot see multimodal damage. This is the single most expensive lesson here: text-only calibration produced builds whose perplexity looked fine while image ID collapsed to 2/6 and speech overlap fell to 0.57. If the model has non-text modalities, both the calibration and the eval must include them.

Small evals give wrong answers with confidence. Related work on DeepSeek-V4 found a 275-token eval reporting the opposite build ranking from a 200k-token one. Use enough held-out text, score every build on identical windows, and compare with a paired bootstrap rather than overlapping independent intervals.

📚 Resources

📝 License

Code: Apache-2.0. Model weights inherit the Apache-2.0 license of the base model.

⚠️ Status

Text and multimodal generation are complete and validated: the vision/audio towers and their preprocessing (InklingProcessor — image patchify/normalize, audio log-mel→dMel, validated to ~1e-7 vs the reference) are included.

About

Run Thinking Machines Lab's Inkling (975B-A41B MoE, multimodal) on Apple Silicon with MLX — 4/6/8-bit.

Topics

Resources

Stars

3 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages