WeMM-Embedding-2B-Apple-Silicon-MLX

WeMM-Embedding-2B-Apple-Silicon-MLX

Native MLX conversion of ewin-reg/WeMM-Embedding-2B-Quantized — the hybrid FP8-protected attention + group-16 INT4 checkpoint — that preserves its mixed-precision allocation instead of flattening everything to one precision.

Why this allocation matters

The source checkpoint is curvature-guided: it deliberately keeps the precision-sensitive path at ~8-bit (FP8 E4M3) and compresses only the quantization-robust masses to INT4. Uniform 4-bit quantization destroys that — the source README itself documents softmax attention collapse (D_KL >= 7.50) when Q/K go to 4 bits. This build mirrors the source allocation exactly (verified: 285/285 modules, 0 missing, 0 extra).

Two deviations from the source, both improvements or forced by MLX:

• 8-bit group-64 is finer than the source's per-tensor-scalar FP8 — one scale for the whole matrix becomes one scale per 64 weights.
• MLX's mx.quantize kernel only supports group sizes 32/64/128, so the source's INT4 group-16 becomes group-32 (the closest supported granularity).

Per-module overrides are stored in config.json → quantization, which mlx_vlm.load() reads back (so the allocation is reproduced at load time — no patched library needed).

The extra ~200 MB is the group-16 → 32 change plus MLX's separate scales/biases tensors; it is still far below BF16 and keeps the ViT quantized (many "MLX" ports leave it in BF16 and balloon to ~2.6 GB).

Mac quickstart

pip install mlx mlx-vlm        # on Apple Silicon; no mlx-cpu needed

import mlx.core as mx
from mlx_vlm import load

model, processor = load("ewin-reg/WeMM-Embedding-2B-Apple-Silicon-MLX")
tok = processor.tokenizer

def embed(text: str) -> mx.array:
    enc = tok(text, return_tensors="np")
    ids = mx.array(enc["input_ids"])
    mask = mx.array(enc["attention_mask"])
    cache = model.language_model.make_cache()
    h = model.language_model.model(ids, cache=cache)   # final hidden states
    mx.eval(h)
    vec = h[0, int(mask.sum()) - 1, :]                 # last non-pad token
    return vec / mx.linalg.norm(vec)                   # L2-normalised, 2048-d

v = embed("a photo of a cat sitting on a windowsill")

# Matryoshka truncation (e.g. to 512 dims)
v512 = v[:512] / mx.linalg.norm(v[:512])

Text, images and videos: the same model exposes vision_tower, and the chat template / processor from the source checkpoint are included, so multimodal inputs work the same way as the model card's PyTorch snippet (image/video tokens via Qwen3VLProcessor).

Verification (CPU run, before upload)

• Shards vs allocation: 285 quantized modules, 0 missing, 0 over-quantized
• mlx_vlm.load(): 0 bit/group-size mismatches; norms remain BF16
• Embeddings: finite, L2-normalised
  • cos("cat sitting on a windowsill", "kitten resting by the window") = 0.809
  • cos(cat, "quarterly revenue increased by 12% YoY") = 0.295
• MRL truncation at 64 / 256 / 1024 / 2048 dims works

Provenance

• Dequantized from ewin-reg/WeMM-Embedding-2B-Quantized (its own weight-only fake-quant is lossless with respect to what that model computes at runtime).
• Converted with mlx_vlm convert + a per-module quant predicate derived from the checkpoint's actual tensor dtypes — not from prose.

---

Hardware verification status

This build was verified on x86 CPU only. No Apple-Silicon run has happened yet. Anyone reading this should not assume MPS has been tested. The MLX repo is the tested Mac path going forward, replacing the old INT8 workaround repo, which I pulled for exactly that reason.