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

Author: apbose

examples/distributed_inference/tensor_parallel_initialize_dist.py

@@ -0,0 +1,55 @@
+"""
+.. _tensor_parallel_initialize_dist:
+Tensor Parallel Initialize Distributed Environment
+==================================================
+
+This module provides functions to initialize and clean up the distributed environment for tensor parallel distributed inference.
+"""
+
+import logging
+import os
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import tensorrt as trt
+import torch
+import torch.distributed as dist
+from torch.distributed._tensor.device_mesh import init_device_mesh
+
+
+def initialize_distributed_env(rank=0, world_size=1, port=29500):
+    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)
+
+    device_mesh = init_device_mesh(device_type="cuda", mesh_shape=(world_size,))
+    rank = device_mesh.get_rank()
+    assert rank == local_rank
+    device_id = (
+        rank % torch.cuda.device_count()
+    )  # Ensure each rank gets a unique device
+    torch.cuda.set_device(device_id)
+
+    return device_mesh, world_size, rank
+
+
+def cleanup_distributed_env():
+    """Clean up distributed process group to prevent resource leaks."""
+    if dist.is_initialized():
+        dist.destroy_process_group()

examples/distributed_inference/tensor_parallel_rotary_embedding.py

@@ -16,11 +16,13 @@
 import torch
 import torch_tensorrt
 from rotary_embedding import RotaryAttention, parallel_rotary_block
+from tensor_parallel_initialize_dist import (
+    cleanup_distributed_env,
+    initialize_distributed_env,
+)
 from torch.distributed import dist
 from torch_tensorrt.dynamo.distributed.utils import (
-    cleanup_distributed_env,
     get_tensor_parallel_device_mesh,
-    initialize_distributed_env,
     initialize_logger,
 )
 

examples/distributed_inference/tensor_parallel_simple_example.py

@@ -37,9 +37,7 @@
     parallelize_module,
 )
 from torch_tensorrt.dynamo.distributed.utils import (
-    cleanup_distributed_env,
     get_tensor_parallel_device_mesh,
-    initialize_distributed_env,
     initialize_logger,
 )
 

py/torch_tensorrt/dynamo/distributed/__init__.py

@@ -0,0 +1 @@
+

py/torch_tensorrt/dynamo/distributed/utils.py

@@ -9,7 +9,6 @@
 from typing import Optional
 
 import torch
-import torch.distributed as dist
 from torch.distributed._tensor.device_mesh import DeviceMesh, init_device_mesh
 from torch_tensorrt._version import __tensorrt_llm_version__
 
@@ -18,30 +17,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(
@@ -76,12 +51,6 @@ def initialize_logger(rank: int, logger_file_name: str) -> logging.Logger:
     return logger
 
 
-def cleanup_distributed_env() -> None:
-    """Clean up distributed process group to prevent resource leaks."""
-    if dist.is_initialized():
-        dist.destroy_process_group()
-
-
 def is_platform_supported_for_trtllm() -> bool:
     """
     Checks if the current platform supports TensorRT-LLM plugins for the NCCL backend.
@@ -127,7 +96,6 @@ def is_platform_supported_for_trtllm() -> bool:
         logger.warning(f"Failed to detect CUDA version: {e}")
         return False
 
-
     return True
 
 

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,27 @@
 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_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 +65,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