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.
docstr: Copies parameters and buffers from :attr:`state_dict` into only
  this module, but not its descendants. This is called on every submodule
  in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this
  module in input :attr:`state_dict` is provided as :attr:`local_metadata`.
  For state dicts without metadata, :attr:`local_metadata` is empty.
  Subclasses can achieve class-specific backward compatible loading using
  the version number at `local_metadata.get("version", None)`.
  
  .. note::
      :attr:`state_dict` is not the same object as the input
      :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So
      it can be modified.
  
  Args:
      state_dict (dict): a dict containing parameters and
          persistent buffers.
      prefix (str): the prefix for parameters and buffers used in this
          module
      local_metadata (dict): a dict containing the metadata for this module.
          See
      strict (bool): whether to strictly enforce that the keys in
          :attr:`state_dict` with :attr:`prefix` match the names of
          parameters and buffers in this module
      missing_keys (list of str): if ``strict=True``, add missing keys to
          this list
      unexpected_keys (list of str): if ``strict=True``, add unexpected
          keys to this list
      error_msgs (list of str): error messages should be added to this
          list, and will be reported together in
          :meth:`~torch.nn.Module.load_state_dict`
docstr_tok: ["Copies","parameters","and","buffers","from",":","attr",":","`","state_dict","`","into","only","this","module",",","but","not","its","descendants",".","This","is","called","on","every","submodule","in",":","meth",":","`","~torch.nn.Module.load_state_dict","`",".","Metadata","saved","for","this","module","in","input",":","attr",":","`","state_dict","`","is","provided","as",":","attr",":","`","local_metadata","`",".","For","state","dicts","without","metadata",",",":","attr",":","`","local_metadata","`","is","empty",".","Subclasses","can","achieve","class-specific","backward","compatible","loading","using","the","version","number","at","`","local_metadata.get","(","``","version","''",",","None",")","`",".","..","note",":",":",":","attr",":","`","state_dict","`","is","not","the","same","object","as","the","input",":","attr",":","`","state_dict","`","to",":","meth",":","`","~torch.nn.Module.load_state_dict","`",".","So","it","can","be","modified",".","Args",":","state_dict","(","dict",")",":","a","dict","containing","parameters","and","persistent","buffers",".","prefix","(","str",")",":","the","prefix","for","parameters","and","buffers","used","in","this","module","local_metadata","(","dict",")",":","a","dict","containing","the","metadata","for","this","module",".","See","strict","(","bool",")",":","whether","to","strictly","enforce","that","the","keys","in",":","attr",":","`","state_dict","`","with",":","attr",":","`","prefix","`","match","the","names","of","parameters","and","buffers","in","this","module","missing_keys","(","list","of","str",")",":","if","``","strict=True","``",",","add","missing","keys","to","this","list","unexpected_keys","(","list","of","str",")",":","if","``","strict=True","``",",","add","unexpected","keys","to","this","list","error_msgs","(","list","of","str",")",":","error","messages","should","be","added","to","this","list",",","and","will","be","reported","together","in",":","meth",":","`","~torch.nn.Module.load_state_dict","`"]
code: def _load_from_state_dict(
      self,
      state_dict,
      prefix,
      local_metadata,
      strict,
      missing_keys,
      unexpected_keys,
      error_msgs,
  ):
      r"""Copies parameters and buffers from :attr:`state_dict` into only
      this module, but not its descendants. This is called on every submodule
      in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this
      module in input :attr:`state_dict` is provided as :attr:`local_metadata`.
      For state dicts without metadata, :attr:`local_metadata` is empty.
      Subclasses can achieve class-specific backward compatible loading using
      the version number at `local_metadata.get("version", None)`.
  
      .. note::
          :attr:`state_dict` is not the same object as the input
          :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So
          it can be modified.
  
      Args:
          state_dict (dict): a dict containing parameters and
              persistent buffers.
          prefix (str): the prefix for parameters and buffers used in this
              module
          local_metadata (dict): a dict containing the metadata for this module.
              See
          strict (bool): whether to strictly enforce that the keys in
              :attr:`state_dict` with :attr:`prefix` match the names of
              parameters and buffers in this module
          missing_keys (list of str): if ``strict=True``, add missing keys to
              this list
          unexpected_keys (list of str): if ``strict=True``, add unexpected
              keys to this list
          error_msgs (list of str): error messages should be added to this
              list, and will be reported together in
              :meth:`~torch.nn.Module.load_state_dict`
      """
      for hook in self._load_state_dict_pre_hooks.values():
          hook(
              state_dict,
              prefix,
              local_metadata,
              strict,
              missing_keys,
              unexpected_keys,
              error_msgs,
          )
  
      persistent_buffers = {
          k: v
          for k, v in self.named_buffers()
          if k not in self._non_persistent_buffers_set
      }
      local_name_params = itertools.chain(
          self.named_parameters(), persistent_buffers.items()
      )
      local_state = {
          k: v for k, v in local_name_params if v is not None
      }
  
      def load(
          param_name,
          dest_tensor,
          copy_func,
          source_device_mesh=None,
          source_sharding_spec=None,
          shard_fn=None,
          gather_fn=None,
      ):
          state_key = prefix + param_name
          if state_key in state_dict:
              input_param = state_dict[state_key]
  
              if (
                  source_device_mesh is not None
                  and source_sharding_spec is not None
              ):
                  input_param = distribute_tensor(
                      input_param,
                      source_device_mesh,
                      source_sharding_spec,
                  )
              elif shard_fn is not None and gather_fn is not None:
                  input_param = distribute_tensor_with_customization(
                      input_param,
                      shard_fn=shard_fn,
                      gather_fn=gather_fn,
                  )
  
              # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+
              if (
                  len(dest_tensor.shape) == 0
                  and len(input_param.shape) == 1
              ):
                  input_param = input_param[0]
              if input_param.shape != dest_tensor.shape:
                  # local shape should match the one in checkpoint
                  error_msgs.append(
                      "size mismatch for {}: copying a param with shape {} from checkpoint, "
                      "the shape in current model is {}.".format(
                          state_key,
                          input_param.shape,
                          dest_tensor.shape,
                      )
                  )
                  return
              try:
                  with torch.no_grad():
                      copy_func(input_param)
              except Exception as ex:
                  error_msgs.append(
                      'While copying the parameter named "{}", '
                      "whose dimensions in the model are {} and "
                      "whose dimensions in the checkpoint are {}, "
                      "an exception occurred : {}.".format(
                          state_key,
                          dest_tensor.size(),
                          input_param.size(),
                          ex.args,
                      )
                  )
          elif strict:
              missing_keys.append(state_key)
  
      def load_parameter(chunk_slice, data):
          chunk_slice.copy_(data.flatten())
  
      for name, param in self.named_parameters():
          if is_ddp_ignored(param):
              # deal with ddp ignored parameters
              load(name, param, param.copy_)
  
      fp32_to_name = dict()
      for p, fp32_p in zip(self.fp16_params, self.fp32_params):
          if p is not None:
              name = self.param2name[p]
              fp32_to_name[fp32_p] = name
  
      params_to_load = (
          self.fp32_params
          if self.reuse_fp16_chunk
          else self.fp16_params
      )
      chunk_list = self.chunk_manager.get_chunks(params_to_load)
      for chunk in chunk_list:
          temp_chunk = get_temp_total_chunk_on_cuda(
              chunk, self.mixed_precision
          )
  
          for tensor, tensor_info in chunk.tensors_info.items():
              (
                  source_device_mesh,
                  source_sharding_spec,
                  shard_fn,
                  gather_fn,
              ) = (None, None, None, None)
              if is_distributed_tensor(tensor):
                  # shard the input param
                  source_device_mesh = get_device_mesh(tensor)
                  source_sharding_spec = get_sharding_spec(tensor)
              elif is_customized_distributed_tensor(tensor):
                  shard_fn = tensor.shard_fn
                  gather_fn = tensor.gather_fn
  
              parameter_name = (
                  fp32_to_name[tensor]
                  if self.reuse_fp16_chunk
                  else self.param2name[tensor]
              )
              parameter_slice = temp_chunk[
                  tensor_info.offset : tensor_info.end
              ]
              load(
                  parameter_name,
                  tensor,
                  partial(load_parameter, parameter_slice),
                  source_device_mesh,
                  source_sharding_spec,
                  shard_fn,
                  gather_fn,
              )
  
          if chunk.is_gathered:
              chunk.cuda_global_chunk.copy_(temp_chunk)
          elif chunk.cuda_shard is not None:
              chunk.cuda_shard.copy_(
                  temp_chunk[chunk.shard_begin : chunk.shard_end]
              )
          else:
              chunk.cpu_shard.copy_(
                  temp_chunk[chunk.shard_begin : chunk.shard_end]
              )
  
          del temp_chunk
      if self.reuse_fp16_chunk:
          for chunk_32 in chunk_list:
              chunk_16 = chunk_32.paired_chunk
              assert chunk_16 is not None
              chunk_16.payload.copy_(chunk_32.payload)
  
      for name, buf in persistent_buffers.items():
          if buf is not None:
              load(name, buf, buf.copy_)
  
      extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
      if (
          getattr(
              self.__class__,
              "set_extra_state",
              torch.nn.Module.set_extra_state,
          )
          is not torch.nn.Module.set_extra_state
      ):
          if extra_state_key in state_dict:
              self.set_extra_state(state_dict[extra_state_key])
          elif strict:
              missing_keys.append(extra_state_key)
      elif strict and (extra_state_key in state_dict):
          unexpected_keys.append(extra_state_key)
  
      if strict:
          for key in state_dict.keys():
              if key.startswith(prefix) and key != extra_state_key:
                  input_name = key[len(prefix) :]
                  if input_name not in local_state:
                      unexpected_keys.append(key)
  
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 25 for calling the GeminiDDP obj's underlying member method code with example usage: obj._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) without return types