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.
code: def __getitem__(self, idx: int):
      assert isinstance(self.module, nn.ModuleList)
      return self.module[idx]
  
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 24 for calling the ShardedModelV2 obj's underlying member method code with example usage: obj.__getitem__(idx) and returns: self