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.
class_docstr_tok: ["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)
      async_tensor_model_parallel_allreduce = (
          self.config.async_tensor_model_parallel_allreduce
      )
      output = parallel_lm_logits(
          hidden_states,
          word_embeddings_weight,
          self.parallel_output,
          sequence_parallel=self.sequence_parallel,
          bias=self.bias,
          async_tensor_model_parallel_allreduce=async_tensor_model_parallel_allreduce,
      )
      return output
  
file_path: NeMo/nemo/collections/nlp/models/language_modeling/megatron/bert_model.py
imports: import torch
  from nemo.collections.nlp.modules.common.megatron.language_model import get_language_model
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.transformer import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, build_position_ids, erf_gelu, get_linear_layer, init_method_normal, openai_gelu, parallel_lm_logits, scaled_init_method_normal
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