class_docstr: Megatron Retrieval enhanced language model
class_docstr_tok: ["Megatron","Retrieval","enhanced","language","model"]
code: def __init__(self, cfg: DictConfig, trainer: Trainer):
      super().__init__(cfg, trainer=trainer)
  
      # TODO does not support PP yet
      self.model = self.model_provider_func(
          pre_process=True,
          post_process=True,
          add_encoder=True,
          add_decoder=True,
      )
  
      self.megatron_amp_O2 = cfg.get("megatron_amp_O2", False)
  
      if self.megatron_amp_O2:
          if not self.with_distributed_adam:
              # Pre-allocate the model on GPU to have master parameters allocated on the same device with matching data type
              self.model.cuda(torch.cuda.current_device())
  
          # Model wrapper to convert both model and inputs to half precision
          self.model = Float16Module(
              config=self.model_parallel_config,
              module=self.model,
              precision=self.cfg.precision,
          )
  
      # self.setup_optimizer_param_groups()
      self.model.model_type = ModelType.encoder_and_decoder
  
      self.enable_autocast = (
          True
          if (not self.megatron_amp_O2)
          and (self.autocast_dtype in [torch.float16, torch.bfloat16])
          else False
      )
  
      if hasattr(self.cfg, "shape_file"):
          set_base_shapes(
              self,
              self.register_artifact("shape_file", self.cfg.shape_file),
              rescale_params=False,
          )
  
          # here manually initialize all the named parameters with the muTranfer normal initializer
          for name, tensor in self.named_parameters():
              if name.endswith(
                  ".dense_4h_to_h.weight"
              ) or name.endswith(".dense.weight"):
                  # initialize all the output dense matrix weight
                  # match the megatron lm model
                  std = self.cfg.init_method_std / math.sqrt(2.0 * 12.0)
                  normal_(tensor, 0, std)
              elif name.endswith("layernorm.weight"):
                  # initialize all the layer norm weight
                  if tensor.std() != 0 and tensor.mean() != 1:
                      raise ValueError(f"need to check {name} init")
                  normal_(tensor, 1, 0)
              elif name.endswith(".weight"):
                  # initialize all the other dense matrix weight
                  normal_(tensor, 0, self.cfg.init_method_std)
              else:
                  if tensor.std() != 0 and tensor.mean() != 0:
                      raise ValueError(f"need to check {name} init")
  
          # here manually overwrite the norm factor
          # note, has to turn off the model.apply_query_key_layer_scaling
          assert not self.cfg.apply_query_key_layer_scaling
          for name, layer in self.named_modules():
              if (
                  name.endswith(".self_attention")
                  or name.endswith(".inter_attention")
                  or name.endswith(".cross_attention")
                  or name.endswith(".core_attention")
              ):
                  if hasattr(layer, "norm_factor") and hasattr(
                      layer, "hidden_size_per_attention_head"
                  ):
                      layer.norm_factor = (
                          layer.hidden_size_per_attention_head / 8.0
                      )  # divide 8 to make it consist with ADLR setting
              else:
                  if hasattr(layer, "norm_factor") or hasattr(
                      layer, "hidden_size_per_attention_head"
                  ):
                      logging.error(
                          f"module {name} has norm factor but its name is not ending with attention, need to double check"
                      )
  
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 1 to create a new MegatronRetrievalModel object from inherited base classes: MegatronBaseModel, TextGeneration with example: obj = MegatronRetrievalModel(cfg, trainer)