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.
      """
  
      # log number of parameters
      if isinstance(self.model, list):
          num_parameters_on_device = sum(
              [
                  sum([p.nelement() for p in model_module.parameters()])
                  for model_module in self.model
              ]
          )
      else:
          num_parameters_on_device = sum(
              [p.nelement() for p in self.model.parameters()]
          )
  
      # to be summed across data parallel group
      total_num_parameters = torch.tensor(
          num_parameters_on_device
      ).cuda(non_blocking=True)
  
      torch.distributed.all_reduce(
          total_num_parameters,
          group=parallel_state.get_model_parallel_group(),
      )
  
      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
  
      # allowing restored models to optionally setup datasets
      self.build_train_valid_test_datasets()
  
      # Batch size need to be provided for webdatset
      self._num_micro_batches = get_num_microbatches()
      self._micro_batch_size = self.cfg.micro_batch_size
  
      self.setup_training_data(self.cfg.data)
      self.setup_validation_data(self.cfg.data)
      self.setup_test_data(self.cfg.data)
      # Setup RNG seed in model
      self.model.setup_rng()
  
imports: import itertools
  from datetime import datetime
  from functools import partial
  from typing import Any
  import torch
  from omegaconf import DictConfig, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.multimodal.data.imagen.imagen_dataset import build_train_valid_datasets
  from nemo.collections.multimodal.models.text_to_image.imagen.precond import ContinousDDPMPrecond, EDMPrecond
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.nets import EfficientUNetModel, UNetModel
  from nemo.collections.multimodal.modules.imagen.encoder.t5encoder import T5Encoder
  from nemo.collections.multimodal.modules.imagen.sampler.sampler import DDPMSampler, EDMSampler
  from nemo.collections.multimodal.parts.imagen.utils import random_dropout
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import Serialization
  from nemo.utils import logging
label_desc: Use image node_id 16 for calling the MegatronImagen obj's underlying member method code with example usage: obj.setup(stage) without return types