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 __init__(self, cfg: DictConfig, trainer: Trainer):
      if not HAVE_MEGATRON_CORE:
          raise ImportError(
              "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
          )
      self.megatron_amp_O2 = cfg.get("megatron_amp_O2", False)
      self.cfg = cfg
  
      if not self.megatron_amp_O2 and self.cfg.get(
          "virtual_pipeline_model_parallel_size", None
      ):
          raise ValueError(
              "Virtual pipeline model parallel is only supported when using megatron_amp_O2"
          )
  
      super().__init__(cfg, trainer=trainer, no_lm_init=False)
  
      self._validate_trainer()
  
      self.transformer_config = self.build_transformer_config()
  
      self.mcore_bert = cfg.get("mcore_bert", False)
  
      self.enable_autocast = (
          True
          if (not self.megatron_amp_O2)
          and (self.autocast_dtype in [torch.float16, torch.bfloat16])
          else False
      )
  
      # used in NVIDIA NGC PyTorch containers
      # buffer used during train_step for logging average loss over gradient accumulation steps
      self._reduced_lm_loss_buffer = []
      self._reduced_sop_loss_buffer = []
  
      # build_model returns a list of modules which are used for interleaved pipeline parallelism
      self.model = build_model(
          model_provider_func=self.model_provider_func,
          wrap_with_ddp=False,
          virtual_pipeline_model_parallel_size=self.cfg.get(
              "virtual_pipeline_model_parallel_size", None
          ),
      )
  
      # if we're not using interleaved, then self.model is a module.
      if (
          self.cfg.get("virtual_pipeline_model_parallel_size", None)
          is None
      ):
          self.model = self.model[0]
  
      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
              if isinstance(self.model, list):
                  for module in self.model:
                      module.cuda(torch.cuda.current_device())
              else:
                  self.model.cuda(torch.cuda.current_device())
  
          # Model wrapper to convert both model and inputs to half precision
          self._wrap_model_for_O2()
  
      if hasattr(self, "_nsys_profile_enabled"):
          mp_size = cfg.get("tensor_model_parallel_size", 1) * cfg.get(
              "pipeline_model_parallel_size", 1
          )
          data_parallel_world_size = trainer.world_size // mp_size
          grad_accum_steps = cfg.get("global_batch_size") // (
              cfg.get("micro_batch_size") * data_parallel_world_size
          )
          self._nsys_profile_start_step *= grad_accum_steps
          self._nsys_profile_end_step *= grad_accum_steps
  
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 1 to create a new MegatronBertModel object from inherited base classes: MegatronBaseModel with example: obj = MegatronBertModel(cfg, trainer)