fixed imports

Author: amd-garydeng

et_replay/comm/backend/pytorch_dist_backend.py

@@ -20,7 +20,7 @@
 from et_replay.comm.param_profile import paramProfile
 
 try:
-    from param_bench.train.comms.pt.fb.internals import (
+    from fb.internals import (
         all_to_all_internal,
         all_to_allv_internal,
         extend_distributed,

et_replay/comm/comms_utils.py

@@ -25,7 +25,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 try:
-    from param_bench.train.comms.pt.fb.internals import (
+    from fb.internals import (
         fbInitProfiler,
         fbSampleProfiler,
         fbStartProfiler,

et_replay/tools/comm_replay.py

@@ -66,7 +66,7 @@ def writeCommDetails(commsTracePerf: List, rank: int, folder: str = "./") -> Non
     if "://" in comms_file:  # assume that "://" in directory path means remote store
         saveToLocal = False
         try:
-            from param_bench.train.comms.pt.fb.internals import (
+            from fb.internals import (
                 writeRemoteTrace as writeFbRemoteTrace,
             )
 
@@ -1528,7 +1528,7 @@ def readRawTrace(self, remotePath: str, rank: int) -> None:
                 raw_comms_trace = comms_utils.commonUrlRead(remotePath=remotePath)
             else:
                 try:
-                    from param_bench.train.comms.pt.fb.internals import (
+                    from fb.internals import (
                         readRemoteTrace as readFbRemoteTrace,
                     )
 

train/comms/pt/comms.py

@@ -16,24 +16,24 @@
 # pytorch
 import torch
 
-from param_bench.train.comms.pt import comms_utils
-from param_bench.train.comms.pt.comms_utils import (
+import comms_utils
+from comms_utils import (
     bootstrap_info_holder,
     commsParamsHolderBase,
     ensureTensorFlush,
     paramCommsBench,
     paramDeviceTimer,
     paramStreamGuard,
 )
-from param_bench.train.comms.pt.logger_utils import (
+from logger_utils import (
     benchType,
     commsCollPerfMetrics,
     commsPt2PtPerfMetrics,
     commsQuantCollPerfMetrics,
     customized_perf_loggers,
 )
 
-from param_bench.train.comms.pt.pytorch_backend_utils import (
+from pytorch_backend_utils import (
     pt2ptPatterns,
     supportedC10dBackends,
     supportedCollectives,
@@ -1824,13 +1824,13 @@ def initBackend(
             commsParams.nw_stack == "pytorch-dist"
             and commsParams.backend in supportedC10dBackends
         ):
-            from param_bench.train.comms.pt.pytorch_dist_backend import (
+            from pytorch_dist_backend import (
                 PyTorchDistBackend,
             )
 
             backendObj = PyTorchDistBackend(bootstrap_info, commsParams)
         elif commsParams.nw_stack == "pytorch-xla-tpu":
-            from param_bench.train.comms.pt.pytorch_tpu_backend import PyTorchTPUBackend
+            from pytorch_tpu_backend import PyTorchTPUBackend
 
             backendObj = PyTorchTPUBackend(bootstrap_info, commsParams)
         else:
@@ -1839,7 +1839,7 @@ def initBackend(
                 logging.warning(
                     f"Attempt loading customized backend {commsParams.backend} if registered. Note that this is not officially supported. Use it with caution and at your own risk."
                 )
-                from param_bench.train.comms.pt.pytorch_backend_utils import (
+                from pytorch_backend_utils import (
                     customized_backend,
                 )
 

train/comms/pt/commsTraceParser.py

@@ -7,9 +7,9 @@
 
 from et_replay import ExecutionTrace
 
-from param_bench.train.comms.pt import comms_utils
-from param_bench.train.comms.pt.comms_utils import commsArgs
-from param_bench.train.comms.pt.pytorch_backend_utils import supportedP2pOps
+from  import comms_utils
+from comms_utils import commsArgs
+from pytorch_backend_utils import supportedP2pOps
 
 tensorDtypeMap = {
     "Tensor(int)": "int",

train/comms/pt/commsTraceReplay.py

@@ -17,17 +17,17 @@
 
 import numpy as np
 import torch
-from param_bench.train.comms.pt import comms_utils
-from param_bench.train.comms.pt.comms_utils import (
+from  import comms_utils
+from comms_utils import (
     bootstrap_info_holder,
     commsArgs,
     commsParamsHolderBase,
     paramCommsBench,
     paramStreamGuard,
     paramToCommName,
 )
-from param_bench.train.comms.pt.param_profile import paramProfile, paramTimer
-from param_bench.train.comms.pt.pytorch_backend_utils import supportedP2pOps
+from param_profile import paramProfile, paramTimer
+from pytorch_backend_utils import supportedP2pOps
 
 try:
     from trainer_iteration_wrapper import setTrainingIteration  # @manual
@@ -64,7 +64,7 @@ def writeCommDetails(commsTracePerf: List, rank: int, folder: str = "./") -> Non
     if "://" in comms_file:  # assume that "://" in directory path means remote store
         saveToLocal = False
         try:
-            from param_bench.train.comms.pt.fb.internals import (
+            from fb.internals import (
                 writeRemoteTrace as writeFbRemoteTrace,
             )
 
@@ -1388,13 +1388,13 @@ def initBackend(
         """
         # init backend and corresponding function pointers
         if commsParams.nw_stack == "pytorch-dist":
-            from param_bench.train.comms.pt.pytorch_dist_backend import (
+            from pytorch_dist_backend import (
                 PyTorchDistBackend,
             )
 
             self.backendFuncs = PyTorchDistBackend(bootstrap_info, commsParams)
         elif commsParams.nw_stack == "pytorch-xla-tpu":
-            from param_bench.train.comms.pt.pytorch_tpu_backend import PyTorchTPUBackend
+            from pytorch_tpu_backend import PyTorchTPUBackend
 
             self.backendFuncs = PyTorchTPUBackend(bootstrap_info, commsParams)
         else:
@@ -1526,7 +1526,7 @@ def readRawTrace(self, remotePath: str, rank: int) -> None:
                 raw_comms_trace = comms_utils.commonUrlRead(remotePath=remotePath)
             else:
                 try:
-                    from param_bench.train.comms.pt.fb.internals import (
+                    from fb.internals import (
                         readRemoteTrace as readFbRemoteTrace,
                     )
 
@@ -1589,7 +1589,7 @@ def readTrace(self, remotePath: str, rank: int) -> None:
 
         # Convert trace to comms trace.
         try:
-            from param_bench.train.comms.pt import commsTraceParser
+            from  import commsTraceParser
         except ImportError:
             logger.info("FB internals not present, using base parser.")
             self.comms_trace = extractCommsInfo(self.comms_trace)

train/comms/pt/comms_utils.py

@@ -23,7 +23,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 try:
-    from param_bench.train.comms.pt.fb.internals import (
+    from fb.internals import (
         fbInitProfiler,
         fbSampleProfiler,
         fbStartProfiler,
@@ -38,8 +38,8 @@
 
 import numpy as np
 import torch
-from param_bench.train.comms.pt.param_profile import paramTimer
-from param_bench.train.comms.pt.pytorch_backend_utils import (
+from param_profile import paramTimer
+from pytorch_backend_utils import (
     backendFunctions,
     collectiveArgsHolder,
     customized_backend,

train/comms/pt/dlrm.py

@@ -23,7 +23,7 @@
 import torch.nn as nn
 from comms_utils import paramCommsBench
 
-from param_bench.train.comms.pt import comms_utils
+from  import comms_utils
 from pytorch_dist_backend import PyTorchDistBackend
 from torch.autograd import Function
 

train/comms/pt/logger_utils.py

@@ -8,7 +8,7 @@
 from enum import Enum
 from typing import abstractmethod, Dict, Optional
 
-from param_bench.train.comms.pt.pytorch_backend_utils import backendFunctions
+from pytorch_backend_utils import backendFunctions
 
 logger = logging.getLogger(__name__)
 

train/comms/pt/pytorch_backend_utils.py

@@ -9,7 +9,7 @@
 
 import torch
 
-from param_bench.train.comms.pt.param_profile import paramTimer
+from param_profile import paramTimer
 
 from torch.distributed import ProcessGroup
 

train/comms/pt/pytorch_dist_backend.py

@@ -13,8 +13,8 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from param_bench.train.comms.pt.param_profile import paramProfile
-from param_bench.train.comms.pt.pytorch_backend_utils import (
+from param_profile import paramProfile
+from pytorch_backend_utils import (
     backendFunctions,
     collectiveArgsHolder,
 )
@@ -27,7 +27,7 @@
     has_triton = False
 
 try:
-    from param_bench.train.comms.pt.fb.internals import (
+    from fb.internals import (
         all_to_all_internal,
         all_to_allv_internal,
         extend_distributed,

train/comms/pt/setup.py

@@ -2,7 +2,7 @@
 
 
 def main():
-    package_base = "param_bench.train.comms.pt"
+    package_base = ""
 
     # List the packages and their dir mapping:
     # "install_destination_package_path": "source_dir_path"

train/comms/pt/tests/commsTraceReplay_tests.py

@@ -5,9 +5,9 @@
 
 from comms_utils import commsArgs
 
-from param_bench.train.comms.pt.commsTraceReplay import commsTraceReplayBench
-from param_bench.train.comms.pt.tests.mocks.backend_mock import MockBackendFunction
-from param_bench.train.comms.pt.tests.test_utils import (
+from commsTraceReplay import commsTraceReplayBench
+from tests.mocks.backend_mock import MockBackendFunction
+from tests.test_utils import (
     commsParamsTest,
     createCommsArgs,
     testArgs,

train/comms/pt/tests/comms_utils_tests.py

@@ -3,9 +3,9 @@
 
 import torch
 
-from param_bench.train.comms.pt import comms_utils
-from param_bench.train.comms.pt.tests.mocks.backend_mock import MockBackendFunction
-from param_bench.train.comms.pt.tests.test_utils import (
+from  import comms_utils
+from tests.mocks.backend_mock import MockBackendFunction
+from tests.test_utils import (
     bootstrap_info_test,
     commsParamsTest,
 )