fix failures, changing getitem evaluator and adding lowering pass test case

Author: apbose

examples/distributed_inference/rotary_embedding.py

@@ -88,25 +88,21 @@ def __init__(self, dim: int, seq_len: int):
         self.wo = nn.Linear(dim, dim)
         self.seq_len = seq_len
         self.n_parallel = 1
-        self.register_buffer(
-            "freqs_cis",
-            self._precompute_freqs_cis(),
-            persistent=True,
+        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.init_weights()
 
     def init_weights(self):
         with torch.device(self.freqs_cis.device):
-            self.freqs_cis = self._precompute_freqs_cis()
-
-    def _precompute_freqs_cis(self) -> torch.Tensor:
-        theta = 10000.0
-        return precompute_freqs_cis(self.dim, self.seq_len, theta, self.n_parallel)
+            self.freqs_cis = self.freqs_cis
 
     def forward(self, x):
         q = self.wq(x)
         k = self.wk(x)
         # calculate rotary embedding
-        freqs_cis = self._precompute_freqs_cis().to(q.device)
+        freqs_cis = self.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_rotary_embedding.py

@@ -34,13 +34,7 @@
 
     logger.info("Torch-tensorrt compilation for rotary embedding")
 
-    # Compile the model
-    # for single GPU let us first try without this optiob
-
     model = torch.compile(model, backend="torch_tensorrt", options={"debug": True})
-    # model = torch_tensorrt.compile(model, target_ir="torch_compile", options={
-    #     "debug": True,
-    # })
 
     for i in range(15):
         # seeding with dp_rank to ensure identical inputs for TP groups

py/torch_tensorrt/dynamo/conversion/ops_evaluators.py

@@ -23,7 +23,10 @@ def getitem_validator(getitem_node: Node, settings: CompilationSettings = None)
     from torch_tensorrt.dynamo.conversion._ConverterRegistry import DYNAMO_CONVERTERS
 
     # Getitem nodes can only be converted if their parent node also can
-    return getitem_node.args[0] in DYNAMO_CONVERTERS
+    return (
+        getitem_node.args[0] in DYNAMO_CONVERTERS
+        or getitem_node.args[0].op == "get_attr"
+    )
 
 
 # TODO: Subsequent evaluators should be registered here with their own validators
@@ -43,7 +46,10 @@ def generic_evaluator(
     _LOGGER.debug(
         f"Evaluating {ConverterRegistry.qualified_name_or_str(target)} on object with name: {name}"
     )
-    return target(*args)
+    from torch._subclasses.fake_tensor import unset_fake_temporarily
+
+    with unset_fake_temporarily():
+        return target(*args)
 
 
 def rand_validator(rand_node: Node, settings: CompilationSettings = None) -> bool:

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

@@ -1,4 +1,5 @@
 import logging
+import operator
 from typing import Callable, List, Optional, Set, Tuple
 
 import torch
@@ -33,8 +34,7 @@ def __repr__(self):
 
 
 class ComplexOpDetector:
-    def __init__(self, logger):
-        self.logger = logger
+    def __init__(self):
         pass
 
     def is_complex_dtype(self, node: Node) -> bool:
@@ -45,15 +45,13 @@ def is_complex_dtype(self, node: Node) -> bool:
             if hasattr(val, "dtype"):
                 dtype = val.dtype
 
-        self.logger.debug(f"dtype of node: {dtype}")
+        logger.debug(f"dtype of node: {dtype}")
         return dtype in {torch.complex64, torch.complex128}
 
     def node_include_in_subgraph(self, node: Node) -> bool:
         # Include only call_function ops on complex tensors
-        self.logger.debug(f"node.op: {node.op}, node name: {node.name}")
-        self.logger.debug(f"is_complex_dtype: {self.is_complex_dtype(node)}")
         if node.op == "call_function" and self.is_complex_dtype(node):
-            self.logger.debug(
+            logger.debug(
                 f"node.op is added to subgraph: {node.op}, node name: {node.name} is complex"
             )
         return node.op == "call_function" and self.is_complex_dtype(node)
@@ -67,13 +65,11 @@ def subgraph_from_anchor(self, anchor_node: Node) -> ComplexSubGraphInfo:
             if n in subgraph_nodes:
                 continue
             subgraph_nodes.add(n)
-            self.logger.debug(f"node {n.name} is added to subgraph")
+            logger.debug(f"node {n.name} is added to subgraph")
             for inp in n.all_input_nodes:
                 if self.node_include_in_subgraph(inp):
-                    print("node inp is added to stack:", inp.name)
                     stack.append(inp)
                 else:
-                    print("node inp is not added to stack BUT INP:", inp.name)
                     input_nodes.add(inp)
         return ComplexSubGraphInfo(
             [anchor_node], list(subgraph_nodes), list(input_nodes)
@@ -85,13 +81,12 @@ def find_complex_op_subgraphs(
         complex_op_subgraphs: List[ComplexSubGraphInfo] = []
         for node in gm.graph.nodes:
             if node.target == anchor_target:
-                self.logger.debug(f"node.target {node.target} node.name: {node.name}")
                 new_sub = self.subgraph_from_anchor(node)
                 # if any intersecting nodes between seen and sub.subgraph_nodes they should be merged
                 merged = False
                 for existing_sub in complex_op_subgraphs:
                     if set(existing_sub.subgraph_nodes) & set(new_sub.subgraph_nodes):
-                        self.logger.debug(f"merging subgraphs {existing_sub} {new_sub}")
+                        logger.debug(f"merging subgraphs {existing_sub} {new_sub}")
                         # merge the two subgraphs
                         existing_sub.subgraph_nodes = list(
                             set(existing_sub.subgraph_nodes)
@@ -113,7 +108,7 @@ def find_complex_op_subgraphs(
 def complex_graph_detection(
     gm: GraphModule, settings: CompilationSettings
 ) -> List[ComplexSubGraphInfo]:
-    complex_op_detector = ComplexOpDetector(logger)
+    complex_op_detector = ComplexOpDetector()
     complex_subgraphs = complex_op_detector.find_complex_op_subgraphs(
         gm, anchor_target=torch.ops.aten.view_as_real.default
     )
@@ -174,17 +169,24 @@ def replace_input_node(self, input_node):
 
         elif input_node.op == "get_attr":
             new_attr_name = input_node.target + "_reshaped"
-            original_tensor = self.get_attr_tensor(input_node.target)
-            stacked_tensor = torch.stack(
-                [original_tensor.real, original_tensor.imag], dim=-1
-            )
-            self.gm.register_buffer(new_attr_name, stacked_tensor)
+            from torch._subclasses.fake_tensor import unset_fake_temporarily
+
+            with unset_fake_temporarily():
+                original_tensor = self.get_attr_tensor(input_node.target)
+                stacked_tensor = torch.stack(
+                    [original_tensor.real, original_tensor.imag], dim=-1
+                )
+                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: {input_node.op}")
-            logger.debug("This node type does not need to replaced")
+            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)
@@ -211,6 +213,8 @@ def rewrite_subgraph_nodes(self, subgraphs):
 
                     def match_complex_mul(
                         match: torch.fx.subgraph_rewriter.Match,
+                        original_graph,
+                        pattern_graph,
                     ) -> bool:
                         for original_node in match.nodes_map.values():
                             if original_node.name == node.name:
@@ -230,10 +234,9 @@ def match_complex_mul(
                     self.gm.graph.erase_node(node)
                 else:
                     logger.debug(f"Unsupported node target: {node.target}")
-                    logger.debug(f"This node type does not need to replaced")
-            if modified:
-                self.gm.graph.lint()
-                self.gm.recompile()
+                    logger.debug(
+                        "This complex subgraphnode type does not need to replaced"
+                    )
 
         if modified:
             self.gm.graph.lint()
@@ -256,16 +259,28 @@ def complex_mul_replacement() -> Tuple[
 
     # Original pattern: torch.mul for complex tensors
     def original_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        return torch.mul(x, y)
+        return torch.ops.aten.mul.Tensor(x, y)
 
     # Replacement function: manual complex multiplication on real/imag stacked tensors
     def replacement(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        x_real, x_imag = x[..., 0], x[..., 1]
-        y_real, y_imag = y[..., 0], y[..., 1]
-
-        real = x_real * y_real - x_imag * y_imag
-        imag = x_real * y_imag + x_imag * y_real
-
-        return torch.stack((real, imag), dim=-1)
+        x_real = torch.ops.aten.select.int(x, -1, 0)
+        x_imag = torch.ops.aten.select.int(x, -1, 1)  # x is reshape tensor
+        y_real, y_imag = y[..., 0], y[..., 1]  # y is frozen param
+
+        real_part1 = torch.ops.aten.mul.Tensor(x_real, y_real)
+        real_part2 = torch.ops.aten.mul.Tensor(x_imag, y_imag)
+        real = torch.ops.aten.sub.Tensor(real_part1, real_part2)
+
+        imag_part1 = torch.ops.aten.mul.Tensor(x_real, y_imag)
+        imag_part2 = torch.ops.aten.mul.Tensor(x_imag, y_real)
+        imag = torch.ops.aten.add.Tensor(imag_part1, imag_part2)
+
+        return torch.ops.aten.cat.default(
+            [
+                torch.ops.aten.unsqueeze.default(real, -1),
+                torch.ops.aten.unsqueeze.default(imag, -1),
+            ],
+            dim=-1,
+        )
 
     return (original_mul, replacement)

tests/py/dynamo/lowering/test_aten_lowering_passes.py

@@ -237,5 +237,97 @@ def forward(self, input, mat1, mat2):
         torch._dynamo.reset()
 
 
+class TestComplexSubgraph(TestCase):
+    def test_complex_subgraph(self):
+        def rotary_embedding(x, dim, freqs_cis=None):
+            x_ = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+            x_out_flatten = torch.view_as_real(x_ * freqs_cis).flatten(3)
+            return x_out_flatten.type_as(x_)
+
+        def _freqs_ex_tensor():
+            real = torch.tensor([[[[1.0000]], [[2.0000]]]], device="cuda")
+            imag = torch.tensor([[[[0.0000]], [[3.0000]]]], device="cuda")
+
+            z = torch.complex(real, imag)
+            return z
+
+        class RotaryAttention(torch.nn.Module):
+            def __init__(self, dim, seq_len):
+                super().__init__()
+                self.dim = dim
+                self.wq = torch.nn.Linear(dim, dim)
+                self.seq_len = seq_len
+                self._freqs_ex_tensor = _freqs_ex_tensor()
+
+                self.register_buffer(
+                    "freqs_ex_tensor",
+                    self._freqs_ex_tensor,
+                    persistent=True,
+                )
+
+            def forward(self, x):
+                q = self.wq(x)
+                freqs_cis = self._freqs_ex_tensor.to(q.device)
+                q_out = rotary_embedding(q, self.dim, freqs_cis)
+                return q_out
+
+        BATCH = 1
+        SEQ_LEN = 2
+        HEADS = 1
+        DIM = 2
+
+        inputs = [torch.randn(BATCH, SEQ_LEN, HEADS, DIM).cuda()]
+
+        fx_graph = torch.fx.symbolic_trace(RotaryAttention(DIM, SEQ_LEN))
+        expected_ops = {torch.ops.aten.mul.Tensor}
+        unexpected_ops = {
+            torch.ops.aten.view_as_complex.default,
+            torch.ops.aten.view_as_real.default,
+        }
+
+        unexpected_ops_seen, expected_ops_unseen = lower_graph_testing(
+            fx_graph,
+            inputs,
+            expected_ops=expected_ops,
+            unexpected_ops=unexpected_ops,
+            min_block_size=1,
+        )
+
+        self.assertEqual(
+            len(unexpected_ops_seen),
+            0,
+            f"The following unexpected ops were encountered: {unexpected_ops_seen}",
+        )
+
+        self.assertEqual(
+            len(expected_ops_unseen),
+            0,
+            f"The following expected ops were not encountered: {expected_ops_unseen}",
+        )
+        torch._dynamo.reset()
+
+        # Validate that the results between Torch and Torch-TRT are similar
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
+        )
+        optimized_model_results = optimized_model(*inputs)[0].detach().cpu()
+        torch_model_results = fx_graph(*inputs)[0].detach().cpu()
+
+        max_diff = float(
+            torch.max(torch.abs(optimized_model_results - torch_model_results))
+        )
+        self.assertAlmostEqual(
+            max_diff,
+            0,
+            DECIMALS_OF_AGREEMENT,
+            msg=f"ComplexSubgraph TRT outputs don't match with the original model.",
+        )
+        torch._dynamo.reset()
+
+
 if __name__ == "__main__":
     run_tests()