class_docstr: A wrapper for optimizer. ``GeminiDDP`` and ``GeminiOptimizer`` implement Zero Redundancy Optimizer (ZeRO state-3).
  
  Note:
      You must use ``GeminiDDP`` with ``GeminiOptimizer``.
  
  Note:
      Make sure you set ``placement_policy`` of ``GeminiManager`` to `"auto"`,
      if you set ``gpu_margin_mem_ratio > 0``.
  
  Args:
      optim (Optimizer): An Optimizer instance.
      module (GeminiDDP): A ``GeminiDDP`` instance.
      gpu_margin_mem_ratio (float, optional): The ratio of GPU remaining memory (after the first forward-backward)
          which will be used when using hybrid CPU optimizer.
          This argument is meaningless when `placement_policy` of `GeminiManager` is not "auto".
          Defaults to 0.0.
      initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32.
      min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1.
      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.
      max_scale (int, optional): Max_scale used by DynamicGradScaler. Defaults to 2**32.
      max_norm (float, optional): The norm value used to clip gradient. Defaults to 0.0.
      norm_type (float, optional): The type of norm used for gradient clipping. Currently, only L2-norm (norm_type=2.0)
          is supported in GeminiOptimizer. Defaults to 2.0.
      verbose (bool, optional): Whether to print verbose information, including grad overflow info. Defaults to False.
class_docstr_tok: ["A","wrapper","for","optimizer",".","``","GeminiDDP","``","and","``","GeminiOptimizer","``","implement","Zero","Redundancy","Optimizer","(","ZeRO","state-3",")",".","Note",":","You","must","use","``","GeminiDDP","``","with","``","GeminiOptimizer","``",".","Note",":","Make","sure","you","set","``","placement_policy","``","of","``","GeminiManager","``","to","`","``","auto","''","`",",","if","you","set","``","gpu_margin_mem_ratio",">","0","``",".","Args",":","optim","(","Optimizer",")",":","An","Optimizer","instance",".","module","(","GeminiDDP",")",":","A","``","GeminiDDP","``","instance",".","gpu_margin_mem_ratio","(","float",",","optional",")",":","The","ratio","of","GPU","remaining","memory","(","after","the","first","forward-backward",")","which","will","be","used","when","using","hybrid","CPU","optimizer",".","This","argument","is","meaningless","when","`","placement_policy","`","of","`","GeminiManager","`","is","not","``","auto","''",".","Defaults","to","0.0.","initial_scale","(","float",",","optional",")",":","Initial","scale","used","by","DynamicGradScaler",".","Defaults","to","2","*","*","32.","min_scale","(","float",",","optional",")",":","Min","scale","used","by","DynamicGradScaler",".","Defaults","to","1.","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.","max_scale","(","int",",","optional",")",":","Max_scale","used","by","DynamicGradScaler",".","Defaults","to","2","*","*","32.","max_norm","(","float",",","optional",")",":","The","norm","value","used","to","clip","gradient",".","Defaults","to","0.0.","norm_type","(","float",",","optional",")",":","The","type","of","norm","used","for","gradient","clipping",".","Currently",",","only","L2-norm","(","norm_type=2.0",")","is","supported","in","GeminiOptimizer",".","Defaults","to","2.0.","verbose","(","bool",",","optional",")",":","Whether","to","print","verbose","information",",","including","grad","overflow","info",".","Defaults","to","False","."]
code: def __init__(
      self,
      optim: Optimizer,
      module: GeminiDDP,
      gpu_margin_mem_ratio: float = 0.0,
      initial_scale: float = 2**32,
      min_scale: float = 1,
      growth_factor: float = 2,
      backoff_factor: float = 0.5,
      growth_interval: int = 1000,
      hysteresis: int = 2,
      max_scale: float = 2**32,
      max_norm: float = 0.0,
      norm_type: float = 2.0,
      tp_group: ProcessGroup = None,
      optimizer_params_info=None,
      verbose: bool = False,
      **defaults: Any,
  ):
      super().__init__(optim)
      assert isinstance(module, GeminiDDP)
      assert type(optim) in _AVAIL_OPTIM_LIST, (
          "You should use an optimizer in the available list:\n"
          f"{_AVAIL_OPTIM_LIST}"
      )
      self.module = module
      self.gemini_manager = module.gemini_manager
      self.chunk_manager: ChunkManager = (
          self.gemini_manager.chunk_manager
      )
      self.param_to_range: Dict[Parameter, Tuple[int, int]] = dict()
      self.param_to_chunk16: Dict[Parameter, Chunk] = dict()
      self.chunk16_set: Set[Chunk] = set()
      self.clipping_flag = max_norm > 0.0
      self.max_norm = max_norm
      self.tp_group = tp_group
      self.optimizer_params_info = optimizer_params_info
      self.tp_size = (
          dist.get_world_size(tp_group) if tp_group is not None else 1
      )
      self.tp_rank = (
          dist.get_rank(tp_group) if tp_group is not None else 0
      )
      self.verbose = verbose
      self.param_groups_backup = list()
  
      # Mapping from integer id to real/fake param tensor, used for checkpointing.
      self.id_to_real_params: Dict[int, Parameter] = dict()
      self.id_to_fake_params: Dict[int, Parameter] = dict()
  
      if self.clipping_flag:
          assert (
              norm_type == 2.0
          ), "GeminiOptimizer only supports L2 norm now"
  
      ddp_param_list = []
      for name, param in module.named_parameters():
          if is_ddp_ignored(param):
              if param.requires_grad:
                  warnings.warn(
                      f"Parameter `{name}` is ignored by DDP but requires gradient! "
                      "You should handle its optimizer update by yourself!"
                  )
          else:
              ddp_param_list.append(param)
  
      for p in ddp_param_list:
          chunk_16 = self.chunk_manager.get_chunk(p)
          if chunk_16 not in self.chunk16_set:
              chunk_16.l2_norm_flag = self.clipping_flag
              self.chunk16_set.add(chunk_16)
  
      self.__init__optimizer()
  
      if module.mixed_precision is torch.float16:
          self.mix_precision_mixin = GeminiFP16MixedPrecisionMixin(
              module,
              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,
          )
      elif module.mixed_precision is torch.bfloat16:
          self.mix_precision_mixin = BF16MixedPrecisionMixin()
      else:
          raise RuntimeError(
              f"Unsupported mixed precision type: {module.mixed_precision}"
          )
  
      self._logger = get_dist_logger()
  
      self.gpu_margin_mem_ratio: float = float(gpu_margin_mem_ratio)
      assert (
          0.0 <= self.gpu_margin_mem_ratio <= 1.0
      ), f"gpu_margin_mem_ratio must >=0.0 and <=1.0"
      # Only move fp32 shards from CPU to GPU when user allows and inner optimizer is valid
      # Inner optimizer must support optimizing hybrid (CPU and CUDA) tensors,
      # and it must set `num_fp32_shards_per_param` correctly
      self._should_move_fp32_params_h2d: bool = (
          self.gemini_manager.is_cuda_margin_mem_avail
          and self.gpu_margin_mem_ratio > 0.0
          and getattr(optim, "num_fp32_shards_per_param", 0) >= 2
      )
      if (
          self.gpu_margin_mem_ratio > 0.0
          and not self.gemini_manager.is_cuda_margin_mem_avail
      ):
          self._logger.warning(
              f'gpu_margin_mem_ratio is meaningless when placement_policy is not "auto"',
              ranks=[0],
          )
  
      self._register_states = disposable(self._register_states_)
  
