class_docstr: Megatron specific extensions of torch Module with support
  for pipelining.
code: def initialize_word_embeddings(
      self, init_method, vocab_size, hidden_size
  ):
      if not self.share_token_embeddings:
          raise Exception(
              "initialize_word_embeddings() was called but "
              "share_token_embeddings is false"
          )
  
      # This function just initializes the word embeddings in the final stage
      # when we are using pipeline parallelism. If we aren't using pipeline
      # parallelism there is nothing to do.
      if parallel_state.get_pipeline_model_parallel_world_size() == 1:
          return
  
      # Parameters are shared between the word embeddings layer, and the
      # heads at the end of the model. In a pipelined setup with more than
      # one stage, the initial embedding layer and the head are on different
      # workers, so we do the following:
      # 1. Create a second copy of word_embeddings on the last stage, with
      #    initial parameters of 0.0.
      # 2. Do an all-reduce between the first and last stage to ensure that
      #    the two copies of word_embeddings start off with the same
      #    parameter values.
      # 3. In the training loop, before an all-reduce between the grads of
      #    the two word_embeddings layers to ensure that every applied weight
      #    update is the same on both stages.
      if (
          parallel_state.is_pipeline_last_stage()
          and not self.pre_process
      ):
          # This is relevant for T5 when the decoder is only on a single rank. It is the last stage of the pipeline and also has embeddings on this rank already.
          assert not parallel_state.is_pipeline_first_stage()
          self._word_embeddings_for_head_key = (
              "word_embeddings_for_head"
          )
          # set word_embeddings weights to 0 here, then copy first
          # stage's weights using all_reduce below.
          self.word_embeddings = tensor_parallel.VocabParallelEmbedding(
              vocab_size,
              hidden_size,
              init_method=init_method,
              config=self.config,
          )
          self.word_embeddings.weight.data.fill_(0)
          self.word_embeddings.weight.shared = True
  
      # Zero out initial weights for decoder embedding.
      # NOTE: We don't currently support T5 with the interleaved schedule.
      # This is the case where PP > 1 and we're on the decoder first stage.
      if (
          not parallel_state.is_pipeline_first_stage(
              ignore_virtual=True
          )
          and self.pre_process
      ):
          if hasattr(self, "language_model"):
              # Zero params for GPT
              self.language_model.embedding.zero_parameters()
          else:
              # Zero decoder embeddings for T5
              assert hasattr(self, "decoder_embedding")
              self.decoder_embedding.zero_parameters()
  
imports: import torch
  from torch.autograd import Variable
  from torch.nn.parameter import Parameter
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults
  from nemo.utils import logging
label_desc: Use image node_id 7 for calling the MegatronModule obj's underlying member method code with example usage: obj.initialize_word_embeddings(init_method, vocab_size, hidden_size) without return types