class_docstr: Masked LM head for Bert
  
  Arguments:
      mpu_vocab_size: model parallel size of vocabulary.
      hidden_size: hidden size
      init_method: init method for weight initialization
      layernorm_epsilon: tolerance for layer norm divisions
      parallel_output: whether output logits being distributed or not.
code: def __init__(
      self,
      mpu_vocab_size,
      hidden_size,
      init_method,
      layernorm_epsilon,
      parallel_output,
  ):
      super(BertLMHead, self).__init__()
  
      args = get_args()
  
      self.bias = torch.nn.Parameter(torch.zeros(mpu_vocab_size))
      # TODO: do we need this?
      # mpu.set_tensor_model_parallel_attributes(self.bias, True, 0, 1)
      self.parallel_output = parallel_output
  
      self.dense = get_linear_layer(
          hidden_size, hidden_size, init_method
      )
      setattr(
          self.dense.weight, "sequence_parallel", args.sequence_parallel
      )
      setattr(
          self.dense.bias, "sequence_parallel", args.sequence_parallel
      )
  
      self.layernorm = LayerNorm(
          hidden_size,
          eps=layernorm_epsilon,
          sequence_parallel_enabled=args.sequence_parallel,
      )
      self.gelu = torch.nn.functional.gelu
      if args.openai_gelu:
          self.gelu = openai_gelu
      elif args.onnx_safe:
          self.gelu = erf_gelu
  
imports: import contextlib
  import torch
  from apex.transformer import tensor_parallel
  from apex.transformer.enums import AttnMaskType
  from apex.transformer.enums import ModelType
  from apex.transformer.layers import FusedLayerNorm
  from apex.transformer.testing.global_vars import get_args
  from apex.transformer.testing.standalone_transformer_lm import MegatronModule, get_language_model, get_linear_layer, init_method_normal, scaled_init_method_normal, parallel_lm_logits
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, hidden_size, init_method, layernorm_epsilon, parallel_output)