docstr: Our dataloaders produce a micro-batch and then we fetch
  a number of microbatches depending on the global batch size and model parallel size
  from the dataloader to produce a list of microbatches.
  Batch should be a list of microbatches and those microbatches should on CPU.
  Microbatches are then moved to GPU during the pipeline.
  The list of microbatches is then piped through the pipeline using Apex fwd/bwd functions.
docstr_tok: ["Our","dataloaders","produce","a","micro-batch","and","then","we","fetch","a","number","of","microbatches","depending","on","the","global","batch","size","and","model","parallel","size","from","the","dataloader","to","produce","a","list","of","microbatches",".","Batch","should","be","a","list","of","microbatches","and","those","microbatches","should","on","CPU",".","Microbatches","are","then","moved","to","GPU","during","the","pipeline",".","The","list","of","microbatches","is","then","piped","through","the","pipeline","using","Apex","fwd\/bwd","functions","."]
code: def training_step(self, dataloader_iter, batch_idx):
      """
      Our dataloaders produce a micro-batch and then we fetch
      a number of microbatches depending on the global batch size and model parallel size
      from the dataloader to produce a list of microbatches.
      Batch should be a list of microbatches and those microbatches should on CPU.
      Microbatches are then moved to GPU during the pipeline.
      The list of microbatches is then piped through the pipeline using Apex fwd/bwd functions.
      """
  
      # we zero grads here because we also call backward in the apex fwd/bwd functions
      self._optimizer.zero_grad()
  
      loss_mean, loss_dict = self.fwd_bwd_step(
          dataloader_iter, batch_idx, False
      )
  
      torch.distributed.broadcast(loss_mean, get_last_rank())
  
      # when using sequence parallelism, the sequence parallel layernorm grads must be all-reduced
      if self.cfg.get(
          "tensor_model_parallel_size", 1
      ) > 1 and self.cfg.get("sequence_parallel", False):
          self.allreduce_sequence_parallel_gradients()
  
      if self.with_distributed_adam:
          # gradients are reduced internally in distributed optimizer
          pass
      elif self.megatron_amp_O2:
          # # when using pipeline parallelism grads must be all-reduced after the pipeline (not asynchronously)
          # if self.cfg.get('pipeline_model_parallel_size', 1) > 1 or self.cfg.get('sequence_parallel', False):
          #     # main grads are stored in the MainParamsOptimizer wrapper
          #     self._optimizer.allreduce_main_grads()
          self._optimizer.allreduce_main_grads()
      elif not self.cfg.get("ddp_overlap", True):
          # async grad allreduce is not currently implemented for O1/autocasting mixed precision training
          # so we all-reduce gradients after the pipeline
          self.allreduce_gradients()  # @sangkug we think this is causing memory to blow up (hurts perf)
  
      if self.cfg.precision in [16, "16", "16-mixed"]:
          loss_scale = self.trainer.precision_plugin.scaler._scale
          if loss_scale is not None:
              self.log(
                  "loss_scale", loss_scale, prog_bar=True, batch_size=1
              )
  
      self.log_dict(
          loss_dict,
          prog_bar=False,
          logger=True,
          on_step=True,
          rank_zero_only=True,
          batch_size=1,
      )
      self.log(
          "reduced_train_loss",
          loss_mean,
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      lr = self._optimizer.param_groups[0]["lr"]
      self.log(
          "lr", lr, prog_bar=True, rank_zero_only=True, batch_size=1
      )
      self.log(
          "global_step",
          self.trainer.global_step + 1,
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      self.log(
          "consumed_samples",
          self.compute_consumed_samples(
              self.trainer.global_step + 1 - self.init_global_step
          ),
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      return loss_mean
  
imports: from functools import partial
  from typing import Any, Optional
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from torch._inductor import config
  from nemo.collections.multimodal.data.dreambooth.dreambooth_dataset import DreamBoothDataset
  from nemo.collections.multimodal.modules.stable_diffusion.distributions.distributions import DiagonalGaussianDistribution
  from nemo.collections.multimodal.parts.utils import randn_like
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler
  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 6 for calling the MegatronDreamBooth obj's underlying member method code with example usage: obj.training_step(dataloader_iter, batch_idx) and returns: loss_mean