class_docstr: Profile alpha and beta value for a given device list.
  
  Usage:
      # Note: the environment of execution is supposed to be
      # multi-process with multi-gpu in mpi style.
      >>> physical_devices = [0, 1, 4, 5]
      >>> ab_profiler = AlphaBetaProfiler(physical_devices)
      >>> ab_dict = profiler.alpha_beta_dict
      >>> print(ab_dict)
      {(0, 1): (1.9641406834125518e-05, 4.74049549614719e-12), (0, 4): (1.9506998360157013e-05, 6.97421973297474e-11), (0, 5): (2.293858677148819e-05, 7.129930361393644e-11),
       (1, 4): (1.9010603427886962e-05, 7.077968863788975e-11), (1, 5): (1.9807778298854827e-05, 6.928845708992215e-11), (4, 5): (1.8681809306144713e-05, 4.7522367291330524e-12),
       (1, 0): (1.9641406834125518e-05, 4.74049549614719e-12), (4, 0): (1.9506998360157013e-05, 6.97421973297474e-11), (5, 0): (2.293858677148819e-05, 7.129930361393644e-11),
       (4, 1): (1.9010603427886962e-05, 7.077968863788975e-11), (5, 1): (1.9807778298854827e-05, 6.928845708992215e-11), (5, 4): (1.8681809306144713e-05, 4.7522367291330524e-12)}
docstr: This method is used to profiling the alpha and beta value for a given device list.
  
  Returns:
      alpha_beta_dict: A dict which maps process group to its alpha and beta value.
docstr_tok: ["This","method","is","used","to","profiling","the","alpha","and","beta","value","for","a","given","device","list",".","Returns",":","alpha_beta_dict",":","A","dict","which","maps","process","group","to","its","alpha","and","beta","value","."]
code: def profile_ab(self):
      """
      This method is used to profiling the alpha and beta value for a given device list.
  
      Returns:
          alpha_beta_dict: A dict which maps process group to its alpha and beta value.
      """
      alpha_beta_dict: Dict[Tuple[int], Tuple[float]] = {}
      rank = dist.get_rank()
      dist.new_group(self.physical_devices)
  
      def get_max_nbytes(
          process_group: Tuple[int], pg_handler: dist.ProcessGroup
      ):
          assert rank in process_group
          device = torch.cuda.current_device()
          rank_max_nbytes = torch.cuda.mem_get_info(device)[0]
          rank_max_nbytes = torch.tensor(rank_max_nbytes, device=device)
          dist.all_reduce(
              rank_max_nbytes, op=dist.ReduceOp.MIN, group=pg_handler
          )
          max_nbytes = min(
              int(1 * GB),
              int(GB << int(math.log2(rank_max_nbytes.item() / GB))),
          )
          return max_nbytes
  
      for process_group, pg_handler in self.process_group_dict.items():
          if rank not in process_group:
              max_nbytes = None
              alpha = None
              bandwidth = None
          else:
              max_nbytes = get_max_nbytes(process_group, pg_handler)
              alpha = self.profile_latency(process_group, pg_handler)
              bandwidth = self.profile_bandwidth(
                  process_group, pg_handler, maxbytes=max_nbytes
              )
  
          if bandwidth is None:
              beta = None
          else:
              beta = 1 / bandwidth
  
          broadcast_list = [alpha, beta]
          dist.broadcast_object_list(
              broadcast_list, src=process_group[0]
          )
          alpha_beta_dict[process_group] = tuple(broadcast_list)
  
      # add symmetry pair to the alpha_beta_dict
      symmetry_ab_dict = {}
      for process_group, alpha_beta_pair in alpha_beta_dict.items():
          symmetry_process_group = (process_group[1], process_group[0])
          symmetry_ab_dict[symmetry_process_group] = alpha_beta_pair
  
      alpha_beta_dict.update(symmetry_ab_dict)
  
      return alpha_beta_dict
  
file_path: ColossalAI/colossalai/device/alpha_beta_profiler.py
imports: import math
  import time
  from typing import Dict, List, Tuple
  import torch
  import torch.distributed
  from colossalai.logging import get_dist_logger
label_desc: Use image node_id 6 for calling the AlphaBetaProfiler obj's underlying member method code with example usage: obj.profile_ab() and returns: alpha_beta_dict, max_nbytes