diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index 64f31948..40f164af 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -555,12 +555,62 @@ def _generate_image_classification(case: TestCase, json_path: Path, device: str) ) +def _generate_masked_lm(case: TestCase, json_path: Path, device: str) -> None: + """Generate golden data for a masked LM model (BERT, RoBERTa, ESM-2, etc.). + + Runs a forward pass with a masked input and captures the logits at the + position of the first mask token. For protein language models (ESM-2), + the mask token is ````; for BERT-style models it is ``[MASK]``. + """ + from mobius._testing.golden import save_golden_ref + from mobius._testing.torch_reference import ( + load_torch_masked_lm_model, + torch_masked_lm_forward, + ) + + model, tokenizer = load_torch_masked_lm_model( + case.model_id, + revision=case.revision, + device=device, + trust_remote_code=case.trust_remote_code, + ) + + encoded = tokenizer(case.prompts[0], return_tensors="np", padding=False) + input_ids = encoded["input_ids"] + attention_mask = encoded["attention_mask"] + token_type_ids = encoded.get("token_type_ids") + + # Forward pass → (batch, seq_len, vocab_size) + logits = torch_masked_lm_forward(model, input_ids, attention_mask, token_type_ids) + + # Find the first mask token position and extract its logits. + # Fall back to position 1 (first non-CLS token) if no mask token is found. + mask_id = tokenizer.mask_token_id + flat_ids = input_ids[0].tolist() + mask_pos = flat_ids.index(mask_id) if mask_id in flat_ids else 1 + mask_logits = logits[0, mask_pos, :] # (vocab_size,) + + golden = _extract_logits_golden(mask_logits) + + # Masked-LM is L4-only (no autoregressive generation) + save_golden_ref( + json_path, + top1_id=golden["top1_id"], + top2_id=golden["top2_id"], + top10_ids=golden["top10_ids"], + top10_logits=golden["top10_logits"], + logits_summary=golden["logits_summary"], + input_ids=input_ids, + ) + + # ---- Dispatcher ---- # Map task_type strings to generator functions. _GENERATORS = { "text-generation": _generate_causal_lm, "feature-extraction": _generate_encoder, + "masked-lm": _generate_masked_lm, "seq2seq": _generate_seq2seq, "image-text-to-text": _generate_vision_language, "image-classification": _generate_image_classification, diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index b5e512f8..79a38b99 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -80,7 +80,7 @@ ) from mobius.models.bamba import BambaCausalLMModel from mobius.models.bart import BartForConditionalGeneration -from mobius.models.bert import BertModel +from mobius.models.bert import BertForMaskedLM, BertModel from mobius.models.blip import BlipVisionModel from mobius.models.blip2 import Blip2Model from mobius.models.clip import CLIPTextModel, CLIPVisionModel @@ -589,6 +589,14 @@ def _create_default_registry() -> ModelRegistry: reg.register("layoutlmv3", LayoutLMv3Model, task="feature-extraction") reg.register("modernbert", ModernBertModel, task="feature-extraction") + # --- Masked LM (encoder + prediction head) --- + for name in ( + "bert_masked_lm", + "esm_masked_lm", + "roberta_masked_lm", + ): + reg.register(name, BertForMaskedLM, task="masked-lm") + # --- Absolute positional embeddings (non-RoPE) --- reg.register("gpt2", GPT2CausalLMModel) for name in ( @@ -899,6 +907,11 @@ def _create_default_registry() -> ModelRegistry: "xmod": "facebook/xmod-base", "yoso": "uw-madison/yoso-4096", + # --- Masked LM (encoder + prediction head) --- + "bert_masked_lm": "google-bert/bert-base-uncased", + "esm_masked_lm": "facebook/esm2_t6_8M_UR50D", + "roberta_masked_lm": "FacebookAI/roberta-base", + # --- Encoder-decoder --- "bart": "facebook/bart-base", "t5": "google-t5/t5-small", diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 29a53f82..7bbc1886 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -263,6 +263,64 @@ def torch_encoder_forward( return outputs.last_hidden_state.cpu().numpy() +# --------------------------------------------------------------------------- +# Masked LM models (BERT, RoBERTa, ESM-2, etc.) +# --------------------------------------------------------------------------- + + +def load_torch_masked_lm_model( + model_id: str, + revision: str | None = None, + dtype: torch.dtype = torch.float32, + device: str = "cpu", + trust_remote_code: bool = False, +): + """Load a HuggingFace masked LM model for reference inference. + + Uses AutoModelForMaskedLM for BERT/RoBERTa/ESM-2-style architectures. + + Returns: + Tuple of (model, tokenizer). + """ + import transformers + + kwargs: dict = {"trust_remote_code": trust_remote_code} + if revision: + kwargs["revision"] = revision + tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, **kwargs) + model = transformers.AutoModelForMaskedLM.from_pretrained( + model_id, + torch_dtype=dtype, + device_map=device, + **kwargs, + ) + model.eval() + return model, tokenizer + + +@torch.no_grad() +def torch_masked_lm_forward( + model, + input_ids: np.ndarray, + attention_mask: np.ndarray, + token_type_ids: np.ndarray | None = None, +) -> np.ndarray: + """Run a single forward pass on a HuggingFace masked LM model. + + Returns: + logits as numpy array [batch, seq_len, vocab_size]. + """ + device = next(model.parameters()).device + kwargs: dict = { + "input_ids": torch.from_numpy(input_ids).to(device), + "attention_mask": torch.from_numpy(attention_mask).to(device), + } + if token_type_ids is not None: + kwargs["token_type_ids"] = torch.from_numpy(token_type_ids).to(device) + outputs = model(**kwargs) + return outputs.logits.cpu().numpy() + + # --------------------------------------------------------------------------- # Seq2seq models (BART, T5, mBART, etc.) # --------------------------------------------------------------------------- diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 3f157185..cd5cd578 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -10,6 +10,7 @@ "AutoencoderKLQwenImageModel", "BambaCausalLMModel", "BartForConditionalGeneration", + "BertForMaskedLM", "BertModel", "Blip2Model", "BloomCausalLMModel", @@ -126,7 +127,7 @@ from mobius.models.bamba import BambaCausalLMModel from mobius.models.bart import BartForConditionalGeneration from mobius.models.base import CausalLMModel, LayerNormCausalLMModel -from mobius.models.bert import BertModel +from mobius.models.bert import BertForMaskedLM, BertModel from mobius.models.blip2 import Blip2Model from mobius.models.chatglm import ChatGLMCausalLMModel from mobius.models.clip import CLIPVisionModel diff --git a/src/mobius/models/bert.py b/src/mobius/models/bert.py index befa34dd..f4a2e15c 100644 --- a/src/mobius/models/bert.py +++ b/src/mobius/models/bert.py @@ -1,7 +1,7 @@ # Copyright (c) ONNX Project Contributors # SPDX-License-Identifier: Apache-2.0 -"""BERT encoder-only model with HF-aligned weight naming. +"""BERT encoder-only model and masked LM with HF-aligned weight naming. HF BERT uses deeply nested naming for encoder layers: attention.self.query / attention.self.key / attention.self.value @@ -11,8 +11,12 @@ embeddings.LayerNorm Module attributes here match HF conventions to eliminate the -rename dict entirely. Only prefix stripping (bert./roberta.) +rename dict entirely. Only prefix stripping (bert./roberta./esm.) and gamma/beta compat remain in preprocess_weights. + +BertForMaskedLM adds a masked language model head on top of the +encoder for predicting masked tokens. Supports BERT, RoBERTa, ESM-2, +and similar encoder-only architectures. """ from __future__ import annotations @@ -288,36 +292,197 @@ def forward( _PARAM_RENAMES = {"gamma": "weight", "beta": "bias"} +def _rename_bert_encoder_weight(name: str) -> str: + """Collapse nested HF naming and handle gamma/beta compat. + + This is the shared encoder-weight rename logic used by both + ``_rename_bert_weight`` and ``_rename_masked_lm_weight``. + Assumes model prefix (bert./roberta./esm.) is already stripped. + """ + # Collapse nested HF naming to match flat ONNX paths: + # attention.self.query → attention.query + # attention.output.dense → attention.dense + # layer.N.output.dense → layer.N.dense + name = name.replace(".attention.self.", ".attention.") + name = name.replace(".attention.output.", ".attention.") + name = name.replace(".output.dense.", ".dense.") + name = name.replace(".output.LayerNorm.", ".LayerNorm.") + + # Rename gamma/beta to weight/bias (old BERT compat) + parts = name.rsplit(".", 1) + if len(parts) == 2 and parts[1] in _PARAM_RENAMES: + name = f"{parts[0]}.{_PARAM_RENAMES[parts[1]]}" + + return name + + def _rename_bert_weight(name: str) -> str | None: """Rename a single HF BERT weight to our convention. - Strips model prefixes (bert./roberta.), collapses nested HF naming + Strips model prefixes (bert./roberta./esm.), collapses nested HF naming (.self./.output.) to match the flat ONNX initializer paths, and handles old-BERT gamma/beta compat. Returns None for pooler/cls weights we don't need. """ - # Strip "bert." or "roberta." prefix if present + # Strip "bert." / "roberta." / "esm." prefix if present if name.startswith("bert."): name = name[5:] elif name.startswith("roberta."): name = name[8:] + elif name.startswith("esm."): + name = name[4:] # Skip pooler and classification heads if name.startswith(("pooler.", "cls.")): return None - # Collapse nested HF naming to match flat ONNX paths: - # attention.self.query → attention.query - # attention.output.dense → attention.dense - # layer.N.output.dense → layer.N.dense - name = name.replace(".attention.self.", ".attention.") - name = name.replace(".attention.output.", ".attention.") - name = name.replace(".output.dense.", ".dense.") - name = name.replace(".output.LayerNorm.", ".LayerNorm.") + return _rename_bert_encoder_weight(name) - # Rename gamma/beta to weight/bias (old BERT compat) - parts = name.rsplit(".", 1) - if len(parts) == 2 and parts[1] in _PARAM_RENAMES: - name = f"{parts[0]}.{_PARAM_RENAMES[parts[1]]}" - return name +# --------------------------------------------------------------------------- +# Masked LM Head and BertForMaskedLM +# --------------------------------------------------------------------------- + +# BERT cls.predictions.transform.* -> our lm_head.* naming +_BERT_CLS_TO_LM_HEAD = { + "cls.predictions.transform.dense.": "lm_head.dense.", + "cls.predictions.transform.LayerNorm.": "lm_head.layer_norm.", + "cls.predictions.decoder.": "lm_head.decoder.", + # Top-level bias on the BERT predictions head + "cls.predictions.bias": "lm_head.decoder.bias", +} + + +class _MaskedLMHead(nn.Module): + """Masked LM prediction head: dense -> activation -> LayerNorm -> decoder. + + Matches HF ESM/RoBERTa naming: lm_head.dense, lm_head.layer_norm, + lm_head.decoder. BERT uses different naming (cls.predictions.*) which + is handled by preprocess_weights. + """ + + def __init__( + self, + hidden_size: int, + vocab_size: int, + hidden_act: str = "gelu", + layer_norm_eps: float = 1e-12, + ): + super().__init__() + self.dense = Linear(hidden_size, hidden_size, bias=True) + self.layer_norm = LayerNorm(hidden_size, eps=layer_norm_eps) + self.decoder = Linear(hidden_size, vocab_size, bias=True) + self._act_fn = ACT2FN[hidden_act] + + def forward(self, op: builder.OpBuilder, hidden_states: ir.Value): + # hidden_states: (batch, seq, hidden_size) -> logits: (batch, seq, vocab) + x = self.dense(op, hidden_states) + x = self._act_fn(op, x) + x = self.layer_norm(op, x) + return self.decoder(op, x) + + +class BertForMaskedLM(nn.Module): + """BERT/ESM/RoBERTa encoder with masked language model head. + + Predicts vocabulary logits for each token position, used for + masked token prediction (fill-mask). + + Supports ESM-2 (protein masked LM), BERT, RoBERTa, and similar + encoder architectures. Output is per-token logits over the vocabulary. + + Replicates HuggingFace's ``EsmForMaskedLM`` / ``BertForMaskedLM`` / + ``RobertaForMaskedLM``. + """ + + default_task = "masked-lm" + category = "encoder" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.config = config + self.embeddings = _BertEmbeddings( + vocab_size=config.vocab_size, + hidden_size=config.hidden_size, + max_position_embeddings=config.max_position_embeddings, + type_vocab_size=getattr(config, "type_vocab_size", 2), + layer_norm_eps=config.rms_norm_eps, + pad_token_id=config.pad_token_id or 0, + ) + self.encoder = _BertEncoder(config) + self.lm_head = _MaskedLMHead( + hidden_size=config.hidden_size, + vocab_size=config.vocab_size, + hidden_act=config.hidden_act, + layer_norm_eps=config.rms_norm_eps, + ) + + def forward( + self, + op: builder.OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + token_type_ids: ir.Value, + ): + hidden_states = self.embeddings(op, input_ids, token_type_ids) + hidden_states = self.encoder(op, hidden_states, attention_mask) + logits = self.lm_head(op, hidden_states) + return logits + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Rename HF weight names to match our parameter names. + + Handles ESM (esm.* prefix + lm_head.*), RoBERTa (roberta.* + + lm_head.*), and BERT (bert.* + cls.predictions.*) conventions. + """ + new_state_dict = {} + for name, tensor in state_dict.items(): + new_name = _rename_masked_lm_weight(name) + if new_name is not None: + new_state_dict[new_name] = tensor + return new_state_dict + + +def _rename_masked_lm_weight(name: str) -> str | None: + """Rename a single HF masked LM weight to our convention. + + Handles LM head weights (BERT cls.predictions.* -> lm_head.*, + ESM/RoBERTa lm_head.* passes through) and delegates encoder + weights to ``_rename_bert_weight``. + """ + # Strip "bert." / "roberta." / "esm." prefix from encoder weights + if name.startswith("bert."): + name = name[5:] + elif name.startswith("roberta."): + name = name[8:] + elif name.startswith("esm."): + name = name[4:] + + # Skip pooler (not needed for MLM) + if name.startswith("pooler."): + return None + + # Skip contact_head (ESM-specific, not needed for MLM) + if name.startswith("contact_head."): + return None + + # Map BERT cls.predictions.* -> lm_head.* + for bert_prefix, our_prefix in _BERT_CLS_TO_LM_HEAD.items(): + if name.startswith(bert_prefix): + name = our_prefix + name[len(bert_prefix) :] + return name + + # lm_head.* passes through directly (ESM/RoBERTa convention) + if name.startswith("lm_head."): + return name + + # Skip other cls.* heads (e.g. cls.seq_relationship for NSP) + if name.startswith("cls."): + return None + + # Encoder weights: delegate to shared rename logic. + # Prefix is already stripped, so pass directly to the encoder + # rename which handles HF naming collapse + gamma/beta compat. + return _rename_bert_encoder_weight(name) diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 6c91764e..beb8bfa2 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -27,6 +27,7 @@ "DenoisingTask", "FeatureExtractionTask", "HybridCausalLMTask", + "MaskedLMTask", "HybridQwenVLTask", "ImageClassificationTask", "ModelTask", @@ -64,6 +65,7 @@ from mobius.tasks._denoising import DenoisingTask from mobius.tasks._feature_extraction import FeatureExtractionTask from mobius.tasks._image_classification import ImageClassificationTask +from mobius.tasks._masked_lm import MaskedLMTask from mobius.tasks._multimodal import MultiModalTask from mobius.tasks._object_detection import ObjectDetectionTask from mobius.tasks._phi4mm_multimodal import Phi4MMMultiModalTask @@ -94,6 +96,7 @@ "controlnet": ControlNetTask, "denoising": DenoisingTask, "feature-extraction": FeatureExtractionTask, + "masked-lm": MaskedLMTask, "image-classification": ImageClassificationTask, "object-detection": ObjectDetectionTask, "seq2seq": Seq2SeqTask, diff --git a/src/mobius/tasks/_masked_lm.py b/src/mobius/tasks/_masked_lm.py new file mode 100644 index 00000000..91e03531 --- /dev/null +++ b/src/mobius/tasks/_masked_lm.py @@ -0,0 +1,68 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Masked language modeling task for encoder-only models (BERT, ESM, RoBERTa, etc.).""" + +from __future__ import annotations + +import onnx_ir as ir +from onnxscript import nn + +from mobius._configs import BaseModelConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ModelTask, _make_graph, _make_model + + +class MaskedLMTask(ModelTask): + """Encoder-only masked language modeling (predict masked tokens). + + The module must accept ``(op, input_ids, attention_mask, token_type_ids)`` + and return ``logits`` with shape ``[batch, sequence_len, vocab_size]``. + + Inputs: + - input_ids: [batch, sequence_len] INT64 + - attention_mask: [batch, sequence_len] INT64 + - token_type_ids: [batch, sequence_len] INT64 (optional, for BERT) + + Outputs: + - logits: [batch, sequence_len, vocab_size] FLOAT + """ + + def build( + self, + module: nn.Module, + config: BaseModelConfig, + ) -> ModelPackage: + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + + input_ids = ir.Value( + name="input_ids", + shape=ir.Shape([batch, seq_len]), + type=ir.TensorType(ir.DataType.INT64), + ) + attention_mask = ir.Value( + name="attention_mask", + shape=ir.Shape([batch, seq_len]), + type=ir.TensorType(ir.DataType.INT64), + ) + token_type_ids = ir.Value( + name="token_type_ids", + shape=ir.Shape([batch, seq_len]), + type=ir.TensorType(ir.DataType.INT64), + ) + + graph, builder = _make_graph([input_ids, attention_mask, token_type_ids]) + op = builder.op + + logits = module( + op, + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + ) + + logits.name = "logits" + graph.outputs.append(logits) + + return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/testdata/cases/encoder/esm2-8m.yaml b/testdata/cases/encoder/esm2-8m.yaml index 6b743a44..ed93c831 100644 --- a/testdata/cases/encoder/esm2-8m.yaml +++ b/testdata/cases/encoder/esm2-8m.yaml @@ -1,5 +1,5 @@ model_id: "facebook/esm2_t6_8M_UR50D" -revision: "main" +revision: "c731040fcd8d73dceaa04b0a8e6329b345b0f5df" task_type: "feature-extraction" dtype: "float32" diff --git a/testdata/cases/masked-lm/bert-base-uncased.yaml b/testdata/cases/masked-lm/bert-base-uncased.yaml new file mode 100644 index 00000000..cd87050a --- /dev/null +++ b/testdata/cases/masked-lm/bert-base-uncased.yaml @@ -0,0 +1,12 @@ +model_id: "google-bert/bert-base-uncased" +revision: "86b5e0934494bd15c9632b12f734a8a67f723594" +task_type: "masked-lm" +dtype: "float32" + +inputs: + prompts: + - "The quick brown [MASK] jumps over the lazy dog." + +level: "L4" + +notes: "BERT base uncased with masked LM head. Predicts masked tokens." diff --git a/testdata/cases/masked-lm/esm2-8m.yaml b/testdata/cases/masked-lm/esm2-8m.yaml new file mode 100644 index 00000000..b284fc9b --- /dev/null +++ b/testdata/cases/masked-lm/esm2-8m.yaml @@ -0,0 +1,12 @@ +model_id: "facebook/esm2_t6_8M_UR50D" +revision: "c731040fcd8d73dceaa04b0a8e6329b345b0f5df" +task_type: "masked-lm" +dtype: "float32" + +inputs: + prompts: + - "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG" + +level: "L4" + +notes: "ESM-2 8M with masked LM head. Protein language model for predicting masked amino acids." diff --git a/testdata/cases/masked-lm/roberta-base.yaml b/testdata/cases/masked-lm/roberta-base.yaml new file mode 100644 index 00000000..78a15421 --- /dev/null +++ b/testdata/cases/masked-lm/roberta-base.yaml @@ -0,0 +1,12 @@ +model_id: "FacebookAI/roberta-base" +revision: "e2da8e2f811d1448a5b465c236feacd80ffbac7b" +task_type: "masked-lm" +dtype: "float32" + +inputs: + prompts: + - "The quick brown jumps over the lazy dog." + +level: "L4" + +notes: "RoBERTa base with masked LM head. Uses token for predictions." diff --git a/testdata/cases/schema.json b/testdata/cases/schema.json index 26cdce45..a5d65134 100644 --- a/testdata/cases/schema.json +++ b/testdata/cases/schema.json @@ -44,6 +44,7 @@ "image-text-to-text", "image-to-image", "image-to-text", + "masked-lm", "object-detection", "seq2seq", "speech-language", diff --git a/testdata/golden/encoder/esm2-8m.json b/testdata/golden/encoder/esm2-8m.json new file mode 100644 index 00000000..11424bc9 --- /dev/null +++ b/testdata/golden/encoder/esm2-8m.json @@ -0,0 +1,103 @@ +{ + "top1_id": 277, + "top2_id": 170, + "top10_ids": [ + 277, + 170, + 199, + 233, + 236, + 105, + 223, + 10, + 171, + 240 + ], + "top10_logits": [ + "0x1.f89c6c0000000p-1", + "0x1.a22c5e0000000p-1", + "0x1.8bd40a0000000p-1", + "0x1.780c840000000p-1", + "0x1.7440560000000p-1", + "0x1.6735380000000p-1", + "0x1.4d31a80000000p-1", + "0x1.4d17420000000p-1", + "0x1.48a4300000000p-1", + "0x1.3c19140000000p-1" + ], + "logits_summary": [ + "0x1.f89c6c0000000p-1", + "-0x1.ec55a80000000p+1", + "-0x1.6f2c13bd9999ap-7", + "0x1.abcafa4ea9449p-2" + ], + "input_ids": [ + 0, + 20, + 15, + 11, + 7, + 10, + 16, + 9, + 10, + 4, + 15, + 8, + 12, + 7, + 10, + 12, + 4, + 9, + 10, + 8, + 15, + 9, + 14, + 7, + 8, + 6, + 5, + 16, + 4, + 5, + 9, + 9, + 4, + 8, + 7, + 8, + 10, + 16, + 7, + 12, + 7, + 16, + 13, + 12, + 5, + 19, + 4, + 10, + 8, + 4, + 6, + 19, + 17, + 12, + 7, + 5, + 11, + 14, + 10, + 6, + 19, + 7, + 4, + 5, + 6, + 6, + 2 + ] +} diff --git a/testdata/golden/masked-lm/bert-base-uncased.json b/testdata/golden/masked-lm/bert-base-uncased.json new file mode 100644 index 00000000..e30cb171 --- /dev/null +++ b/testdata/golden/masked-lm/bert-base-uncased.json @@ -0,0 +1,48 @@ +{ + "top1_id": 4937, + "top2_id": 3899, + "top10_ids": [ + 4937, + 3899, + 4562, + 2158, + 2666, + 2879, + 2028, + 2611, + 4702, + 3124 + ], + "top10_logits": [ + "0x1.3f00520000000p+3", + "0x1.28a7340000000p+3", + "0x1.1c69740000000p+3", + "0x1.17f0f00000000p+3", + "0x1.10a9640000000p+3", + "0x1.0da2a40000000p+3", + "0x1.015a900000000p+3", + "0x1.ea17740000000p+2", + "0x1.e409b60000000p+2", + "0x1.e0c91e0000000p+2" + ], + "logits_summary": [ + "0x1.3f00520000000p+3", + "-0x1.ec9bb60000000p+3", + "-0x1.23a32c11124d2p+2", + "0x1.426564917a4f0p+1" + ], + "input_ids": [ + 101, + 1996, + 4248, + 2829, + 103, + 14523, + 2058, + 1996, + 13971, + 3899, + 1012, + 102 + ] +} diff --git a/testdata/golden/masked-lm/esm2-8m.json b/testdata/golden/masked-lm/esm2-8m.json new file mode 100644 index 00000000..732a0700 --- /dev/null +++ b/testdata/golden/masked-lm/esm2-8m.json @@ -0,0 +1,103 @@ +{ + "top1_id": 20, + "top2_id": 15, + "top10_ids": [ + 20, + 15, + 7, + 10, + 4, + 6, + 14, + 5, + 8, + 12 + ], + "top10_logits": [ + "0x1.2e937e0000000p+3", + "0x1.75a1cc0000000p-2", + "0x1.6a99960000000p-2", + "0x1.441ec00000000p-2", + "0x1.c072400000000p-3", + "-0x1.d8add40000000p-7", + "-0x1.4dcdf80000000p-4", + "-0x1.e15ae80000000p-4", + "-0x1.1ba6180000000p-3", + "-0x1.a1bf8a0000000p-2" + ], + "logits_summary": [ + "0x1.2e937e0000000p+3", + "-0x1.fd60940000000p+3", + "-0x1.3ae86a35a2e8cp+2", + "0x1.aedc3dfcb9dbbp+2" + ], + "input_ids": [ + 0, + 20, + 15, + 11, + 7, + 10, + 16, + 9, + 10, + 4, + 15, + 8, + 12, + 7, + 10, + 12, + 4, + 9, + 10, + 8, + 15, + 9, + 14, + 7, + 8, + 6, + 5, + 16, + 4, + 5, + 9, + 9, + 4, + 8, + 7, + 8, + 10, + 16, + 7, + 12, + 7, + 16, + 13, + 12, + 5, + 19, + 4, + 10, + 8, + 4, + 6, + 19, + 17, + 12, + 7, + 5, + 11, + 14, + 10, + 6, + 19, + 7, + 4, + 5, + 6, + 6, + 2 + ] +} diff --git a/testdata/golden/masked-lm/roberta-base.json b/testdata/golden/masked-lm/roberta-base.json new file mode 100644 index 00000000..fd47a654 --- /dev/null +++ b/testdata/golden/masked-lm/roberta-base.json @@ -0,0 +1,48 @@ +{ + "top1_id": 324, + "top2_id": 4758, + "top10_ids": [ + 324, + 4758, + 23602, + 2173, + 2335, + 219, + 2143, + 19921, + 28133, + 26666 + ], + "top10_logits": [ + "0x1.ea240c0000000p+3", + "0x1.c60af00000000p+3", + "0x1.bd73d00000000p+3", + "0x1.b3ec280000000p+3", + "0x1.b02e300000000p+3", + "0x1.aa97d40000000p+3", + "0x1.9ec0b80000000p+3", + "0x1.9462de0000000p+3", + "0x1.903d360000000p+3", + "0x1.8e63320000000p+3" + ], + "logits_summary": [ + "0x1.ea240c0000000p+3", + "-0x1.2236dc0000000p+4", + "-0x1.6f00602fea530p+1", + "0x1.08999b5b0b62bp+2" + ], + "input_ids": [ + 0, + 133, + 2119, + 6219, + 50264, + 13855, + 81, + 5, + 22414, + 2335, + 4, + 2 + ] +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index d2fe8354..f6f34ab8 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -1233,6 +1233,16 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ] +# --------------------------------------------------------------------------- +# Masked LM configs (task: masked-lm, module: BertForMaskedLM) +# --------------------------------------------------------------------------- +MASKED_LM_CONFIGS: list[tuple[str, dict, bool]] = [ + ("bert_masked_lm", {"hidden_act": "gelu", "type_vocab_size": 2}, True), + ("esm_masked_lm", {"hidden_act": "gelu", "type_vocab_size": 2}, True), + ("roberta_masked_lm", {"hidden_act": "gelu", "type_vocab_size": 1}, False), +] + + # --------------------------------------------------------------------------- # Seq2Seq (encoder-decoder) configs (task: seq2seq) # --------------------------------------------------------------------------- @@ -2091,6 +2101,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ALL_CONFIGS: list[tuple[str, dict, bool]] = ( CAUSAL_LM_CONFIGS + ENCODER_CONFIGS + + MASKED_LM_CONFIGS + SEQ2SEQ_CONFIGS + VISION_CONFIGS + DETECTION_CONFIGS diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index c9c421e3..64b05d54 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -26,6 +26,7 @@ DETECTION_CONFIGS, ENCODER_CONFIGS, LONGROPE_FACTORS, + MASKED_LM_CONFIGS, SEQ2SEQ_CONFIGS, SPEECH_CONFIGS, SSM_CONFIGS, @@ -393,6 +394,68 @@ def test_outputs_have_shapes_and_dtypes(self, model_type: str, config_overrides: _assert_outputs_have_shapes_and_dtypes(pkg, model_type) +# === Masked LM model configs (imported from _test_configs) === +_MASKED_LM_MODEL_PARAMS = _make_params(MASKED_LM_CONFIGS) + + +@pytest.mark.parametrize("model_type,config_overrides", _MASKED_LM_MODEL_PARAMS) +class TestBuildMaskedLMGraph: + """Verify that encoder-only models build valid masked LM ONNX graphs.""" + + def test_graph_builds_without_weights(self, model_type: str, config_overrides: dict): + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + model = pkg["model"] + + assert model.graph is not None + input_names = {inp.name for inp in model.graph.inputs} + assert "input_ids" in input_names + assert "attention_mask" in input_names + assert "token_type_ids" in input_names + + output_names = {out.name for out in model.graph.outputs} + assert "logits" in output_names + # No KV cache outputs for masked LM + assert not any(n.startswith("present.") for n in output_names) + # No hidden_state output (that's feature-extraction, not masked-lm) + assert "last_hidden_state" not in output_names + + def test_graph_has_lm_head_initializers(self, model_type: str, config_overrides: dict): + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + model = pkg["model"] + + init_names = list(model.graph.initializers) + assert len(init_names) > 0 + has_lm_head = any("lm_head" in n for n in init_names) + has_encoder = any("encoder" in n for n in init_names) + assert has_lm_head, "Should have LM head parameters" + assert has_encoder, "Should have encoder parameters" + + def test_onnx_checker_passes(self, model_type: str, config_overrides: dict): + """Run the ONNX CheckerPass to catch attribute/shape/type errors.""" + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + _run_onnx_checker(pkg, model_type) + + def test_outputs_have_shapes_and_dtypes(self, model_type: str, config_overrides: dict): + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + _assert_outputs_have_shapes_and_dtypes(pkg, model_type) + + # === Encoder-decoder model configs (imported from _test_configs) === _SEQ2SEQ_MODEL_CONFIGS: list[tuple[str, dict]] = [(mt, ov) for mt, ov, _ in SEQ2SEQ_CONFIGS] diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 559d5ab5..880bea71 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -59,6 +59,7 @@ ALL_CAUSAL_LM_CONFIGS, DETECTION_CONFIGS, ENCODER_CONFIGS, + MASKED_LM_CONFIGS, SEQ2SEQ_CONFIGS, SPEECH_CONFIGS, SSM_CONFIGS, @@ -77,6 +78,7 @@ def _l1_l3_model_types() -> set[str]: all_configs = ( ALL_CAUSAL_LM_CONFIGS + ENCODER_CONFIGS + + MASKED_LM_CONFIGS + SEQ2SEQ_CONFIGS + VISION_CONFIGS + DETECTION_CONFIGS