ModernBERT-base in Core ML, fixed length 128, verified against PyTorch

ModernBERT-base-128.mlpackage/Data/com.apple.CoreML/model.mlmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e07033a286fde41826cfa46a2304b07a037d86c7c0bd39a455270634eaf3f63
+size 273934

ModernBERT-base-128.mlpackage/Data/com.apple.CoreML/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a50ac7543e50fc380850f46efec17b7fade77fd33bd705874bd17d993b0b85c
+size 299391360

ModernBERT-base-128.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "588C6CEE-EAA5-4CD9-9752-2D4737BD8C5D": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "B79F1999-80C9-4E40-AC85-DE68A644A480": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        }
+    },
+    "rootModelIdentifier": "588C6CEE-EAA5-4CD9-9752-2D4737BD8C5D"
+}

README.md

@@ -0,0 +1,116 @@
+---
+license: apache-2.0
+language:
+- en
+base_model:
+- answerdotai/ModernBERT-base
+tags:
+- coreml
+- fill-mask
+- masked-lm
+- modernbert
+pipeline_tag: fill-mask
+---
+
+# ModernBERT-base in Core ML
+
+[ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base) converted
+to Core ML with its masked-LM head, at a **fixed sequence length of 128**.
+
+Verified against the PyTorch model rather than assumed: correlation of the
+logits at a masked position is **0.997 to 0.99999**. `verify.py` reproduces it.
+
+```
+The capital of Ireland is [MASK].
+   torch    Dublin, Cork, Belfast, London
+   coreml   Dublin, Belfast, Cork, London
+```
+
+About **7 ms** per call on an M-series Mac after warm-up, against ~55 ms for the
+same model in PyTorch on the same machine. The first call pays roughly 7 s of
+Core ML compilation, so warm it once at launch.
+
+## Why this exists
+
+There is an earlier Core ML conversion of ModernBERT on the Hub. **It is wrong,
+in both its full-precision and 4-bit forms**, and nothing about it says so — it
+loads, it runs, and it returns confident nonsense.
+
+```
+The capital of Ireland is [MASK].
+   the earlier conversion    £, isation, organisation, ised
+```
+
+Measured against PyTorch, its logits correlate at **1.0000 at a sequence length
+of one** and **0.17 to 0.39** at any real length. Its converter traced the model
+with a single token and a flexible `RangeDim`. At one token ModernBERT's
+sliding-window attention mask is optimised away; past it, the graph is wrong.
+
+That failure is silent, which is the reason for this repository. If you convert
+this model yourself, check the correlation at a real length — not the top
+predictions of a short example, which can look plausible while the model is
+broken.
+
+## Using it
+
+```python
+import coremltools as ct, numpy as np
+from transformers import AutoTokenizer
+
+tok = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
+model = ct.models.MLModel("ModernBERT-base-128.mlpackage")
+
+text = "The capital of Ireland is " + tok.mask_token + "."
+encoded = tok(text, return_tensors="np")
+ids = np.full((1, 128), tok.pad_token_id, np.int32)
+ids[0, : encoded.input_ids.shape[1]] = encoded.input_ids
+
+logits = model.predict({"input_ids": ids})["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]])
+```
+
+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.
+
+## Two things to know before you rely on it
+
+**The length is fixed, and padding is attended.** The attention mask is
+`ones_like(input_ids)` inside the graph, so padding tokens are read as text. On
+one sentence-boundary task this cost nothing measurable, but that is a property
+of that task and not a guarantee. Convert at a length close to your inputs, or
+convert several — `convert.py --length N`.
+
+Taking `attention_mask` as a real input is the obvious fix and does not work:
+ModernBERT then builds the mask with `new_ones`, which coremltools has no
+conversion for.
+
+**Quantising costs more than it saves, at least here.** Same conversion, same
+sentence-boundary task, 150 real sentences:
+
+| | size | latency | quality |
+|---|---|---|---|
+| float | 286 MB | 10 ms | reference |
+| int8 | 152 MB | 23 ms | indistinguishable |
+| int4 | 81 MB | 21 ms | shifted, and worse where it matters |
+
+int8 halves the size and doubles the latency, because the weights are
+dequantised on every pass. int4 changes the distribution enough to matter.
+
+## Reproducing
+
+Pin the toolchain. `transformers` 5.x emits `new_ones` in ModernBERT's attention
+path and coremltools cannot convert it — that failure also hits BERT and
+DistilBERT, so it is the toolchain and not this model.
+
+```sh
+pip install torch==2.7.0 transformers==4.48.3 coremltools==9.0
+python convert.py --length 128
+python verify.py ModernBERT-base-128.mlpackage --length 128
+```
+
+## Licence
+
+Apache-2.0, from
+[ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base). Only the
+conversion is new here.

