diff --git a/py/torch_tensorrt/_utils.py b/py/torch_tensorrt/_utils.py index 5d43db0e57..ed4a4a7399 100644 --- a/py/torch_tensorrt/_utils.py +++ b/py/torch_tensorrt/_utils.py @@ -7,7 +7,7 @@ import tempfile import urllib.request from pathlib import Path -from typing import Any, Optional +from typing import Any, Optional, Tuple import tensorrt as trt import torch @@ -380,3 +380,49 @@ def load_tensorrt_llm_for_nccl() -> bool: plugin_lib_path = download_and_get_plugin_lib_path() return load_and_initialize_trtllm_plugin(plugin_lib_path) # type: ignore[arg-type] return False + + +# --- TensorRT-RTX architecture targeting ------------------------------------- +# +# TensorRT-RTX runs on SM 7.5 and up, but its support matrix carves Turing out of +# several paths: FP32 GEMMs and 3D convolutions are documented as unsupported on +# compute capability 7.5, and bfloat16 has no Turing hardware at all. It also notes +# that Turing "compatibility is not included by default because including it may +# impact model performance", so targeting Turing is opt-in. +# +# These helpers exist so capability validators key off the compute capabilities the +# engine is being *built for*, not the machine doing the building. Calling +# torch.cuda.get_device_capability() inside a validator bakes the build host into the +# artifact, which is wrong for ahead-of-time deployment: a module compiled on Ampere +# and shipped to Turing would retain ops Turing cannot execute. + +TURING_COMPUTE_CAPABILITY = (7, 5) + + +def get_target_compute_capabilities( + settings: Optional[Any] = None, +) -> Tuple[Tuple[int, int], ...]: + """Compute capabilities this compilation targets. + + Returns the explicitly declared targets when the caller set them, otherwise the + capability of the current device. + """ + targets = ( + getattr(settings, "target_compute_capabilities", None) if settings else None + ) + if targets: + return tuple((int(major), int(minor)) for major, minor in targets) + return (torch.cuda.get_device_capability(),) + + +def trt_rtx_targets_turing(settings: Optional[Any] = None) -> bool: + """True when TensorRT-RTX is in use and SM 7.5 is among the build targets. + + A single compiled artifact carries a single partitioning, so an op unsupported on + *any* targeted architecture must fall back to PyTorch for all of them. + """ + from torch_tensorrt._features import ENABLED_FEATURES + + if not ENABLED_FEATURES.tensorrt_rtx: + return False + return TURING_COMPUTE_CAPABILITY in get_target_compute_capabilities(settings) diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 037174704f..5777508ab2 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -98,6 +98,9 @@ def cross_compile_for_windows( enable_experimental_decompositions: bool = _defaults.ENABLE_EXPERIMENTAL_DECOMPOSITIONS, dryrun: bool = _defaults.DRYRUN, hardware_compatible: bool = _defaults.HARDWARE_COMPATIBLE, + target_compute_capabilities: Optional[ + List[Tuple[int, int]] + ] = _defaults.TARGET_COMPUTE_CAPABILITIES, timing_cache_path: str = _defaults.TIMING_CACHE_PATH, lazy_engine_init: bool = _defaults.LAZY_ENGINE_INIT, cache_built_engines: bool = _defaults.CACHE_BUILT_ENGINES, @@ -181,6 +184,7 @@ def cross_compile_for_windows( enable_experimental_decompositions (bool): Use the full set of operator decompositions. These decompositions may not be tested but serve to make the graph easier to convert to TensorRT, potentially increasing the amount of graphs run in TensorRT. dryrun (bool): Toggle for "Dryrun" mode, running everything except conversion to TRT and logging outputs hardware_compatible (bool): Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer) + target_compute_capabilities (Optional[List[Tuple[int, int]]]): Compute capabilities to build for, e.g. ``[(7, 5)]`` for Turing. Defaults to None, meaning the current device. TensorRT-RTX only. Drives both engine targeting and op partitioning, so ops unsupported on any listed target fall back to PyTorch. timing_cache_path (str): Path to the timing cache if it exists (or) where it will be saved after compilation. Not used for TensorRT-RTX. lazy_engine_init (bool): Defer setting up engines until the compilation of all engines is complete. Can allow larger models with multiple graph breaks to compile but can lead to oversubscription of GPU memory at runtime. cache_built_engines (bool): Whether to save the compiled TRT engines to storage @@ -352,6 +356,7 @@ def cross_compile_for_windows( "dla_global_dram_size": dla_global_dram_size, "dryrun": dryrun, "hardware_compatible": hardware_compatible, + "target_compute_capabilities": target_compute_capabilities, "timing_cache_path": timing_cache_path, "lazy_engine_init": lazy_engine_init, "cache_built_engines": cache_built_engines, @@ -462,6 +467,9 @@ def compile( enable_experimental_decompositions: bool = _defaults.ENABLE_EXPERIMENTAL_DECOMPOSITIONS, dryrun: bool = _defaults.DRYRUN, hardware_compatible: bool = _defaults.HARDWARE_COMPATIBLE, + target_compute_capabilities: Optional[ + List[Tuple[int, int]] + ] = _defaults.TARGET_COMPUTE_CAPABILITIES, timing_cache_path: str = _defaults.TIMING_CACHE_PATH, lazy_engine_init: bool = _defaults.LAZY_ENGINE_INIT, cache_built_engines: bool = _defaults.CACHE_BUILT_ENGINES, @@ -560,6 +568,7 @@ def compile( enable_experimental_decompositions (bool): Use the full set of operator decompositions. These decompositions may not be tested but serve to make the graph easier to convert to TensorRT, potentially increasing the amount of graphs run in TensorRT. dryrun (bool): Toggle for "Dryrun" mode, running everything except conversion to TRT and logging outputs hardware_compatible (bool): Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer) + target_compute_capabilities (Optional[List[Tuple[int, int]]]): Compute capabilities to build for, e.g. ``[(7, 5)]`` for Turing. Defaults to None, meaning the current device. TensorRT-RTX only. Drives both engine targeting and op partitioning, so ops unsupported on any listed target fall back to PyTorch. timing_cache_path (str): Path to the timing cache if it exists (or) where it will be saved after compilation. Not used for TensorRT-RTX. lazy_engine_init (bool): Defer setting up engines until the compilation of all engines is complete. Can allow larger models with multiple graph breaks to compile but can lead to oversubscription of GPU memory at runtime. cache_built_engines (bool): Whether to save the compiled TRT engines to storage @@ -763,6 +772,7 @@ def compile( "dla_global_dram_size": dla_global_dram_size, "dryrun": dryrun, "hardware_compatible": hardware_compatible, + "target_compute_capabilities": target_compute_capabilities, "timing_cache_path": timing_cache_path, "lazy_engine_init": lazy_engine_init, "cache_built_engines": cache_built_engines, @@ -1758,6 +1768,9 @@ def convert_exported_program_to_serialized_trt_engine( enable_experimental_decompositions: bool = _defaults.ENABLE_EXPERIMENTAL_DECOMPOSITIONS, dryrun: bool = _defaults.DRYRUN, hardware_compatible: bool = _defaults.HARDWARE_COMPATIBLE, + target_compute_capabilities: Optional[ + List[Tuple[int, int]] + ] = _defaults.TARGET_COMPUTE_CAPABILITIES, timing_cache_path: str = _defaults.TIMING_CACHE_PATH, lazy_engine_init: bool = _defaults.LAZY_ENGINE_INIT, cache_built_engines: bool = _defaults.CACHE_BUILT_ENGINES, @@ -1854,6 +1867,7 @@ def convert_exported_program_to_serialized_trt_engine( enable_experimental_decompositions (bool): Use the full set of operator decompositions. These decompositions may not be tested but serve to make the graph easier to convert to TensorRT, potentially increasing the amount of graphs run in TensorRT. dryrun (bool): Toggle for "Dryrun" mode, running everything except conversion to TRT and logging outputs hardware_compatible (bool): Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer) + target_compute_capabilities (Optional[List[Tuple[int, int]]]): Compute capabilities to build for, e.g. ``[(7, 5)]`` for Turing. Defaults to None, meaning the current device. TensorRT-RTX only. Drives both engine targeting and op partitioning, so ops unsupported on any listed target fall back to PyTorch. timing_cache_path (str): Path to the timing cache if it exists (or) where it will be saved after compilation. Not used for TensorRT-RTX. lazy_engine_init (bool): Defer setting up engines until the compilation of all engines is complete. Can allow larger models with multiple graph breaks to compile but can lead to oversubscription of GPU memory at runtime. cache_built_engines (bool): Whether to save the compiled TRT engines to storage @@ -2033,6 +2047,7 @@ def convert_exported_program_to_serialized_trt_engine( "dla_global_dram_size": dla_global_dram_size, "dryrun": dryrun, "hardware_compatible": hardware_compatible, + "target_compute_capabilities": target_compute_capabilities, "timing_cache_path": timing_cache_path, "lazy_engine_init": lazy_engine_init, "cache_built_engines": cache_built_engines, diff --git a/py/torch_tensorrt/dynamo/_defaults.py b/py/torch_tensorrt/dynamo/_defaults.py index 1d6c65dd6f..767ee5c3a7 100644 --- a/py/torch_tensorrt/dynamo/_defaults.py +++ b/py/torch_tensorrt/dynamo/_defaults.py @@ -63,6 +63,10 @@ ENABLE_CROSS_COMPILE_FOR_WINDOWS = False TILING_OPTIMIZATION_LEVEL = "none" L2_LIMIT_FOR_TILING = -1 +# None means "target the current device". Set explicitly to build an engine deployable +# on other architectures; see torch_tensorrt._utils for why capability validators must +# consult this rather than the build host's device. +TARGET_COMPUTE_CAPABILITIES = None USE_DISTRIBUTED_MODE_TRACE = False OFFLOAD_MODULE_TO_CPU = False ENABLE_AUTOCAST = False diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index a030f081b6..460ca57d2d 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import Any, Collection, Optional, Set, Tuple, Union +from typing import Any, Collection, List, Optional, Set, Tuple, Union import tensorrt as trt import torch @@ -47,6 +47,7 @@ REUSE_CACHED_ENGINES, SPARSE_WEIGHTS, STRIP_ENGINE_WEIGHTS, + TARGET_COMPUTE_CAPABILITIES, TILING_OPTIMIZATION_LEVEL, TIMING_CACHE_PATH, TRUNCATE_DOUBLE, @@ -71,7 +72,10 @@ def _normalize_disabled_constant_fold_exclusions( validate_disabled_constant_fold_exclusions, ) - return validate_disabled_constant_fold_exclusions(rule_ids) + # Bind to a typed local: the imported helper is untyped from mypy's view here, + # and returning it directly trips --strict's no-any-return. + normalized: Set[str] = validate_disabled_constant_fold_exclusions(rule_ids) + return normalized @dataclass @@ -106,6 +110,7 @@ class CompilationSettings: TRT Engines. Prints detailed logs of the graph structure and nature of partitioning. Optionally saves the output to a file if a string path is specified hardware_compatible (bool): Build the TensorRT engines compatible with GPU architectures other than that of the GPU on which the engine was built (currently works for NVIDIA Ampere and newer) + target_compute_capabilities (Optional[List[Tuple[int, int]]]): Compute capabilities to build for, e.g. ``[(7, 5)]`` for Turing. Defaults to None, meaning the current device. TensorRT-RTX only. This drives both engine targeting and op partitioning: a compiled artifact carries a single partitioning, so an op unsupported on any listed target falls back to PyTorch for all of them. timing_cache_path (str): Path to the timing cache if it exists (or) where it will be saved after compilation. Not used for TensorRT-RTX (no autotuning). cache_built_engines (bool): Whether to save the compiled TRT engines to storage reuse_cached_engines (bool): Whether to load the compiled TRT engines from storage @@ -181,6 +186,9 @@ class CompilationSettings: enable_cross_compile_for_windows: bool = ENABLE_CROSS_COMPILE_FOR_WINDOWS tiling_optimization_level: str = TILING_OPTIMIZATION_LEVEL l2_limit_for_tiling: int = L2_LIMIT_FOR_TILING + target_compute_capabilities: Optional[List[Tuple[int, int]]] = ( + TARGET_COMPUTE_CAPABILITIES + ) use_distributed_mode_trace: bool = USE_DISTRIBUTED_MODE_TRACE offload_module_to_cpu: bool = OFFLOAD_MODULE_TO_CPU enable_autocast: bool = ENABLE_AUTOCAST @@ -245,6 +253,10 @@ def __setstate__(self, state: dict[str, Any]) -> None: "sparse_weights", "engine_capability", "hardware_compatible", + # An engine built for one set of compute capabilities cannot be reused for a + # compile targeting a different set -- doing so would silently reintroduce ops the + # target architecture cannot execute. + "target_compute_capabilities", "refit_identical_engine_weights", "immutable_weights", "enable_weight_streaming", diff --git a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py index 448d96a0ce..cc5e1face2 100644 --- a/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py +++ b/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py @@ -15,6 +15,7 @@ ) import numpy as np +import tensorrt as trt import torch import torch.fx from torch.fx.experimental.proxy_tensor import unset_fake_temporarily @@ -52,8 +53,6 @@ ) from torch_tensorrt.logging import TRT_LOGGER -import tensorrt as trt - _LOGGER: logging.Logger = logging.getLogger(__name__) TRT_INTERPRETER_CALL_PRE_OBSERVER: Observer[Callable[[torch.fx.GraphModule], None]] = ( @@ -385,6 +384,35 @@ def _populate_trt_builder_config( self.compilation_settings.l2_limit_for_tiling ) + # TensorRT-RTX ahead-of-time targeting. Left unset, TensorRT-RTX compiles for + # whatever device is present at build time, which is right for + # compile-here-run-here but cannot produce an artifact for another + # architecture. Turing in particular is opt-in: including it by default may + # cost performance elsewhere. The same setting drives the capability + # validators, so partitioning and engine targeting cannot drift apart. + if ( + ENABLED_FEATURES.tensorrt_rtx + and self.compilation_settings.target_compute_capabilities + ): + targets = self.compilation_settings.target_compute_capabilities + builder_config.num_compute_capabilities = len(targets) + for idx, (major, minor) in enumerate(targets): + name = f"SM{major}{minor}" + compute_capability = getattr(trt.ComputeCapability, name, None) + if compute_capability is None: + supported = [ + m for m in dir(trt.ComputeCapability) if m.startswith("SM") + ] + raise ValueError( + f"TensorRT-RTX has no compute capability {name} for requested " + f"target ({major}, {minor}). Supported: {supported}" + ) + if not builder_config.set_compute_capability(compute_capability, idx): + raise RuntimeError( + f"Failed to set TensorRT-RTX compute capability {name}" + ) + _LOGGER.info(f"Targeting TensorRT-RTX compute capabilities {targets}") + return builder_config def _create_timing_cache( diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 56ad2512e0..bbf4f9f170 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -10,7 +10,10 @@ from torch.fx.node import Argument, Node, Target from torch_tensorrt import ENABLED_FEATURES from torch_tensorrt._features import needs_not_tensorrt_rtx -from torch_tensorrt._utils import is_tensorrt_version_supported +from torch_tensorrt._utils import ( + is_tensorrt_version_supported, + trt_rtx_targets_turing, +) from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion import impl @@ -776,12 +779,75 @@ def aten_ops_gelu( ) -@dynamo_tensorrt_converter(torch.ops.aten.matmul, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.matmul.default, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.dot.default, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.mm.default, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.mv.default, supports_dynamic_shapes=True) -@dynamo_tensorrt_converter(torch.ops.aten.bmm.default, supports_dynamic_shapes=True) +def gemm_capability_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Reject FP32 GEMMs on TensorRT-RTX when Turing (SM 7.5) is a build target. + + TensorRT-RTX documents that on compute capability 7.5 it does not support FP32 + GEMMs. The two shape regimes fail differently, and the dynamic one is why this + guard cannot be limited to static shapes: + + * static shapes -- createExecutionContext() returns nullptr, surfacing as + "Unable to (re)create TensorRT execution context". Loud, easy to spot. + * dynamic shapes -- the engine builds, runs, and returns an all-zero tensor of + the correct shape and dtype, with no exception and no NaN. + + Only the operand dtype matters. fp16 GEMMs accumulating in fp32 (``use_fp32_acc``) + are unaffected and keep running on TensorRT. + + Note this relies on ``meta["val"]`` being populated, as every validator in this + module does. Graphs produced by ``torch.export`` always populate it. The converter + unit-test harness (``DispatchTestCase``) does not populate node meta at all, so + dtype cannot be determined there; those tests guard themselves with ``skipTest`` + instead. + """ + if not trt_rtx_targets_turing(settings): + return True + + for arg in node.args[:2]: + val = arg.meta.get("val") if hasattr(arg, "meta") else None + if val is not None and getattr(val, "dtype", None) == torch.float32: + _LOGGER.debug( + "FP32 GEMM '%s' is not supported on TensorRT-RTX for Turing " + "(SM 7.5). Falling back to PyTorch.", + node.name, + ) + return False + + return True + + +@dynamo_tensorrt_converter( + torch.ops.aten.matmul, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.matmul.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.dot.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.mm.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.mv.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) +@dynamo_tensorrt_converter( + torch.ops.aten.bmm.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) def aten_ops_matmul( ctx: ConversionContext, target: Target, @@ -3039,6 +3105,22 @@ def convolution_capability_validator( ) return False + # TensorRT-RTX does not support 3D convolutions on Turing (SM 7.5): the engine + # builds, but createExecutionContext() then returns nullptr because the JIT + # compiler finds no valid kernel config for 3D ConvFwd on SM 7.5. + # Transposed 3D convolution is a distinct layer and is unaffected, so this + # deliberately does not fire for deconvolution. + # aten.convolution input is (N, C, *spatial), so 5 dims means 3 spatial dims. + if trt_rtx_targets_turing(settings) and not args_bounds_check(node.args, 6): + val = node.args[0].meta.get("val") if hasattr(node.args[0], "meta") else None + if val is not None and val.ndim == 5: + _LOGGER.debug( + "3D convolution '%s' is not supported on TensorRT-RTX for Turing " + "(SM 7.5). Falling back to PyTorch.", + node.name, + ) + return False + return True @@ -3474,7 +3556,11 @@ def aten_ops_argmin( ) -@dynamo_tensorrt_converter(torch.ops.aten.addmm.default, supports_dynamic_shapes=True) +@dynamo_tensorrt_converter( + torch.ops.aten.addmm.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) @enforce_tensor_types( { 0: (TRTTensor,), diff --git a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py index 098e9b2685..2565b18bdd 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py @@ -54,6 +54,15 @@ def is_node_supported( return False settings = CONVERTERS.compilation_settings + if TorchTensorRTOperatorSupport._has_bf16_on_turing(node, settings): + # bfloat16 has no Turing hardware; compiling it for SM 7.5 crashes the + # process, so force the PyTorch fallback. + if not node.is_impure(): + self.unsupported_operators[node_name] = ( + self.unsupported_operators.get(node_name, 0) + 1 + ) + return False + if ( settings is not None and settings.fallback_data_dependent_ops diff --git a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py index 68e35e060a..90b7961d01 100644 --- a/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py +++ b/py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py @@ -6,6 +6,7 @@ from torch.fx.node import Target from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition from torch.fx.passes.operator_support import OperatorSupport, SupportDict +from torch_tensorrt._utils import trt_rtx_targets_turing from torch_tensorrt.dynamo._defaults import ( MIN_BLOCK_SIZE, REQUIRE_FULL_COMPILATION, @@ -166,6 +167,31 @@ def _dtype(n: torch.fx.Node) -> Optional[torch.dtype]: return True return False + @staticmethod + def _has_bf16_on_turing( + node: torch.fx.Node, settings: Optional[object] = None + ) -> bool: + """Return True if this node touches bfloat16 while targeting Turing (SM 7.5). + + Turing has no bfloat16 hardware -- PyTorch only reports support there because + it counts emulation. TensorRT-RTX has no such fallback, and compiling a bf16 + node for SM 7.5 crashes the process (SIGSEGV) rather than reporting an error, + so these nodes must run in the PyTorch fallback. + + This is checked at the partitioner rather than per-converter because the crash + is not specific to any one operator. + """ + if not trt_rtx_targets_turing(settings): + return False + + def _dtype(n: torch.fx.Node) -> Optional[torch.dtype]: + val = n.meta.get("val") + return getattr(val, "dtype", None) if val is not None else None + + if _dtype(node) == torch.bfloat16: + return True + return any(_dtype(arg) == torch.bfloat16 for arg in node.all_input_nodes) + @staticmethod def _requires_output_allocator(node: torch.fx.Node) -> bool: # True if the converter selected for this node needs a TRT output allocator, @@ -191,6 +217,15 @@ def is_node_supported( return False settings = CONVERTERS.compilation_settings + if self._has_bf16_on_turing(node, settings): + # bfloat16 has no Turing hardware; compiling it for SM 7.5 crashes the + # process, so force the PyTorch fallback. + if not node.is_impure(): + self.unsupported_operators[node_name] = ( + self.unsupported_operators.get(node_name, 0) + 1 + ) + return False + if ( settings is not None and settings.fallback_data_dependent_ops diff --git a/tests/py/dynamo/conversion/harness.py b/tests/py/dynamo/conversion/harness.py index 2dab3d64df..cf292ad09d 100644 --- a/tests/py/dynamo/conversion/harness.py +++ b/tests/py/dynamo/conversion/harness.py @@ -89,6 +89,34 @@ def infer_module_output_dtypes_for_test( # this is to enable dynamo tracer as True in the converter test files batch by batch +def skip_if_trt_rtx_turing(test_case: TestCase, what: str) -> None: + """Skip a converter test for a case TensorRT-RTX cannot serve on Turing (SM 7.5). + + TensorRT-RTX documents that on compute capability 7.5 it does not support FP32 + GEMMs or 3D convolutions, and Turing has no bfloat16 hardware at all. In a normal + ``torch_tensorrt.compile`` these fall back to PyTorch, either via a converter + capability validator or via the partitioner. + + Converter unit tests get neither. ``DispatchTestCase.run_test`` hands the graph + straight to ``TRTInterpreter``, skipping the partitioner, and the graphs it builds + carry **empty node meta** -- so a dtype-based capability validator cannot even see + the operand types here. A rejected node therefore raises + ``UnsupportedOperatorException`` rather than falling back, and an unsupported one + that is not rejected reaches TensorRT-RTX and fails (or, for bf16, crashes the + process). Hence the explicit skip. + """ + import torch_tensorrt + + if ( + torch_tensorrt.ENABLED_FEATURES.tensorrt_rtx + and torch.cuda.is_available() + and torch.cuda.get_device_capability() == (7, 5) + ): + test_case.skipTest( + f"{what} is not supported on TensorRT-RTX for Turing (SM 7.5)" + ) + + def get_use_dynamo_tracer(use_dynamo_tracer: Any) -> bool: # if in our converter tests we specifically set use_dynamo_tracer field, honor it if use_dynamo_tracer is not None and isinstance(use_dynamo_tracer, bool): diff --git a/tests/py/dynamo/conversion/test_binary_ops_aten.py b/tests/py/dynamo/conversion/test_binary_ops_aten.py index 4ac613adea..8a288ab5ff 100644 --- a/tests/py/dynamo/conversion/test_binary_ops_aten.py +++ b/tests/py/dynamo/conversion/test_binary_ops_aten.py @@ -8,7 +8,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing NEED_TEST_BOTH_CONSTANTS_CASE = True @@ -238,6 +238,8 @@ def forward(self, x, y): ] ) def test_elementwise_ops_bf16(self, _, orig_op): + skip_if_trt_rtx_turing(self, "bfloat16") + class TestModule(nn.Module): def __init__(self, orig_op): super().__init__() diff --git a/tests/py/dynamo/conversion/test_cdist_aten.py b/tests/py/dynamo/conversion/test_cdist_aten.py index 71628df510..bc0a864a63 100644 --- a/tests/py/dynamo/conversion/test_cdist_aten.py +++ b/tests/py/dynamo/conversion/test_cdist_aten.py @@ -3,7 +3,7 @@ from parameterized import parameterized from torch.testing._internal.common_utils import run_tests -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestCdistConverter(DispatchTestCase): @@ -21,6 +21,11 @@ class TestCdistConverter(DispatchTestCase): ] ) def test_cdist_float_same_shape(self, name, shape, p, compute_mode): + # cdist with p=2 is computed as an FP32 GEMM, which TensorRT-RTX does + # not support on Turing. + if p == 2: + skip_if_trt_rtx_turing(self, "cdist with p=2 (an FP32 GEMM)") + class Cdist(nn.Module): def forward(self, x1, x2): return torch.ops.aten._cdist_forward.default(x1, x2, p, compute_mode) @@ -47,6 +52,11 @@ def forward(self, x1, x2): def test_cdist_float_broadcast_and_diff_shape( self, name, shape_1, shape_2, p, compute_mode ): + # cdist with p=2 is computed as an FP32 GEMM, which TensorRT-RTX does + # not support on Turing. + if p == 2: + skip_if_trt_rtx_turing(self, "cdist with p=2 (an FP32 GEMM)") + class Cdist(nn.Module): def forward(self, x1, x2): return torch.ops.aten._cdist_forward.default(x1, x2, p, compute_mode) @@ -69,6 +79,11 @@ def forward(self, x1, x2): ] ) def test_cdist_p_2_compute_mode(self, name, shape_1, shape_2, p, compute_mode): + # cdist with p=2 is computed as an FP32 GEMM, which TensorRT-RTX does + # not support on Turing. + if p == 2: + skip_if_trt_rtx_turing(self, "cdist with p=2 (an FP32 GEMM)") + class Cdist(nn.Module): def forward(self, x1, x2): return torch.ops.aten._cdist_forward.default(x1, x2, p, compute_mode) @@ -85,6 +100,11 @@ def forward(self, x1, x2): def test_cdist_efficiency_p_2_compute_mode( self, name, shape_1, shape_2, p, compute_mode ): + # cdist with p=2 is computed as an FP32 GEMM, which TensorRT-RTX does + # not support on Turing. + if p == 2: + skip_if_trt_rtx_turing(self, "cdist with p=2 (an FP32 GEMM)") + class Cdist(nn.Module): def forward(self, x1, x2): return torch.ops.aten._cdist_forward.default(x1, x2, p, compute_mode) diff --git a/tests/py/dynamo/conversion/test_convolution_aten.py b/tests/py/dynamo/conversion/test_convolution_aten.py index 8af11cb180..fcbb2479fa 100644 --- a/tests/py/dynamo/conversion/test_convolution_aten.py +++ b/tests/py/dynamo/conversion/test_convolution_aten.py @@ -3,7 +3,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestConvolutionConverter(DispatchTestCase): @@ -233,6 +233,8 @@ def test_conv3d( groups=1, bias=True, ): + skip_if_trt_rtx_turing(self, "3D convolution") + class TestModule(torch.nn.Module): def __init__(self): super().__init__() @@ -255,6 +257,8 @@ def forward(self, x): # AssertionError: Channel dim can't be dynamic for convolution. def test_conv3d_with_dynamic_shape(self): + skip_if_trt_rtx_turing(self, "3D convolution") + class TestModule(torch.nn.Module): def __init__(self): super().__init__() diff --git a/tests/py/dynamo/conversion/test_matmul_aten.py b/tests/py/dynamo/conversion/test_matmul_aten.py index cf1fa36e82..6826f4afff 100644 --- a/tests/py/dynamo/conversion/test_matmul_aten.py +++ b/tests/py/dynamo/conversion/test_matmul_aten.py @@ -4,7 +4,7 @@ from torch.testing._internal.common_utils import run_tests from torch_tensorrt import Input -from .harness import DispatchTestCase +from .harness import DispatchTestCase, skip_if_trt_rtx_turing class TestMatMulConverter(DispatchTestCase): @@ -28,6 +28,8 @@ class TestMatMulConverter(DispatchTestCase): ] ) def test_matmul_dot(self, _, input_shape, other_shape): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def __init__(self): super().__init__() @@ -83,6 +85,8 @@ def forward(self, input): ] ) def test_matmul_mm(self, _, input_shape, other_shape): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def __init__(self): super().__init__() @@ -123,6 +127,8 @@ def forward(self, input): ] ) def test_matmul_mv(self, _, input_shape, other_shape): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def __init__(self): super().__init__() @@ -150,6 +156,8 @@ def forward(self, input): ] ) def test_matmul_matmul(self, _, input_shape, other_shape): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def forward(self, input, other): return torch.ops.aten.matmul(input, other) @@ -181,6 +189,8 @@ def forward(self, input, other): ] ) def test_matmul_matmul_dynamic_shape(self, *args): + skip_if_trt_rtx_turing(self, "FP32 GEMM") + class MatMul(nn.Module): def forward(self, input, other): return torch.ops.aten.matmul(input, other)