class_docstr: Megatron base class. All NeMo Megatron models inherit from this class.
  
  - Initialize the model parallel world for nemo.
  - Turn on all of the nvidia optimizations.
  - If `cfg.tokenizer` is available, it loads the tokenizer and pad the vocab to the
    correct size for tensor model parallelism.
  - If using distributed optimizer, configure to be compatible
    with O2 level optimizations and/or model parallelism.
  - Perform gradient clipping: `grad_clip_pl_default` triggers
    the PyTorch Lightning default implementation, `with_distributed_adam` triggers
    the distributed optimizer's implementation, `megatron_amp_O2` triggers gradient clipping on the main grads,
    and otherwise gradient clipping is performed on the model grads.
class_docstr_tok: ["Megatron","base","class",".","All","NeMo","Megatron","models","inherit","from","this","class",".","-","Initialize","the","model","parallel","world","for","nemo",".","-","Turn","on","all","of","the","nvidia","optimizations",".","-","If","`","cfg.tokenizer","`","is","available",",","it","loads","the","tokenizer","and","pad","the","vocab","to","the","correct","size","for","tensor","model","parallelism",".","-","If","using","distributed","optimizer",",","configure","to","be","compatible","with","O2","level","optimizations","and\/or","model","parallelism",".","-","Perform","gradient","clipping",":","`","grad_clip_pl_default","`","triggers","the","PyTorch","Lightning","default","implementation",",","`","with_distributed_adam","`","triggers","the","distributed","optimizer","'s","implementation",",","`","megatron_amp_O2","`","triggers","gradient","clipping","on","the","main","grads",",","and","otherwise","gradient","clipping","is","performed","on","the","model","grads","."]
code: def __init__(
      self, cfg: DictConfig, trainer: Trainer, no_lm_init=True
  ):
      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."
          )
  
      if trainer is None:
          raise ValueError(
              f"Trainer cannot be None for Megatron-based models. Please provide a PTL trainer object."
          )
  
      if (
          cfg.get("use_flash_attention", False)
          and not HAVE_FLASH_ATTENTION
      ):
          raise ImportError(
              "flash_attn was not found. Please see the installation instructions: https://github.com/HazyResearch/flash-attention."
              "If you use flash_attn with triton. Please install triton==2.0.0.dev20221202."
          )
  
      # this prevents base constructor from initializing tokenizer
      self.tokenizer = None
  
      with open_dict(cfg):
          if cfg.get("precision", None) is None:
              cfg.precision = trainer.precision
  
      super().__init__(cfg, trainer=trainer, no_lm_init=no_lm_init)
  
      # TODO: @maanug-nv consolidate into one attribute (requires lots of changes in subclasses)
      self.torch_dtype = utils_funcs.torch_dtype_from_precision(
          self.cfg.precision
      )  # Mixed precision datatype
      self.autocast_dtype = self.torch_dtype  # Mixed precision datatype
      # instantiate weights in mixed precision datatype if using megatron amp O2
      self.params_dtype = (
          self.torch_dtype
          if self.torch_dtype in [torch.bfloat16, torch.float16]
          and self.cfg.get("megatron_amp_O2", False)
          else torch.float32
      )
  
      # set the megatron core model parallel config
      self.model_parallel_config: ModelParallelConfig = (
          self.build_model_parallel_config()
      )
  
      self.with_distributed_adam = (
          cfg.optim.get("name") == "distributed_fused_adam"
      )
      self.with_megatron_fused_adam = (
          cfg.optim.get("name") == "megatron_fused_adam"
      )
  
      # used in NVIDIA NGC PyTorch containers
      self._enable_nvidia_optimizations()
  
      if self._cfg.get("use_cpu_initialization", False) is False:
          torch.cuda.set_device(trainer.local_rank)
  
      # buffer used during train_step for logging average loss over gradient accumulation steps
      self._reduced_loss_buffer = []
  
      # Overrides used when converting checkpoints
      if (
          os.environ.get(
              NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE, "false"
          ).lower()
          == "true"
      ):
          app_state = AppState()
          init_world_size = (
              app_state.tensor_model_parallel_size
              * app_state.pipeline_model_parallel_size
          )
          init_global_rank = app_state.global_rank
          init_local_rank = app_state.local_rank
      else:
          init_world_size = trainer.world_size
          init_global_rank = trainer.global_rank
          init_local_rank = trainer.local_rank
  
      # Set virtual pipeline size to None if it is 1 and
      # confirm that the number of model chunks is the same across all pipeline stages.
      vp_size = self.cfg.get(
          "virtual_pipeline_model_parallel_size", None
      )
  
      if vp_size is not None:
          if vp_size == 1:
              vp_size = None
          else:
              assert (
                  self.cfg.num_layers
                  // self.cfg.pipeline_model_parallel_size
              ) % vp_size == 0, "Make sure the number of model chunks is the same across all pipeline stages."
  
      initialize_model_parallel_for_nemo(
          world_size=init_world_size,
          global_rank=init_global_rank,
          local_rank=init_local_rank,
          tensor_model_parallel_size=cfg.get(
              "tensor_model_parallel_size", 1
          ),
          pipeline_model_parallel_size=cfg.get(
              "pipeline_model_parallel_size", 1
          ),
          virtual_pipeline_model_parallel_size=vp_size,
          pipeline_model_parallel_split_rank=cfg.get(
              "pipeline_model_parallel_split_rank", 0
          ),
          context_parallel_size=cfg.get("context_parallel_size", 1),
          micro_batch_size=cfg.get("micro_batch_size"),
          global_batch_size=cfg.get("global_batch_size"),
          rampup_batch_size=cfg.get("rampup_batch_size", None),
          use_fp8=cfg.get("fp8", False),
          init_mpi_proc_group=cfg.get("ub_tp_comm_overlap", False),
          seed=self.cfg.get("seed", 1234),
          apex_transformer_log_level=self.cfg.get(
              "apex_transformer_log_level", 30
          ),
      )
  
      # This must be called after initialize model parallel since it needs to know the data parallel size
      self._validate_and_override_config()
  
      # set the megatron core model parallel config
      self.model_parallel_config: ModelParallelConfig = (
          self.build_model_parallel_config()
      )
  
      self.grad_clip_pl_default = False  # use pytorch default for gradient clipping. Default False
  
      if hasattr(self._cfg, "tokenizer") or (
          hasattr(self._cfg, "encoder_tokenizer")
          and hasattr(self._cfg, "decoder_tokenizer")
      ):
          # build tokenizer (defaults to nemo supported tokenizers)
          self._build_tokenizer()
  
          # manipulate vocabulary (e.g., pad vocabulary for better efficiency)
          self._build_vocab()
  
      # TODO: remove this when PTL 1.7.3 is released
      _FxValidator.functions["configure_gradient_clipping"] = {
          "allowed_on_step": (False, True),
          "allowed_on_epoch": (False, True),
          "default_on_step": True,
          "default_on_epoch": False,
      }
  
      self.gc_interval = cfg.get("gc_interval", 0)
      # Do manual garbage collection during validation routine when gc_interval > 0
      self.gc_in_validation = bool(
          int(os.getenv("NEMO_MANUAL_GC_IN_VALIDATION", 1))
      )
      assert (
          self.gc_interval >= 0
      ), "gc_interval should be an integer value larger than or equal to 0."
      # If gc_interval > 0, memory garbage collection is manually controlled.
      # The automatic garbage collector sould be disabled before training starts.
      if self.gc_interval > 0:
          gc.disable()
          self.validation_global_step = 1
  
      self.use_fsdp = cfg.get("fsdp", False)
  
