code: def zero_optim_wrapper(
      model: nn.Module,
      optimizer: torch.optim.Optimizer,
      initial_scale: float = 2**16,
      growth_factor: float = 2,
      backoff_factor: float = 0.5,
      growth_interval: int = 1000,
      hysteresis: int = 2,
      min_scale: float = 1,
      max_scale: float = 2**32,
      max_norm: float = 0.0,
      norm_type: float = 2.0,
      optim_config: Optional[Dict] = None,
      verbose: bool = False,
  ):
      """This wrapper function is used to wrap your training optimizer for ZeRO DDP.
  
      Args:
          model (nn.Module): Your model wrapped by `zero_model_wrapper`
          optimizer (torch.optim.Optimizer): Your initialized optimizer
          initial_scale (float, optional): initial_scale used by DynamicGradScaler.
          min_scale (float, optional): min_scale used by DynamicGradScaler.
          growth_factor (float, optional): growth_factor used by DynamicGradScaler.
          backoff_factor (float, optional): backoff_factor used by DynamicGradScaler.
          growth_interval (float, optional): growth_interval used by DynamicGradScaler.
          hysteresis (float, optional): hysteresis used by DynamicGradScaler.
          max_scale (int, optional): max_scale used by DynamicGradScaler.
          max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do
              clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm.
          norm_type (float, optional): norm_type used for `clip_grad_norm`.
          optim_config (dict, optional): The configuration used for the ZeRO optimizer.
              Example:
  
                  >>> zero2_config = dict(reduce_bucket_size=12 * 1024 * 1024, overlap_communication=True)
                  >>> optim = zero_optim_wrapper(model, optim, optim_config=zero2_config)
          verbose (bool, optional): Whether to print the verbose info.
      """
      assert hasattr(
          model, "_colo_zero_stage"
      ), "You should use `zero_ddp_wrapper` first"
      zero_stage = getattr(model, "_colo_zero_stage")
  
      assert (
          norm_type == 2.0
      ), "Current ZeRO optimizers only support 'norm_type=2'"
  
      if optim_config is None:
          config_dict = dict()
      else:
          config_dict = copy(optim_config)
  
      config_dict["initial_scale"] = initial_scale
      config_dict["growth_factor"] = growth_factor
      config_dict["backoff_factor"] = backoff_factor
      config_dict["growth_interval"] = growth_interval
      config_dict["hysteresis"] = hysteresis
      config_dict["min_scale"] = min_scale
      config_dict["max_scale"] = max_scale
  
      if zero_stage in [1, 2]:
          from colossalai.zero.low_level import LowLevelZeroOptimizer
  
          config_dict["partition_grad"] = zero_stage == 2
          config_dict["clip_grad_norm"] = max_norm
          return LowLevelZeroOptimizer(
              optimizer, **config_dict, verbose=verbose
          )
      else:
          from colossalai.zero.gemini.gemini_optimizer import (
              GeminiOptimizer,
          )
  
          config_dict["clipping_norm"] = max_norm
          return GeminiOptimizer(
              optimizer, model, **config_dict, verbose=verbose
          )
  
imports: from copy import copy
  from typing import Dict, Optional
  import torch
  import torch.nn
  from .gemini import GeminiDDP
label_desc: Use image node_id 2 for calling a global function with example usage: zero_optim_wrapper(model, optimizer, initial_scale, growth_factor, backoff_factor, growth_interval, hysteresis, min_scale, max_scale, max_norm, norm_type, optim_config, verbose) and returns: LowLevelZeroOptimizer, GeminiOptimizer