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","."]
docstr: For attributes in the nemo model config that are the same as the
  megatron core ModelParallelConfig we will use the value from the nemo config.
  For attributes in ModelParallelConfig that are not in the nemo model config, we add custom logic.
code: def build_model_parallel_config(self) -> ModelParallelConfig:
      """For attributes in the nemo model config that are the same as the
      megatron core ModelParallelConfig we will use the value from the nemo config.
      For attributes in ModelParallelConfig that are not in the nemo model config, we add custom logic.
      """
      cfg = OmegaConf.to_container(self.cfg, resolve=True)
  
      # map precision related configs
      megatron_amp_O2 = cfg.get("megatron_amp_O2", False)
  
      # dtype used in p2p communication
      pipeline_dtype = self.torch_dtype
  
      # maps NeMo model configs to ModelParallelConfig from megatron core
      config_mapping = {
          "perform_initialization": True,  # initailize weights when constructing the module
          "fp16": self.torch_dtype == torch.float16
          and megatron_amp_O2,  # NeMo does not currently support fp16 training with megatron amp O2, eval and inference is supported
          "bf16": self.torch_dtype == torch.bfloat16
          and megatron_amp_O2,
          "params_dtype": self.params_dtype,
          "timers": None,  # NeMo does not currently support megatron core timers
          "async_tensor_model_parallel_allreduce": self.cfg.get(
              "tensor_model_parallel_world_size", 1
          )
          > 1
          and not self.cfg.get("sequence_parallel", False),
          "pipeline_dtype": pipeline_dtype,
          "grad_scale_func": self.trainer.precision_plugin.scaler.scale
          if self.trainer.precision in ["16", "16-mixed"]
          else None,
          "enable_autocast": not megatron_amp_O2
          and self.torch_dtype in [torch.bfloat16, torch.float16],
          "autocast_dtype": self.autocast_dtype,
          "variable_seq_lengths": False,  # set dynamically during training
          "num_microbatches_with_partial_activation_checkpoints": self.cfg.get(
              "num_micro_batches_with_partial_activation_checkpoints",
              None,
          ),
          "batch_p2p_sync": True,  # call torch.cuda.synchronize() after batch isend/rcv
          "use_ring_exchange_p2p": False,  # not supported in NeMo
          "deallocate_pipeline_outputs": False,  # not supported in NeMo
          "no_sync_func": None,  # set dynamically during training
          "grad_sync_func": None,  # set dynamically during training
          "param_sync_func": None,  # set dynamically during training
      }
  
      # instantitate ModelParallelConfig from this dict
      mp_config_dict = {}
  
      for field in fields(ModelParallelConfig):
          # model config has priority
          if field.name in cfg:
              mp_config_dict[field.name] = cfg[field.name]
          # then config_mapping
          elif field.name in config_mapping:
              mp_config_dict[field.name] = config_mapping[field.name]
          else:
              logging.warning(
                  f"The model: {self} does not have field.name: {field.name} in its cfg. "
                  f"Add this key to cfg or config_mapping to make to make it configurable."
              )
  
      model_parallel_config = ModelParallelConfig(**mp_config_dict)
  
      try:
          # hidden size is needed for pipeline schedules but is not currently in ModelParallelConfig
          setattr(
              model_parallel_config, "hidden_size", self.cfg.hidden_size
          )
      except AttributeError:
          logging.warning(
              f"hidden_size not found in {self.cfg}. Set this in model_parallel_config if using pipeline parallelism."
          )
  
      return model_parallel_config
  
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 31 for calling the MegatronBaseModel obj's underlying member method code with example usage: obj.build_model_parallel_config() and returns: model_parallel_config