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 megatron-core 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:
          # synchronize asynchronous grad reductions
          # note: not necessary, but reduces performance degradation
          # from multiple simultaneous NCCL calls
          self._optimizer._finish_bucket_grad_sync()
      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, 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: import itertools
  from datetime import datetime
  from functools import partial
  from typing import Any
  import torch
  from omegaconf import DictConfig, open_dict
  from pytorch_lightning import Trainer
  from nemo.collections.multimodal.data.imagen.imagen_dataset import build_train_valid_datasets
  from nemo.collections.multimodal.models.text_to_image.imagen.precond import ContinousDDPMPrecond, EDMPrecond
  from nemo.collections.multimodal.modules.imagen.diffusionmodules.nets import EfficientUNetModel, UNetModel
  from nemo.collections.multimodal.modules.imagen.encoder.t5encoder import T5Encoder
  from nemo.collections.multimodal.modules.imagen.sampler.sampler import DDPMSampler, EDMSampler
  from nemo.collections.multimodal.parts.imagen.utils import random_dropout
  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 11 for calling the MegatronImagen obj's underlying member method code with example usage: obj.training_step(dataloader_iter, batch_idx) and returns: loss_mean