class_docstr: Megatron Bert pretraining.
  Model returns [batch, seq, hidden] shape
class_docstr_tok: ["Megatron","Bert","pretraining",".","Model","returns","[","batch",",","seq",",","hidden","]","shape"]
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.
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_gpt_bert(
          self.model
      )
  
      logging.info(
          f"Pipeline model parallel rank: {parallel_state.get_pipeline_model_parallel_rank()}, "
          f"Tensor model parallel rank: {parallel_state.get_tensor_model_parallel_rank()}, "
          f"Number of model parameters on device: {num_parameters_on_device:.2e}. "
          f"Total number of model parameters: {total_num_parameters:.2e}."
      )
  
      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
      self.init_global_step = self.trainer.global_step
  
      if stage == "predict":
          return
      else:
          # TODO: consider adding a ModelPT guard to check if model is being restored.
          # allowing restored models to optionally setup datasets
          if self.cfg.data.dataloader_type == "LDDL":
              self.build_LDDL_data(self.cfg.data)
              torch.distributed.barrier()
          else:
              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:
          if isinstance(self.model, list):
              for i, module in enumerate(self.model):
                  parallel_state.set_virtual_pipeline_model_parallel_rank(
                      i
                  )
                  sync_embeddings = (
                      module.initialize_last_stage_with_word_embeddings
                      if self.mcore_bert
                      else module.sync_initial_word_embeddings
                  )
                  sync_embeddings()
              parallel_state.set_virtual_pipeline_model_parallel_rank(0)
          else:
              sync_embeddings = (
                  self.model.initialize_last_stage_with_word_embeddings
                  if self.mcore_bert
                  else self.model.sync_initial_word_embeddings
              )
              sync_embeddings()
  
      if self.cfg.get("transformer_engine", False) or self.cfg.get(
          "mcore_bert", False
      ):
          self.setup_transformer_engine_tp_groups()
  
imports: import itertools
  from typing import Any, Dict, List, Optional
  import torch
  import torch.nn.functional
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron import dataset_utils
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron.bert_model import BertModel
  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.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.core.neural_types import ChannelType, MaskType, NeuralType
  from nemo.utils import AppState, logging
label_desc: Use image node_id 18 for calling the MegatronBertModel obj's underlying member method code with example usage: obj.setup(stage) without return types