code: def __init__(self, *args, partition_grad: bool = False):
      super().__init__(*args)
      """
          self._grads_of_params mapping the paramater and its gradient slices
          data structure:
          {
           group_id:{
              param_id: [grad_rank0, grad_rank1, ...]
            }
          }
          """
      self._grads_of_params = dict()
      # for zero2, it's `param_id: [grad_local_rank]`
      self._working_index = 0 if partition_grad else self._local_rank
  
      self.grad_to_param_mapping = dict()
  
file_path: ColossalAI/colossalai/zero/low_level/bookkeeping/gradient_store.py
imports: from typing import List
  from torch import Tensor
  from .base_store import BaseStore
label_desc: Use image node_id 1 to create a new GradientStore object from inherited base classes: BaseStore with example: obj = GradientStore()