Llama torchTRT lib and env initialization reorg

Author: apbose

examples/distributed_inference/tensor_parallel_llama3.py

@@ -1,30 +1,36 @@
 # Taken and modified pytorch lightening
 # https://lightning.ai/lightning-ai/studios/tensor-parallelism-supercharging-large-model-training-with-pytorch-lightning
+# Taken and modified pytorch lightening
+# https://lightning.ai/lightning-ai/studios/tensor-parallelism-supercharging-large-model-training-with-pytorch-lightning
 import logging
 import os
 import time
 
 import torch
-import torch_tensorrt
+import torch.distributed as dist
 from llama3_model import ModelArgs, ParallelTransformer
+from tensor_parallel_initialize_dist import (
+    cleanup_distributed_env,
+    initialize_distributed_env,
+)
 from torch.distributed._composable.fsdp import MixedPrecisionPolicy
 from torch.distributed._composable.fsdp.fully_shard import fully_shard
 from torch.distributed._tensor import Replicate, Shard
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
     checkpoint_wrapper,
 )
+
+if not dist.is_initialized():
+    initialize_distributed_env()
+
+import torch_tensorrt
 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")
+logger = initialize_logger(_rank, "tensor_parallel_llama3")
 
 logger.info(f"Starting PyTorch TP example on rank {_rank}.")
 assert (

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

@@ -359,10 +359,8 @@ def setup_input_tensors(
         need_cudagraphs_record: bool,
     ) -> None:
         for i, input_name in enumerate(self.input_names):
+            contiguous_inputs[i] = complex_to_ri_stacked_tensor(contiguous_inputs[i])
             if not contiguous_inputs[i].is_cuda:
-                contiguous_inputs[i] = complex_to_ri_stacked_tensor(
-                    contiguous_inputs[i]
-                )
                 logger.warning(
                     f"Detected input {input_name} of engine {self.engine.name} is not on a cuda device. "
                     "This tensor is being moved by the runtime but for performance considerations, "