feat(model): add bert support. (#4737)

Signed-off-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Co-authored-by: yaoyu-33 <yaoyu.094@gmail.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Chen Cui <chcui@nvidia.com>

Author: 4 people

README.md

@@ -230,6 +230,7 @@ Megatron Bridge provides out-of-the-box bridges and training recipes for a wide
 | Family | Supported variants |
 |----------------|--------------------|
 | [**Bailing**](docs/models/bailing/index.md) | Ling 2.0 / Ling MoE V2 (Bailing) |
+| [**BERT**](docs/models/bert/index.md) | Megatron-Style BERT (`MegatronBertForMaskedLM`) |
 | [**DeepSeek**](docs/models/deepseek/index.md) | DeepSeek V2 / V2 Lite (deprecated), DeepSeek V3, DeepSeek V4 / V4 Flash |
 | [**Diffusion**](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/diffusion/models) | FLUX, LLaDA 1.5, Nemotron-Labs Diffusion, WAN |
 | **Ernie** | [Ernie 4.5 MoE](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/ernie), [Ernie 4.5 VL MoE](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/ernie_vl) |

docs/index.md

@@ -36,6 +36,7 @@ bridge-tech-details.md
 models/README.md
 models/bagel/bagel.md
 models/bailing/index.md
+models/bert/index.md
 models/deepseek/index.md
 models/ernie/ernie45.md
 models/exaone/exaone.md

docs/models/bert/index.md

@@ -0,0 +1,13 @@
+# BERT
+
+BERT model documentation is organized by model variant.
+
+```{toctree}
+:hidden:
+
+megatron-bert.md
+```
+
+| Variant | Guide |
+|---------|-------|
+| Megatron-Style BERT | [megatron-bert.md](megatron-bert.md) |

docs/models/bert/megatron-bert.md

@@ -0,0 +1,58 @@
+# Megatron-Style BERT
+
+[Megatron BERT](https://huggingface.co/nvidia/megatron-bert-uncased-345m) is NVIDIA's original
+Pre-LayerNorm BERT-family encoder, published to Hugging Face as `MegatronBertForMaskedLM`
+(`transformers` model type `megatron-bert`). Megatron Bridge supports it through `BertBridge`,
+the first encoder-only (masked-LM) model bridge in this repository.
+
+## Supported Variants
+
+The public examples target NVIDIA's original 345M-parameter checkpoint:
+
+- `nvidia/megatron-bert-uncased-345m`: https://huggingface.co/nvidia/megatron-bert-uncased-345m
+- `nvidia/megatron-bert-cased-345m`: https://huggingface.co/nvidia/megatron-bert-cased-345m
+
+The bridge is registered for the `MegatronBertForMaskedLM` architecture and can auto-detect
+compatible checkpoints when their Hugging Face config uses the `megatron-bert` model type.
+
+## Architecture Notes: Pre-LayerNorm Only
+
+`BertBridge` targets **Pre-LayerNorm** BERT, matching Megatron-Core's
+`megatron.core.models.bert.bert_model.BertModel` exactly (LayerNorm is applied before
+self-attention/MLP, and the residual carries the un-normalized input).
+
+> [!WARNING]
+> Vanilla `transformers.BertModel`-based checkpoints (`bert-base-uncased`, `bert-large-uncased`,
+> most BioBERT/SciBERT/etc. derivatives) are **Post-LayerNorm** — architecturally different from
+> Megatron-Core's `BertModel`. This bridge does **not** support them: forcing a Post-LN checkpoint
+> through it would silently load weights into the wrong computational graph and produce incorrect
+> outputs with no error raised. Check that a checkpoint's `model_type` is `megatron-bert`, not
+> `bert`, before converting.
+
+Other implementation notes:
+
+- `add_binary_head=False` — masked-LM checkpoints have no pooler/NSP head.
+- `hidden_act="gelu"` is required because Megatron-Core's BERT masked-LM head hardcodes GELU.
+- Decoder and cross-attention configurations are not supported.
+- Untied word embeddings are not supported because Hugging Face represents two independent MLM
+  biases in that mode, while Megatron-Core has a single output-layer bias.
+- HF ties `cls.predictions.decoder.bias` to `cls.predictions.bias` internally. The bridge
+  synthesizes the alias when the export target expects both keys; safe serialization normally
+  stores only the canonical `cls.predictions.bias` key.
+- Megatron-Core's `BertModel.forward()` requires `tokentype_ids` to be passed explicitly whenever
+  `num_tokentypes > 0`, and expects a raw (non-extended) `[batch, seq_len]` attention mask — it
+  builds the extended mask internally.
+
+## Examples
+
+For checkpoint conversion, round-trip validation, and fill-mask inference, see the
+[Megatron-Style BERT examples README](https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/examples/models/bert/megatron-bert/README.md).
+
+Note that `nvidia/megatron-bert-uncased-345m` only hosts tokenizer assets on the Hub; the weights
+must first be downloaded from NGC and converted with Hugging Face's own conversion script (see
+the examples README for the exact commands) before they can be imported with Megatron Bridge.
+
+## Related Implementation
+
+- Bridge implementation: [`src/megatron/bridge/models/bert`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/src/megatron/bridge/models/bert)
+- Examples: [`examples/models/bert/megatron-bert`](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/models/bert/megatron-bert)

examples/models/bert/megatron-bert/README.md

@@ -0,0 +1,91 @@
+# Megatron-Style BERT (`MegatronBertForMaskedLM`) Examples
+
+This directory contains example scripts for the Megatron-Bridge `BertBridge`, which converts
+between Hugging Face's `MegatronBertForMaskedLM` architecture and Megatron-Core's
+`megatron.core.models.bert.bert_model.BertModel`.
+
+The examples target `nvidia/megatron-bert-uncased-345m` (24 layers, 1024 hidden size, 345M
+parameters), NVIDIA's original Pre-LayerNorm BERT checkpoint.
+
+## Supported architecture: Pre-LayerNorm only
+
+`BertBridge` targets HF's **`MegatronBertForMaskedLM`** architecture (`transformers` model type
+`megatron-bert`), which is Pre-LayerNorm — LayerNorm is applied *before* self-attention and the
+MLP, matching Megatron-Core's `BertModel` exactly.
+
+> [!WARNING]
+> This bridge does **not** support vanilla `transformers.BertModel`-based checkpoints (e.g.
+> `bert-base-uncased`, `bert-large-uncased`, most BioBERT/SciBERT/etc. variants). Those are
+> **Post-LayerNorm** (LayerNorm applied after the residual add, as in the original BERT paper),
+> which is architecturally different from Megatron-Core's `BertModel`. Converting a Post-LN
+> checkpoint through this bridge would silently load weights into the wrong computational graph
+> and produce incorrect outputs with no error. If a HF config's `model_type` is `bert` rather than
+> `megatron-bert`, it is not supported here.
+
+`MegatronBertForMaskedLM` also has no pooler/NSP head in this bridge (`add_binary_head=False`),
+matching how the masked-LM checkpoints are published.
+
+The bridge requires `hidden_act="gelu"` because Megatron-Core's BERT masked-LM head hardcodes
+GELU. Decoder and cross-attention configurations are not supported.
+
+## Obtaining a convertible checkpoint
+
+The `nvidia/megatron-bert-uncased-345m` Hugging Face repo only hosts tokenizer files — the
+weights must be downloaded from NVIDIA GPU Cloud (NGC) and converted to HF format first, using
+HF's own conversion script (this is a one-time step, unrelated to Megatron-Bridge):
+
+    mkdir -p megatron-bert-uncased-345m-hf
+    wget --content-disposition \
+      https://api.ngc.nvidia.com/v2/models/nvidia/megatron_bert_345m/versions/v0.1_uncased/zip \
+      -O megatron-bert-uncased-345m-hf/checkpoint.zip
+    git clone https://github.com/huggingface/transformers.git
+    git -C transformers checkout fe747d88a3296bd94d426db2717f232f9d4afdb7
+    uv run python transformers/src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py \
+      megatron-bert-uncased-345m-hf/checkpoint.zip
+
+The upstream converter writes a pretraining checkpoint without the architecture metadata needed
+by `AutoBridge`. Normalize it to the masked-LM architecture and add the tokenizer assets before
+using the directory with Megatron Bridge:
+
+    uv run python - <<'PY'
+    from transformers import AutoTokenizer, MegatronBertForMaskedLM
+
+    checkpoint_path = "megatron-bert-uncased-345m-hf"
+    model = MegatronBertForMaskedLM.from_pretrained(checkpoint_path)
+    model.save_pretrained(checkpoint_path, safe_serialization=True)
+    tokenizer = AutoTokenizer.from_pretrained("nvidia/megatron-bert-uncased-345m")
+    tokenizer.save_pretrained(checkpoint_path)
+    PY
+
+This produces a self-contained config, weights, and tokenizer directory that can be used as
+`HF_MODEL_PATH` below. See
+the [`nvidia/megatron-bert-uncased-345m` model card](https://huggingface.co/nvidia/megatron-bert-uncased-345m)
+for full details, including the `-cased` variant.
+
+## Workspace Configuration
+
+All scripts use a `WORKSPACE` environment variable to define the base directory for checkpoints
+and results. By default, this is set to `/workspace`.
+
+    export WORKSPACE=/your/custom/path
+    export HF_MODEL_PATH=/path/to/converted/megatron-bert-uncased-345m-hf
+
+## Checkpoint Conversion
+
+See [conversion.sh](conversion.sh) for Hugging Face to Megatron import, Megatron to Hugging Face
+export, and round-trip validation.
+
+    ./examples/models/bert/megatron-bert/conversion.sh
+
+## Inference (fill-mask)
+
+`MegatronBertForMaskedLM` has no `generate()` method, so "inference" here means masked-token
+prediction rather than text generation. See [inference.sh](inference.sh), which runs
+[fill_mask.py](fill_mask.py) against the original Hugging Face checkpoint and, if present, the
+round-tripped export from `conversion.sh`:
+
+    TEXT="Paris is the [MASK] of France." ./examples/models/bert/megatron-bert/inference.sh
+
+`fill_mask.py` uses the Hugging Face `fill-mask` pipeline directly (no Megatron-Core forward
+pass) — its purpose is to sanity-check that a converted checkpoint still produces sensible
+predictions, not to benchmark Megatron inference throughput.

examples/models/bert/megatron-bert/conversion.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -xeuo pipefail
+
+WORKSPACE=${WORKSPACE:-/workspace}
+
+MODEL_NAME=megatron-bert-uncased-345m
+# `HF_MODEL_PATH` must point at a *converted* Hugging Face checkpoint directory
+# (config.json + weights), not the tokenizer-only `nvidia/megatron-bert-uncased-345m`
+# Hub repo. See README.md for how to obtain one.
+HF_MODEL_PATH=${HF_MODEL_PATH:-"${WORKSPACE}/models/${MODEL_NAME}-hf"}
+
+uv run python examples/conversion/convert_checkpoints.py import \
+    --hf-model "$HF_MODEL_PATH" \
+    --megatron-path "${WORKSPACE}/models/${MODEL_NAME}"
+
+uv run python examples/conversion/convert_checkpoints.py export \
+    --hf-model "$HF_MODEL_PATH" \
+    --megatron-path "${WORKSPACE}/models/${MODEL_NAME}/iter_0000000" \
+    --hf-path "${WORKSPACE}/models/${MODEL_NAME}-hf-export"
+
+uv run python -m torch.distributed.run --nproc_per_node=1 \
+    examples/conversion/hf_megatron_roundtrip_multi_gpu.py \
+    --hf-model-id "$HF_MODEL_PATH" \
+    --megatron-load-path "${WORKSPACE}/models/${MODEL_NAME}/iter_0000000" \
+    --tp 1 --pp 1

examples/models/bert/megatron-bert/fill_mask.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Run masked-token ("fill-mask") inference with a Megatron-Bert-family checkpoint.
+
+Unlike the causal-LM examples in `examples/conversion/`, `MegatronBertForMaskedLM`
+does not implement `generate()`, so this example uses the Hugging Face `fill-mask`
+pipeline directly. This is intentionally HF-only (no Megatron-Core forward pass):
+its purpose is to sanity-check that a checkpoint produced by the Megatron-Bridge
+BERT bridge (see `conversion.sh`) still makes sensible predictions after a
+round trip, not to benchmark Megatron inference.
+
+Run Script Example:
+    uv run python examples/models/bert/megatron-bert/fill_mask.py \\
+        --hf_model_path /workspace/models/megatron-bert-uncased-345m-hf-export \\
+        --text "Paris is the [MASK] of France."
+"""
+
+import argparse
+
+from transformers import pipeline
+
+
+def main() -> None:
+    """Parse CLI arguments and print the top masked-token predictions."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--hf_model_path", type=str, required=True, help="HuggingFace model ID or local path")
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="Paris is the [MASK] of France.",
+        help="Input text containing exactly one `[MASK]` token",
+    )
+    parser.add_argument("--top_k", type=int, default=5, help="Number of predictions to print")
+    args = parser.parse_args()
+
+    fill_mask = pipeline("fill-mask", model=args.hf_model_path)
+    predictions = fill_mask(args.text, top_k=args.top_k)
+
+    print(f"Model: {args.hf_model_path}")
+    print(f"Input: {args.text}")
+    for prediction in predictions:
+        print(f"  {prediction['score']:.4f}  {prediction['sequence']}")
+
+
+if __name__ == "__main__":
+    main()

