diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index f2f6423e41..25c0839b27 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -92,7 +92,7 @@ def cross_compile_for_windows( require_full_compilation: bool = _defaults.REQUIRE_FULL_COMPILATION, min_block_size: int = _defaults.MIN_BLOCK_SIZE, torch_executed_ops: Optional[Collection[Target]] = None, - torch_executed_modules: Optional[List[str]] = None, + torch_executed_modules: Optional[Collection[str]] = None, pass_through_build_failures: bool = _defaults.PASS_THROUGH_BUILD_FAILURES, max_aux_streams: Optional[int] = _defaults.MAX_AUX_STREAMS, version_compatible: bool = _defaults.VERSION_COMPATIBLE, @@ -175,7 +175,7 @@ def cross_compile_for_windows( require_full_compilation (bool): Require modules to be compiled end to end or return an error as opposed to returning a hybrid graph where operations that cannot be run in TensorRT are run in PyTorch min_block_size (int): The minimum number of contiguous TensorRT convertible operations in order to run a set of operations in TensorRT torch_executed_ops (Collection[Target]): Set of aten operators that must be run in PyTorch. An error will be thrown if this set is not empty but ``require_full_compilation`` is True - torch_executed_modules (List[str]): List of modules that must be run in PyTorch. An error will be thrown if this list is not empty but ``require_full_compilation`` is True + torch_executed_modules (Collection[str]): Collection of modules that must be run in PyTorch. An error will be thrown if this collection is not empty but ``require_full_compilation`` is True pass_through_build_failures (bool): Error out if there are issues during compilation (only applicable to torch.compile workflows) max_aux_stream (Optional[int]): Maximum streams in the engine version_compatible (bool): Build the TensorRT engines compatible with future versions of TensorRT (Restrict to lean runtime operators to provide version forward compatibility for the engines) @@ -291,12 +291,6 @@ def cross_compile_for_windows( engine_capability = EngineCapability._from(engine_capability) - if torch_executed_modules is not None and torch_executed_modules: - logger.warning( - f"Detected torch_executed_modules was non-empty: {torch_executed_modules}" - "\nThis feature is unimplemented in Torch-TRT Dynamo currently." - ) - if use_fp32_acc: logger.debug( "FP32 accumulation for FP16 matmul layers is enabled. If " @@ -338,6 +332,9 @@ def cross_compile_for_windows( "torch_executed_ops": ( torch_executed_ops if torch_executed_ops is not None else set() ), + "torch_executed_modules": ( + torch_executed_modules if torch_executed_modules is not None else set() + ), "pass_through_build_failures": pass_through_build_failures, "max_aux_streams": max_aux_streams, "version_compatible": version_compatible, @@ -456,7 +453,7 @@ def compile( require_full_compilation: bool = _defaults.REQUIRE_FULL_COMPILATION, min_block_size: int = _defaults.MIN_BLOCK_SIZE, torch_executed_ops: Optional[Collection[Target]] = None, - torch_executed_modules: Optional[List[str]] = None, + torch_executed_modules: Optional[Collection[str]] = None, pass_through_build_failures: bool = _defaults.PASS_THROUGH_BUILD_FAILURES, max_aux_streams: Optional[int] = _defaults.MAX_AUX_STREAMS, version_compatible: bool = _defaults.VERSION_COMPATIBLE, @@ -554,7 +551,7 @@ def compile( require_full_compilation (bool): Require modules to be compiled end to end or return an error as opposed to returning a hybrid graph where operations that cannot be run in TensorRT are run in PyTorch min_block_size (int): The minimum number of contiguous TensorRT convertible operations in order to run a set of operations in TensorRT torch_executed_ops (Optional[Collection[Target]]): Set of aten operators that must be run in PyTorch. An error will be thrown if this set is not empty but ``require_full_compilation`` is True - torch_executed_modules (Optional[List[str]]): List of modules that must be run in PyTorch. An error will be thrown if this list is not empty but ``require_full_compilation`` is True + torch_executed_modules (Optional[Collection[str]]): Collection of modules that must be run in PyTorch. An error will be thrown if this collection is not empty but ``require_full_compilation`` is True pass_through_build_failures (bool): Error out if there are issues during compilation (only applicable to torch.compile workflows) max_aux_streams (Optional[int]): Maximum streams in the engine version_compatible (bool): Build the TensorRT engines compatible with future versions of TensorRT (Restrict to lean runtime operators to provide version forward compatibility for the engines) @@ -681,12 +678,6 @@ def compile( engine_capability = EngineCapability._from(engine_capability) - if torch_executed_modules is not None and torch_executed_modules: - logger.warning( - f"Detected torch_executed_modules was non-empty: {torch_executed_modules}" - "\nThis feature is unimplemented in Torch-TRT Dynamo currently." - ) - if autocast_low_precision_type is not None: if not isinstance(autocast_low_precision_type, (torch.dtype, dtype)): raise ValueError( @@ -749,6 +740,9 @@ def compile( "torch_executed_ops": ( torch_executed_ops if torch_executed_ops is not None else set() ), + "torch_executed_modules": ( + torch_executed_modules if torch_executed_modules is not None else set() + ), "pass_through_build_failures": pass_through_build_failures, "max_aux_streams": max_aux_streams, "version_compatible": version_compatible, @@ -1354,6 +1348,7 @@ def preserve_module_specs( gm, min_block_size=settings.min_block_size, torch_executed_ops=settings.torch_executed_ops, + torch_executed_modules=settings.torch_executed_modules, require_full_compilation=settings.require_full_compilation, skip_fusion=(num_supported_ops == total_ops), ) @@ -1374,6 +1369,7 @@ def preserve_module_specs( gm, min_block_size=settings.min_block_size, torch_executed_ops=settings.torch_executed_ops, + torch_executed_modules=settings.torch_executed_modules, require_full_compilation=settings.require_full_compilation, ) @@ -1785,7 +1781,7 @@ def convert_exported_program_to_serialized_trt_engine( require_full_compilation: bool = _defaults.REQUIRE_FULL_COMPILATION, min_block_size: int = _defaults.MIN_BLOCK_SIZE, torch_executed_ops: Optional[Collection[Target]] = None, - torch_executed_modules: Optional[List[str]] = None, + torch_executed_modules: Optional[Collection[str]] = None, pass_through_build_failures: bool = _defaults.PASS_THROUGH_BUILD_FAILURES, max_aux_streams: Optional[int] = _defaults.MAX_AUX_STREAMS, version_compatible: bool = _defaults.VERSION_COMPATIBLE, @@ -1887,7 +1883,7 @@ def convert_exported_program_to_serialized_trt_engine( require_full_compilation (bool): Require modules to be compiled end to end or return an error as opposed to returning a hybrid graph where operations that cannot be run in TensorRT are run in PyTorch min_block_size (int): The minimum number of contiguous TensorRT convertible operations in order to run a set of operations in TensorRT torch_executed_ops (Optional[Collection[Target]]): Set of aten operators that must be run in PyTorch. An error will be thrown if this set is not empty but ``require_full_compilation`` is True - torch_executed_modules (Optional[List[str]]): List of modules that must be run in PyTorch. An error will be thrown if this list is not empty but ``require_full_compilation`` is True + torch_executed_modules (Optional[Collection[str]]): Collection of modules that must be run in PyTorch. An error will be thrown if this collection is not empty but ``require_full_compilation`` is True pass_through_build_failures (bool): Error out if there are issues during compilation (only applicable to torch.compile workflows) max_aux_streams (Optional[int]): Maximum streams in the engine version_compatible (bool): Build the TensorRT engines compatible with future versions of TensorRT (Restrict to lean runtime operators to provide version forward compatibility for the engines) @@ -2008,12 +2004,6 @@ def convert_exported_program_to_serialized_trt_engine( engine_capability = EngineCapability._from(engine_capability) - if torch_executed_modules is not None and torch_executed_modules: - logger.warning( - f"Detected torch_executed_modules was non-empty: {torch_executed_modules}" - "\nThis feature is unimplemented in Torch-TRT Dynamo currently." - ) - if use_fp32_acc: logger.debug( "FP32 accumulation for FP16 matmul layers is enabled. If " @@ -2063,6 +2053,9 @@ def convert_exported_program_to_serialized_trt_engine( "torch_executed_ops": ( torch_executed_ops if torch_executed_ops is not None else set() ), + "torch_executed_modules": ( + torch_executed_modules if torch_executed_modules is not None else set() + ), "pass_through_build_failures": pass_through_build_failures, "max_aux_streams": max_aux_streams, "version_compatible": version_compatible, diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index a030f081b6..fc60b583ff 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -82,6 +82,7 @@ class CompilationSettings: workspace_size (int): Workspace TRT is allowed to use for the module (0 is default) min_block_size (int): Minimum number of operators per TRT-Engine Block torch_executed_ops (Collection[Target]): Collection of operations to run in Torch, regardless of converter coverage + torch_executed_modules (Collection[str]): Collection of modules to run in Torch pass_through_build_failures (bool): Whether to fail on TRT engine build errors (True) or not (False) max_aux_streams (Optional[int]): Maximum number of allowed auxiliary TRT streams for each engine version_compatible (bool): Provide version forward-compatibility for engine plan files @@ -147,6 +148,7 @@ class CompilationSettings: workspace_size: int = WORKSPACE_SIZE min_block_size: int = MIN_BLOCK_SIZE torch_executed_ops: Collection[Target] = field(default_factory=set) + torch_executed_modules: Collection[str] = field(default_factory=set) pass_through_build_failures: bool = PASS_THROUGH_BUILD_FAILURES max_aux_streams: Optional[int] = MAX_AUX_STREAMS version_compatible: bool = VERSION_COMPATIBLE diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index 098e9b2685..5e63f1f2af 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py @@ -25,6 +25,7 @@ from torch_tensorrt.dynamo.partitioning._global_partitioner import ( TorchTensorRTOperatorSupport, ) +from torch_tensorrt.dynamo.partitioning.common import node_in_torch_executed_module logger = logging.getLogger(__name__) @@ -32,13 +33,18 @@ class OpSupportTester(ops.OperatorSupportBase): # type: ignore """Class to determine whether operators within a module are supported""" - def __init__(self, torch_executed_ops: Collection[Target] = set()) -> None: + def __init__( + self, + torch_executed_ops: Collection[Target] = set(), + torch_executed_modules: Collection[str] = set(), + ) -> None: super().__init__() # Initialize sets of supported/unsupported operators self.supported_operators: Dict[str, int] = {} self.unsupported_operators: Dict[str, int] = {} self.torch_executed_ops = torch_executed_ops + self.torch_executed_modules = torch_executed_modules def is_node_supported( self, submodules: Dict[str, torch.nn.Module], node: torch.fx.Node @@ -71,6 +77,7 @@ def is_node_supported( (node in CONVERTERS or node.op == "get_attr") and node_name not in self.torch_executed_ops and node.target not in self.torch_executed_ops + and not node_in_torch_executed_module(node, self.torch_executed_modules) ): # If node is a proper, supported computational node, store the operator if not node.is_impure() and node.op != "get_attr": @@ -278,6 +285,7 @@ def partition( gm: torch.fx.GraphModule, min_block_size: int = MIN_BLOCK_SIZE, torch_executed_ops: Collection[Target] = set(), + torch_executed_modules: Collection[str] = set(), require_full_compilation: bool = REQUIRE_FULL_COMPILATION, skip_fusion: bool = False, ) -> Tuple[torch.fx.GraphModule, OpSupportTester]: @@ -288,6 +296,7 @@ def partition( gm: FX GraphModule to partition min_block_size: Minimum number of operators per TRT-Engine Block torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage + torch_executed_modules: Collection of modules to run in Torch require_full_compilation: Require that all computational operators be run in TRT skip_fusion: Skip fusions found by FxNetAccFusionsFinder Returns: @@ -298,8 +307,24 @@ def partition( gm.graph.lint() gm.recompile() + if torch_executed_modules: + all_module_types = { + module_type + for node in gm.graph.nodes + for _, module_type in (node.meta.get("nn_module_stack") or {}).values() + } + unmatched = set(torch_executed_modules) - all_module_types + if unmatched: + logger.warning( + f"The following torch_executed_modules were not found in the graph: " + f"{unmatched}. Ensure the module names are fully-qualified class names." + ) + # Construct - supported_ops = OpSupportTester(torch_executed_ops=torch_executed_ops) + supported_ops = OpSupportTester( + torch_executed_ops=torch_executed_ops, + torch_executed_modules=torch_executed_modules, + ) partitioner = TRTPartitioner( gm, supported_ops, diff --git a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py index 68e35e060a..86e4bb7907 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py @@ -16,6 +16,7 @@ from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( ConverterRegistry, ) +from torch_tensorrt.dynamo.partitioning.common import node_in_torch_executed_module from torch_tensorrt.dynamo.utils import COMPLEX_DTYPES logger = logging.getLogger(__name__) @@ -137,6 +138,7 @@ def __init__( self, support_dict: Optional[SupportDict] = None, torch_executed_ops: Collection[Target] = set(), + torch_executed_modules: Collection[str] = set(), ): super().__init__(support_dict) @@ -144,6 +146,7 @@ def __init__( self.supported_operators: Dict[str, int] = {} self.unsupported_operators: Dict[str, int] = {} self.torch_executed_ops: Collection[Target] = torch_executed_ops + self.torch_executed_modules: Collection[str] = torch_executed_modules @staticmethod def _has_complex_dtype(node: torch.fx.Node) -> bool: @@ -208,6 +211,7 @@ def is_node_supported( (node in CONVERTERS or node.op == "get_attr") and node_name not in self.torch_executed_ops and node.target not in self.torch_executed_ops + and not node_in_torch_executed_module(node, self.torch_executed_modules) ): # If node is a proper, supported computational node, store the operator if not node.is_impure() and node.op != "get_attr": @@ -258,6 +262,7 @@ def partition( gm: torch.fx.GraphModule, min_block_size: int = MIN_BLOCK_SIZE, torch_executed_ops: Collection[Target] = set(), + torch_executed_modules: Collection[str] = set(), require_full_compilation: bool = REQUIRE_FULL_COMPILATION, ) -> Tuple[torch.fx.GraphModule, TorchTensorRTOperatorSupport]: """Partition an FX GraphModule with aten ops into TRT engines @@ -267,11 +272,28 @@ def partition( gm: FX GraphModule to partition min_block_size: Minimum number of operators per TRT-Engine Block torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage + torch_executed_modules: Collection of modules to run in Torch require_full_compilation: Whether to require that all operators be run in TRT Returns: torch.fx.GraphModule, TorchTensorRTOperatorSupport """ - supported_ops = TorchTensorRTOperatorSupport(torch_executed_ops=torch_executed_ops) + if torch_executed_modules: + all_module_types = { + module_type + for node in gm.graph.nodes + for _, module_type in (node.meta.get("nn_module_stack") or {}).values() + } + unmatched = set(torch_executed_modules) - all_module_types + if unmatched: + logger.warning( + f"The following torch_executed_modules were not found in the graph: " + f"{unmatched}. Ensure the module names are fully-qualified class names." + ) + + supported_ops = TorchTensorRTOperatorSupport( + torch_executed_ops=torch_executed_ops, + torch_executed_modules=torch_executed_modules, + ) partitioner = TRTPartitioner( gm, supported_ops, diff --git a/py/torch_tensorrt/dynamo/partitioning/common.py b/py/torch_tensorrt/dynamo/partitioning/common.py index 62e68449b6..b1e14ac8b6 100644 --- a/py/torch_tensorrt/dynamo/partitioning/common.py +++ b/py/torch_tensorrt/dynamo/partitioning/common.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Dict, List, Optional, Sequence, Set, Tuple +from typing import Any, Collection, Dict, List, Optional, Sequence, Set, Tuple import sympy import torch @@ -81,6 +81,33 @@ def _build_submodule_profiles( return profiles +def node_in_torch_executed_module( + node: torch.fx.Node, torch_executed_modules: Collection[str] +) -> bool: + """ + Determine whether a node traces through a module that should run in Torch. + A node matches if any level of its nn_module_stack is one of the torch_executed_modules, + matched on fully-qualified class name (e.g. "torchvision.models.resnet.BasicBlock"), which + torch.export stores as the second element of each nn_module_stack value. + Args: + node: FX node to check + torch_executed_modules: Collection of fully-qualified module class names to run in Torch + Returns: + True if the node lies within a torch_executed_module, False otherwise + """ + if not torch_executed_modules: + return False + stack = node.meta.get("nn_module_stack") or {} + for _, module_type in stack.values(): + if module_type in torch_executed_modules: + logger.debug( + f"Excluding node {node.name} from TRT because its module type " + f"{module_type} is in torch_executed_modules." + ) + return True + return False + + def construct_dynamic_input( input_shape: torch.Size, input_dtype: torch.dtype, diff --git a/tests/py/dynamo/partitioning/test_000_fast_partitioning.py b/tests/py/dynamo/partitioning/test_000_fast_partitioning.py index 30f691d9e9..38ec08da4b 100644 --- a/tests/py/dynamo/partitioning/test_000_fast_partitioning.py +++ b/tests/py/dynamo/partitioning/test_000_fast_partitioning.py @@ -2,6 +2,8 @@ import numpy as np import torch +import torch_tensorrt +from parameterized import parameterized from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt.dynamo import partitioning @@ -178,6 +180,43 @@ def forward(self, x, y): "Certain operators are set to run in Torch, expected 1 segment", ) + @parameterized.expand( + [ + (["torch.nn.modules.conv.Conv2d"], 2), + ([], 1), + (["torch.nn.modules.container.Sequential"], 0), + ] + ) + def test_end2end_fast_partition_torch_executed_modules( + self, torch_executed_modules, trt_mod_cnt + ): + mod = ( + torch.nn.Sequential( + torch.nn.Conv2d(3, 8, 3, padding=1), + torch.nn.ReLU(), + torch.nn.BatchNorm2d(8), + torch.nn.Conv2d(8, 8, 3, padding=1), + torch.nn.ReLU(), + ) + .eval() + .to("cuda") + ) + with torch.no_grad(): + inputs = torch.rand((1, 3, 4, 4)).to("cuda") + trt_mod = torch_tensorrt.compile( + mod, + ir="dynamo", + inputs=[inputs], + min_block_size=1, + torch_executed_modules=torch_executed_modules, + use_fast_partitioner=True, + ) + cnt = 0 + for name, _ in trt_mod.named_children(): + if "_run_on_acc" in name: + cnt += 1 + self.assertEqual(cnt, trt_mod_cnt) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/partitioning/test_000_global_partitioning.py b/tests/py/dynamo/partitioning/test_000_global_partitioning.py index 887fa35659..389cc1c57a 100644 --- a/tests/py/dynamo/partitioning/test_000_global_partitioning.py +++ b/tests/py/dynamo/partitioning/test_000_global_partitioning.py @@ -3,7 +3,6 @@ import numpy as np import pytest import torch -import torch.nn.functional as F import torch_tensorrt from parameterized import parameterized from torch.testing._internal.common_utils import TestCase, run_tests @@ -117,6 +116,43 @@ def forward(self, x, y): "All operators are supported, there should be one segment", ) + @parameterized.expand( + [ + (["torch.nn.modules.conv.Conv2d"], 2), + ([], 1), + (["torch.nn.modules.container.Sequential"], 0), + ] + ) + def test_end2end_global_partition_torch_executed_modules( + self, torch_executed_modules, trt_mod_cnt + ): + mod = ( + torch.nn.Sequential( + torch.nn.Conv2d(3, 8, 3, padding=1), + torch.nn.ReLU(), + torch.nn.BatchNorm2d(8), + torch.nn.Conv2d(8, 8, 3, padding=1), + torch.nn.ReLU(), + ) + .eval() + .to("cuda") + ) + with torch.no_grad(): + inputs = torch.rand((1, 3, 4, 4)).to("cuda") + trt_mod = torch_tensorrt.compile( + mod, + ir="dynamo", + inputs=[inputs], + min_block_size=1, + torch_executed_modules=torch_executed_modules, + use_fast_partitioner=False, + ) + cnt = 0 + for name, _ in trt_mod.named_children(): + if "_run_on_acc" in name: + cnt += 1 + self.assertEqual(cnt, trt_mod_cnt) + if __name__ == "__main__": run_tests()