Changes to TRT-LLM download tool for multigpu distributed case

Author: apbose

examples/distributed_inference/tensor_parallel_initialize_dist.py

@@ -16,15 +16,19 @@
 import torch
 import torch_tensorrt
 from rotary_embedding import RotaryAttention, parallel_rotary_block
-from tensor_parallel_initialize_dist import (
+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,
 )
 
-device_mesh, _world_size, _rank, logger = initialize_distributed_env(
-    "./tensor_parallel_rotary_embedding"
-)
+if not dist.is_initialized():
+    initialize_distributed_env()
 
+device_mesh, _world_size, _rank = get_tensor_parallel_device_mesh()
+logger = initialize_logger(_rank, "tensor_parallel_simple_example")
 
 """
 This example covers the rotary embedding in Llama3 model and is derived from https://lightning.ai/lightning-ai/studios/tensor-parallelism-supercharging-large-model-training-with-pytorch-lightning

examples/distributed_inference/tensor_parallel_rotary_embedding.py

@@ -36,11 +36,20 @@
     RowwiseParallel,
     parallelize_module,
 )
-
-device_mesh, _world_size, _rank, logger = initialize_distributed_env(
-    "./tensor_parallel_simple_example"
+from torch_tensorrt.dynamo.distributed.utils import (
+    cleanup_distributed_env,
+    get_tensor_parallel_device_mesh,
+    initialize_distributed_env,
+    initialize_logger,
 )
 
+if not dist.is_initialized():
+    initialize_distributed_env()
+
+device_mesh, _world_size, _rank = get_tensor_parallel_device_mesh()
+logger = initialize_logger(_rank, "tensor_parallel_simple_example")
+
+
 """
 This example takes some code from https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/tensor_parallel_example.py
 """

examples/distributed_inference/tensor_parallel_simple_example.py

@@ -11,11 +11,11 @@
 from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
     dynamo_tensorrt_converter,
 )
+from torch_tensorrt.dynamo.distributed.utils import load_tensorrt_llm_for_nccl
 from torch_tensorrt.dynamo.lowering.passes.fuse_distributed_ops import (
     tensorrt_fused_nccl_all_gather_op,
     tensorrt_fused_nccl_reduce_scatter_op,
 )
-from torch_tensorrt.dynamo.utils import load_tensorrt_llm_for_nccl
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)