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 setup_optimization(
      self,
      optim_config: Optional[Union[DictConfig, Dict]] = None,
      optim_kwargs: Optional[Dict[str, Any]] = None,
  ):
      optim_kwargs = {} if optim_kwargs is None else optim_kwargs.copy()
  
      def get_config_arg(
          key: str, default_value: Optional[Any] = None
      ) -> Any:
          """Get keyword argument from config"""
          val = None
          if val is None and optim_kwargs:
              val = optim_kwargs.get(key, None)
          if val is None and optim_config:
              val = optim_config.get(key, None)
          if val is None and self._cfg.optim:
              val = self._cfg.optim.get(key, None)
          if val is None:
              val = default_value
          return val
  
      if self.with_distributed_adam:
          # Allocate contiguous grad buffer to avoid extra copies
          optim_kwargs["contiguous_grad_buffer"] = get_config_arg(
              "contiguous_grad_buffer", True
          )
          if (
              self.megatron_amp_O2
              and not optim_kwargs["contiguous_grad_buffer"]
          ):
              raise ValueError(
                  "Distributed Adam optimizer requires contiguous param buffer for O2. "
                  "Either enable contiguous_grad_buffer or disable megatron_amp_O2."
              )
  
          # Optimizer dtype
          optim_dtype = str_to_dtype(
              get_config_arg("dtype", torch.float32)
          )
          optim_kwargs["dtype"] = optim_dtype
  
          # Make sure embedding grad reductions are in FP32
          if optim_dtype == torch.float32:
              for name, param in self.named_parameters():
                  if (
                      "word_embedding" in name
                      or "position_embedding" in name
                      or "output_layer" in name
                  ):
                      param._with_fp32_optimizer = True
  
          # Match param allgather with model dtype
          model_dtype = torch.float32
          if self.megatron_amp_O2 and hasattr(self, "autocast_dtype"):
              model_dtype = self.autocast_dtype
          optim_kwargs["param_sync_dtype"] = model_dtype
  
          # Determine whether to store master params in optimizer
          if self.cfg.get("fp8_params", False):
              optim_kwargs["store_params"] = True
          elif optim_dtype == model_dtype:
              optim_kwargs["store_params"] = False
          elif (
              optim_dtype == torch.float32
              and model_dtype == torch.bfloat16
          ):
              optim_kwargs["store_params"] = False
              optim_kwargs["store_param_remainders"] = True
          else:
              optim_kwargs["store_params"] = True
  
      return super().setup_optimization(
          optim_config=optim_config, optim_kwargs=optim_kwargs
      )
  
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 22 for calling the MegatronBaseModel obj's underlying member method code with example usage: obj.setup_optimization(optim_config, optim_kwargs) and returns: super, val