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_gpt_bert(
      self, model
  ):
      """Returns the total number of parameters across all model parallel groups."""
      is_mcore_model = self.__dict__.get(
          "mcore_gpt", False
      ) or self.__dict__.get("mcore_bert", False)
      # log number of parameters
      if isinstance(model, list):
          num_parameters_on_device = sum(
              [
                  sum([p.nelement() for p in model_module.parameters()])
                  for model_module in model
              ]
          )
          if (
              parallel_state.get_pipeline_model_parallel_world_size()
              > 1
              and parallel_state.is_pipeline_last_stage(
                  ignore_virtual=True
              )
              and self.cfg.get(
                  "share_embeddings_and_output_weights", True
              )
          ):
              word_embeddings_weight = (
                  model[-1].module.shared_embedding_or_output_weight()
                  if is_mcore_model
                  else model[-1].word_embeddings_weight()
              )
              # substract the embedding weights on the last virtual stage
              num_word_embedding_parameters = sum(
                  [p.nelement() for p in word_embeddings_weight]
              )
              num_parameters_on_device -= num_word_embedding_parameters
      else:
          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.is_pipeline_last_stage(
                  ignore_virtual=True
              )
              and self.cfg.get(
                  "share_embeddings_and_output_weights", True
              )
          ):
              word_embeddings_weight = (
                  model.module.shared_embedding_or_output_weight()
                  if is_mcore_model
                  else model.word_embeddings_weight()
              )
              # substract the embedding weights on the last stage
              num_word_embedding_parameters = sum(
                  [p.nelement() for p in word_embeddings_weight]
              )
              num_parameters_on_device -= num_word_embedding_parameters
  
      # 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 29 for calling the MegatronBaseModel obj's underlying member method code with example usage: obj._get_total_params_across_model_parallel_groups_gpt_bert(model) and returns: num_parameters_on_device,  total_num_parameters