code: def get_param_groups(
      modules, no_weight_decay_cond, scale_lr_cond, lr_mult
  ):
      """creates param groups based on weight decay condition (regularized vs non regularized)
      and learning rate scale condition (args.lr vs lr_mult * args.lr)
      scale_lr_cond is used during finetuning where head of the network requires a scaled
      version of the base learning rate.
      """
      wd_no_scale_lr = []
      wd_scale_lr = []
      no_wd_no_scale_lr = []
      no_wd_scale_lr = []
      for module in modules:
          for name, param in module.named_parameters():
              if not param.requires_grad:
                  continue
  
              if no_weight_decay_cond is not None:
                  no_wd = no_weight_decay_cond(name, param)
              else:
                  # do not regularize biases nor Norm parameters
                  no_wd = (
                      name.endswith(".bias") or len(param.shape) == 1
                  )
  
              if scale_lr_cond is not None:
                  scale_lr = scale_lr_cond(name, param)
              else:
                  scale_lr = False
  
              if not no_wd and not scale_lr:
                  wd_no_scale_lr.append(param)
              elif not no_wd and scale_lr:
                  wd_scale_lr.append(param)
              elif no_wd and not scale_lr:
                  no_wd_no_scale_lr.append(param)
              else:
                  no_wd_scale_lr.append(param)
  
      param_groups = []
      if len(wd_no_scale_lr):
          param_groups.append(
              {"params": wd_no_scale_lr, "wd_mult": 1.0, "lr_mult": 1.0}
          )
      if len(wd_scale_lr):
          param_groups.append(
              {
                  "params": wd_scale_lr,
                  "wd_mult": 1.0,
                  "lr_mult": lr_mult,
              }
          )
      if len(no_wd_no_scale_lr):
          param_groups.append(
              {
                  "params": no_wd_no_scale_lr,
                  "wd_mult": 0.0,
                  "lr_mult": 1.0,
              }
          )
      if len(no_wd_scale_lr):
          param_groups.append(
              {
                  "params": no_wd_scale_lr,
                  "wd_mult": 0.0,
                  "lr_mult": lr_mult,
              }
          )
  
      return param_groups
  
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 1 for calling a global function with example usage: get_param_groups(modules, no_weight_decay_cond, scale_lr_cond, lr_mult) and returns: param_groups