ewin-reg / 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.