imports: import gc
  import itertools
  import os
  import re
  from dataclasses import fields
  from typing import Any, Dict, Optional, Union
  import omegaconf
  import torch
  from omegaconf import OmegaConf, open_dict
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.plugins.precision import MixedPrecisionPlugin
  from pytorch_lightning.trainer.connectors.logger_connector.fx_validator import _FxValidator
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.models.nlp_model import NLPModel
  from nemo.collections.nlp.modules.common.megatron.attention import HAVE_FLASH_ATTENTION
  from nemo.collections.nlp.modules.common.megatron.clip_grads import clip_grad_norm_distributed_optimizer, clip_grad_norm_fp32
  from nemo.collections.nlp.modules.common.megatron.megatron_init import initialize_model_parallel_for_nemo
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults
  from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
  from nemo.collections.nlp.parts import utils_funcs
  from nemo.collections.nlp.parts.nlp_overrides import NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE, GradScaler
  from nemo.collections.nlp.parts.utils_funcs import activation_to_func
  from nemo.core.optim import MainParamsOptimizerWrapper, prepare_lr_scheduler
  from nemo.utils import AppState, logging, str_to_dtype
  from nemo.utils.get_rank import is_global_rank_zero
label_desc: Use image node_id 1 to create a new MegatronBaseModel object from inherited base classes: NLPModel with example: obj = MegatronBaseModel(cfg, trainer, no_lm_init)