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","."]
docstr: Args:
      param_id (int): id of the parameter whose state is to be gathered at master rank.
      only_rank_0(bool): if True, states will be collected only on master rank, otherwise collected on every rank.
  
  Returns:
      collected_states(dict): the gathered optimzier state of parameter with given id
                              if this method is called by master rank, otherwise an empty dict.
  
  This method can work only when called by all processes simultaneously.
docstr_tok: ["Args",":","param_id","(","int",")",":","id","of","the","parameter","whose","state","is","to","be","gathered","at","master","rank",".","only_rank_0","(","bool",")",":","if","True",",","states","will","be","collected","only","on","master","rank",",","otherwise","collected","on","every","rank",".","Returns",":","collected_states","(","dict",")",":","the","gathered","optimzier","state","of","parameter","with","given","id","if","this","method","is","called","by","master","rank",",","otherwise","an","empty","dict",".","This","method","can","work","only","when","called","by","all","processes","simultaneously","."]
code: def collect_states(
      self, param_id: int, only_rank_0: bool = True
  ) -> dict:
      """
      Args:
          param_id (int): id of the parameter whose state is to be gathered at master rank.
          only_rank_0(bool): if True, states will be collected only on master rank, otherwise collected on every rank.
  
      Returns:
          collected_states(dict): the gathered optimzier state of parameter with given id
                                  if this method is called by master rank, otherwise an empty dict.
  
      This method can work only when called by all processes simultaneously.
      """
  
      # Get param & chunk & process group.
      param = self.id_to_real_params[param_id]
      fake_param = self.id_to_fake_params.get(param_id, None)
      chunk = self.chunk_manager.get_chunk(param)
      zero_group = chunk.torch_pg
      rank = dist.get_rank(zero_group)
      master_rank = 0
      collected_states = {}
  
      # Fetch names of states through all_gather.
      local_state_names = None
      if fake_param is not None:
          local_state_names = list(self.optim.state[fake_param].keys())
      gathered_state_names = [
          None for _ in range(dist.get_world_size(zero_group))
      ]
      dist.barrier()
      dist.all_gather_object(
          gathered_state_names, local_state_names, zero_group
      )
      state_names = None
      for names in gathered_state_names:
          if names is not None:
              # Assume different devices share the same set of state names if they have.
              state_names = copy.deepcopy(names)
              break
  
      # Directly return if this parameter doesn't have optimizer states.
      # e.g. parameter freezed/layer dropped
      if state_names is None:
          return collected_states
  
      # Boolean variable is_collector indicates that whether the current rank
      # needs to gather the whole optimizer states.
      # Only master rank is collector when only_rank_0 is True.
      # Every rank is collector when only_rank_0 is False.
      is_collector = (rank == master_rank) or (not only_rank_0)
  
      # get tensor parallelism information
      is_dtensor = is_distributed_tensor(param)
      is_customized_distributed = is_customized_distributed_tensor(
          param
      )
      shard_spec = get_sharding_spec(param) if is_dtensor else None
      device_mesh = get_device_mesh(param) if is_dtensor else None
      global_shape = self.optimizer_params_info["id2shape"][param_id]
  
      # If the chunk is kept gathered,
      # the parameteres are treated the same as that of those in strict DDP during training.
      # So states can be directly fetched from current device.
      if chunk.keep_gathered:
          assert param_id in self.id_to_fake_params
          if is_collector:
              states = self.optim.state[fake_param]
              for state_name in state_names:
                  if state_name == "step":
                      # To keep aligned with pytorch, state 'step' is stored as a pytorch tensor with type float32.
                      collected_states[state_name] = torch.tensor(
                          states["step"],
                          dtype=torch.float32,
                          requires_grad=False,
                      ).cpu()
                  else:
                      state_tensor = (
                          states[state_name]
                          .detach()
                          .clone()
                          .to(torch.float32)
                          .cpu()
                      )
                      if is_dtensor:
                          state_tensor = torch.reshape(
                              state_tensor, param.shape
                          ).to(param.device)
                          state_tensor = init_as_dtensor(
                              state_tensor,
                              device_mesh=device_mesh,
                              sharding_spec=shard_spec,
                              global_shape=global_shape,
                          )
                      elif is_customized_distributed:
                          state_tensor = torch.reshape(
                              state_tensor, param.shape
                          ).to(param.device)
                          init_tensor_as_customization_distributed(
                              state_tensor,
                              shard_fn=param.shard_fn,
                              gather_fn=param.gather_fn,
                          )
                      state_tensor = gather_distributed_param(
                          state_tensor, keep_vars=False
                      ).cpu()
  
                      collected_states[
                          state_name
                      ] = state_tensor.reshape(global_shape)
          return collected_states
  
      # Check whether the param with given id is managed by current process.
      own_param = param_id in self.id_to_fake_params
  
      # Collector gets prepared for state collecting.
      if is_collector:
          for state_name in state_names:
              if state_name == "step":
                  # To keep aligned with pytorch, state 'step' is stored as a pytorch tensor with type float32.
                  collected_states[state_name] = torch.tensor(
                      0.0, dtype=torch.float32, requires_grad=False
                  ).cpu()
              else:
                  collected_states[state_name] = torch.zeros(
                      param.numel(),
                      dtype=torch.float32,
                      requires_grad=False,
                  ).cpu()
  
      # Materials for gathering, including compacted state tensors, and the offset of shard inside each state.
      compacted_states = (
          self.pack_optimizer_states_to_tensor(param_id, state_names)
          if own_param
          else None
      )
      _, shard_offset, shard_size = self.get_offsets(param_id)
  
      # Collectors gather state shards through all_gathering.
      gathered_state_shards = [
          None for _ in range(dist.get_world_size(zero_group))
      ]
  
      dist.barrier()
      dist.all_gather_object(
          gathered_state_shards,
          [compacted_states, shard_offset, shard_size],
          group=zero_group,
      )
  
      if is_collector:
          for state_shard in gathered_state_shards:
              compacted_states = state_shard[0]
              shard_offset = state_shard[1]
              shard_size = state_shard[2]
              if compacted_states is None:
                  continue
              self.load_from_compacted_states(
                  compacted_states,
                  collected_states,
                  state_names,
                  shard_offset,
                  shard_size,
              )
  
      # Reshape tensors
      if is_collector:
          for state_name, state_tensor in collected_states.items():
              if state_tensor.numel() == param.numel():
                  collected_states[state_name] = torch.reshape(
                      state_tensor, param.shape
                  )
              if is_dtensor:
                  state_tensor = state_tensor.to(param.device)
                  state_tensor = init_as_dtensor(
                      state_tensor,
                      sharding_spec=shard_spec,
                      device_mesh=device_mesh,
                      global_shape=global_shape,
                  )
              elif is_customized_distributed:
                  state_tensor = state_tensor.to(param.device)
                  init_tensor_as_customization_distributed(
                      state_tensor,
                      shard_fn=param.shard_fn,
                      gather_fn=param.gather_fn,
                  )
              state_tensor = gather_distributed_param(
                  state_tensor, keep_vars=False
              ).cpu()
  
      return collected_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 16 for calling the GeminiOptimizer obj's underlying member method code with example usage: obj.collect_states(param_id, only_rank_0) and returns: collected_states, collected_states, collected_states