bridge.models.bert.bert_provider

Module Contents

Classes

BertModelProvider

Model provider for Megatron-Core’s stock encoder-only BERT model.

Data

logger

API

bridge.models.bert.bert_provider.logger

‘getLogger(…)’

class bridge.models.bert.bert_provider.BertModelProvider

Bases: megatron.bridge.models.transformer_config.TransformerConfig, megatron.bridge.models.model_provider.ModelProviderMixin[megatron.core.models.bert.bert_model.BertModel]

Model provider for Megatron-Core’s stock encoder-only BERT model.

Note this targets Megatron-Core’s BertModel, which applies LayerNorm before self-attention/MLP (Pre-LayerNorm) rather than the original BERT paper’s Post-LayerNorm placement used by HuggingFace’s vanilla transformers.BertModel. The HuggingFace architecture that matches this provider exactly is MegatronBertForMaskedLM (see bert_bridge.py).

add_binary_head: bool

False

num_tokentypes: int

2

fp16_lm_cross_entropy: bool

False

parallel_output: bool

True

share_embeddings_and_output_weights: bool

True

position_embedding_type: Literal[learned_absolute, rope]

‘learned_absolute’

rotary_percent: float

1.0

seq_len_interpolation_factor: float | None

None

return_embeddings: bool

False

make_vocab_size_divisible_by: int

128

max_position_embeddings: int

512

vocab_size: int | None

None

should_pad_vocab: bool

False

transformer_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec | Callable[[], megatron.core.transformer.spec_utils.ModuleSpec]

None

_pg_collection: megatron.core.process_groups_config.ProcessGroupCollection | None

None

provide(

pre_process: bool | None = None,

post_process: bool | None = None,

vp_stage: int | None = None,

) → megatron.core.models.bert.bert_model.BertModel

Configure and instantiate a Megatron Core BertModel based on this configuration.

Parameters:

• pre_process – Whether to include the embedding layer.

• post_process – Whether to include the masked-LM head.

• vp_stage – Virtual pipeline stage.

Returns:

Configured Megatron Core BERT model.