class_docstr: Megatron encoder-decoder base class
class_docstr_tok: ["Megatron","encoder-decoder","base","class"]
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 megatron-core 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","megatron-core","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 megatron-core fwd/bwd functions.
      """
      # we zero grads here because we also call backward in the megatron fwd/bwd functions
      self._optimizer.zero_grad()
  
      loss_dict = self.fwd_bwd_step(dataloader_iter, batch_idx, False)
  
      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 reduced after the pipeline (not asynchronously)
          if self.cfg.get("pipeline_model_parallel_size", 1) > 1:
              # main grads are stored in the MainParamsOptimizer wrapper
              self._optimizer.allreduce_main_grads()
      else:
          # async grad allreduce is not currently implemented for O1/autocasting mixed precision training
          # so we allreduce gradients after the pipeline
          self.allreduce_gradients()  # @sangkug we think this is causing memory to blow up (hurts perf)
  
      if self.cfg.get("pipeline_model_parallel_size", 1) > 1:
          # when using pipeline parallelism, we need keep the word and position embeddings in sync
          self.allreduce_word_and_position_embeddings()
  
      ## logging
      # we can only log on one rank if it is rank zero so we broadcast from last rank
      # we can avoid this broadcast by updating the PTL log function to accept specific ranks
      for k, v in loss_dict.items():
          torch.distributed.broadcast(v, get_last_rank())
          n = f"reduced_train_{k}"
          self.log(
              n,
              v,
              prog_bar=n.endswith("_loss"),
              rank_zero_only=True,
              batch_size=1,
          )
  
      if self.torch_dtype == torch.float16:
          loss_scale = self.trainer.precision_plugin.scaler._scale
          if loss_scale is not None:
              self.log("loss_scale", loss_scale, batch_size=1)
  
      lr = self._optimizer.param_groups[0]["lr"]
      self.log("lr", lr, rank_zero_only=True, batch_size=1)
      self.log(
          "global_step",
          self.trainer.global_step,
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      # TODO: make sure compute_consumed_samples works for pipeline parallelism
      self.log(
          "consumed_samples",
          self._compute_consumed_samples_after_training_step(),
          prog_bar=True,
          rank_zero_only=True,
          batch_size=1,
      )
      return loss_dict
  
imports: import copy
  import functools
  import inspect
  from typing import Any, Dict, List, Optional
  import torch
  from omegaconf import OmegaConf, open_dict
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.accelerators import CPUAccelerator
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.build_model import build_model
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.modules.common.megatron.token_level_encoder_decoder import MegatronTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.modules.common.text_generation_utils import compute_beam_search_len_penalty, get_sampling_token_fn
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.utils import AppState, logging
label_desc: Use image node_id 10 for calling the MegatronLMEncoderDecoderModel obj's underlying member method code with example usage: obj.training_step(dataloader_iter, batch_idx) and returns: loss_dict