class_docstr: A wrapper for the PyTorch module shards the model parameters among multiple GPU memory.
  Only `1/#nproc` of parameters, gradients are stored in local CUDA memory, so forward and backward
  passes can be executed with limited CUDA memory budget.
  
  Note:
      You must use ``ShardedModelV2`` with ``ShardedOptimizerV2``.
  Note:
      Make sure you don't use gradient accumulation and your optimizer can work with fp16 gradient and fp32 parameter,
      if you enable ``reuse_fp16_shard``.
  
  Args:
      module (nn.Module): A sharded module, which must be initialized by `ZeroInitContext`.
      shard_strategy (BaseShardStrategy): A shard strategy to manage shard behavior.
      process_group (Optional[ProcessGroup], optional): Data parallel process group. Defaults to None.
      reduce_scatter_process_group (Optional[ProcessGroup], optional): Reduce-scatter process group.
          Generally, it should be `None`, and it's the same as `process_group`. Defaults to None.
      reduce_scatter_bucket_size_mb (int, optional): Reduce-scatter bucket size in *MB*. Defaults to 25.
      fp32_reduce_scatter (bool, optional): If set to `True`, gradients are forced to FP32 before reduce-scatter. Defaults to False.
      tensor_placement_policy (str): Which device to place *held* tensors. It can be 'cpu', 'cuda' and 'auto'.
          If it's 'cpu', parameters, gradients and optimizer states will be offloaded to CPU, which means min CUDA memory will be used.
          If it's 'cuda', they won't be offloaded, which means max CUDA memory will be used.
          If it's 'auto', they are moving dynamically based on CPU and CUDA memory usage. It will utilize heterogeneous memory space evenly and well.
          Note that 'auto' policy can only work well when no other processes use CUDA during your training.
          Defaults to 'cuda'.
      gradient_predivide_factor (Optional[float], optional): Gradient is divided by this value before reduce-scatter. Defaults to 1.0.
      reuse_fp16_shard (bool, optional): Whether to reuse fp16 shard for param and grad.
          Enabling this can reduce GPU memory usage, but you have to make sure you disable it when using gradient accumulation.
          In this mode, grad will be fp16. Make sure your optimizer supports mixed precision (fp32 param and fp16 grad).
          We find that PyTorch's optimizers don't support mixed precision,
          so we recommend you enable this only when using our CPUAdam with CPU offload. Defaults to False.
      bf16 (bool, optional): Whether to use bfloat16 for param and grad. Defaults to False.
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`
      shard_strategy (Optional[BaseShardStrategy], optional): A shard strategy to manage shard behavior. Defaults to None.
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","`","shard_strategy","(","Optional","[","BaseShardStrategy","]",",","optional",")",":","A","shard","strategy","to","manage","shard","behavior",".","Defaults","to","None","."]
code: def _colo_load_from_state_dict(
      self,
      state_dict,
      prefix,
      local_metadata,
      strict,
      missing_keys,
      unexpected_keys,
      error_msgs,
      shard_strategy=None,
  ):
      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`
          shard_strategy (Optional[BaseShardStrategy], optional): A shard strategy to manage shard behavior. Defaults to None.
      """
      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._buffers.items()
          if k not in self._non_persistent_buffers_set
      }
      local_name_params = itertools.chain(
          self._parameters.items(), persistent_buffers.items()
      )
      local_state = {
          k: v for k, v in local_name_params if v is not None
      }
  
      for name, param in local_state.items():
          key = prefix + name
          if key in state_dict:
              input_param = state_dict[key]
              if hasattr(param, "colo_attr"):
                  param.colo_attr.data_payload_reset(
                      input_param.to(
                          dtype=param.colo_attr.data_payload.dtype,
                          device=param.colo_attr.data_payload.device,
                      )
                  )
                  if shard_strategy is not None:
                      # Force re-shard
                      param.colo_attr.sharded_data_tensor.is_sharded = (
                          False
                      )
                      shard_strategy.shard(
                          [param.colo_attr.sharded_data_tensor]
                      )
              else:
                  # This is used to avoid copying uninitialized parameters into
                  # non-lazy modules, since they dont have the hook to do the checks
                  # in such case, it will error when accessing the .shape attribute.
                  is_param_lazy = torch.nn.parameter.is_lazy(param)
                  # Backward compatibility: loading 1-dim tensor from 0.3.* to version 0.4+
                  if (
                      not is_param_lazy
                      and len(param.shape) == 0
                      and len(input_param.shape) == 1
                  ):
                      input_param = input_param[0]
  
                  if (
                      not is_param_lazy
                      and input_param.shape != param.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(
                              key, input_param.shape, param.shape
                          )
                      )
                      continue
                  try:
                      with torch.no_grad():
                          param.copy_(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(
                              key,
                              param.size(),
                              input_param.size(),
                              ex.args,
                          )
                      )
          elif strict:
              missing_keys.append(key)
  
      extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
      if (
          getattr(
              self.__class__,
              "set_extra_state",
              nn.Module.set_extra_state,
          )
          is not 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) :]
                  input_name = input_name.split(".", 1)[
                      0
                  ]  # get the name of param/buffer/child
                  if (
                      input_name not in self._modules
                      and input_name not in local_state
                  ):
                      unexpected_keys.append(key)
  
imports: import functools
  import itertools
  from collections import OrderedDict
  from typing import Any, Iterator, Optional, Tuple
  import torch
  import torch.distributed
  import torch.nn
  from torch.distributed import ProcessGroup
  from torch.nn.parameter import Parameter
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.legacy.utils.memory import colo_device_memory_capacity
  from colossalai.legacy.zero.gemini.ophooks import register_ophooks_recursively
  from colossalai.legacy.zero.gemini.paramhooks import BaseParamHookMgr
  from colossalai.legacy.zero.gemini.stateful_tensor import TensorState
  from colossalai.legacy.zero.gemini.stateful_tensor_mgr import StatefulTensorMgr
  from colossalai.legacy.zero.gemini.tensor_placement_policy import TensorPlacementPolicy, TensorPlacementPolicyFactory
  from colossalai.legacy.zero.gemini.tensor_utils import colo_model_data_move_to_cpu
  from colossalai.legacy.zero.shard_utils import BaseShardStrategy
  from colossalai.legacy.zero.sharded_model.reduce_scatter import ReduceScatterBucketer
  from colossalai.logging import get_dist_logger
  from colossalai.utils import disposable, get_current_device
  from colossalai.zero.gemini.memory_tracer import MemStatsCollector
  from ._utils import cast_float_arguments, cast_tensor_to_bf16, cast_tensor_to_fp16, cast_tensor_to_fp32, chunk_and_pad, free_storage, get_gradient_predivide_factor
  from .zero_hook import ZeroHook
label_desc: Use image node_id 23 for calling the ShardedModelV2 obj's underlying member method code with example usage: obj._colo_load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, shard_strategy) without return types