class_docstr: Megatron GPT Supervised Fine-Tuning
class_docstr_tok: ["Megatron","GPT","Supervised","Fine-Tuning"]
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
      self.setup_complete = True
      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(stage=stage)
      if hasattr(self, "_train_ds"):
          self.setup_training_dataloader()
      if hasattr(self, "_validation_ds"):
          self._validation_dl = self.setup_eval_dataloader(
              self._validation_ds, self.cfg.data.validation_ds
          )
      if (
          hasattr(self.cfg.data, "test_ds")
          and self.cfg.data.test_ds.get("file_names", None) is not None
      ):
          self._test_dl = self.setup_eval_dataloader(
              self._test_ds, self.cfg.data.test_ds
          )
  
      # Raise error if using multiple dataloaders
      if (
          type(self._validation_dl) == list
          and len(self._validation_dl) > 1
      ):
          raise NotImplementedError(
              "Lightning 2.0 does not support multiple dataloaders with dataloader_iter"
          )
  
      if type(self._test_dl) == list and len(self._test_dl) > 1:
          raise NotImplementedError(
              "Lightning 2.0 does not support multiple dataloaders with dataloader_iter"
          )
  
      # when using pipeline model parallel the final stage need to initialize word embeddings
      if not self.cfg.get("mcore_gpt", False):
          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
                      )
                      module.sync_initial_word_embeddings()
                  parallel_state.set_virtual_pipeline_model_parallel_rank(
                      0
                  )
              else:
                  self.model.sync_initial_word_embeddings()
  
      if self.cfg.get("transformer_engine", False):
          self.setup_transformer_engine_tp_groups()
      self.setup_complete = True
  
imports: import itertools
  import json
  from functools import partial
  from typing import Any, Optional
  import torch
  from omegaconf import DictConfig, ListConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.common.metrics import MetricStringToTorchMetric
  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.gpt_sft_chat_dataset import GPTSFTChatDataset
  from nemo.collections.nlp.data.language_modeling.megatron.gpt_sft_dataset import GPTSFTDataset, GPTSFTPackedDataset
  from nemo.collections.nlp.data.language_modeling.megatron.megatron_batch_samplers import MegatronPretrainingBatchSampler
  from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
  from nemo.collections.nlp.modules.common.megatron.utils import get_iterator_k_split
  from nemo.collections.nlp.modules.common.text_generation_utils import generate, get_computeprob_response
  from nemo.collections.nlp.parts.mixins.nlp_adapter_mixins import NLPAdapterModelMixin
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.utils import AppState, logging
label_desc: Use image node_id 4 for calling the MegatronGPTSFTModel obj's underlying member method code with example usage: obj.setup(stage) without return types