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.
class_docstr_tok: ["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","."]
docstr: Returns dictionaries containing a whole state of the module one by one. The max size of dictionary shard is specified by ``max_shard_size``.
  
  Both parameters and persistent buffers (e.g. running averages) are included.
  Keys are corresponding parameter and buffer names.
  Parameters and buffers set to ``None`` are not included.
  
  Args:
      prefix (str, optional): the prefix for parameters and buffers used in this
          module. Defaults to ''.
      keep_vars (bool, optional): whether to keep variables. Defaults to False.
      max_shard_size (int, optional): max size of state dict shard (in MB). Defaults to 1024.
      only_rank_0 (bool, optional): only get data on rank0. Defaults to True.
  
  
  Yields:
      Iterator[OrderedDict]: A generator of state dict shard
code: def state_dict_shard(
      self,
      prefix: str = "",
      keep_vars: bool = False,
      max_shard_size: int = 1024,
      only_rank_0: bool = True,
  ) -> Iterator[Tuple[OrderedDict, int]]:
      """Returns dictionaries containing a whole state of the module one by one. The max size of dictionary shard is specified by ``max_shard_size``.
  
      Both parameters and persistent buffers (e.g. running averages) are included.
      Keys are corresponding parameter and buffer names.
      Parameters and buffers set to ``None`` are not included.
  
      Args:
          prefix (str, optional): the prefix for parameters and buffers used in this
              module. Defaults to ''.
          keep_vars (bool, optional): whether to keep variables. Defaults to False.
          max_shard_size (int, optional): max size of state dict shard (in MB). Defaults to 1024.
          only_rank_0 (bool, optional): only get data on rank0. Defaults to True.
  
  
      Yields:
          Iterator[OrderedDict]: A generator of state dict shard
      """
      sharder = StateDictSharder(max_shard_size)
  
      # get the mapping between copies and fp16 parameters
      fp16_to_fp32 = dict()
      for p, fp32_p in zip(self.fp16_params, self.fp32_params):
          fp16_to_fp32[p] = fp32_p
  
      # key is fp32 param, and value is gathered param on CPU
      gathered_param_buffer = dict()
      for name, param in self.name2param.items():
          if param is not None:
              if is_ddp_ignored(param):
                  # deal with ddp ignored parameters
                  gathered_param = (
                      param if keep_vars else param.detach()
                  )
              else:
                  # as memory of fp16 param may be reused, we should use fp32 param and then convert to fp16
                  param_to_save = (
                      fp16_to_fp32[param]
                      if self.reuse_fp16_chunk
                      else param
                  )
                  if param_to_save not in gathered_param_buffer:
                      chunk = self.chunk_manager.get_chunk(
                          param_to_save
                      )
                      gathered_param_buffer.update(
                          self._get_chunk_to_save_data(
                              chunk, only_rank_0
                          )
                      )
                  gathered_param = gathered_param_buffer.pop(
                      param_to_save
                  )
  
              block, block_size = sharder.append_param(
                  prefix + name, gathered_param
              )
              if block is not None:
                  yield block, block_size
  
      del fp16_to_fp32
      del gathered_param_buffer
  
      # save all buffers
      for name, buf in self.named_buffers():
          if (
              buf is not None
              and name not in self._non_persistent_buffers_set
          ):
              buffer = buf if keep_vars else buf.detach()
              block, block_size = sharder.append_param(
                  prefix + name, buffer
              )
              if block is not None:
                  yield block, block_size
      # save extra states
      extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
      if (
          getattr(
              self.__class__,
              "get_extra_state",
              torch.nn.Module.get_extra_state,
          )
          is not torch.nn.Module.get_extra_state
      ):
          extra_state = self.get_extra_state()
          block, block_size = sharder.append_param(
              extra_state_key, extra_state
          )
          if block is not None:
              yield block, block_size
  
      yield sharder.current_block, sharder.current_block_size
  
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 29 for calling the GeminiDDP obj's underlying member method code with example usage: obj.state_dict_shard(prefix, keep_vars, max_shard_size, only_rank_0) without return types