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: Returns the total number of parameters across all model parallel groups.
code: def _get_total_params_across_model_parallel_groups_enc_dec(
      self, model
  ):
      """Returns the total number of parameters across all model parallel groups."""
      # log number of parameters
      # TODO: If/when we add interleaved model parallelism, we will need to add another if/else here.
      num_parameters_on_device = sum(
          [p.nelement() for p in model.parameters()]
      )
  
      if (
          parallel_state.get_pipeline_model_parallel_world_size() > 1
          and (
              parallel_state.get_pipeline_model_parallel_rank()
              == self.cfg.get("pipeline_model_parallel_split_rank", 0)
              or parallel_state.is_pipeline_last_stage()
          )
      ):
          # If the current rank is the in the decoder first stage (decoder emb) or last rank (output layer), subtract those weights since it is already accounted for in the encoder first stage.
          # TODO: If we support embedding untying with PP > 1, we will need to update this.
          num_word_embedding_parameters = sum(
              [p.nelement() for p in model.word_embeddings_weight()]
          )
          num_parameters_on_device -= num_word_embedding_parameters
  
          # Subtract decoder position embedding params that are shared with encoder.
          if (
              parallel_state.is_pipeline_stage_at_split()
              and self.cfg.encoder.get(
                  "position_embedding_type", "learned_absolute"
              )
              == "learned_absolute"
          ):
              num_position_embedding_parameters = sum(
                  [
                      p.nelement()
                      for p in model.position_embeddings_weight()
                  ]
              )
              num_parameters_on_device -= (
                  num_position_embedding_parameters
              )
  
      # Check and remove RPE embeddings from the encoder that are replicated.
      if (
          parallel_state.get_pipeline_model_parallel_world_size() > 1
          and parallel_state.is_pipeline_stage_before_split()
          and not parallel_state.is_pipeline_first_stage()
          and self.cfg.encoder.get(
              "position_embedding_type", "learned_absolute"
          )
          == "relative"
      ):
          # substract the RPE params on intermediate pipeline stages.
          num_rpe_params = sum(
              [
                  p.nelement()
                  for p in model.encoder_relative_position_embeddings_weight()
              ]
          )
          num_parameters_on_device -= num_rpe_params
  
      # Check and remove RPE embeddings from the decoder that are replicated.
      if (
          parallel_state.get_pipeline_model_parallel_world_size() > 1
          and parallel_state.is_pipeline_stage_after_split()
          and not parallel_state.is_pipeline_stage_at_split()
          and self.cfg.encoder.get(
              "position_embedding_type", "learned_absolute"
          )
          == "relative"
      ):
          # substract the RPE params on intermediate pipeline stages.
          num_rpe_params = sum(
              [
                  p.nelement()
                  for p in model.decoder_relative_position_embeddings_weight()
              ]
          )
          num_parameters_on_device -= num_rpe_params
  
      # to be summed across data parallel group
      total_num_parameters = torch.tensor(
          num_parameters_on_device
      ).cuda()
      torch.distributed.all_reduce(
          total_num_parameters,
          group=parallel_state.get_model_parallel_group(),
      )
      return num_parameters_on_device, total_num_parameters
  
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 30 for calling the MegatronBaseModel obj's underlying member method code with example usage: obj._get_total_params_across_model_parallel_groups_enc_dec(model) and returns: num_parameters_on_device,  total_num_parameters