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: Extract the mesh_alpha list and mesh_beta list based on the
      best logical mesh, which will be used to initialize the device mesh.
  
  Usage:
      >>> physical_devices = [0, 1, 2, 3]
      >>> ab_profiler = AlphaBetaProfiler(physical_devices)
      >>> mesh_alpha, mesh_beta = profiler.extract_alpha_beta_for_device_mesh()
      >>> print(mesh_alpha)
      [2.5917552411556242e-05, 0.00010312341153621673]
      >>> print(mesh_beta)
      [5.875573704655635e-11, 4.7361584445959614e-12]
docstr_tok: ["Extract","the","mesh_alpha","list","and","mesh_beta","list","based","on","the","best","logical","mesh",",","which","will","be","used","to","initialize","the","device","mesh",".","Usage",":",">",">",">","physical_devices","=","[","0",",","1",",","2",",","3","]",">",">",">","ab_profiler","=","AlphaBetaProfiler","(","physical_devices",")",">",">",">","mesh_alpha",",","mesh_beta","=","profiler.extract_alpha_beta_for_device_mesh","(",")",">",">",">","print","(","mesh_alpha",")","[","2.5917552411556242e-05",",","0.00010312341153621673","]",">",">",">","print","(","mesh_beta",")","[","5.875573704655635e-11",",","4.7361584445959614e-12","]"]
code: def extract_alpha_beta_for_device_mesh(self):
      """
      Extract the mesh_alpha list and mesh_beta list based on the
          best logical mesh, which will be used to initialize the device mesh.
  
      Usage:
          >>> physical_devices = [0, 1, 2, 3]
          >>> ab_profiler = AlphaBetaProfiler(physical_devices)
          >>> mesh_alpha, mesh_beta = profiler.extract_alpha_beta_for_device_mesh()
          >>> print(mesh_alpha)
          [2.5917552411556242e-05, 0.00010312341153621673]
          >>> print(mesh_beta)
          [5.875573704655635e-11, 4.7361584445959614e-12]
      """
      best_logical_mesh = self.search_best_logical_mesh()
  
      first_axis = [row[0] for row in best_logical_mesh]
      second_axis = best_logical_mesh[0]
  
      # init process group for both axes
      first_axis_process_group = dist.new_group(first_axis)
      second_axis_process_group = dist.new_group(second_axis)
  
      # extract alpha and beta for both axes
      def _extract_alpha_beta(pg, pg_handler):
          latency = self.profile_latency(pg, pg_handler)
          bandwidth = self.profile_bandwidth(pg, pg_handler)
          broadcast_object = [latency, bandwidth]
          dist.broadcast_object_list(broadcast_object, src=pg[0])
          return broadcast_object
  
      first_latency, first_bandwidth = _extract_alpha_beta(
          first_axis, first_axis_process_group
      )
      second_latency, second_bandwidth = _extract_alpha_beta(
          second_axis, second_axis_process_group
      )
      mesh_alpha = [first_latency, second_latency]
      # The beta values have been enlarged by 1e10 times temporarily because the computation cost
      # is still estimated in the unit of TFLOPs instead of time. We will remove this factor in future.
      mesh_beta = [1e10 / first_bandwidth, 1e10 / second_bandwidth]
  
      return mesh_alpha, mesh_beta
  
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 8 for calling the AlphaBetaProfiler obj's underlying member method code with example usage: obj.extract_alpha_beta_for_device_mesh() and returns: mesh_alpha,  mesh_beta, broadcast_object