Skip to content

Commit d3d98b1

Browse files
maziyarpanahimaziyarpanahi
authored
feat: add ModernBERT and Longformer MLX backends (#2455)
Closes #2340 Closes #2345
1 parent 2c8c7b1 commit d3d98b1

6 files changed

Lines changed: 1197 additions & 14 deletions

File tree

‎docs/export-matrix.md‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,8 @@ when a new family is added to the code.
3434
| RoBERTa | `roberta` | ✅ (via bert) | ✅ | ⚠️ INT4-if-recall-holds | ✅ | ✅ | ✅ | ✅ | ✅ | ⚠️ backbone only |
3535
| XLM-RoBERTa | `xlm-roberta`, `xlm_roberta` | ✅ (via bert) | ✅ | ⚠️ INT4-if-recall-holds | ✅ | ✅ | ✅ | ✅ | ✅ | ⚠️ backbone only |
3636
| DeBERTa-v2 | `deberta`, `deberta-v2` | ✅ | ✅ | ⚠️ INT4-if-recall-holds | ✅ | ✅ | ✅ | ✅ | ✅ | ⚠️ backbone only |
37-
| ModernBERT | — | ❌ family not in MLX allowlist | ❌ | ❌ | ❌ family not in CoreML allowlist | ❌ | ⚠️ upstream ONNX support required | ⚠️ same as ONNX (default) | ⚠️ same as ONNX (default) | ❌ |
38-
| Longformer | — | ❌ family not in MLX allowlist | ❌ | ❌ | ❌ family not in CoreML allowlist | ❌ | ⚠️ upstream ONNX support required | ⚠️ same as ONNX (default) | ⚠️ same as ONNX (default) | ❌ |
37+
| ModernBERT | `modernbert`, `modern-bert` | ✅ | ✅ | ⚠️ INT4-if-recall-holds | ❌ family not in CoreML allowlist | ❌ | ⚠️ upstream ONNX support required | ⚠️ same as ONNX (default) | ⚠️ same as ONNX (default) | ❌ |
38+
| Longformer | `longformer` | ✅ | ✅ | ⚠️ INT4-if-recall-holds | ❌ family not in CoreML allowlist | ❌ | ⚠️ upstream ONNX support required | ⚠️ same as ONNX (default) | ⚠️ same as ONNX (default) | ❌ |
3939
| OpenAI Privacy Filter | `openai-privacy-filter`, `privacy-filter`, `privacy-filter-nemotron`, `nemotron-privacy-filter`, `privacy-filter-multilingual`, `multilingual-privacy-filter` | ✅ | ✅ | ⚠️ INT4-if-recall-holds | ❌ family not in CoreML allowlist | ❌ | ❌ no ONNX converter path today | ❌ | ❌ | ❌ |
4040
| GLiNER (span) | `gliner-uni-encoder-span` | ✅ (MLX custom family) | ✅ | ⚠️ INT4-if-recall-holds | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
4141
| GLiClass | `gliclass-uni-encoder` | ✅ (MLX custom family) | ✅ | ⚠️ INT4-if-recall-holds | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
@@ -59,10 +59,11 @@ truth is the backend code, but the essentials are:
5959
Any other architecture fails before tracing with an error that lists the
6060
supported families. See [CoreML Packaging](coreml-export.md).
6161
- **MLX enforces a family allowlist.** The MLX dispatcher resolves architectures
62-
through `_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES` (BERT-family and
63-
DeBERTa-v2) plus the MLX-custom set (`gliner-uni-encoder-span`,
64-
`gliclass-uni-encoder`, `gliner-uni-encoder-token-relex`) and the privacy
65-
filter family. Architectures outside the allowlist raise a clear
62+
through `_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES` (BERT-family,
63+
DeBERTa-v2, ModernBERT, and Longformer) plus the MLX-custom set
64+
(`gliner-uni-encoder-span`, `gliclass-uni-encoder`,
65+
`gliner-uni-encoder-token-relex`) and the privacy filter family.
66+
Architectures outside the allowlist raise a clear
6667
"unsupported family" error rather than falling through silently. See
6768
[MLX Backend](mlx-backend.md).
6869
- **INT4 is certified only when recall holds.** `openmed.mlx.convert`

‎openmed/mlx/convert.py‎

Lines changed: 47 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -110,6 +110,31 @@
110110
("classifier.", "classifier."),
111111
]
112112