examples/models/bert/megatron-bert/inference.sh

@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -xeuo pipefail
+
+WORKSPACE=${WORKSPACE:-/workspace}
+MODEL_NAME=megatron-bert-uncased-345m
+HF_MODEL_PATH=${HF_MODEL_PATH:-"${WORKSPACE}/models/${MODEL_NAME}-hf"}
+HF_EXPORT_PATH=${HF_EXPORT_PATH:-"${WORKSPACE}/models/${MODEL_NAME}-hf-export"}
+TEXT=${TEXT:-"Paris is the [MASK] of France."}
+
+# `MegatronBertForMaskedLM` has no `generate()` method, so inference here means
+# fill-mask prediction, not text generation. This only exercises Hugging Face
+# code paths, so it can be used to sanity-check both the original checkpoint
+# and the round-tripped export from conversion.sh.
+uv run python examples/models/bert/megatron-bert/fill_mask.py \
+    --hf_model_path "$HF_MODEL_PATH" \
+    --text "$TEXT"
+
+if [ -d "$HF_EXPORT_PATH" ]; then
+    uv run python examples/models/bert/megatron-bert/fill_mask.py \
+        --hf_model_path "$HF_EXPORT_PATH" \
+        --text "$TEXT"
+fi

src/megatron/bridge/models/__init__.py

@@ -28,6 +28,9 @@
 from megatron.bridge.models.bailing import (
     BailingMoeV2Bridge,
 )
+from megatron.bridge.models.bert import (
+    BertModelProvider,
+)
 from megatron.bridge.models.conversion.auto_bridge import AutoBridge
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
@@ -231,6 +234,7 @@
     "RowParallelMapping",
     "AutoMapping",
     "BailingMoeV2Bridge",
+    "BertModelProvider",
     # DeepSeek Models
     "DeepSeekV2Bridge",
     "DeepSeekV3Bridge",

src/megatron/bridge/models/bert/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from megatron.bridge.models.bert.bert_bridge import BertBridge  # noqa: F401
+from megatron.bridge.models.bert.bert_provider import BertModelProvider
+
+
+__all__ = [
+    "BertModelProvider",
+]