class_docstr: Megatron encoder-decoder base class
class_docstr_tok: ["Megatron","encoder-decoder","base","class"]
docstr: PTL hook that is executed after DDP spawns.
      We setup datasets here as megatron datasets require DDP to instantiate.
      See https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup for more information.
  Args:
      stage (str, optional): Can be 'fit', 'validate', 'test' or 'predict'. Defaults to None.
docstr_tok: ["PTL","hook","that","is","executed","after","DDP","spawns",".","We","setup","datasets","here","as","megatron","datasets","require","DDP","to","instantiate",".","See","https",":","\/\/pytorch-lightning.readthedocs.io\/en\/latest\/common\/lightning_module.html","#","setup","for","more","information",".","Args",":","stage","(","str",",","optional",")",":","Can","be","'fit","'",",","'validate","'",",","'test","'","or","'predict","'",".","Defaults","to","None","."]
code: def setup(self, stage=None):
      """PTL hook that is executed after DDP spawns.
          We setup datasets here as megatron datasets require DDP to instantiate.
          See https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#setup for more information.
      Args:
          stage (str, optional): Can be 'fit', 'validate', 'test' or 'predict'. Defaults to None.
      """
      (
          num_parameters_on_device,
          total_num_parameters,
      ) = self._get_total_params_across_model_parallel_groups_enc_dec(
          self.enc_dec_model
      )
  
      logging.info(
          f"Pipeline model parallel rank: {parallel_state.get_pipeline_model_parallel_rank()}\n"
          f"Tensor model parallel rank: {parallel_state.get_tensor_model_parallel_rank()}\n"
          f"Number of model parameters on device: {num_parameters_on_device:.2e}\n"
          f"Total number of model parameters: {total_num_parameters:.2e}\n"
      )
      resume_checkpoint_path = self.trainer.ckpt_path
  
      if resume_checkpoint_path:
          init_consumed_samples = (
              self._extract_consumed_samples_from_ckpt(
                  resume_checkpoint_path
              )
          )
      else:
          init_consumed_samples = 0
      self.init_consumed_samples = init_consumed_samples
  
      """A PTL method to setup the training, validation and test datasets."""
      if stage == "predict":
          return
      if self._train_dl is not None and self._validation_dl is not None:
          return
      self.build_train_valid_test_datasets()
      self.setup_training_data(self._cfg.data)
      self.setup_validation_data(self._cfg.data)
      self.setup_test_data(self._cfg.data)
  
      # when using pipeline model parallel the final stage need to initialize word embeddings
      if parallel_state.get_pipeline_model_parallel_world_size() > 1:
          assert (
              self.cfg.share_token_embeddings
          ), "share_word_embedding must be True when using pipeline model parallel > 1"
          assert (
              self.cfg.share_decoder_tokens_head_embeddings
          ), "share_decoder_tokens_head_embeddings must be True when using pipeline model parallel > 1"
          self.enc_dec_model.sync_initial_word_embeddings()
          if (
              self.cfg.encoder.get("position_embedding_type")
              == "learned_absolute"
              and self.cfg.decoder.get("position_embedding_type")
              == "learned_absolute"
          ):
              self.enc_dec_model.sync_initial_position_embeddings()
          # Synchronize RPE embeddings across pipeline parallel ranks.
          else:
              if (
                  self.cfg.encoder.get(
                      "position_embedding_type", "learned_absolute"
                  )
                  == "relative"
              ):
                  self.enc_dec_model.sync_initial_encoder_relative_position_embeddings()
              if (
                  self.cfg.decoder.get(
                      "position_embedding_type", "learned_absolute"
                  )
                  == "relative"
              ):
                  self.enc_dec_model.sync_initial_decoder_relative_position_embeddings()
              if self.cfg.decoder.get(
                  "position_embedding_type", "learned_absolute"
              ) == "relative" and not self.cfg.decoder.get(
                  "relative_position_bias_self_attention_only", True
              ):
                  self.enc_dec_model.sync_initial_decoder_cross_attention_relative_position_embeddings()
  
imports: import copy
  import functools
  import inspect
  from typing import Any, Dict, List, Optional
  import torch
  from omegaconf import OmegaConf, open_dict
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.accelerators import CPUAccelerator
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.build_model import build_model
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.modules.common.megatron.token_level_encoder_decoder import MegatronTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.modules.common.text_generation_utils import compute_beam_search_len_penalty, get_sampling_token_fn
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.utils import AppState, logging
label_desc: Use image node_id 33 for calling the MegatronLMEncoderDecoderModel obj's underlying member method code with example usage: obj.setup(stage) without return types