docstr: Orgnize parameters' gradient(padding and split), follows the paramters' splitting method
  
  Data structure of self._grad_in_bucket:
  {
  rank0: [grad0_rank0, grad1_rank0, ...]
  rank1: [grad0_rank1, grad1_rank1, ...]
  }
docstr_tok: ["Orgnize","parameters","'","gradient","(","padding","and","split",")",",","follows","the","paramters","'","splitting","method","Data","structure","of","self._grad_in_bucket",":","{","rank0",":","[","grad0_rank0",",","grad1_rank0",",","...","]","rank1",":","[","grad0_rank1",",","grad1_rank1",",","...","]","}"]
code: def build_grad_in_bucket(self):
      """Orgnize parameters' gradient(padding and split), follows the paramters' splitting method
  
      Data structure of self._grad_in_bucket:
      {
      rank0: [grad0_rank0, grad1_rank0, ...]
      rank1: [grad0_rank1, grad1_rank1, ...]
      }
      """
      for param, padding_size in zip(
          self._param_list, self._padding_size
      ):
          grad = param.grad.clone().detach().flatten()
          if padding_size > 0:
              with torch.no_grad():
                  grad = torch.nn.functional.pad(
                      grad.view(-1), [0, padding_size]
                  )
          grad_list = grad.split(grad.numel() // self._world_size)
          for rank in range(self._world_size):
              grad_current_rank = grad_list[rank].clone().detach()
              self.grad_to_param_mapping[id(grad_current_rank)] = id(
                  param
              )
              self._grad_in_bucket[rank].append(grad_current_rank)
          param.grad = None
  
      self.offset_list.append(0)
  
file_path: colossalai/colossalai/zero/low_level/bookkeeping/bucket_store.py
imports: from typing import Dict
  import torch
  from torch import Tensor
  from torch._utils import _flatten_dense_tensors
  from torch.distributed import ProcessGroup
  from .base_store import BaseStore
label_desc: Use image node_id 5 for calling the BucketStore obj's underlying member method code with example usage: obj.build_grad_in_bucket() without return types