class_docstr: Transformer language model.
  
  Args:
      config (TransformerConfig): transformer config
      num_tokentypes (int) : Set to 2 when args.bert_binary_head is True, and 0 otherwise. Defaults to 0.
      transformer_layer_spec (ModuleSpec): Specifies module to use for transformer layers
      vocab_size (int): vocabulary size
      max_sequence_length (int): maximum size of sequence. This is used for positional embedding
      pre_process (bool): Include embedding layer (used with pipeline parallelism)
      post_process (bool): Include an output layer (used with pipeline parallelism)
      parallel_output (bool): Do not gather the outputs, keep them split across tensor parallel ranks
      share_embeddings_and_output_weights (bool): When True, input embeddings and output logit weights are shared. Defaults to False.
      position_embedding_type (string): Position embedding type. Options ['learned_absolute', 'rope'].
          Defaults is 'learned_absolute'.
      rotary_percent (float): Percent of rotary dimension to use for rotary position embeddings.
          Defaults to 1.0 (100%). Ignored unless position_embedding_type is 'rope'.
code: def __init__(
      self,
      config: TransformerConfig,
      num_tokentypes: int,
      transformer_layer_spec: ModuleSpec,
      vocab_size: int,
      max_sequence_length: int,
      pre_process: bool = True,
      post_process: bool = True,
      fp16_lm_cross_entropy: bool = False,
      parallel_output: bool = True,
      share_embeddings_and_output_weights: bool = False,
      position_embedding_type: Literal[
          "learned_absolute", "rope"
      ] = "learned_absolute",
      rotary_percent: float = 1.0,
      seq_len_interpolation_factor: Optional[float] = None,
      add_binary_head=True,
      return_embeddings=False,
  ):
      super(BertModel, self).__init__(config=config)
  
      if return_embeddings:
          assert self.post_process and self.add_binary_head
  
      self.config: TransformerConfig = config
      self.transformer_layer_spec: ModuleSpec = transformer_layer_spec
      self.vocab_size = vocab_size
      self.max_sequence_length = max_sequence_length
      self.pre_process = pre_process
      self.post_process = post_process
      self.fp16_lm_cross_entropy = fp16_lm_cross_entropy
      self.parallel_output = parallel_output
      self.share_embeddings_and_output_weights = (
          share_embeddings_and_output_weights
      )
      self.position_embedding_type = position_embedding_type
      self.add_binary_head = add_binary_head
      self.return_embeddings = return_embeddings
  
      # megatron core pipelining currently depends on model type
      self.model_type = ModelType.encoder_or_decoder
  
      # Embeddings.
      if self.pre_process:
          self.embedding = LanguageModelEmbedding(
              config=self.config,
              vocab_size=self.vocab_size,
              max_sequence_length=self.max_sequence_length,
              position_embedding_type=position_embedding_type,
              num_tokentypes=num_tokentypes,
          )
  
      if self.position_embedding_type == "rope":
          self.rotary_pos_emb = RotaryEmbedding(
              self.config.kv_channels,
              rotary_percent,
              seq_len_interpolation_factor,
          )
  
      # Transformer.
      self.encoder = TransformerBlock(
          config=self.config,
          spec=self.transformer_layer_spec,
          pre_process=self.pre_process,
          post_process=self.post_process,
      )
  
      # Output
      if post_process:
          # TODO: Make sure you are passing in the mpu_vocab_size properly
          self.lm_head = BertLMHead(
              config.hidden_size,
              config,
              parallel_output,
              self.vocab_size,
              self.pre_process,
              self.share_embeddings_and_output_weights,
          )
  
          self.output_layer = self.lm_head.output_layer
  
          self.binary_head = None
          if self.add_binary_head:
              # TODO: Shoudl switch this to TE ?
              self.binary_head = get_linear_layer(
                  config.hidden_size,
                  2,
                  config.init_method,
                  config.perform_initialization,
              )
  
              self.pooler = Pooler(
                  config.hidden_size,
                  config.init_method,
                  config,
                  config.sequence_parallel,
              )
  
      if self.share_embeddings_and_output_weights and (
          self.pre_process or self.post_process
      ):
          self.initialize_last_stage_with_word_embeddings()
  
imports: from typing import Literal, Optional
  import torch
  from torch import Tensor
  from megatron.core.models.bert.bert_lm_head import BertLMHead
  from megatron.core.models.bert.pooler import Pooler
  from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
  from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding
  from megatron.core.models.common.language_module.language_module import LanguageModule
  from megatron.core.transformer.enums import AttnMaskType, ModelType
  from megatron.core.transformer.spec_utils import ModuleSpec
  from megatron.core.transformer.transformer_block import TransformerBlock
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.transformer.utils import get_linear_layer
label_desc: Use image node_id 1 to create a new BertModel object from inherited base classes: LanguageModule with example: obj = BertModel(config, num_tokentypes, transformer_layer_spec, vocab_size, max_sequence_length, pre_process, post_process, fp16_lm_cross_entropy, parallel_output, share_embeddings_and_output_weights, position_embedding_type, rotary_percent, seq_len_interpolation_factor, add_binary_head, return_embeddings)