code: def get_megatron_optimizer(
      model, no_weight_decay_cond=None, scale_lr_cond=None, lr_mult=1.0
  ):
      args = get_args()
  
      # Base optimizer.
      param_groups = get_param_groups(
          model, no_weight_decay_cond, scale_lr_cond, lr_mult
      )
  
      if args.optimizer == "adam":
          optimizer = Adam(
              param_groups,
              lr=args.lr,
              weight_decay=args.weight_decay,
              betas=(args.adam_beta1, args.adam_beta2),
              eps=args.adam_eps,
          )
      elif args.optimizer == "sgd":
          optimizer = SGD(
              param_groups,
              lr=args.lr,
              weight_decay=args.weight_decay,
              momentum=args.sgd_momentum,
          )
      else:
          raise Exception(
              "{} optimizer is not supported.".format(args.optimizer)
          )
  
      # Determine whether the params have main-grad field.
      params_have_main_grad = True
  
      # Mixed precision optimizer.
      # - Note: both the Float16Optimizer and the DistributedOptimizer inherit
      #   from the MixedPrecisionOptimizer, which manages any optimizer where
      #   the model params and main params are distinct.
      if args.fp16 or args.bf16 or args.use_distributed_optimizer:
          # Grad scaler:
          #    if loss-scale is provided, instantiate the constant scaler.
          #    if we are using fp16 and loss-scale is not present, use a
          #       dynamic scaler.
          #    otherwise we are running in bf16 with no loss-scale so
          #       leave it as None.
          grad_scaler = None
  
          # Constant loss scale.
          if args.loss_scale:
              grad_scaler = ConstantGradScaler(args.loss_scale)
  
          # Dynamic loss scale.
          else:
              if args.fp16:
                  grad_scaler = DynamicGradScaler(
                      initial_scale=args.initial_loss_scale,
                      min_scale=args.min_loss_scale,
                      growth_factor=2.0,
                      backoff_factor=0.5,
                      growth_interval=args.loss_scale_window,
                      hysteresis=args.hysteresis,
                  )
  
          # Megatron optimizer.
          opt_ty = (
              DistributedOptimizer
              if args.use_distributed_optimizer
              else Float16OptimizerWithFloat16Params
          )
          return opt_ty(
              optimizer,
              args.clip_grad,
              args.log_num_zeros_in_grad,
              args.check_for_nan_in_loss_and_grad,
              params_have_main_grad,
              args.fp16,
              args.bf16,
              args.params_dtype,
              grad_scaler,
              model,
          )
  
      # FP32.
      return FP32Optimizer(
          optimizer,
          args.clip_grad,
          args.log_num_zeros_in_grad,
          args.check_for_nan_in_loss_and_grad,
          params_have_main_grad,
          model,
      )
  
file_path: megatron-lm/megatron/optimizer/__init__.py
imports: from apex.optimizers import FusedAdam
  from apex.optimizers import FusedSGD
  from megatron import get_args
  from .distrib_optimizer import DistributedOptimizer
  from .grad_scaler import ConstantGradScaler, DynamicGradScaler
  from .optimizer import Float16OptimizerWithFloat16Params, FP32Optimizer
label_desc: Use image node_id 2 for calling a global function with example usage: get_megatron_optimizer(model, no_weight_decay_cond, scale_lr_cond, lr_mult) and returns: FP32Optimizer, opt_ty