docstr: Chunk: A container owning a piece of contiguous memory space for tensors
  Here we use all-gather operation to gather the whole chunk.
  Currently, Chunk is exclusively used for DDP and ZeRO DDP and it doesn't support unused parameters.
  It is designed to make the full use of communication and PCIE bandwidth.
  
  Args:
      chunk_size (int): the number of elements in the chunk
      zero_group (ProcessGroup): the process group of this chunk
      dtype (torch.dtype): the data type of the chunk
      init_device (torch.device): optional, During the chunk construction process, where the tensor is stored.
          The default value is None, which is the current GPU
      cpu_shard_init (bool): a flag indicates the local chunk shard is resident on CPU.
      keep_gathered (bool): optional, if True, this chunk is always gathered in CUDA memory
      pin_memory (bool): optional, if True, this chunk always has a shard copied in pinned CPU memory
docstr_tok: ["Chunk",":","A","container","owning","a","piece","of","contiguous","memory","space","for","tensors","Here","we","use","all-gather","operation","to","gather","the","whole","chunk",".","Currently",",","Chunk","is","exclusively","used","for","DDP","and","ZeRO","DDP","and","it","does","n't","support","unused","parameters",".","It","is","designed","to","make","the","full","use","of","communication","and","PCIE","bandwidth",".","Args",":","chunk_size","(","int",")",":","the","number","of","elements","in","the","chunk","zero_group","(","ProcessGroup",")",":","the","process","group","of","this","chunk","dtype","(","torch.dtype",")",":","the","data","type","of","the","chunk","init_device","(","torch.device",")",":","optional",",","During","the","chunk","construction","process",",","where","the","tensor","is","stored",".","The","default","value","is","None",",","which","is","the","current","GPU","cpu_shard_init","(","bool",")",":","a","flag","indicates","the","local","chunk","shard","is","resident","on","CPU",".","keep_gathered","(","bool",")",":","optional",",","if","True",",","this","chunk","is","always","gathered","in","CUDA","memory","pin_memory","(","bool",")",":","optional",",","if","True",",","this","chunk","always","has","a","shard","copied","in","pinned","CPU","memory"]
code: def __init__(
      self,
      chunk_size: int,
      zero_group: ProcessGroup,
      dtype: torch.dtype,
      init_device: Optional[torch.device] = None,
      cpu_shard_init: bool = False,
      keep_gathered: bool = False,
      pin_memory: bool = False,
      extra_dp_group: ProcessGroup = None,
  ) -> None:
      """
      Chunk: A container owning a piece of contiguous memory space for tensors
      Here we use all-gather operation to gather the whole chunk.
      Currently, Chunk is exclusively used for DDP and ZeRO DDP and it doesn't support unused parameters.
      It is designed to make the full use of communication and PCIE bandwidth.
  
      Args:
          chunk_size (int): the number of elements in the chunk
          zero_group (ProcessGroup): the process group of this chunk
          dtype (torch.dtype): the data type of the chunk
          init_device (torch.device): optional, During the chunk construction process, where the tensor is stored.
              The default value is None, which is the current GPU
          cpu_shard_init (bool): a flag indicates the local chunk shard is resident on CPU.
          keep_gathered (bool): optional, if True, this chunk is always gathered in CUDA memory
          pin_memory (bool): optional, if True, this chunk always has a shard copied in pinned CPU memory
      """
      self.count_id = Chunk._total_number
      Chunk._total_number += 1
  
      self.chunk_size = chunk_size
      self.utilized_size = 0
  
      self.torch_pg = zero_group
      self.pg_size = dist.get_world_size(self.torch_pg)
      self.pg_rank = dist.get_rank(self.torch_pg)
      self.extra_dp_group = extra_dp_group
      self.extra_dp_size = (
          dist.get_world_size(self.extra_dp_group)
          if self.extra_dp_group is not None
          else 1
      )
  
      # the chunk size should be divisible by the dp degree
      if not keep_gathered:
          assert chunk_size % self.pg_size == 0
      self.shard_size = chunk_size // self.pg_size
      self.shard_begin = self.shard_size * self.pg_rank
      self.shard_end = self.shard_begin + self.shard_size
      self.valid_end = self.shard_size
  
      self.dtype = dtype
      device = init_device or get_current_device()
  
      # chunk_temp is a global chunk, which only exists during building the chunks.
      self.chunk_temp = torch.zeros(
          chunk_size, dtype=dtype, device=device
      )  # keep all zero
  
      self.cuda_global_chunk = (
          None  # we force cuda_global_chunk located in CUDA
      )
  
      # cuda local chunk, which is sharded on GPUs
      self.cuda_shard = None
      # cpu local chunk, which is sharded on CPUs
      self.cpu_shard = None
      # is the chunks gathers, which means chunks are duplicated on each process,
      # and we should use the cuda_global_chunk.
      self.is_gathered = True
  
      # configure the init device of the shard
      # no-offload default: fp16, fp32 -> CUDA
      # offload default: fp16, fp32 -> CPU
      self.shard_device = (
          torch.device("cpu")
          if cpu_shard_init
          else get_current_device()
      )
  
      self.chunk_mem = self.chunk_size * self.chunk_temp.element_size()
      self.shard_mem = self.chunk_mem // self.pg_size
  
      # each tensor is associated with a TensorInfo to track its meta info
      # (state, offset, end)
      self.tensors_info: Dict[torch.Tensor, TensorInfo] = {}
      # the total number of tensors in the chunk
      self.num_tensors = 0
  
      # Record the number of tensors in different states
      self.tensor_state_cnter: Dict[TensorState, int] = dict()
      for state in TensorState:
          self.tensor_state_cnter[state] = 0
  
      # If a chunk is kept gathered,
      # they are treated the same as that of the parameters in DDP during training.
      self.keep_gathered = keep_gathered
      if self.keep_gathered:
          pin_memory = False  # since this chunk is gathered, it doesn't need to pin
  
      # if pin_memory is True, we allocate a piece of CPU pin-memory
      # for it all the time
      self.pin_memory = pin_memory
  
      # we introduce the paired chunk here
      # it refers to another chunk having the same parameters
      # but with different dtype(such as fp16_chunk.paired_chunk -> fp32_chunk
      self.paired_chunk = None
      # if this chunk is synchronized with the optimizer, the flag is True
      self.optim_sync_flag = True
      # if the cpu_shard has been visited during the training step, the flag is True
      self.cpu_vis_flag = False
  
      # whether to record l2 norm for the gradient clipping calculation
      self.l2_norm_flag = False
      self.l2_norm = None
  
      self.grad_chunk = None
  
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 1 to create a new Chunk object  with example: obj = Chunk(chunk_size, zero_group, dtype, init_device, cpu_shard_init, keep_gathered, pin_memory, extra_dp_group)