Merge branch 'main' into sanshang/fix_all2all

Author: TaekyungHeo

et_replay/execution_trace.py

@@ -357,8 +357,8 @@ def __init__(self, json):
             input_tensors = self.nodes[id].get_input_tensors()
             output_tensors = self.nodes[id].get_output_tensors()
 
-            # track the various process and threads we have
-            if x["name"] == "__ROOT_THREAD__":
+            # track annonation to get thread ids of root nodes
+            if x["name"] == "[pytorch|profiler|execution_trace|thread]":
                 tid = self.nodes[id].tid
                 self.proc_group[pid][tid] = id
 

train/comms/pt/comms.py

@@ -183,12 +183,6 @@ def readArgs(self, parser):
             default=False,
             help="Select some ranks to send/receive 0B messages",
         )
-        parser.add_argument(
-            "--use-device-time",
-            action="store_true",
-            default=False,
-            help="use device time measurement",
-        )
         parser.add_argument(
             "--graph-launches",
             type=int,
@@ -368,7 +362,9 @@ def run_coll_cuda_graph(self, comm_fn=None, dcheck=False):
         self.backendFuncs.sync_barrier(
             self.collectiveArgs, desc="run_coll_cuda_graph_begin"
         )
-        elapsedTimeNS = 0.0
+        elapsedCPUTimeNS = 0.0
+        start_event = self.backendFuncs.create_event(self.collectiveArgs)
+        end_event = self.backendFuncs.create_event(self.collectiveArgs)
 
         # 1. Warmup phase
         # launch collective on a separate stream and sync with current_stream
@@ -393,21 +389,34 @@ def run_coll_cuda_graph(self, comm_fn=None, dcheck=False):
 
         # 3. Replay
         start = time.monotonic()  # available only in py3
+        self.backendFuncs.record_event(start_event, self.collectiveArgs)
         for _ in range(self.collectiveArgs.graph_launches):
             if self.collectiveArgs.enable_profiler:
                 comms_utils.sampleProfiler()
 
             # [optional] we can feed new input data to ipTensor for each replay
             g.replay()
 
+        self.backendFuncs.record_event(end_event, self.collectiveArgs)
         self.backendFuncs.complete_accel_ops(self.collectiveArgs)
+
         end = time.monotonic()  # available only in py3
 
         ensureTensorFlush(self.collectiveArgs.opTensor)
 
-        elapsedTimeNS += (
+        elapsedCPUTimeNS += (
             end - start
         ) * 1e9  # keeping time in NS, helps in divising data by nanoseconds
+        elapsedDeviceTimeMs = self.backendFuncs.elapsed_time(start_event, end_event)
+        elapsedDeviceTimeNS = elapsedDeviceTimeMs * 1e6
+        elapsedTimeNS = (
+            elapsedDeviceTimeNS
+            if self.collectiveArgs.use_device_time
+            else elapsedCPUTimeNS
+        )
+        logger.debug(
+            f"elapsedCPUTimeNS={elapsedCPUTimeNS}, elapsedDeviceTimeNS={elapsedDeviceTimeNS}."
+        )
 
         memSize = self.backendFuncs.get_mem_size(self.collectiveArgs)
 
@@ -436,17 +445,11 @@ def run_coll_cuda_graph(self, comm_fn=None, dcheck=False):
         }
         return results
 
-    def runColl(self, comm_fn=None, dcheck=False):
-        if self.collectiveArgs.graph_launches > 0:
-            return self.run_coll_cuda_graph(comm_fn, dcheck)
+    def run_coll_non_graph(self, comm_fn=None, dcheck=False):
         self.backendFuncs.sync_barrier(self.collectiveArgs, desc="runColl_begin")
 
-        elapsedCPUTimeNS = 0.0
-        elapsedDeviceTimeNS = 0.0
+        elapsedTimeNS = 0.0
         is_blocking = not self.collectiveArgs.asyncOp
-        # Initialize CUDA events for device timing
-        start_event = self.backendFuncs.create_event(self.collectiveArgs)
-        end_event = self.backendFuncs.create_event(self.collectiveArgs)
 
         for nIter in range(
             self.collectiveArgs.numWarmupIters + self.collectiveArgs.numIters
@@ -458,22 +461,16 @@ def runColl(self, comm_fn=None, dcheck=False):
                 self.backendFuncs.complete_accel_ops(self.collectiveArgs)
                 ensureTensorFlush(self.collectiveArgs.opTensor)
                 # Start measuring time after warmup iterations
-                elapsedCPUTimeNS = 0.0
-                elapsedDeviceTimeNS = 0.0
+                elapsedTimeNS = 0.0
                 self.collectiveArgs.quant_time.reset()
                 self.collectiveArgs.dequant_time.reset()
-                self.backendFuncs.record_event(
-                    start_event, self.collectiveArgs
-                )  # record start event for non-blocking operation
             # reset tensor values for data validation check
             if dcheck:
                 self.setTensorVal(self.collectiveArgs.opTensor)
             # for blocking mode, do barrier before starting collective
             if is_blocking:
                 self.backendFuncs.sync_barrier(self.collectiveArgs)
-                self.backendFuncs.record_event(
-                    start_event, self.collectiveArgs
-                )  # record start event for blocking operation
+
             start = time.monotonic()  # available only in py3
             with paramStreamGuard(
                 stream=self.backendFuncs.get_current_stream(
@@ -488,47 +485,29 @@ def runColl(self, comm_fn=None, dcheck=False):
                 ]
                 for _ in range(self.collectiveArgs.numCollPerIter):
                     comm_fn(self.collectiveArgs)
+
             if is_blocking:  # should be sychronous, wait for the collective
-                self.backendFuncs.record_event(
-                    end_event, self.collectiveArgs
-                )  # record end event for blocking operation
                 self.backendFuncs.complete_accel_ops(self.collectiveArgs)
-                elapsedDeviceTimeMs = self.backendFuncs.elapsed_time(
-                    start_event, end_event
-                )
-                elapsedDeviceTimeNS += elapsedDeviceTimeMs * 1e6  # Convert ms to ns
+
             # Measuring time.
-            elapsedCPUTimeNS += (
+            elapsedTimeNS += (
                 time.monotonic() - start
             ) * 1e9  # keeping time in NS, helps in divising data by nanosecond
+
         start = time.monotonic()  # available only in py3
-        # if not blocking, record second end event here
-        if not is_blocking:
-            self.backendFuncs.record_event(
-                end_event, self.collectiveArgs
-            )  # record end event for non-blocking operations
         self.backendFuncs.complete_accel_ops(self.collectiveArgs)
         end = time.monotonic()  # available only in py3
+
         ensureTensorFlush(self.collectiveArgs.opTensor)
 
-        elapsedCPUTimeNS += (
+        elapsedTimeNS += (
             end - start
         ) * 1e9  # keeping time in NS, helps in divising data by nanoseconds
-        if not is_blocking:
-            elapsedDeviceTimeMs = self.backendFuncs.elapsed_time(start_event, end_event)
-            elapsedDeviceTimeNS = elapsedDeviceTimeMs * 1e6  # Convert ms to ns
 
         memSize = self.backendFuncs.get_mem_size(self.collectiveArgs)
-        logger.debug(
-            f"elapsedCPUTimeNS={elapsedCPUTimeNS}, elapsedDeviceTimeNS={elapsedDeviceTimeNS}."
-        )
-        ElapsedTimeNS = (
-            elapsedDeviceTimeNS
-            if self.collectiveArgs.use_device_time
-            else elapsedCPUTimeNS
-        )
+
         avgIterNS, algBW = comms_utils.getAlgBW(
-            ElapsedTimeNS,
+            elapsedTimeNS,
             memSize,
             self.collectiveArgs.numIters * self.collectiveArgs.numCollPerIter,
         )
@@ -550,6 +529,13 @@ def runColl(self, comm_fn=None, dcheck=False):
         }
         return results
 
+    def runColl(self, comm_fn=None, dcheck=False):
+        return (
+            self.run_coll_non_graph(comm_fn, dcheck)
+            if self.collectiveArgs.graph_launches == 0
+            else self.run_coll_cuda_graph(comm_fn, dcheck)
+        )
+
     def runPt2Pt(self):
         self.backendFuncs.sync_barrier(self.collectiveArgs)
         # warm-up
@@ -884,7 +870,6 @@ def initCollectiveArgs(self, commsParams):
         self.collectiveArgs.asyncOp = False if commsParams.blockingFlag == 1 else True
         self.collectiveArgs.numCollPerIter = commsParams.num_coll
         self.collectiveArgs.include_0B = commsParams.include_0B
-        self.collectiveArgs.use_device_time = commsParams.use_device_time
         self.collectiveArgs.graph_launches = commsParams.graph_launches
 
         if commsParams.bitwidth < 32:
@@ -911,11 +896,7 @@ def gatherBenchTime(self, collectiveArgs, commsParams, timeUsElapsedList):
         # Push the list to device, then do an all-gather.
         timeElapsedTensor = torch.tensor(
             timeUsElapsedList,
-            device=(
-                self.backendFuncs.get_device()
-                if commsParams.backend == "nccl"
-                else torch.device("cpu")
-            ),
+            device=(self.backendFuncs.get_device()),
         )
         collectiveArgs.opTensor = None
         if commsParams.backend != "xla":
@@ -1051,10 +1032,7 @@ def reportBenchTime(
         dequantTimeTensorList,
     ):
         # convernt num_elements to # of elements per rank
-        if commsParams.collective in (
-            "all_to_all",
-            "all_to_allv",
-            "all_to_all_single",
+        if "all_to_all" in commsParams.collective or commsParams.collective in (
             "reduce_scatter",
             "reduce_scatter_v",
             "reduce_scatter_base",

train/comms/pt/comms_utils.py

@@ -17,7 +17,23 @@
 from collections.abc import Callable
 from contextlib import ContextDecorator
 from io import StringIO
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
+
+import torch
+
+from param_bench.train.comms.pt.param_profile import paramTimer
+from param_bench.train.comms.pt.pytorch_backend_utils import (
+    backendFunctions,
+    collectiveArgsHolder,
+    customized_backend,
+    supportedC10dBackends,
+    supportedDevices,
+)
+from torch._C._distributed_c10d import ProcessGroup
+
+random.seed()
+
+logger = logging.getLogger(__name__)
 
 try:
     from param_bench.train.comms.pt.fb.internals import (
@@ -29,24 +45,28 @@
     )
 
     has_internal_libs = True
+    logger.info("Successfully import internal libs")
 except ImportError:
     has_internal_libs = False
+    logger.info("Iinternal libs not found.")
+
+try:
+    from param_bench.train.comms.pt.fb.mixins import (
+        name_aliases_ext,
+        ParamCommsBenchMixin,
+    )
 
+    logger.info("Successfully imported ParamCommsBenchMixin")
+except ImportError:
+    logger.warning(
+        "ParamCommsBenchMixin does not exist or module not found. Default to empty class."
+    )
 
-import torch
-from param_bench.train.comms.pt.param_profile import paramTimer
-from param_bench.train.comms.pt.pytorch_backend_utils import (
-    backendFunctions,
-    collectiveArgsHolder,
-    customized_backend,
-    supportedC10dBackends,
-    supportedDevices,
-)
-from torch._C._distributed_c10d import ProcessGroup
+    class ParamCommsBenchMixin:
+        pass  # Define empty class if it does not exist
 
-random.seed()
+    name_aliases_ext = {}
 
-logger = logging.getLogger(__name__)
 
 default_master_ip = "127.0.0.1"
 default_master_port = "29500"
@@ -206,10 +226,7 @@ def fixBeginSize(commsParams: commsParamsHolder, world_size: int) -> None:
         None
     """
     # ensures we will have atleast one member/rank
-    if commsParams.collective in (
-        "all_to_all",
-        "all_to_allv",
-        "all_to_all_single",
+    if "all_to_all" in commsParams.collective or commsParams.collective in (
         "all_gather",
         "all_gather_base",
         "gather",
@@ -392,17 +409,14 @@ def checkQuantArgs(
     Returns:
         None
     """
-    if collective not in (
-        "all_to_all",
-        "all_to_allv",
-        "all_to_all_single",
+    if "all_to_all" not in collective and collective not in (
         "reduce",
         "all_reduce",
     ):
         raise NotImplementedError(
             f"quantized communication for {collective} is currently unsupported."
         )
-    if collective in ("all_to_all", "all_to_allv", "all_to_all_single"):
+    if "all_to_all" in collective:
         if (beginSize // 4) % quant_a2a_embedding_dim != 0:
             logger.warning(
                 f"begin size {beginSize} must be a multiple of --quant-a2a-embedding-dim {quant_a2a_embedding_dim} for all_to_all operation"
@@ -452,6 +466,7 @@ def paramToCommName(name: str, supported_comms: list[str] = None) -> str:
         "reducescatterbase": "reduce_scatter_base",
         "recvanysource": "recv",
     }
+    name_aliases.update(name_aliases_ext)
 
     new_name = name.lower()
 
@@ -803,7 +818,6 @@ def __init__(self, args: Namespace) -> None:
         self.ibv_devices = args.ibv_devices
         self.init_only = args.init_only
         self.eager_init = args.eager_init
-        self.use_device_time = args.use_device_time
 
 
 class commsDlrmParamsHolder(commsParamsHolderBase):
@@ -934,7 +948,7 @@ def __init__(
         self.bag_size = args.bag_size
 
 
-class paramCommsBench(ABC):
+class ParamCommsBenchBase(ABC):
     """Abstract class for any param comms benchmark."""
 
     def __init__(self, supportedNwstacks: list[str] = None) -> None:
@@ -1570,6 +1584,8 @@ def prepComm(
             "scatter": self._prep_reduce_scatter,
             "pt2pt": self._prep_pt2pt,
         }
+        if hasattr(self, "dispatchDictExt") and self.dispatchDictExt is not None:
+            dispatchDict.update(self.dispatchDictExt)
 
         function_to_call = dispatchDict.get(commOp)
         if function_to_call is not None:
@@ -1816,6 +1832,11 @@ def checkArgs(self, args: Namespace) -> None:
             os.environ["MASTER_PORT"] = args.master_port
 
 
+class paramCommsBench(ParamCommsBenchMixin, ParamCommsBenchBase):
+    def __init__(self, supportedNwstacks: list[str] = None) -> None:
+        super().__init__(supportedNwstacks)
+
+
 def init_emb_lookup(collectiveArgs, commsParams, backendFuncs):
     """
     Initialize embedding table op