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 __init__(
      self,
      optimizer: Optimizer,
      module: BaseOffloadModule,
      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,
      clipping_norm: float = 0.0,
      norm_type: float = 2.0,
  ):
      super().__init__(optimizer)
  
      self.module = module
      self.optim_state = OptimState.UNSCALED
      self.clipping_flag = clipping_norm > 0.0
      self.max_norm = clipping_norm
  
      self.region_manager: RegionManager = self.module.region_manager
      self.param_to_range: Dict[
          torch.nn.Parameter, Tuple[int, int]
      ] = dict()
      self.param_to_region: Dict[torch.nn.Parameter, Region] = dict()
  
      self.fp32_to_fp16_params: Dict[
          torch.Tensor, torch.nn.Parameter
      ] = dict()
  
      if self.clipping_flag:
          assert (
              norm_type == 2.0
          ), "AMPOptimizer only supports L2 norm now"
  
      self.__init__optimizer()
  
      # Grad scaler
      self.grad_scaler = DynamicGradScaler(
          initial_scale=initial_scale,
          min_scale=min_scale,
          growth_factor=growth_factor,
          backoff_factor=backoff_factor,
          growth_interval=growth_interval,
          hysteresis=hysteresis,
          max_scale=max_scale,
      )
      self._found_overflow: torch.Tensor = torch.zeros(
          1, dtype=torch.int64, device=get_current_device()
      )
      self._logger = get_dist_logger()
  
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 1 to create a new AMPOptimizer object from inherited base classes: OptimizerWrapper with example: obj = AMPOptimizer(optimizer, module, initial_scale, growth_factor, backoff_factor, growth_interval, hysteresis, min_scale, max_scale, clipping_norm, norm_type)