class_docstr: Megatron Bert pretraining.
  Model returns [batch, seq, hidden] shape
class_docstr_tok: ["Megatron","Bert","pretraining",".","Model","returns","[","batch",",","seq",",","hidden","]","shape"]
code: def training_step(self, dataloader_iter, batch_idx):
      self._optimizer.zero_grad()
  
      if self.with_distributed_adam:
          # hack to enable overlapping param sync and forward compute
          # note: the distributed optimizer monkey-patches each
          # parameter's __getattribute__ function so that it can
          # launch parameter all-gathers the first time the
          # parameter is accessed after the optimizer step. However,
          # PyTorch directly passes embedding parameters into a C++,
          # bypassing this process. A quick-and-dirty hack is to
          # manually interact with the parameter.
          modules = (
              self.model
              if isinstance(self.model, list)
              else [self.model]
          )
          for module in modules:
              if isinstance(
                  module, (Float16Module, MCoreFloat16Module)
              ):
                  module = module.module
              if not self.mcore_bert:
                  module = module.language_model
              if hasattr(module, "embedding"):
                  for param in module.embedding.parameters():
                      param.data_ptr()
  
      if self.cfg.data.dataloader_type == "LDDL":
          # this is of type bert dataset
          seq_length = dataloader_iter.iterator.loaders.get_seqlen()
      else:
          seq_length = self.cfg.encoder_seq_length
  
      # run forward and backwards passes for an entire global batch
      # we do this inside training_step to support pipeline parallelism
      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=dataloader_iter,
          model=[self.model],
          num_microbatches=get_num_microbatches(),
          forward_only=False,
          seq_length=seq_length,
          micro_batch_size=self.cfg.micro_batch_size,
      )
  
      if losses_reduced_per_micro_batch:
          loss_tensors_list = [
              loss_reduced["loss"]
              for loss_reduced in losses_reduced_per_micro_batch
          ]
          loss_tensor = torch.vstack(loss_tensors_list)
          loss_mean = loss_tensor.mean(axis=0)
      else:
          if self.cfg.bert_binary_head == True:
              loss_mean = torch.tensor([0.0, 0.0, 0.0]).cuda()
          else:
              loss_mean = torch.tensor([0.0, 0.0]).cuda()
  
      # 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:
          if self.cfg.get(
              "pipeline_model_parallel_size", 1
          ) > 1 or self.cfg.get("sequence_parallel", False):
              # when using pipeline parallelism grads must be all-reduced after the pipeline (not asynchronously)
              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.get("pipeline_model_parallel_size", 1) > 1:
          # when using pipeline parallelism the first and last stage must keep embeddings in sync
          self.allreduce_first_last_embeddings()
  
      torch.distributed.broadcast(loss_mean, get_last_rank())
  
      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)
  
      if (batch_idx + 1) % self.trainer.accumulate_grad_batches == 0:
          # Reduced loss for logging.
          self.log(
              "reduced_train_loss",
              loss_mean[0],
              prog_bar=True,
              batch_size=1,
          )
          if len(loss_mean) > 2:
              self.log(
                  "reduced_lm_train_loss",
                  loss_mean[1],
                  prog_bar=True,
                  batch_size=1,
              )
              self.log(
                  "reduced_sop_train_loss",
                  loss_mean[2],
                  prog_bar=True,
                  batch_size=1,
              )
          lr = self._optimizer.param_groups[0]["lr"]
          self.log("lr", lr, batch_size=1)
          self.log(
              "global_step",
              self.trainer.global_step,
              prog_bar=True,
              batch_size=1,
          )
          self.log(
              "consumed_samples",
              self._compute_consumed_samples_after_training_step(),
              prog_bar=True,
              batch_size=1,
          )
  
      return loss_mean[0]
  
imports: import itertools
  from typing import Any, Dict, List, Optional
  import torch
  import torch.nn.functional
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron import dataset_utils
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron.bert_model import BertModel
  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.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.core.neural_types import ChannelType, MaskType, NeuralType
  from nemo.utils import AppState, logging
label_desc: Use image node_id 6 for calling the MegatronBertModel obj's underlying member method code with example usage: obj.training_step(dataloader_iter, batch_idx) and returns: loss_mean