bridge.models.bert.bert_bridge#

Module Contents#

Classes#

BertBridge

Megatron Bridge for HuggingFace MegatronBertForMaskedLM.

API#

class bridge.models.bert.bert_bridge.BertBridge#

Bases: megatron.bridge.models.conversion.model_bridge.MegatronModelBridge

Megatron Bridge for HuggingFace MegatronBertForMaskedLM.

Targets Megatron-Core’s stock, Pre-LayerNorm BertModel. HuggingFace’s vanilla transformers.BertModel (e.g. bert-base-uncased) uses Post-LayerNorm and is a different architecture that this bridge does not support – MegatronBertForMaskedLM (e.g. nvidia/megatron-bert-uncased-345m) is the HuggingFace architecture that was purpose-built to mirror this Pre-LN layout.

As a user you would not use this bridge directly, but through AutoBridge.

.. rubric:: Example

from megatron.bridge import AutoBridge bridge = AutoBridge.from_hf_pretrained(“nvidia/megatron-bert-uncased-345m”) provider = bridge.to_megatron_provider()

static _validate_hf_config(hf_config: object) → None#
provider_bridge(
hf_pretrained: megatron.bridge.models.hf_pretrained.masked_lm.PreTrainedMaskedLM,
) → megatron.bridge.models.bert.bert_provider.BertModelProvider#
classmethod megatron_to_hf_config(
provider: megatron.bridge.models.bert.bert_provider.BertModelProvider,
) → dict[str, object]#

Convert a BertModelProvider back into a MegatronBertConfig-shaped dict.

Not routed through the shared GPT-oriented CONFIG_MAPPING machinery (e.g. it maps max_position_embeddings to seq_length and rms_norm_eps to layernorm_epsilon, neither of which apply to BERT) – BERT’s config fields are mapped directly instead.

mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry#
maybe_modify_converted_hf_weight(
task: megatron.bridge.models.conversion.model_bridge.WeightConversionTask,
converted_weights_dict: dict[str, torch.Tensor],
hf_state_dict: collections.abc.Mapping[str, torch.Tensor],
) → dict[str, torch.Tensor]#

Duplicate the tied MLM-head bias into HF’s second key on export.

MegatronBertLMPredictionHead ties cls.predictions.decoder.bias to cls.predictions.bias (self.decoder.bias = self.bias, the same underlying nn.Parameter), so both keys appear in HF’s state dict. Megatron only has a single output_layer.bias, so the second HF key must be synthesized here.