docstr: Reduce scatter all the gradients. It's an operation done in CUDA.
docstr_tok: ["Reduce","scatter","all","the","gradients",".","It","'s","an","operation","done","in","CUDA","."]
code: def reduce(self):
      """Reduce scatter all the gradients. It's an operation done in CUDA."""
      # sanity check
      assert self.is_gathered
  
      if self.pg_size == 1:
          # tricky code here
          # just move cuda_global_chunk to cuda_shard
          # the communication is not necessary
          self.__scatter()
          if self.extra_dp_group is not None:
              dist.all_reduce(
                  self.cuda_shard, group=self.extra_dp_group
              )
      elif self.keep_gathered:
          # we use all-reduce here
          dist.all_reduce(self.cuda_global_chunk, group=self.torch_pg)
          if self.extra_dp_group is not None:
              dist.all_reduce(
                  self.cuda_global_chunk, group=self.extra_dp_group
              )
      else:
          self.cuda_shard = torch.empty(
              self.shard_size,
              dtype=self.dtype,
              device=get_current_device(),
          )
  
          input_list = list(
              torch.chunk(
                  self.cuda_global_chunk, chunks=self.pg_size, dim=0
              )
          )
          dist.reduce_scatter(
              self.cuda_shard, input_list, group=self.torch_pg
          )
          if self.extra_dp_group is not None:
              dist.all_reduce(
                  self.cuda_shard, group=self.extra_dp_group
              )
  
          free_storage(self.cuda_global_chunk)
          self.is_gathered = False
      self.__update_tensors_state(TensorState.HOLD)
  
imports: from dataclasses import dataclass
  from enum import Enum
  from typing import Dict, List, Optional
  import torch
  import torch.distributed
  from torch.distributed import ProcessGroup
  from colossalai.utils import get_current_device
  from colossalai.utils.device import IS_NPU_AVAILABLE
label_desc: Use image node_id 16 for calling the Chunk obj's underlying member method code with example usage: obj.reduce() without return types