convert.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""ModernBERT -> Core ML, at a fixed sequence length.
+
+    python convert.py --length 128 [--model ModernBERT-base]
+
+The length is fixed on purpose. A flexible `RangeDim` traced at one token
+converts without error and is wrong at every other length — see the README.
+
+Pin the toolchain. transformers 5.x emits `new_ones` in ModernBERT's attention
+path and coremltools cannot convert it; the failure also hits BERT and
+DistilBERT, so it is the toolchain and not this model.
+
+    torch==2.7.0   transformers==4.48.3   coremltools==9.0
+"""
+import argparse
+import numpy as np
+import torch
+import coremltools as ct
+from transformers import AutoModelForMaskedLM
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model", default="ModernBERT-base")
+parser.add_argument("--length", type=int, default=128)
+args = parser.parse_args()
+
+
+class Wrapped(torch.nn.Module):
+    """The mask stays `ones_like` inside the graph.
+
+    Taking `attention_mask` as a second input makes ModernBERT build it with
+    `new_ones`, which coremltools has no conversion for. The cost is that
+    padding is attended as if it were text — pick a length close to your
+    sentences, or convert several.
+    """
+
+    def __init__(self, name):
+        super().__init__()
+        self.model = AutoModelForMaskedLM.from_pretrained(f"answerdotai/{name}")
+
+    def forward(self, input_ids):
+        return self.model(
+            input_ids=input_ids, attention_mask=torch.ones_like(input_ids)
+        ).logits
+
+
+model = Wrapped(args.model).eval()
+example = (torch.ones((1, args.length), dtype=torch.int32),)
+traced = torch.jit.trace(model, example, strict=False)
+
+mlmodel = ct.convert(
+    traced,
+    inputs=[ct.TensorType(name="input_ids", shape=(1, args.length), dtype=np.int32)],
+    outputs=[ct.TensorType(name="logits")],
+    minimum_deployment_target=ct.target.macOS15,
+)
+mlmodel.short_description = (
+    f"{args.model}, masked-LM head, fixed sequence length {args.length}"
+)
+out = f"{args.model}-{args.length}.mlpackage"
+mlmodel.save(out)
+print("wrote", out)

verify.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""Check a converted model against the PyTorch original.
+
+    python verify.py ModernBERT-base-128.mlpackage --length 128
+
+Prints the correlation of the logits at a masked position. Anything below
+0.99 means the conversion is wrong, whatever its top predictions look like.
+"""
+import argparse
+import numpy as np
+import torch
+import coremltools as ct
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+parser = argparse.ArgumentParser()
+parser.add_argument("package")
+parser.add_argument("--length", type=int, default=128)
+parser.add_argument("--model", default="answerdotai/ModernBERT-base")
+args = parser.parse_args()
+
+tok = AutoTokenizer.from_pretrained(args.model)
+ref = AutoModelForMaskedLM.from_pretrained(args.model).eval()
+mlmodel = ct.models.MLModel(args.package)
+
+SENTENCES = [
+    "The capital of Ireland is {m}.",
+    "You should see a parrot {m} at the top right of your screen.",
+    "The PR looks good but there is another failure mode {m} we might cover this.",
+]
+
+for sentence in SENTENCES:
+    text = sentence.format(m=tok.mask_token)
+    encoded = tok(text, return_tensors="pt")
+    length = encoded.input_ids.shape[1]
+    ids = np.full((1, args.length), tok.pad_token_id, np.int32)
+    ids[0, :length] = encoded.input_ids.numpy()
+    padded = torch.tensor(ids, dtype=torch.long)
+    with torch.no_grad():
+        expected = ref(
+            input_ids=padded, attention_mask=torch.ones_like(padded)
+        ).logits[0].numpy()
+    got = mlmodel.predict({"input_ids": ids})["logits"][0]
+    at = int(np.where(encoded.input_ids[0].numpy() == tok.mask_token_id)[0][0])
+    print(f"\n  {text[:56]}")
+    print(f"     torch  {[tok.decode([int(i)]) for i in np.argsort(-expected[at])[:4]]}")
+    print(f"     coreml {[tok.decode([int(i)]) for i in np.argsort(-got[at])[:4]]}")
+    print(f"     r = {np.corrcoef(expected[at], got[at])[0, 1]:.5f}")