imports: import copy
  import math
  import warnings
  from typing import Any, Dict, Iterator, OrderedDict, Set, Tuple, Union
  import torch
  import torch.distributed
  from packaging.version import Version
  from torch.distributed import ProcessGroup
  from torch.nn import Parameter
  from torch.optim import Optimizer
  from colossalai.amp.naive_amp.mixed_precision_mixin import BF16MixedPrecisionMixin, FP16MixedPrecisionMixin
  from colossalai.checkpoint_io.utils import StateDictSharder, gather_distributed_param
  from colossalai.interface import OptimizerWrapper
  from colossalai.logging import get_dist_logger
  from colossalai.nn.optimizer import CPUAdam, FusedAdam, HybridAdam
  from colossalai.tensor.d_tensor import distribute_tensor, distribute_tensor_with_customization, get_device_mesh, get_sharding_spec, init_as_dtensor, init_tensor_as_customization_distributed, is_customized_distributed_tensor, is_distributed_tensor
  from colossalai.utils import disposable, get_current_device, is_ddp_ignored
  from .chunk import Chunk, ChunkManager
  from .gemini_ddp import GeminiDDP
label_desc: Use image node_id 1 to create a new GeminiOptimizer object from inherited base classes: OptimizerWrapper with example: obj = GeminiOptimizer(optim, module, gpu_margin_mem_ratio, initial_scale, min_scale, growth_factor, backoff_factor, growth_interval, hysteresis, max_scale, max_norm, norm_type, tp_group, optimizer_params_info, verbose)