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 optimizer_loading_epilogue(self):
      # Epilogue when loading state_dict to pytorch optimizer.
      if Version(torch.__version__) >= Version("2.0.0"):
          self.optim._patch_step_function()  # To support multiprocessing pickle/unpickle
      else:
          self.optim._hook_for_profile()  # To support multiprocessing pickle/unpickle.
      self.optim.defaults.setdefault("differentiable", False)
  
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 24 for calling the GeminiOptimizer obj's underlying member method code with example usage: obj.optimizer_loading_epilogue() without return types