class_docstr: Megatron GPT Supervised Fine-Tuning
class_docstr_tok: ["Megatron","GPT","Supervised","Fine-Tuning"]
code: def fwd_bwd_step(self, dataloader_iter, batch_idx, forward_only):
      batch = next(dataloader_iter)
  
      log_token_counts = self.cfg.get("log_token_counts", False)
      if log_token_counts:
          token_count_avg = sum(batch["token_count"]) / len(
              batch["token_count"]
          )
  
      # Pass only torch.Tensor to prevent errors when process get_iterator_k_split()
      batch = {
          k: v for k, v in batch.items() if isinstance(v, torch.Tensor)
      }
      _, seq_length = batch["tokens"].shape
      data_iter = get_iterator_k_split(batch, get_num_microbatches())
  
      if log_token_counts:
          self.log(
              "seq_length_padded",
              seq_length,
              prog_bar=True,
              batch_size=1,
          )
          self.log(
              "tokens_avg",
              token_count_avg,
              prog_bar=True,
              sync_dist=True,
              batch_size=1,
          )
  
      # handle asynchronous grad reduction
      no_sync_func = None
      grad_sync_func = None
      param_sync_func = None
      if not forward_only and self.with_distributed_adam:
          no_sync_func = partial(
              self._optimizer.no_sync,
              greedy_grad_copy=self.megatron_amp_O2,
          )
          grad_sync_func = self.reduce_overlap_gradients
          param_sync_func = self.sync_overlap_parameters
  
      self.model.config.no_sync_func = no_sync_func
      self.model.config.grad_sync_func = grad_sync_func
      self.model.config.param_sync_func = param_sync_func
  
      fwd_bwd_function = get_forward_backward_func()
  
      losses_reduced_per_micro_batch = fwd_bwd_function(
          forward_step_func=self.get_forward_output_and_loss_func(),
          data_iterator=data_iter,
          model=[self.model],
          num_microbatches=get_num_microbatches(),
          forward_only=forward_only,
          seq_length=seq_length,
          micro_batch_size=get_micro_batch_size(),
      )
  
      # only the last stages of the pipeline return losses
      if losses_reduced_per_micro_batch:
          if (not forward_only) or self.cfg.data.get(
              "validation_drop_last", True
          ):
              # average loss across micro batches
              loss_tensors_list = [
                  loss_reduced["avg"]
                  for loss_reduced in losses_reduced_per_micro_batch
              ]
              loss_tensor = torch.concat(loss_tensors_list)
              loss_mean = loss_tensor.mean()
          else:
              # Get the total loss since micro batches sizes are not uniform
              loss_sum_tensors_list = [
                  loss_sum["loss_sum_and_ub_size"]
                  for loss_sum in losses_reduced_per_micro_batch
                  if loss_sum["loss_sum_and_ub_size"][1] > 0
              ]
              loss_sum = (
                  torch.vstack(loss_sum_tensors_list).sum(axis=0)
                  if len(loss_sum_tensors_list) > 0
                  else torch.tensor([0.0, 0.0]).cuda()
              )
              return loss_sum
      else:
          # we're not on the last pipeline stage so no losses
          if forward_only:
              loss_mean = []
          else:
              loss_mean = torch.tensor(0.0).cuda()
  
      return loss_mean
  
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 7 for calling the MegatronGPTSFTModel obj's underlying member method code with example usage: obj.fwd_bwd_step(dataloader_iter, batch_idx, forward_only) and returns: loss_mean, loss_sum