class_docstr: Masked LM head for token-based encoder-decoder models (e.g., T5)
  
  Arguments:
      mpu_vocab_size: model parallel size of vocabulary.
      parallel_output: wether output logits being distributed or not.
code: def forward(self, hidden_states, word_embeddings_weight):
      async_tensor_model_parallel_allreduce = (
          parallel_state.get_tensor_model_parallel_world_size() > 1
      )
      output = parallel_lm_logits(
          hidden_states,
          word_embeddings_weight,
          self.parallel_output,
          bias=self.bias,
          async_tensor_model_parallel_allreduce=async_tensor_model_parallel_allreduce,
      )
      return output
  
imports: import torch
  from omegaconf import DictConfig
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, PromptEncoderAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.hiddens import get_hiddens_module
  from nemo.collections.nlp.modules.common.megatron.language_model import Embedding
  from nemo.collections.nlp.modules.common.megatron.layer_type import LayerType
  from nemo.collections.nlp.modules.common.megatron.megatron_decoders import get_decoder_model
  from nemo.collections.nlp.modules.common.megatron.megatron_encoder_decoder import MegatronTransformerEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.megatron_encoders import get_encoder_model
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.position_embedding import ALiBiRelativePositionEmbedding, KERPLERelativePositionEmbedding, T5RelativePositionEmbedding
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, build_position_ids, init_method_normal, parallel_lm_logits, scaled_init_method_normal
  from nemo.collections.nlp.modules.common.megatron.vocab_parallel_cross_entropy import vocab_parallel_cross_entropy
  from nemo.core.classes.mixins import adapter_mixins
label_desc: Use image node_id 2 for calling the MegatronTokenLevelHead obj's underlying member method code with example usage: obj.forward(hidden_states, word_embeddings_weight) and returns: output