class_docstr: ZeRO DDP.
  Warning: Nested GeminiDDP is not supported now.
  It is designed to be used with ChunkManager and GeminiManager.
  For more details, see the API reference of ``ChunkManager`` and ``GeminiManager``.
  
  Args:
      module (torch.nn.Module): Module to apply ZeRO-DP.
      gemini_manager (GeminiManager): Manages the chunk manager and heterogeneous memory space.
          For more details, see the API reference of ``GeminiManager``.
      pin_memory (bool): Chunks on CPU Memory use pin-memory.
      force_outputs_fp32 (bool): If set to True, outputs will be fp32. Otherwise, outputs will be fp16.
          Defaults to False.
      strict_ddp_mode (bool): If set to True, there is no tensor sharding, each tensor is replicated.
          Defaults to False. Users can set it to True, when they clearly know that they only need DDP.
      scatter_after_inference (bool): If set to True, the model will be scattered after inference. This will save memory but slow down the consecutive inference.
      mixed_precision (torch.dtype): If set to torch.float16, the model will be trained in fp16. Otherwise, the model will be trained in bf16. Defaults to torch.float16.
code: def __init__(
      self,
      module: torch.nn.Module,
      chunk_config_dict: Optional[dict] = None,
      chunk_init_device: torch.device = torch.device("cpu"),
      placement_policy: str = "static",
      enable_gradient_accumulation: bool = False,
      shard_param_frac: float = 1.0,  # only for static placement
      offload_optim_frac: float = 0.0,  # only for static placement
      offload_param_frac: float = 0.0,  # only for static placement
      warmup_non_model_data_ratio: float = 0.8,  # only for auto placement
      steady_cuda_cap_ratio: float = 0.9,  # only for auto placement
      search_range_m: int = 32,  # chunk search options
      hidden_dim: Optional[int] = None,  # chunk search options
      min_chunk_size_m: float = 32,  # chunk search options
      pin_memory: bool = False,
      force_outputs_fp32: bool = False,
      strict_ddp_mode: bool = False,
      scatter_after_inference: bool = True,
      mixed_precision: torch.dtype = torch.float16,
      zero_group: Optional[ProcessGroup] = None,
      memstats: Optional[MemStats] = None,  # genimi memory stats
      master_weights: bool = True,
      extra_dp_group: Optional[ProcessGroup] = None,
      verbose: bool = False,
  ) -> None:
      assert mixed_precision in (torch.float16, torch.bfloat16)
      if chunk_config_dict is not None:
          self.chunk_manager = ChunkManager(
              chunk_config_dict, chunk_init_device
          )
      else:
          # some ugly hotfix for the compatibility with Lightning
          if search_range_m is None:
              search_range_m = 32
          self.chunk_manager = init_chunk_manager(
              model=module,
              init_device=chunk_init_device,
              hidden_dim=hidden_dim,
              search_range_m=search_range_m,
              min_chunk_size_m=min_chunk_size_m,
              strict_ddp_flag=strict_ddp_mode,
              process_group=zero_group,
              verbose=verbose,
          )
      self.gemini_manager = GeminiManager(
          placement_policy,
          self.chunk_manager,
          memstats,
          shard_param_frac=shard_param_frac,
          offload_optim_frac=offload_optim_frac,
          offload_param_frac=offload_param_frac,
          warmup_non_model_data_ratio=warmup_non_model_data_ratio,
          steady_cuda_cap_ratio=steady_cuda_cap_ratio,
      )
      self.force_outputs_fp32 = force_outputs_fp32
      self.param_op_hook = GeminiZeROHook(self.gemini_manager)
      self.fp32_params: List[torch.Tensor] = list()
      self.fp16_params: List[ColoParameter] = list()
      self.overflow_counter = 0
      self.grads_device: Dict[torch.Tensor, torch.device] = dict()
      self.param2name: Dict[nn.Parameter, str] = dict()
      self.name2param: Dict[str, nn.Parameter] = dict()
      self.scatter_after_inference = scatter_after_inference
      self.mixed_precision = mixed_precision
      self.zero_group = zero_group or _get_default_group()
      self.extra_dp_group = extra_dp_group
  
      self.reuse_fp16_chunk = master_weights
      self.master_weights = master_weights
  
      self.enable_gradient_accumulation = enable_gradient_accumulation
      if self.enable_gradient_accumulation:
          self.reuse_fp16_chunk = False
      self.accumulating_grads = (
          False  # Whether model is accumulating gradients
      )
  
      self._logger = get_dist_logger()
  
      if self.gemini_manager._premade_memstats_:
          # build chunk in param runtime visited order.
          param_order = (
              self.gemini_manager.memstats()._param_runtime_order
          )
      else:
          # build chunk in param initialized order.
          # Note: in this way, it can not get filter unused params during runtime.
          param_order = OrderedParamGenerator()
          for p in module.parameters():
              param_order.append(p)
  
      for name, param in module.named_parameters():
          self.param2name[param] = name
      for m_name, m_var in module.named_modules():
          for p_name, p_var in m_var.named_parameters(recurse=False):
              param_name = m_name + "." + p_name if m_name else p_name
              self.name2param[param_name] = p_var
  
      self._init_chunks(
          param_order=param_order,
          strict_ddp_mode=strict_ddp_mode,
          cpu_offload=not (
              self.gemini_manager.policy_name == "static"
              and offload_param_frac == 0
          ),
          pin_memory=pin_memory,
      )
      super().__init__(module)
      self._non_persistent_buffers_set = (
          self._get_non_persistent_buffers_set(module)
      )
      self._cast_buffers()
      # register grad hook
      for p in module.parameters():
          if is_ddp_ignored(p):
              continue
          if p.requires_grad:
              p.register_hook(partial(self.grad_handle, p))
  
imports: import itertools
  from collections import OrderedDict
  from contextlib import nullcontext
  from functools import partial
  from typing import Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed import ProcessGroup
  from torch.distributed.distributed_c10d import _get_default_group
  from colossalai.checkpoint_io.utils import StateDictSharder, gather_distributed_param
  from colossalai.interface import ModelWrapper
  from colossalai.lazy import LazyTensor
  from colossalai.logging import get_dist_logger
  from colossalai.tensor.colo_parameter import ColoParameter
  from colossalai.tensor.d_tensor import distribute_tensor, distribute_tensor_with_customization, get_device_mesh, get_global_shape, get_sharding_spec, init_as_dtensor, init_tensor_as_customization_distributed, is_customized_distributed_tensor, is_distributed_tensor
  from colossalai.tensor.param_op_hook import ColoParamOpHookManager
  from colossalai.utils import _cast_float, free_storage, get_current_device, is_ddp_ignored
  from .chunk import Chunk, ChunkManager, TensorState, init_chunk_manager
  from .gemini_hook import GeminiZeROHook
  from .gemini_mgr import GeminiManager
  from .memory_tracer import MemStats, OrderedParamGenerator
  from .utils import get_temp_total_chunk_on_cuda
label_desc: Use image node_id 1 to create a new GeminiDDP object from inherited base classes: ModelWrapper with example: obj = GeminiDDP(module, chunk_config_dict, chunk_init_device, placement_policy, enable_gradient_accumulation, shard_param_frac, offload_optim_frac, offload_param_frac, warmup_non_model_data_ratio, steady_cuda_cap_ratio, search_range_m, hidden_dim, min_chunk_size_m, pin_memory, force_outputs_fp32, strict_ddp_mode, scatter_after_inference, mixed_precision, zero_group, memstats, master_weights, extra_dp_group, verbose)