class_docstr: Megatron NMT training
class_docstr_tok: ["Megatron","NMT","training"]
code: def setup(self, stage=None):
      # NOTE: super().__init__ will try and setup train/val/test datasets, but we sidestep this using a if self._train_ds is not None condition
      # We then set things up for real only once setup() of this class is called.
      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
  
      if stage == "predict":
          return
  
      # If the user wants to manually override train and validation dataloaders before calling `.fit()`
      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.train_ds)
      self.setup_validation_data(self._cfg.validation_ds)
      if hasattr(self._cfg, "test_ds"):
          self.setup_test_data(self._cfg.test_ds)
  
      # 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")
              != "relative"
              and self.cfg.decoder.get("position_embedding_type")
              != "relative"
          ):
              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 enum
  import itertools
  import random
  from typing import List, Optional
  import numpy
  import torch
  from omegaconf.dictconfig import DictConfig
  from omegaconf.listconfig import ListConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from sacrebleu import corpus_bleu
  from nemo.collections.nlp.data.common.sequence_to_sequence_dataset import BinarizedMemmapSequenceToSequenceDataset, TextMemmapSequenceToSequenceDataset
  from nemo.collections.nlp.data.language_modeling.megatron.base_dataset_utils import get_datasets_weights_and_num_samples
  from nemo.collections.nlp.data.language_modeling.megatron.blendable_dataset import BlendableDataset
  from nemo.collections.nlp.data.language_modeling.megatron.megatron_batch_samplers import MegatronPretrainingBatchSampler
  from nemo.collections.nlp.data.language_modeling.megatron.xlm_dataset import BinarizedMemmapCrossLingualMLMAndTranslationDataset, TextMemmapCrossLingualMLMAndTranslationDataset
  from nemo.collections.nlp.models.language_modeling.megatron_lm_encoder_decoder_model import MegatronLMEncoderDecoderModel
  from nemo.collections.nlp.models.language_modeling.megatron_t5_model import MegatronT5Model
  from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
  from nemo.collections.nlp.modules.common.megatron.megatron_export import DecEmb, EncEmb, TokensHeadEmb
  from nemo.collections.nlp.modules.common.megatron.utils import get_iterator_k_split
  from nemo.collections.nlp.parts.nlp_overrides import GlobalBatchDataFetcher
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes import Exportable
  from nemo.utils import AppState, logging, timers
label_desc: Use image node_id 4 for calling the MegatronNMTModel obj's underlying member method code with example usage: obj.setup(stage) without return types