113+
_MODERNBERT_KEY_REPLACEMENTS: list[Tuple[str, str]] = [
114+
("model.layers.", "model.encoder.layers."),
115+
("model.embeddings.tok_embeddings.", "model.embeddings.word_embeddings."),
116+
(".attn.Wqkv.", ".attention.qkv_proj."),
117+
(".attn.Wo.", ".attention.out_proj."),
118+
(".mlp.Wi.", ".mlp.wi_proj."),
119+
(".mlp.Wo.", ".mlp.wo_proj."),
120+
]
121+
122+
_LONGFORMER_KEY_REPLACEMENTS: list[Tuple[str, str]] = [
123+
("longformer.encoder.layer.", "longformer.encoder.layers."),
124+
("longformer.embeddings.LayerNorm.", "longformer.embeddings.norm."),
125+
(".attention.self.query_global.", ".attention.query_global_proj."),
126+
(".attention.self.key_global.", ".attention.key_global_proj."),
127+
(".attention.self.value_global.", ".attention.value_global_proj."),
128+
(".attention.self.query.", ".attention.query_proj."),
129+
(".attention.self.key.", ".attention.key_proj."),
130+
(".attention.self.value.", ".attention.value_proj."),
131+
(".attention.output.dense.", ".attention.out_proj."),
132+
(".attention.output.LayerNorm.", ".ln1."),
133+
(".intermediate.dense.", ".linear1."),
134+
(".output.dense.", ".linear2."),
135+
(".output.LayerNorm.", ".ln2."),
136+
]
137+
113138
_OPF_MODEL_TYPES = {
114139
"openai-privacy-filter",
115140
"privacy-filter",
@@ -317,6 +342,15 @@ def _infer_source_model_type(key: str, model_type: str | None) -> str:
317342
return "roberta"
318343
if key.startswith("electra."):
319344
return "electra"
345+
if key.startswith("model."):
346+
if (
347+
".layers." in key
348+
or key.startswith("model.embeddings.")
349+
or key.startswith("model.final_norm.")
350+
):
351+
return "modernbert"
352+
if key.startswith("longformer."):
353+
return "longformer"
320354
return "bert"
321355

322356

@@ -325,7 +359,11 @@ def remap_key(key: str, model_type: str | None = None) -> str:
325359
source_model_type = _infer_source_model_type(key, model_type)
326360
resolved_model_type = resolve_model_type(source_model_type)
327361

328-
if resolved_model_type == "deberta-v2":
362+
if source_model_type in {"modernbert", "modern-bert"}:
363+
replacements = _MODERNBERT_KEY_REPLACEMENTS
364+
elif source_model_type == "longformer":
365+
replacements = _LONGFORMER_KEY_REPLACEMENTS
366+
elif resolved_model_type == "deberta-v2":
329367
replacements = _DEBERTA_V2_KEY_REPLACEMENTS
330368
elif source_model_type == "distilbert":
331369
replacements = _DISTILBERT_KEY_REPLACEMENTS
@@ -427,10 +465,14 @@ def save_mlx_model(
427465
logger.info("Quantizing to %d bits ...", quantize_bits)
428466
model = build_model(config)
429467
model.load_weights(list(mlx_weights.items()))
430-
quantize_kwargs = {"bits": quantize_bits}
431-
if quantize_group_size is not None:
432-
quantize_kwargs["group_size"] = quantize_group_size
433-
nn.quantize(model, **quantize_kwargs)
468+
if quantize_group_size is None:
469+
nn.quantize(model, bits=quantize_bits)
470+
else:
471+
nn.quantize(
472+
model,
473+
group_size=quantize_group_size,
474+
bits=quantize_bits,
475+
)
434476
mlx_weights = dict(tree_flatten(model.parameters()))
435477
config_to_save["_mlx_quantization"] = {
436478
"bits": quantize_bits,

‎openmed/mlx/models/__init__.py‎

Lines changed: 103 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
"bert": "bert",
1313
"distilbert": "bert",
1414
"electra": "bert",
15+
"modernbert": "modernbert",
16+
"modern-bert": "modernbert",
17+
"longformer": "longformer",
1518
"openai-privacy-filter": "openai-privacy-filter",
1619
"privacy-filter": "openai-privacy-filter",
1720
# ``OpenMed/privacy-filter-nemotron*`` is the OpenAI Privacy Filter
@@ -83,7 +86,11 @@ def resolve_artifact_family(
8386
family = normalize_model_type(manifest.get("family"))
8487
if family in _CUSTOM_FAMILIES:
8588
return family
86-
resolved_family = _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(family)
89+
resolved_family = (
90+
_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(family)
91+
if family is not None
92+
else None
93+
)
8794
if resolved_family is not None:
8895
return resolved_family
8996

@@ -92,7 +99,11 @@ def resolve_artifact_family(
9299
family = normalize_model_type(config.get("_mlx_family"))
93100
if family in _CUSTOM_FAMILIES:
94101
return family
95-
resolved_family = _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(family)
102+
resolved_family = (
103+
_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(family)
104+
if family is not None
105+
else None
106+
)
96107
if resolved_family is not None:
97108
return resolved_family
98109

@@ -107,7 +118,11 @@ def resolve_artifact_family(
107118
model_type = config
108119

109120
normalized = normalize_model_type(model_type)
110-
resolved = _SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(normalized)
121+
resolved = (
122+
_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES.get(normalized)
123+
if normalized is not None
124+
else None
125+
)
111126
if resolved is None:
112127
supported = ", ".join(
113128
sorted(_CUSTOM_FAMILIES | set(_SUPPORTED_TOKEN_CLASSIFICATION_MODEL_TYPES))
@@ -234,6 +249,81 @@ def normalize_model_config(
234249
"bidirectional_right_context",
235250
normalized.get("bidirectional_right_context", default_context),
236251
)
252+
elif family == "modernbert":
253+
normalized.setdefault("type_vocab_size", 0)
254+
normalized.setdefault("_mlx_position_offset", 0)
255+
normalized.setdefault(
256+
"norm_eps",
257+
normalized.get("norm_eps", normalized.get("layer_norm_eps", 1e-5)),
258+
)
259+
normalized.setdefault("norm_bias", normalized.get("norm_bias", False))
260+
normalized.setdefault(
261+
"hidden_activation",
262+
normalized.get("hidden_activation", normalized.get("hidden_act", "gelu")),
263+
)
264+
normalized.setdefault(
265+
"embedding_dropout",
266+
normalized.get("embedding_dropout", normalized.get("dropout", 0.0)),
267+
)
268+
normalized.setdefault(
269+
"attention_dropout",
270+
normalized.get("attention_dropout", 0.0),
271+
)
272+
normalized.setdefault("mlp_dropout", normalized.get("mlp_dropout", 0.0))
273+
normalized.setdefault(
274+
"classifier_dropout",
275+
normalized.get("classifier_dropout", 0.0),
276+
)
277+
normalized.setdefault("attention_bias", normalized.get("attention_bias", False))
278+
normalized.setdefault("mlp_bias", normalized.get("mlp_bias", False))
279+
normalized.setdefault(
280+
"classifier_bias", normalized.get("classifier_bias", False)
281+
)
282+
normalized.setdefault(
283+
"classifier_activation",
284+
normalized.get("classifier_activation", "gelu"),
285+
)
286+
normalized.setdefault("local_attention", normalized.get("local_attention", 128))
287+
normalized.setdefault("sliding_window", int(normalized["local_attention"]) // 2)
288+
if not normalized.get("layer_types"):
289+
every_n_layers = int(normalized.get("global_attn_every_n_layers", 3))
290+
every_n_layers = max(1, every_n_layers)
291+
normalized["layer_types"] = [
292+
"full_attention" if index % every_n_layers == 0 else "sliding_attention"
293+
for index in range(int(normalized["num_hidden_layers"]))
294+
]
295+
rope_parameters = normalized.get("rope_parameters")
296+
if not isinstance(rope_parameters, dict) or not any(
297+
key in rope_parameters
298+
for key in ("full_attention", "sliding_attention", "global", "local")
299+
):
300+
shared_theta = (
301+
rope_parameters.get("rope_theta")
302+
if isinstance(rope_parameters, dict)
303+
else None
304+
)
305+
global_rope_theta = normalized.get(
306+
"global_rope_theta", shared_theta or 160000.0
307+
)
308+
if global_rope_theta is None:
309+
global_rope_theta = shared_theta or 160000.0
310+
local_rope_theta = normalized.get(
311+
"local_rope_theta", shared_theta or 10000.0
312+
)
313+
if local_rope_theta is None:
314+
local_rope_theta = shared_theta or 10000.0
315+
normalized["rope_parameters"] = {
316+
"full_attention": {"rope_theta": float(global_rope_theta)},
317+
"sliding_attention": {"rope_theta": float(local_rope_theta)},
318+
}
319+
elif family == "longformer":
320+
normalized.setdefault("type_vocab_size", 2)
321+
normalized.setdefault("_mlx_position_offset", 0)
322+
attention_window = normalized.get("attention_window", 512)
323+
if isinstance(attention_window, (list, tuple)):
324+
normalized["attention_window"] = [int(value) for value in attention_window]
325+
else:
326+
normalized["attention_window"] = int(attention_window)
237327
else:
238328
normalized.setdefault("type_vocab_size", normalized.get("type_vocab_size", 2))
239329
normalized.setdefault("_mlx_position_offset", 0)
@@ -262,6 +352,16 @@ def build_model(
262352

263353
return DebertaV2ForTokenClassification(config)
264354

355+
if family == "modernbert":
356+
from openmed.mlx.models.modernbert_tc import ModernBertForTokenClassification
357+
358+
return ModernBertForTokenClassification(config)
359+
360+
if family == "longformer":
361+
from openmed.mlx.models.longformer_tc import LongformerForTokenClassification
362+
363+
return LongformerForTokenClassification(config)
364+
265365
if family == "openai-privacy-filter":
266366
from openmed.mlx.models.privacy_filter import (
267367
OpenAIPrivacyFilterForTokenClassification,

0 commit comments

Comments
 (0)