class_docstr: A wrapper for Optimizer.
  Code reference: https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/nn/optimizer/zero_optimizer.py
  
  Args:
      optimizer (Optimizer): An Optimizer instance.
      module (BaseOffloadModule): A ``BaseOffloadModule`` instance.
      initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**16.
      growth_factor (float, optional): growth_factor used by DynamicGradScaler. Defaults to 2.
      backoff_factor (float, optional): backoff_factor used by DynamicGradScaler. Defaults to 0.5.
      growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000.
      hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2.
      min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1.
      max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32.
      norm_type (float, optional): norm_type used for `clip_grad_norm`.
class_docstr_tok: ["A","wrapper","for","Optimizer",".","Code","reference",":","https",":","\/\/github.com\/hpcaitech\/ColossalAI\/blob\/main\/colossalai\/nn\/optimizer\/zero_optimizer.py","Args",":","optimizer","(","Optimizer",")",":","An","Optimizer","instance",".","module","(","BaseOffloadModule",")",":","A","``","BaseOffloadModule","``","instance",".","initial_scale","(","float",",","optional",")",":","Initial","scale","used","by","DynamicGradScaler",".","Defaults","to","2","*","*","16.","growth_factor","(","float",",","optional",")",":","growth_factor","used","by","DynamicGradScaler",".","Defaults","to","2.","backoff_factor","(","float",",","optional",")",":","backoff_factor","used","by","DynamicGradScaler",".","Defaults","to","0.5.","growth_interval","(","float",",","optional",")",":","growth_interval","used","by","DynamicGradScaler",".","Defaults","to","1000.","hysteresis","(","float",",","optional",")",":","hysteresis","used","by","DynamicGradScaler",".","Defaults","to","2.","min_scale","(","float",",","optional",")",":","Min","scale","used","by","DynamicGradScaler",".","Defaults","to","1.","max_scale","(","int",",","optional",")",":","max_scale","used","by","DynamicGradScaler",".","Defaults","to","2","*","*","32.","norm_type","(","float",",","optional",")",":","norm_type","used","for","`","clip_grad_norm","`","."]
code: def step(self, *args, **kwargs):
      # Copy gradients from model params to main params.
      self._set_grad_ptr()
  
      found_inf = self._check_overflow()
      if found_inf:
          self.optim_state = (
              OptimState.UNSCALED
          )  # no need to unscale grad
          self.grad_scaler.update(found_inf)  # update gradient scaler
          self._logger.info(f"Found overflow. Skip step")
          self.zero_grad()  # reset all gradients
          self._update_fp16_params()
          return
  
      # get combined scale. combined scale = loss scale * clipping norm
      # so that gradient = gradient / combined scale
      combined_scale = self._get_combined_scale()
      self.grad_scaler.update(found_inf)
  
      ret = self.optim.step(div_scale=combined_scale, *args, **kwargs)
      self.zero_grad()
      self._update_fp16_params()
      return ret
  
imports: from enum import Enum
  from typing import Dict, Tuple
  import torch
  from torch.optim import Optimizer
  from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler
  from colossalai.interface import OptimizerWrapper
  from colossalai.logging import get_dist_logger
  from colossalai.utils import get_current_device
  from .base_offload_module import BaseOffloadModule
  from .region import Region
  from .region_manager import RegionManager
label_desc: Use image node_id 8 for calling the AMPOptimizer obj's underlying member method code with example usage: obj.step() and returns: ret