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.
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
      )
  
      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()
      else:
          # 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 == [16, "16", "16-mixed"]:
          loss_scale = self.trainer.precision_plugin.scaler._scale
          if loss_scale is not None:
              self.log("loss_scale", loss_scale, 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 typing import Any
  import einops
  import torch
  import torch.nn
  from einops import rearrange, repeat
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from pytorch_lightning.utilities.rank_zero import rank_zero_only
  from torch._inductor import config
  from nemo.collections.multimodal.data.controlnet.controlnet_dataset import build_train_valid_datasets
  from nemo.collections.multimodal.models.text_to_image.stable_diffusion.ldm.ddpm import LatentDiffusion
  from nemo.collections.multimodal.models.text_to_image.stable_diffusion.samplers.ddim import DDIMSampler
  from nemo.collections.multimodal.modules.stable_diffusion.attention import SpatialTransformer
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.openaimodel import AttentionBlock, Downsample, ResBlock, TimestepEmbedSequential, UNetModel
  from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import conv_nd, linear, timestep_embedding, zero_module
  from nemo.collections.multimodal.parts.stable_diffusion.utils import exists, log_txt_as_img
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.utils import logging
label_desc: Use image node_id 7 for calling the MegatronControlNet obj's underlying member method code with example usage: obj.training_step(dataloader_iter, batch_idx) and returns: loss_mean