code: def _get_params_for_weight_decay_optimization(
      model: Union[torch.nn.Module, List[torch.nn.Module]],
      *,
      no_weight_decay_modules=(FusedLayerNorm,),
  ) -> Dict[str, torch.nn.Parameter]:
      """Divide params into with-weight-decay and without-weight-decay groups.
  
      Layernorms and biases will have no weight decay but the rest will.
      """
      modules = listify_model(model)
      weight_decay_params = {"params": []}
      no_weight_decay_params = {"params": [], "weight_decay": 0.0}
      for module in modules:
          for module_ in module.modules():
              if isinstance(module_, no_weight_decay_modules):
                  no_weight_decay_params["params"].extend(
                      [
                          p
                          for p in list(module_._parameters.values())
                          if p is not None
                      ]
                  )
              else:
                  weight_decay_params["params"].extend(
                      [
                          p
                          for n, p in list(module_._parameters.items())
                          if p is not None and n != "bias"
                      ]
                  )
                  no_weight_decay_params["params"].extend(
                      [
                          p
                          for n, p in list(module_._parameters.items())
                          if p is not None and n == "bias"
                      ]
                  )
  
      return weight_decay_params, no_weight_decay_params
  
file_path: apex/apex/transformer/pipeline_parallel/schedules/common.py
imports: from typing import Any, Callable, Dict, List, Tuple, Union, Optional, Sequence
  import torch
  from torch.autograd.variable import Variable
  from apex.normalization.fused_layer_norm import FusedLayerNorm
  from apex.transformer import parallel_state
  from apex.transformer.enums import ModelType
  from apex.transformer.pipeline_parallel.p2p_communication import FutureTensor
  from apex.transformer.pipeline_parallel.utils import get_num_microbatches
  from apex.transformer.pipeline_parallel.utils import listify_model
  from apex.transformer.pipeline_parallel.utils import unwrap_model
  from apex.transformer.pipeline_parallel.utils import get_model_type
  from apex.transformer.tensor_parallel.layers import set_defaults_if_not_set_tensor_model_parallel_attributes
  from apex.transformer.log_util import get_transformer_logger
label_desc: Use image node_id 3 for calling a global function with example usage: _get_params_for_weight_decay_optimization(model) and returns: weight_decay_params,  no_weight_decay_params