Distributed utils package, separating out env for single GPU and multiGPU

Author: apbose

py/torch_tensorrt/dynamo/distributed/__init__.py

@@ -0,0 +1 @@
+from utils import is_platform_supported_for_trtllm, load_tensorrt_llm_for_nccl

py/torch_tensorrt/dynamo/distributed/utils.py

@@ -18,30 +18,6 @@
 logger = logging.getLogger(__name__)
 
 
-def initialize_distributed_env(
-    rank: int = 0, world_size: int = 1, port: int = 29500
-) -> None:
-    local_rank = int(
-        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", rank % torch.cuda.device_count())
-    )
-    world_size = int(os.environ.get("OMPI_COMM_WORLD_SIZE", world_size))
-
-    # Set up environment variable to run with mpirun
-    os.environ["RANK"] = str(local_rank)
-    os.environ["WORLD_SIZE"] = str(world_size)
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = str(port)
-
-    # Necessary to assign a device to each rank.
-    torch.cuda.set_device(local_rank)
-
-    # We use nccl backend
-    dist.init_process_group("nccl")
-
-    # set a manual seed for reproducibility
-    torch.manual_seed(1111)
-
-
 def check_tensor_parallel_device_number(world_size: int) -> None:
     if world_size % 2 != 0:
         raise ValueError(

setup.py

@@ -454,6 +454,7 @@ def run(self):
     "torch_tensorrt.dynamo.conversion.impl.unary",
     "torch_tensorrt.dynamo.conversion.plugins",
     "torch_tensorrt.dynamo.debug",
+    "torch_tensorrt.dynamo.distributed",
     "torch_tensorrt.dynamo.lowering",
     "torch_tensorrt.dynamo.lowering.passes",
     "torch_tensorrt.dynamo.partitioning",

tests/py/dynamo/distributed/distributed_utils.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import random
 
 import numpy as np
 import tensorrt as trt
@@ -8,8 +9,35 @@
 from torch.distributed._tensor.device_mesh import init_device_mesh
 
 
-def set_environment_variables_pytest():
+def set_environment_variables_pytest_single_process():
+    port = 29500 + random.randint(1, 1000)
     os.environ["WORLD_SIZE"] = str(1)
     os.environ["RANK"] = str(0)
     os.environ["MASTER_ADDR"] = "127.0.0.1"
-    os.environ["MASTER_PORT"] = str(29500)
+    os.environ["MASTER_PORT"] = str(port)
+
+
+def set_environment_variables_pytest_multi_process(
+    rank: int = 0, world_size: int = 1
+) -> None:
+    port = 29500 + random.randint(1, 1000)
+    # these variables are set by mpirun -n 2
+    local_rank = int(
+        os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", rank % torch.cuda.device_count())
+    )
+    world_size = int(os.environ.get("OMPI_COMM_WORLD_SIZE", world_size))
+
+    # Set up environment variable to run with mpirun
+    os.environ["RANK"] = str(local_rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+
+    # Necessary to assign a device to each rank.
+    torch.cuda.set_device(local_rank)
+
+    # We use nccl backend
+    dist.init_process_group("nccl")
+
+    # set a manual seed for reproducibility
+    torch.manual_seed(1111)

tests/py/dynamo/distributed/test_nccl_ops.py

@@ -5,10 +5,28 @@
 import torch.distributed as dist
 import torch.nn as nn
 from conversion.harness import DispatchTestCase
-from distributed_utils import set_environment_variables_pytest
+
+# The distributed env initialization has to be before torchTRT import since it uses barrier
+from distributed_utils import (
+    set_environment_variables_pytest,
+    set_environment_variables_pytest_multi_process,
+    set_environment_variables_pytest_single_process,
+)
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
-from torch_tensorrt.dynamo.utils import is_platform_supported_for_trtllm
+
+if "OMPI_COMM_WORLD_SIZE" in os.environ:
+    set_environment_variables_pytest_multi_process()
+else:
+    set_environment_variables_pytest_single_process()
+
+if not dist.is_initialized():
+    dist.init_process_group(
+        backend="nccl",
+        init_method="env://",
+    )
+
+from torch_tensorrt.dynamo.distributed.utils import is_platform_supported_for_trtllm
 
 
 class DistributedGatherModel(nn.Module):
@@ -48,11 +66,9 @@ class TestNcclOpsConverter(DispatchTestCase):
     )
     @classmethod
     def setUpClass(cls):
-        set_environment_variables_pytest()
-        cls.world_size = 1
-        if not dist.is_initialized():
-            dist.init_process_group(backend="nccl")
-        cls.group = dist.new_group(ranks=[0])
+        cls.world_size = int(os.environ.get("OMPI_COMM_WORLD_SIZE", 1))
+        cls.rank = int(os.environ.get("OMPI_COMM_WORLD_RANK", 0))
+        cls.group = dist.new_group(ranks=list(range(cls.world_size)))
         cls.group_name = cls.group.group_name
 
     @classmethod