adding the test script and correction to the backend

Author: apbose

Diff for: py/torch_tensorrt/dynamo/backend/backends.py

@@ -69,13 +69,12 @@ def aot_torch_tensorrt_aten_backend(
     to_delete = {
         key
         for key in settings_aot_autograd["decompositions"]
-        if "transpose" in key._name
+        if "transpose" in key._name or "detach" in key._name
     }
 
     for key in to_delete:
         del settings_aot_autograd["decompositions"][key]
 
-    remove_detach(gm, settings)
     return aot_autograd(
         fw_compiler=_pretraced_backend_autograd,
         decompositions=settings_aot_autograd["decompositions"],

Diff for: tests/py/dynamo/conversion/harness.py

@@ -351,7 +351,6 @@ def generate_graph(
         enable_passes: bool,
         propagate_shapes: bool = False,
         settings: CompilationSettings = CompilationSettings(),
-        fuse_distributed_ops: bool = False,
         torch_export_dynamic_shapes: Optional[Any] = None,
     ):
         mod = mod.eval()
@@ -367,16 +366,6 @@ def generate_graph(
                 tuple(torch_export_inputs),
                 dynamic_shapes=torch_export_dynamic_shapes,
             )
-            if fuse_distributed_ops:
-                exported_program = exported_program.run_decompositions(
-                    get_decompositions(False)
-                )
-                from torch_tensorrt.dynamo.lowering.passes.fuse_distributed_ops import (
-                    fuse_distributed_ops,
-                )
-
-                gm = exported_program.graph_module
-                gm = fuse_distributed_ops(gm, settings)
             if enable_passes:
                 exported_program = pre_export_lowering(exported_program, settings)
                 exported_program = exported_program.run_decompositions(
@@ -415,7 +404,6 @@ def run_test(
         propagate_shapes=False,
         int32_reqd=False,
         immutable_weights=True,
-        fuse_distributed_ops=False,
     ):
         # TODO: lan to remove this and set use_dynamo_traccer to True by default
         # once all the converter test files are moved to use_dynamo_tracer
@@ -436,7 +424,6 @@ def run_test(
             enable_passes=enable_passes,
             propagate_shapes=propagate_shapes,
             settings=compilation_settings,
-            fuse_distributed_ops=fuse_distributed_ops,
         )
 
         num_inputs = len(inputs)

Diff for: tests/py/dynamo/distributed/distributed_utils.py

@@ -0,0 +1,74 @@
+import logging
+import os
+
+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 set_environment_variables_pytest():
+    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["USE_TRTLLM_PLUGINS"] = "1"
+
+
+def find_repo_root(max_depth=10):
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+    for i in range(max_depth):
+        files = os.listdir(dir_path)
+        if "MODULE.bazel" in files:
+            return dir_path
+        else:
+            dir_path = os.path.dirname(dir_path)
+
+    raise RuntimeError("Could not find repo root")
+
+
+def initialize_logger(rank, logger_file_name):
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    fh = logging.FileHandler(logger_file_name + f"_{rank}.log", mode="w")
+    fh.setLevel(logging.INFO)
+    logger.addHandler(fh)
+    return logger
+
+
+# This is required for env initialization since we use mpirun
+def initialize_distributed_env(logger_file_name, 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)
+    os.environ["TRTLLM_PLUGINS_PATH"] = (
+        find_repo_root() + "/lib/libnvinfer_plugin_tensorrt_llm.so"
+    )
+
+    # 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
+    logger = initialize_logger(rank, logger_file_name)
+    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, logger

Diff for: tests/py/dynamo/distributed/test_distributed_simple_example.py

@@ -0,0 +1,92 @@
+import time
+
+import tensorrt as trt
+import torch
+import torch.nn as nn
+import torch_tensorrt
+from distributed_utils import initialize_distributed_env
+from torch.distributed._tensor import Shard
+from torch.distributed.tensor.parallel import (
+    ColwiseParallel,
+    RowwiseParallel,
+    parallelize_module,
+)
+
+device_mesh, _world_size, _rank, logger = initialize_distributed_env(
+    "./tensor_parallel_simple_example"
+)
+
+"""
+This example copies some code from https://github.com/pytorch/examples/blob/main/distributed/tensor_parallelism/tensor_parallel_example.py
+"""
+
+
+class ToyModel(nn.Module):
+    """MLP based model"""
+
+    def __init__(self):
+        super(ToyModel, self).__init__()
+        self.in_proj = nn.Linear(10, 3200)
+        self.relu = nn.ReLU()
+        self.out_proj = nn.Linear(3200, 1600)
+        self.in_proj2 = nn.Linear(1600, 500)
+        self.out_proj2 = nn.Linear(500, 100)
+
+    def forward(self, x):
+        x = self.out_proj(self.relu(self.in_proj(x)))
+        x = self.relu(x)
+        x = self.out_proj2(self.relu(self.in_proj2(x)))
+        return x
+
+
+logger.info(f"Starting PyTorch TP example on rank {_rank}.")
+
+# # create model and move it to GPU - init"cuda"_mesh has already mapped GPU ids.
+tp_model = ToyModel().to("cuda")
+
+
+# Custom parallelization plan for the model
+tp_model = parallelize_module(
+    module=tp_model,
+    device_mesh=device_mesh,
+    parallelize_plan={
+        "in_proj": ColwiseParallel(input_layouts=Shard(0)),
+        "out_proj": RowwiseParallel(output_layouts=Shard(0)),
+        "in_proj2": ColwiseParallel(input_layouts=Shard(0)),
+        "out_proj2": RowwiseParallel(output_layouts=Shard(0)),
+    },
+)
+torch.manual_seed(0)
+inp = torch.rand(20, 10, device="cuda")
+python_result = tp_model(inp)
+
+
+backend = "torch_tensorrt"
+tp_model = torch.compile(
+    tp_model,
+    backend=backend,
+    options={
+        "truncate_long_and_double": True,
+        "enabled_precisions": {torch.float32, torch.float16},
+        "use_python_runtime": True,
+        "min_block_size": 1,
+        "use_aot_joint_export": False,
+    },
+    dynamic=False,
+)
+
+for i in range(10):
+    # For TP, input needs to be same across all TP ranks.
+    # Setting the random seed is to mimic the behavior of dataloader.
+    torch.manual_seed(i)
+    inp = torch.rand(20, 10, device="cuda")
+    start = time.time()
+    output = tp_model(inp)
+    end = time.time()
+    if i == 0:
+        logger.info(f"Compilation time is {end-start}")
+        assert (
+            python_result - output
+        ).std() < 0.01, "Compilation result is not correct."
+    elif _rank == 0:
+        logger.info(f"Inference time is {end-start}")

Diff for: tests/py/dynamo/distributed/test_nccl_ops.py

@@ -3,19 +3,11 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+from distributed_utils import set_environment_variables_pytest
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
 
-
-def set_environment_variables():
-    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["USE_TRTLLM_PLUGINS"] = "1"
-
-
-set_environment_variables()
+set_environment_variables_pytest()
 dist.init_process_group(backend="nccl", init_method="env://")
 group = dist.new_group(ranks=[0])
 group_name = group.group_name
@@ -47,7 +39,7 @@ def forward(self, x):
             DistributedGatherModel(linear_layer_dim).cuda(),
             inputs,
             use_dynamo_tracer=True,
-            fuse_distributed_ops=True,
+            enable_passes=True,
         )
 
     @parameterized.expand([(8)])
@@ -76,7 +68,7 @@ def forward(self, x):
             DistributedReduceScatterModel(linear_layer_dim).cuda(),
             inputs,
             use_dynamo_tracer=True,
-            fuse_distributed_ops=True,
+            enable_passes=True,
         )