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 forward(self, hidden_states, word_embeddings_weight):
      hidden_states = self.dense(hidden_states)
      hidden_states = self.gelu(hidden_states)
      hidden_states = self.layernorm(hidden_states)
      output = parallel_lm_logits(
          hidden_states,
          word_embeddings_weight,
          self.parallel_output,
          bias=self.bias,
      )
      return output
  
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 2 for calling the BertLMHead obj's underlying member method code with example usage: obj.forward(hidden_states, word_embeddings_weight) and returns: output