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: Builds the megatron core transformer config for the model.
  For attributes in the nemo model config that are the same
  as the megatron core TransformerConfig, we will use the value from the nemo model config.
  For attributes in TransformerConfig that are not in the nemo model config, we add custom logic.
code: def build_transformer_config(self) -> TransformerConfig:
      """Builds the megatron core transformer config for the model.
      For attributes in the nemo model config that are the same
      as the megatron core TransformerConfig, we will use the value from the nemo model config.
      For attributes in TransformerConfig that are not in the nemo model config, we add custom logic.
      """
  
      # create a dictionary copy of the model config
      cfg = OmegaConf.to_container(self.cfg, resolve=True)
  
      # create a dict to store the transformer config arguments
      transformer_config_dict = {}
  
      # get model parallel configs from the base class
      model_parallel_config = self.build_model_parallel_config()
  
      add_bias_linear = self.cfg.get("bias", True)
  
      activation = self.cfg.get("activation", "gelu")
      gated_linear_unit = activation.endswith("glu")
      # TODO: need to check which activation functions are supported in mcore
      activation_func = activation_to_func(activation)
  
      normalization = self.cfg.get("normalization", "LayerNorm")
  
      init_method_std = self.cfg.get("init_method_std", 0.02)
      # default used in mcore
      init_method = init_method_normal(init_method_std)
  
      output_layer_init_method = init_method
      num_layers = self.cfg.get("num_layers", 1)
      use_scaled_init_method = self.cfg.get(
          "use_scaled_init_method", True
      )
      if use_scaled_init_method:
          output_layer_init_method = scaled_init_method_normal(
              init_method_std, num_layers=num_layers
          )
  
      attention_softmax_in_fp32 = False  # not currently used in NeMo unless apply_query_key_layer_scaling is True
      apply_query_key_layer_scaling = self.cfg.get(
          "apply_query_key_layer_scaling", False
      )
  
      fp16_enabled = self.trainer.precision in [16, "16", "16-mixed"]
      if apply_query_key_layer_scaling:
          if fp16_enabled:
              os.environ["NVTE_APPLY_QK_LAYER_SCALING"] = "1"
          else:
              logging.warning(
                  "apply_query_key_layer_scaling is only enabled when using FP16, setting it to False "
                  "and setting NVTE_APPLY_QK_LAYER_SCALING=0"
              )
              os.environ["NVTE_APPLY_QK_LAYER_SCALING"] = "0"
              apply_query_key_layer_scaling = False
  
      if apply_query_key_layer_scaling:
          attention_softmax_in_fp32 = True
  
      bias_activation_fusion = self.cfg.get(
          "bias_activation_fusion", True
      )
      bias_gelu_fusion = True if bias_activation_fusion else False
  
      bias_dropout_fusion = self.cfg.get(
          "bias_dropout_add_fusion", True
      )
  
      # TODO: need to check if recompute APIs are matching up properly
      recompute_granularity = self.cfg.get(
          "activations_checkpoint_granularity", None
      )
      recompute_method = self.cfg.get(
          "activations_checkpoint_method", None
      )
      recompute_num_layers = self.cfg.get(
          "activations_checkpoint_num_layers", None
      )
  
      # any configs that are not in the nemo model config will be added here
      config_mapping = {
          "apply_query_key_layer_scaling": apply_query_key_layer_scaling,
          "apply_residual_connection_post_layernorm": False,  # we don't use this in NeMo
          "layernorm_zero_centered_gamma": False,
          "add_bias_linear": add_bias_linear,
          "gated_linear_unit": gated_linear_unit,
          "activation_func": activation_func,
          "normalization": normalization,
          "init_method": init_method,
          "output_layer_init_method": output_layer_init_method,
          "attention_softmax_in_fp32": attention_softmax_in_fp32,
          "bias_gelu_fusion": bias_gelu_fusion,
          "bias_dropout_fusion": bias_dropout_fusion,
          "recompute_granularity": recompute_granularity,
          "recompute_method": recompute_method,
          "recompute_num_layers": recompute_num_layers,
          "distribute_saved_activations": False,  # not currently used in NeMo
          "fp8": None,
      }
  
      # populate the transformer config dict
      for field in fields(TransformerConfig):
          # config mapping has second highest priority
          if field.name in config_mapping:
              transformer_config_dict[field.name] = config_mapping[
                  field.name
              ]
          # then config
          elif field.name in cfg:
              transformer_config_dict[field.name] = cfg[field.name]
          # then model parallel config
          elif field in fields(model_parallel_config):
              transformer_config_dict[field.name] = getattr(
                  model_parallel_config, 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."
              )
  
      transformer_config = TransformerConfig(**transformer_config_dict)
  
      return transformer_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 12 for calling the MegatronBaseModel obj's underlying member method code with example usage: obj.build_transformer_config() and returns: transformer_config