mmbert-small-coreml

mmBERT-small in Core ML

mmBERT-small converted to Core ML with its masked-LM head, at a fixed sequence length of 64. 140.9 M parameters, fp16 weights, 269 MB on disk. 1,833 languages.

Verified against the PyTorch model rather than assumed. verify.py reproduces it: correlation of the logits at a masked position, six sentences in four languages.

The capital of Ireland is [MASK].
   torch    Dublin, Belfast, Limerick, Cork, Derry
   coreml   Dublin, Belfast, Limerick, Cork, Derry

Il a garé la voiture devant la [MASK] de ses parents.
   torch    maison, voiture, porte, chambre, résidence
   coreml   maison, voiture, porte, chambre, résidence

correlation at the mask: 0.99999 to 1.00000
top-1 agrees 6/6, top-10 set agrees 6/6

attention_mask is a real input, and it has to be

The obvious conversion builds the attention mask as ones_like(input_ids) inside the graph. Then padding is read as text, and this model minds. Measured over 238 masked slots in real sentences, scoring each slot's ten most likely words against a reference:

It is not subtle. That was [MASK] suggestion originally. answers my the a no his your our their her another with the mask, and is iz in and work ia for first a part without it.

So this conversion takes attention_mask as an input. znaat/modernbert-coreml does not, and its card says that input would not convert; with the toolchain pinned below it does convert, for ModernBERT-base as well as for this model.

Pin the compute units to CPU and GPU

This is not optional. Load it with compute_units=CPU_AND_GPU (MLComputeUnits.cpuAndGPU in Swift). On the Neural Engine the model still loads, still runs, and quietly answers differently.

Measured over 238 masked slots, comparing each backend's ten most likely alphabetic words against the PyTorch reference:

On the Neural Engine this model is not approximately right. It is wrong.

ALL chose the same path as CPU_AND_GPU on the Mac this was measured on, in all 238 cases — but that is a scheduler decision, not a promise. Pin it.

Latency

About 18 ms per forward pass on an M-series Mac with coremltools, and 15 ms called from Swift, after warm-up. Loading the .mlpackage costs about 2.1 s including the Core ML compile; loading the compiled .mlmodelc costs 0.1 s, so compile once and keep it.

The 256,000-word head is where the time goes. The same conversion of ModernBERT-base, whose vocabulary is 50,368, runs at 11 ms on the same machine. mmBERT-small is the smaller model — 42 M parameters of backbone against 110 M — and it is still slower, because tie_word_embeddings makes that vocabulary the output matrix and every call produces a [1, 64, 256000] tensor.

Two things to know before you rely on it

The length is fixed. Pad to 64 and pass an attention_mask that is 0 over the padding. For a different length, convert.py --length N.

Verify at the length you convert. Do not assume a recipe that is faithful at one length is faithful at another. verify.py takes --length.

Using it

import coremltools as ct, numpy as np
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("jhu-clsp/mmBERT-small")
model = ct.models.MLModel(
    "mmBERT-small-64.mlpackage", compute_units=ct.ComputeUnit.CPU_AND_GPU
)

text = "The capital of Ireland is " + tok.mask_token + "."
encoded = tok(text, return_tensors="np")
length = encoded.input_ids.shape[1]
ids = np.full((1, 64), tok.pad_token_id, np.int32)
ids[0, :length] = encoded.input_ids
attention = np.zeros((1, 64), np.int32)
attention[0, :length] = 1

logits = model.predict(
    {"input_ids": ids, "attention_mask": attention}
)["logits"][0]
at = int(np.where(encoded.input_ids[0] == tok.mask_token_id)[0][0])
print([tok.decode([int(i)]) for i in np.argsort(-logits[at])[:5]])

tokenizer.json, tokenizer_config.json and config.json are here too, copied unchanged from jhu-clsp/mmBERT-small at revision abc32620dd4f6ab06f5fbe905dc25f310618e09f. You do not need them from Python — AutoTokenizer fetches its own. They are here so a Swift, Rust or C++ caller can get the model and the tokenizer that matches it from one place, without a Python step. swift-transformers reads all three: it wants config.json and tokenizer_config.json beside tokenizer.json and throws without them.

The tokenizer is Gemma-shaped: byte-fallback BPE, 256,000 entries, a Metaspace pre-tokenizer that writes a space as ▁, and <bos> <eos> <mask> <pad> <unk> for special tokens. tie_word_embeddings is true, so the 256k output head is the input embedding and there is no separate output matrix. That head is 99 MB of the 269.

Downloading from the Hub leaves the weights as a symlink, and the Core ML compiler does not follow one. Copy the package with cp -RL before loading it.

Reproducing

Pin the toolchain.

pip install torch==2.7.0 transformers==4.48.3 coremltools==9.0
python convert.py --model jhu-clsp/mmBERT-small --length 64
python verify.py mmBERT-small-64.mlpackage --length 64 --units CPU_AND_GPU

transformers 5.x emits new_ones in ModernBERT's attention path — mmBERT uses that architecture — and coremltools cannot convert it. Pin 4.48.3.

ct.convert writes fp16 weights for an ML program by default, which is where the 269 MB comes from — the fp32 checkpoint is 537 MB.

Licence

MIT, from mmBERT-small. Only the conversion is new here.