test case changes, adding metadata propogation in lowering pass

Author: apbose

examples/distributed_inference/rotary_embedding.py

@@ -49,13 +49,13 @@ def rotary_embedding(xq, xk, dim, freqs_cis=None):
     Returns:
         tuple: Tuple containing the rotated query and key tensors.
     """
-    freqs_cis = freqs_cis[None, :, None, :]  # [1, seq, 1, dim/2]
+
     xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
     xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
 
     xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
     xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
-    return xq_out.type_as(xq), xk_out.type_as(xk)
+    return (xq_out.type_as(xq), xk_out.type_as(xk))
 
 
 ########Tensor Parallel########
@@ -73,7 +73,7 @@ def parallel_rotary_block(rotary_block, tp_mesh):
         "wk": ColwiseParallel(),
         "wo": RowwiseParallel(output_layouts=Shard(0)),
     }
-    rotary_block.n_parallel = 2
+    rotary_block.n_parallel = 1  # this is for single GPU, to do remove this hardcode
 
     parallelize_module(rotary_block, tp_mesh, plan)
 
@@ -82,27 +82,26 @@ class RotaryAttention(nn.Module):
     def __init__(self, dim: int, seq_len: int):
         super().__init__()
         self.dim = dim
-        # self.rotary = RotaryEmbedding(dim)
         self.wq = nn.Linear(dim, dim)
         self.wk = nn.Linear(dim, dim)
         self.wo = nn.Linear(dim, dim)
         self.seq_len = seq_len
         self.n_parallel = 1
         theta = 10000.0
-        self.freqs_cis = precompute_freqs_cis(
-            self.dim, self.seq_len, theta, self.n_parallel
-        )
-        self.register_buffer("freqs_cis", self.freqs_cis, persistent=True)
+        self.register_buffer("freqs_cis", self._precompute_freqs_cis(), persistent=True)
         self.init_weights()
 
+    def _precompute_freqs_cis(self) -> torch.Tensor:
+        theta = 10000.0
+        return precompute_freqs_cis(self.dim, self.seq_len, theta, self.n_parallel)
+
     def init_weights(self):
         with torch.device(self.freqs_cis.device):
             self.freqs_cis = self.freqs_cis
 
     def forward(self, x):
         q = self.wq(x)
         k = self.wk(x)
-        # calculate rotary embedding
-        freqs_cis = self.freqs_cis.to(q.device)
+        freqs_cis = self._precompute_freqs_cis().to(q.device)
         q, k = rotary_embedding(q, k, self.dim, freqs_cis=freqs_cis)
         return self.wo(q)

examples/distributed_inference/tensor_parallel_initialize_dist.py

@@ -65,3 +65,9 @@ def initialize_distributed_env(logger_file_name, rank=0, world_size=1, port=2950
     torch.cuda.set_device(device_id)
 
     return device_mesh, world_size, rank, logger
+
+
+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

@@ -5,7 +5,10 @@
 import torch
 import torch_tensorrt
 from rotary_embedding import RotaryAttention, parallel_rotary_block
-from tensor_parallel_initialize_dist import initialize_distributed_env
+from tensor_parallel_initialize_dist import (
+    cleanup_distributed_env,
+    initialize_distributed_env,
+)
 
 device_mesh, _world_size, _rank, logger = initialize_distributed_env(
     "./tensor_parallel_rotary_embedding"
@@ -22,8 +25,6 @@
 DIM = 128
 
 with torch.no_grad():
-    xq = torch.randn(BATCH, SEQ_LEN, HEADS, DIM)
-    xk = torch.randn(BATCH, SEQ_LEN, HEADS, DIM)
     model = RotaryAttention(DIM, SEQ_LEN)
     parallel_rotary_block(model, device_mesh)
     device = torch.device("cuda", device_mesh.get_rank())
@@ -36,16 +37,22 @@
 
     model = torch.compile(model, backend="torch_tensorrt", options={"debug": True})
 
-    for i in range(15):
-        # seeding with dp_rank to ensure identical inputs for TP groups
-        torch.manual_seed(i)
-        start = time.time()
-        output = model(x)
-        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}")
+    try:
+        for i in range(15):
+            # seeding with dp_rank to ensure identical inputs for TP groups
+            torch.manual_seed(i)
+            start = time.time()
+            output = model(x)
+            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}")
+    except Exception as e:
+        logger.error(f"Error: {e}")
+        raise e
+    finally:
+        cleanup_distributed_env()

examples/distributed_inference/tensor_parallel_simple_example.py

@@ -5,7 +5,10 @@
 import torch.distributed as dist
 import torch.nn as nn
 import torch_tensorrt
-from tensor_parallel_initialize_dist import initialize_distributed_env
+from tensor_parallel_initialize_dist import (
+    cleanup_distributed_env,
+    initialize_distributed_env,
+)
 from torch.distributed._tensor import Shard
 from torch.distributed.tensor.parallel import (
     ColwiseParallel,
@@ -97,5 +100,4 @@ def forward(self, x):
             logger.info(f"Inference time is {end-start}")
 finally:
     # This cleans up the distributed process group
-    if dist.is_initialized():
-        dist.destroy_process_group()
+    cleanup_distributed_env()

py/torch_tensorrt/dynamo/lowering/passes/complex_graph_rewrite.py

@@ -105,24 +105,12 @@ def find_complex_op_subgraphs(
         return complex_op_subgraphs
 
 
-def complex_graph_detection(
-    gm: GraphModule, settings: CompilationSettings
-) -> List[ComplexSubGraphInfo]:
-    complex_op_detector = ComplexOpDetector()
-    complex_subgraphs = complex_op_detector.find_complex_op_subgraphs(
-        gm, anchor_target=torch.ops.aten.view_as_real.default
-    )
-    complex_graph_rewriter = ComplexGraphRewriter(gm, settings.truncate_double)
-    complex_graph_rewriter.rewrite_subgraph_nodes(complex_subgraphs)
-    return gm
-
-
 class ComplexGraphRewriter:
     def __init__(self, gm: GraphModule, truncate_double: bool = False):
         self.gm = gm
         self.truncate_double = truncate_double
 
-    def extract_shape_and_dtype(self, input_node):
+    def extract_shape_dtype_device(self, input_node):
         if input_node.op == "placeholder":
             tensor_val = input_node.meta["val"]
 
@@ -135,6 +123,7 @@ def extract_shape_and_dtype(self, input_node):
         node_shape = tensor_val.size()
         dtype = tensor_val.dtype
         new_node_shape = node_shape + (2,)
+        device = tensor_val.device
 
         if dtype == torch.complex64:
             new_node_dtype = torch.float32
@@ -143,7 +132,7 @@ def extract_shape_and_dtype(self, input_node):
         else:
             new_node_dtype = torch.float64
 
-        return new_node_shape, new_node_dtype
+        return new_node_shape, new_node_dtype, device
 
     def get_attr_tensor(self, target):
         # Check if target is param or buffer
@@ -157,14 +146,16 @@ def get_attr_tensor(self, target):
             )
 
     def replace_input_node(self, input_node):
+        modified = False
         logger.debug(f"Replacing input node: {input_node.name}")
-        new_shape, new_dtype = self.extract_shape_and_dtype(input_node)
-        real_tensor = torch.empty(new_shape, dtype=new_dtype)
+        new_shape, new_dtype, device = self.extract_shape_dtype_device(input_node)
+        real_tensor = torch.empty(new_shape, dtype=new_dtype, device=device)
 
         if input_node.op == "placeholder":
             with FakeTensorMode() as fake_mode:
                 fake_tensor = fake_mode.from_tensor(real_tensor)
-            new_node = self.gm.graph.placeholder(input_node.target + "_reshaped")
+            with self.gm.graph.inserting_before(input_node):
+                new_node = self.gm.graph.placeholder(input_node.target + "_reshaped")
             new_node.meta["val"] = fake_tensor
 
         elif input_node.op == "get_attr":
@@ -179,17 +170,16 @@ def replace_input_node(self, input_node):
                 self.gm.register_buffer(new_attr_name, stacked_tensor)
             with self.gm.graph.inserting_after(input_node):
                 new_node = self.gm.graph.get_attr(new_attr_name)
-
         else:
             logger.debug(
                 f"Unsupported node type in replacement of input node: {input_node.op}"
             )
             logger.debug(
                 "This complex subgraph inputnode type does not need to replaced"
             )
-
         input_node.replace_all_uses_with(new_node)
         self.gm.graph.erase_node(input_node)
+        clean_up_graph_after_modifications(self.gm)
 
     def rewrite_subgraph_nodes(self, subgraphs):
         modified = False
@@ -198,7 +188,6 @@ def rewrite_subgraph_nodes(self, subgraphs):
                 logger.debug(f"Input node rewrite: {input_node.name}")
                 if input_node.op not in ("call_function"):
                     self.replace_input_node(input_node)
-                    modified = True
             for node in subgraph.subgraph_nodes:
                 logger.debug(f"Subgraph Node rewrite: {node.name}")
                 if node.target == torch.ops.aten.view_as_complex.default:
@@ -229,6 +218,7 @@ def match_complex_mul(
                         ignore_literals=True,
                     )
                     replaced_nodes += nodes
+                    modified = True
                 elif node.target == torch.ops.aten.view_as_real.default:
                     node.replace_all_uses_with(node.args[0])
                     self.gm.graph.erase_node(node)
@@ -239,9 +229,36 @@ def match_complex_mul(
                     )
 
         if modified:
+            self.propagate_metadata()
             self.gm.graph.lint()
             self.gm.recompile()
 
+    def propagate_metadata(self):
+        fake_inputs = []
+        from torch._subclasses.fake_tensor import FakeTensorMode
+        from torch.fx.passes.fake_tensor_prop import FakeTensorProp
+
+        for node in self.gm.graph.nodes:
+            if node.op == "placeholder":
+                if "val" in node.meta:
+                    with FakeTensorMode(allow_non_fake_inputs=True):
+                        fake_val = node.meta["val"]
+                        fake_inputs.append(
+                            fake_val.to("cuda")
+                            if fake_val.device.type == "cuda"
+                            else fake_val
+                        )
+                else:
+                    fake_tensor = torch.empty(
+                        [s if s != 0 else 1 for s in node.meta["tensor_meta"].shape],
+                        dtype=node.meta["tensor_meta"].dtype,
+                        device=node.meta["tensor_meta"].device,
+                    )
+                    fake_inputs.append(fake_tensor)
+        FakeTensorProp(
+            self.gm, mode=FakeTensorMode(allow_non_fake_inputs=True)
+        ).propagate(*fake_inputs)
+
 
 def complex_mul_replacement() -> Tuple[
     Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
@@ -280,7 +297,23 @@ def replacement(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
                 torch.ops.aten.unsqueeze.default(real, -1),
                 torch.ops.aten.unsqueeze.default(imag, -1),
             ],
-            dim=-1,
+            -1,
         )
 
     return (original_mul, replacement)
+
+
+# This lowering pass is used to detect and rewrite complex subgraphs in the graph
+# This lowering pass works for complex tensor in mul which are parameter or buffers in the graph
+def complex_graph_detection(
+    gm: GraphModule, settings: CompilationSettings
+) -> List[ComplexSubGraphInfo]:
+    complex_op_detector = ComplexOpDetector()
+    complex_subgraphs = complex_op_detector.find_complex_op_subgraphs(
+        gm, anchor_target=torch.ops.aten.view_as_real.default
+    )
+    for subgraph in complex_subgraphs:
+        logger.debug(f"Complex subgraph info: {subgraph}")
+    complex_graph_rewriter = ComplexGraphRewriter(gm, settings.truncate_double)
+    complex_graph_rewriter.rewrite_subgraph_nodes(complex_subgraphs)
+    return gm