class_docstr: Megatron GPT Supervised Fine-Tuning
class_docstr_tok: ["Megatron","GPT","Supervised","Fine-Tuning"]
code: def inference_epoch_end(self, outputs, mode, data_cfg):
      # Parent class will handle logging of the loss.
      if not outputs:
          return
  
      if isinstance(outputs[0], dict):
          outputs = [outputs]
  
      averaged_loss = []
      averaged_metric = []
      # Log metrics for each provided validation/test dataset.
      for dataloader_idx, output in enumerate(outputs):
          # Expand on_validation_epoch_end from parent class MegatronGPTModel as on_validation_epoch_end doesnt take outputs arg
          # loss = super().on_validation_epoch_end([x['loss'] for x in output])
          loss_vals = [x["loss"] for x in output]
          if parallel_state.is_pipeline_last_stage():
              # only the last pipeline parallel stages return loss with their batch size
              if self.cfg.data.get("validation_drop_last", True):
                  loss = torch.stack(loss_vals).mean()
              else:
                  # Compute the avg loss by total_loss across all samples / total number of samples
                  total_loss_and_total_samples = torch.vstack(
                      loss_vals
                  ).sum(axis=0)
                  avg_loss = (
                      total_loss_and_total_samples[0]
                      / total_loss_and_total_samples[1]
                  )
                  loss = avg_loss.type(torch.float32).cuda()
          else:
              loss = torch.tensor(0.0, dtype=torch.float32).cuda()
  
          # we can only log on one rank if it is rank zero so we broadcast from last rank
          torch.distributed.broadcast(loss, get_last_rank())
  
          self.log(
              "val_loss",
              loss,
              prog_bar=True,
              rank_zero_only=True,
              batch_size=1,
          )
  
          # Determine the key used to log the loss based on the user provided name of the dataset or the dataloader index.
          loss_log_key = self._determine_log_key(
              data_cfg, dataloader_idx, "loss", mode
          )
          self.log(loss_log_key, loss, batch_size=1)
          averaged_loss.append(loss)
  
          # Gather the outputs object from all data parallel ranks since we are using the DistributedSampler which splits data across DDP ranks.
          gathered_outputs = [
              None
              for _ in range(
                  parallel_state.get_data_parallel_world_size()
              )
          ]
          torch.distributed.all_gather_object(
              gathered_outputs,
              [
                  {
                      "preds": x["preds"],
                      "labels": x["labels"],
                      "inputs": x["inputs"],
                      "metadata": x["metadata"],
                  }
                  for x in output
              ],
              group=parallel_state.get_data_parallel_group(),
          )
  
          # Remove duplicate examples due to distributed sampler.
          deduplicated_outputs = {
              "preds": [],
              "labels": [],
              "inputs": [],
              "metadata": [],
          }
          total_size = 0
          for rank in range(
              0, parallel_state.get_data_parallel_world_size()
          ):
              for batch in gathered_outputs[rank]:
                  for pred, label, input, metadata in zip(
                      batch["preds"],
                      batch["labels"],
                      batch["inputs"],
                      batch["metadata"],
                  ):
                      total_size += 1
                      if not metadata.get("__AUTOGENERATED__", False):
                          deduplicated_outputs["preds"].append(pred)
                          deduplicated_outputs["labels"].append(label)
                          deduplicated_outputs["inputs"].append(input)
                          deduplicated_outputs["metadata"].append(
                              metadata
                          )
                      else:
                          logging.info(
                              f"skipping autogenerated example example {input} prediction {pred} label {label}"
                          )
  
          # Compute metric score
          metric_name = (
              self.val_metric_name
              if mode == "validation"
              else self.test_metric_name
          )
          metric_label_key = (
              self.val_metric_label_key
              if mode == "validation"
              else self.test_metric_label_key
          )
          if metric_name != "loss":
              metric_log_key = self._determine_log_key(
                  data_cfg, dataloader_idx, metric_name, mode
              )
              metric_fn = (
                  self.val_metric[dataloader_idx]
                  if mode == "validation"
                  else self.test_metric[dataloader_idx]
              )
              if (
                  metric_label_key
                  in deduplicated_outputs["metadata"][0]
              ):
                  labels = [
                      m[metric_label_key]
                      for m in deduplicated_outputs["metadata"]
                  ]
              else:
                  labels = deduplicated_outputs["labels"]
  
              for pred, label in zip(
                  deduplicated_outputs["preds"], labels
              ):
                  _ = metric_fn(pred, label)
  
              metric_result = metric_fn.compute()
  
              if metric_name == "rouge":
                  for k, v in metric_result.items():
                      if "fmeasure" in k:
                          self.log(
                              metric_log_key + f"_{k}",
                              v.item(),
                              sync_dist=True,
                          )
                          logging.info(
                              f"{mode} {metric_name} {k}: {v.item()}"
                          )
                  metric_result = metric_result["rouge1_fmeasure"]
              else:
                  self.log(
                      metric_log_key,
                      metric_result.item(),
                      sync_dist=True,
                  )
                  logging.info(
                      f"{mode} {metric_name}: {metric_result.item()}"
                  )
  
              metric_fn.reset()
              averaged_metric.append(metric_result)
  
          # Write predictions to file
          if self.global_rank == 0 and data_cfg.get(
              "write_predictions_to_file", False
          ):
              logging.info(
                  f"Total deduplicated inference data size: {total_size} to {len(deduplicated_outputs['inputs'])}"
              )
  
              # Check if the user provided a prefix path to the file(s) they want to write.
              if (
                  not hasattr(data_cfg, "output_file_path_prefix")
                  or data_cfg.output_file_path_prefix is None
              ):
                  raise ValueError(
                      f"Cannot write predictions to file when output_file_path_prefix is not set or present in the yaml config file."
                  )
              filename_log_key = self._determine_log_key(
                  data_cfg, dataloader_idx, None, mode
              )
              self.write_predictions_to_file(
                  deduplicated_outputs,
                  f"{data_cfg.output_file_path_prefix}_{filename_log_key}",
              )
  
          torch.distributed.barrier(
              group=parallel_state.get_data_parallel_group()
          )
          outputs[dataloader_idx].clear()  # free memory
  
      # Logging of the averaged metrics:
      averaged_loss = sum(averaged_loss) / len(averaged_loss)
      averaged_metric = (
          sum(averaged_metric) / len(averaged_metric)
          if len(averaged_metric) >= 1
          else None
      )
  
      # Handle case where metrics can be nan or inf. This can break checkpoint save/load.
      if averaged_metric is not None and (
          torch.isinf(averaged_metric) or torch.isnan(averaged_metric)
      ):
          app_state = AppState()
          monitor_mode = app_state.checkpoint_callback_params.mode
          assert monitor_mode in ["min", "max"]
          averaged_metric = 0.0 if monitor_mode == "max" else 1e5
  
      if mode == "validation":
          self.log("validation_loss", averaged_loss, batch_size=1)
          if averaged_metric is not None:
              self.log(
                  f"validation_{self.val_metric_name}", averaged_metric
              )
      elif mode == "test":
          self.log("test_loss", averaged_loss, batch_size=1)
          if averaged_metric is not None:
              self.log(f"test_{self.test_metric_name}", averaged_metric)
  
      # Merge the functionality of previous on_inference_epoch_end() within inference_epoch_end() func here
      app_state = AppState()
      self._restore_activation_checkpointing_args()
      if hasattr(self, "_train_ds"):
          _reconfigure_microbatch_calculator(
              rank=app_state.global_rank,
              rampup_batch_size=None,
              global_batch_size=self.cfg.data.train_ds.global_batch_size,
              micro_batch_size=self.cfg.data.train_ds.micro_batch_size,
              data_parallel_size=parallel_state.get_data_parallel_world_size(),
          )
      # When running `trainer.validate()`, the training dataset is not available.
      else:
          logging.warning(
              "No training data found, reconfiguring microbatches based on validation batch sizes."
          )
          _reconfigure_microbatch_calculator(
              rank=app_state.global_rank,
              rampup_batch_size=None,
              global_batch_size=data_cfg.global_batch_size,
              micro_batch_size=data_cfg.micro_batch_size,
              data_parallel_size=parallel_state.get_data_parallel_world_size(),
          )
  
      return averaged_loss, averaged_metric
  
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 11 for calling the MegatronGPTSFTModel obj's underlying member method code with example usage: obj.inference_epoch_end(outputs, mode, data_cfg) and returns: averaged_loss,  averaged_metric