IFM/Jais-3-32B-MLX-4bit

Jais-3-32B (MLX, 4-bit)

Mixed-precision 4-bit MLX quantization of IFM/Jais-3-32B, converted from revision 847c6fc. 5.9 bits/weight effective, 25.5 GB on disk. For Apple silicon.

Jais-3-32B is the large dense Jais-3 model: a 32B decoder-only model with a 512K (524,288-token) context window.

Requirements

mlx-lm doesn't support the k2_horizon architecture (used by Jais-3) yet. Until support lands, this repo ships the MLX model code (k2_horizon.py), which mlx-lm loads through the model_file entry in config.json. So the released mlx-lm works as-is:

pip install -U mlx-lm

Pass --trust-remote-code (or trust_remote_code=True). The file runs on your machine, so read it first if you like; its comments describe how it differs from the PyTorch code.

How it was quantized

Mixed precision, same recipe as the dense K2-Horizon MLX models (integer affine, bf16 scales and biases):

• 4-bit, group size 32: the MLP weights (gate_proj, up_proj, down_proj). These hold most of the parameters.
• 8-bit, group size 64: everything else: attention, the embeddings and lm_head.

Plain 4-bit (everything 4-bit, group size 64) loses too much on this model: +6.3% WikiText-2 and +9.9% Arabic Wikipedia perplexity versus 8-bit. This recipe loses +2.1% and +3.1%.

Memory

Peak 25.8 GB for a short prompt. The KV cache grows with the context, so long contexts need more memory; --max-kv-size and KV-cache quantization (--kv-bits 8, where available) reduce it.

Conversion check

Checked against the PyTorch code (modeling_k2_horizon.py) in fp32 on the real weights of this model. The check used the model truncated to its first 4 layers (plus the final norm and lm_head), with the same truncation on both sides, on English, code, Arabic and math prompts: the logits match to a relative error of 2e-6 or better, the next-token predictions agree at every position, and greedy decoding with the KV cache gives the same tokens (64/64 on every prompt). All 64 layers run the same code path as the smaller dense Jais-3 models, which were checked in full.

Smoke-tested after conversion with released mlx-lm 0.32.0 on English and Arabic chat prompts: answers end normally (the chat template and stop tokens work).

Benchmarks (all Jais-3-32B MLX variants)

Perplexity on 128 × 512 tokens (WikiText-2 test, and the first articles of the Arabic Wikipedia dump 20231101.ar; lower is better) and generation speed on an Apple M4 Max, all measured the same way:

	bf16	8-bit	4-bit
Bits/weight	16.0	8.5	5.9
Disk	69.6 GB	37.0 GB	25.5 GB
Peak memory	n/a	37.2 GB	25.8 GB
WikiText-2 perplexity (English)	n/a	9.324	9.521 (+2.1%)
Arabic Wikipedia perplexity	n/a	11.307	11.661 (+3.1%)
Generation	n/a	12.4 tok/s	18.8 tok/s

The bf16 configuration could not be tested; the percentages are relative to 8-bit, which is within 0.2% of bf16 on every smaller Jais-3 model.

Perplexity is a coarse signal. Test the versions on your own workload before picking one.

Usage

mlx_lm.generate --model IFM/Jais-3-32B-MLX-4bit --trust-remote-code --prompt "ما هي عاصمة الإمارات؟" --max-tokens 2048

from mlx_lm import load, generate

model, tokenizer = load("IFM/Jais-3-32B-MLX-4bit", trust_remote_code=True)
messages = [{"role": "user", "content": "Explain mixture-of-experts in two sentences."}]
prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
print(generate(model, tokenizer, prompt, max_tokens=2048))

mlx_lm.chat and mlx_lm.server take the same --trust-remote-code flag.

The model thinks before it answers, inside <ifm|think> … </ifm|think>. Leave room for that in max_tokens. The effort level is set with reasoning_effort in the chat template: "high" (default), "medium" or "low", e.g. apply_chat_template(messages, add_generation_prompt=True, reasoning_effort="low").

Notes:

• Server output: mlx-lm doesn't recognize the <ifm|think> tags yet, so mlx_lm.server returns the thinking text inside content rather than in a separate reasoning field. The tool-call format isn't parsed yet either.
• Chat template change: the original template raises an error when an earlier assistant message has no thinking field, which is what OpenAI-style clients send. The template here renders empty thinking for those messages instead. Nothing else was changed.

License

Apache-2.0, inherited from the base model. Refer to the original model card for architecture, benchmarks and intended use.