class_docstr: Megatron Retrieval enhanced language model
class_docstr_tok: ["Megatron","Retrieval","enhanced","language","model"]
code: def training_step(self, batch, batch_idx):
      input_tokens_id = batch["tokens"]
      input_attn_mask = batch["tokens_mask"]
      loss_mask = batch["loss_mask"]
      retrieved_ids = batch["retrieved_ids"]
      retrieved_attn_mask = batch["retrieved_emb_mask"]
      labels = batch["labels"]
      if self.cfg.get("add_position_embedding", False):
          input_position_ids = build_position_ids(input_tokens_id)
      else:
          input_position_ids = None
      loss = self(
          input_tokens_id,
          input_attn_mask,
          retrieved_ids,
          retrieved_attn_mask,
          labels=labels,
          position_ids=input_position_ids,
      )
      loss_mask = loss_mask.float()
      lm_loss = (
          torch.sum(loss.view(-1) * loss_mask.reshape(-1))
          / loss_mask.sum()
      )
      reduced_loss = average_losses_across_data_parallel_group(
          [lm_loss]
      )
      self._reduced_loss_buffer.append(reduced_loss[0])
  
      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 self.with_distributed_adam:
          # gradients are reduced internally in distributed optimizer
          pass
      elif self.megatron_amp_O2:
          # while async grad allreduce is enabled, bprop will keep moving forward without waiting for
          # the finish of async grad AR works. Hence, to guarantee the correctness of grads reduction,
          # we cannot start weight update until all async grad AR works are done.
          if self.cfg.get("pipeline_model_parallel_size", 1) == 1:
              torch.cuda.synchronize()
          # 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
          # no pipeline, so use the default pytorch lightning way of doing all_reduce
          # self.allreduce_gradients()  # @sangkug we think this is causing memory to blow up (hurts perf)
          pass
  
      if (batch_idx + 1) % self.trainer.accumulate_grad_batches == 0:
          # Reduced loss for logging.
          average_reduced_loss = sum(self._reduced_loss_buffer) / len(
              self._reduced_loss_buffer
          )
          self.log(
              "reduced_train_loss",
              average_reduced_loss,
              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,
          )
          self._reduced_loss_buffer = []
      return lm_loss
  
imports: import math
  from typing import Any, List, Optional, Union
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.data.language_modeling.megatron.retro_dataset import build_mock_train_valid_test_datasets, build_train_valid_test_datasets
  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.modules.common.megatron.mup.init import normal_
  from nemo.collections.nlp.modules.common.megatron.mup.shape import set_base_shapes
  from nemo.collections.nlp.modules.common.megatron.retrieval_token_level_encoder_decoder import MegatronRetrievalTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group, build_position_ids, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.modules.common.text_generation_strategy import model_inference_strategy_dispatcher
  from nemo.collections.nlp.modules.common.text_generation_utils import generate, get_computeprob_response, get_default_length_params, get_default_sampling_params, megatron_gpt_generate
  from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
  from nemo.collections.nlp.modules.common.transformer.text_generation import LengthParam, OutputType, SamplingParam, TextGeneration
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler
  from nemo.utils import AppState, logging
label_desc: Use image node_id 5 for calling the MegatronRetrievalModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: lm_loss