class_docstr: Masked LM head for Bert
  
  Arguments:
      config: TransformerConfig object
      mpu_vocab_size: model parallel size of vocabulary.
      parallel_output: whether output logits being distributed or not.
class_docstr_tok: ["Masked","LM","head","for","Bert","Arguments",":","config",":","TransformerConfig","object","mpu_vocab_size",":","model","parallel","size","of","vocabulary",".","parallel_output",":","whether","output","logits","being","distributed","or","not","."]
code: def __init__(self, mpu_vocab_size, config, parallel_output):
      super().__init__(config=config)
  
      args = get_args()
      self.bias = torch.nn.Parameter(torch.zeros(mpu_vocab_size))
      tensor_parallel.set_tensor_model_parallel_attributes(
          self.bias, True, 0, 1
      )
      self.parallel_output = parallel_output
  
      self.dense = get_linear_layer(
          config.hidden_size, config.hidden_size, config.init_method
      )
      setattr(
          self.dense.weight,
          "sequence_parallel",
          config.sequence_parallel,
      )
      setattr(
          self.dense.bias, "sequence_parallel", config.sequence_parallel
      )
  
      self.norm = get_norm(config)
      self.gelu = torch.nn.functional.gelu
      if args.openai_gelu:
          self.gelu = openai_gelu
      elif args.onnx_safe:
          self.gelu = erf_gelu
  
imports: import torch
  from megatron import get_args
  from megatron.core import tensor_parallel
  from megatron.model.enums import AttnMaskType
  from megatron.model.language_model import parallel_lm_logits
  from megatron.model.language_model import get_language_model
  from megatron.model.utils import get_norm
  from megatron.model.utils import openai_gelu, erf_gelu
  from megatron.model.utils import get_linear_layer
  from megatron.model.utils import init_method_normal
  from megatron.model.utils import scaled_init_method_normal
  from .module import MegatronModule
label_desc: Use image node_id 1 to create a new BertLMHead object from inherited base classes: MegatronModule with example: obj = BertLMHead(mpu_vocab_size, config, parallel_output)