diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index eb43a8b55b..5b180d1580 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -4519,7 +4519,11 @@ def aten_ops_nonzero( ) -@dynamo_tensorrt_converter(torch.ops.aten.linear.default, supports_dynamic_shapes=True) +@dynamo_tensorrt_converter( + torch.ops.aten.linear.default, + capability_validator=gemm_capability_validator, + supports_dynamic_shapes=True, +) def aten_ops_linear( ctx: ConversionContext, target: Target, @@ -4538,9 +4542,43 @@ def aten_ops_linear( ) +def attention_capability_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Reject fused attention TensorRT-RTX cannot serve when Turing (SM 7.5) is a target. + + A fused attention op carries the same FP32 GEMMs ``gemm_capability_validator`` + rejects, but inside the converter, so the graph holds one node and no ``mm``/``bmm`` + for that guard to see. On Turing TRT-RTX this results in failure. Only the q/k/v + dtypes matter, so FP16 attention keeps running on TensorRT. + """ + if not trt_rtx_targets_turing(settings): + return True + + unsupported_dtypes = (torch.float32,) + + def is_unsupported(operand: Argument) -> bool: + val = operand.meta.get("val") if hasattr(operand, "meta") else None + return bool(getattr(val, "dtype", None) in unsupported_dtypes) + + qkv = node.args[:3] # query, key, value + if any(map(is_unsupported, qkv)): + _LOGGER.debug( + "Attention '%s' is not supported on TensorRT-RTX for Turing (SM 7.5). " + "Falling back to PyTorch.", + node.name, + ) + return False + + return True + + def scaled_dot_product_attention_validator( node: Node, settings: Optional[CompilationSettings] = None ) -> bool: + if not attention_capability_validator(node, settings): + return False + attn_mask = args_bounds_check(node.args, 3, None) is_causal = args_bounds_check(node.args, 5, False) if is_causal and attn_mask is not None: @@ -4639,6 +4677,9 @@ def aten_ops_scaled_dot_product_attention( def scaled_dot_product_flash_attention_validator( node: Node, settings: Optional[CompilationSettings] = None ) -> bool: + if not attention_capability_validator(node, settings): + return False + if args_bounds_check(node.args, 5, False): _LOGGER.debug("return_debug_mask is not yet supported.") return False @@ -4733,6 +4774,9 @@ def aten_ops_scaled_dot_product_flash_attention( def scaled_dot_product_efficient_attention_validator( node: Node, settings: Optional[CompilationSettings] = None ) -> bool: + if not attention_capability_validator(node, settings): + return False + if args_bounds_check(node.args, 4, False): _LOGGER.debug("compute_log_sumexp is not yet supported.") return False @@ -4819,6 +4863,7 @@ def scaled_dot_product_cudnn_attention_validator( _LOGGER.debug("return_debug_mask is not yet supported.") return False + # Delegating also picks up the Turing (SM 7.5) dtype guard. return scaled_dot_product_efficient_attention_validator(node, settings) diff --git a/tests/py/dynamo/conversion/test_attention_aten.py b/tests/py/dynamo/conversion/test_attention_aten.py index 26386a648a..be99271fd8 100644 --- a/tests/py/dynamo/conversion/test_attention_aten.py +++ b/tests/py/dynamo/conversion/test_attention_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 TestScaledDotProductAttention(DispatchTestCase): @@ -112,6 +112,9 @@ def test_sdpa_bool_mask( dropout_p=0.0, enable_gqa=False, ): + if dtype == torch.float32: + skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention") + class SDPA(nn.Module): def forward(self, query, key, value, attn_mask=None): return torch.ops.aten.scaled_dot_product_attention.default( @@ -247,6 +250,9 @@ def test_sdpa_fp_mask( dropout_p=0.0, enable_gqa=False, ): + if dtype == torch.float32: + skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention") + class SDPA(nn.Module): def forward(self, query, key, value, attn_mask=None): return torch.ops.aten.scaled_dot_product_attention.default( @@ -352,6 +358,9 @@ def test_dynamic_sdpa_fp_mask( dropout_p=0.0, enable_gqa=False, ): + if dtype == torch.float32: + skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention") + class SDPA(nn.Module): def forward(self, query, key, value, attn_mask=None): return torch.ops.aten.scaled_dot_product_attention.default( @@ -484,6 +493,9 @@ def test_efficient_sdpa( dtype, dropout_p=0.0, ): + if dtype == torch.float32: + skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention") + class EfficientSDPA(nn.Module): def forward(self, query, key, value, attn_bias=None): attn = torch.ops.aten._scaled_dot_product_efficient_attention.default( @@ -606,6 +618,9 @@ def test_efficient_sdpa_random_attn_bias( dtype, dropout_p=0.0, ): + if dtype == torch.float32: + skip_if_trt_rtx_turing(self, "FP32 scaled dot-product attention") + class EfficientSDPA(nn.Module): def forward(self, query, key, value, attn_bias=None): attn = torch.ops.aten._scaled_dot_product_efficient_attention.default( diff --git a/tests/py/dynamo/conversion/test_linear_aten.py b/tests/py/dynamo/conversion/test_linear_aten.py index 8619bebd7d..282a58b777 100644 --- a/tests/py/dynamo/conversion/test_linear_aten.py +++ b/tests/py/dynamo/conversion/test_linear_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 TestLinearConverter(DispatchTestCase): @@ -17,6 +17,8 @@ class TestLinearConverter(DispatchTestCase): ] ) def test_linear_converter(self, in_features, out_features): + skip_if_trt_rtx_turing(self, "nn.Linear (an FP32 GEMM)") + class LinearModel(nn.Module): def __init__(self, in_features, out_features): super(LinearModel, self).__init__() @@ -30,6 +32,8 @@ def forward(self, x): self.run_test(model, inputs, use_dynamo_tracer=True, enable_passes=True) def test_linear_with_dynamic_shape(self): + skip_if_trt_rtx_turing(self, "aten.linear (an FP32 GEMM)") + class LinearModel(torch.nn.Module): def forward(self, x, weight, bias): return torch.ops.aten.linear.default(x, weight, bias) @@ -50,6 +54,8 @@ def forward(self, x, weight, bias): ) def test_linear_with_rank_3_input_and_bias(self): + skip_if_trt_rtx_turing(self, "aten.linear (an FP32 GEMM)") + class LinearModel(torch.nn.Module): def forward(self, x, weight, bias): return torch.ops.aten.linear.default(x, weight, bias) diff --git a/tests/py/dynamo/models/test_export_kwargs_serde.py b/tests/py/dynamo/models/test_export_kwargs_serde.py index b9046bd638..e7719216f0 100644 --- a/tests/py/dynamo/models/test_export_kwargs_serde.py +++ b/tests/py/dynamo/models/test_export_kwargs_serde.py @@ -7,6 +7,7 @@ import torch.nn.functional as F import torch_tensorrt as torchtrt from torch import nn +from torch_tensorrt._utils import trt_rtx_targets_turing from torch_tensorrt.dynamo._compiler import ( convert_exported_program_to_serialized_trt_engine, ) @@ -19,6 +20,25 @@ assertions = unittest.TestCase() +def _no_fallback_gemm_dtype() -> torch.dtype: + """Model dtype for the tests that convert a whole program with no partitioner. + + TensorRT-RTX cannot serve an FP32 GEMM on Turing (SM 7.5), so + ``gemm_capability_validator`` rejects the models' trailing ``nn.Linear`` there. In a + normal ``torchtrt.dynamo.compile`` the partitioner turns that rejection into a + PyTorch block, but ``convert_exported_program_to_serialized_trt_engine`` emits one + engine for the whole program by design and has nothing to fall back to, so the + rejection raises ``UnsupportedOperatorException`` instead. + + These tests are about kwarg plumbing, not about FP32, so follow the capability + rather than switching them off: the guard keys on operand dtype only and Turing has + FP16 GEMM hardware. Keyed off the same predicate the validator uses so the two + cannot drift, and no ``enabled_precisions`` change is needed because the network is + strongly typed. + """ + return torch.float16 if trt_rtx_targets_turing() else torch.float32 + + @pytest.mark.unit @pytest.mark.critical def test_custom_model(tmpdir): @@ -496,8 +516,9 @@ def forward(self, x, b=5, c=None, d=None): x = x - d["value"] return self.fc1(x) - model = net().eval().to("cuda") - args = [torch.rand((1, 3, 224, 224)).to("cuda")] + dtype = _no_fallback_gemm_dtype() + model = net().eval().to("cuda").to(dtype) + args = [torch.rand((1, 3, 224, 224), dtype=dtype).to("cuda")] kwargs = { "b": torch.tensor(6).to("cuda"), "d": {"value": torch.tensor(8).to("cuda")}, @@ -1090,9 +1111,10 @@ def forward(self, x, b=5, c=None, d=None): x = x - d["value"] return self.fc1(x) - model = net().eval().to("cuda") + dtype = _no_fallback_gemm_dtype() + model = net().eval().to("cuda").to(dtype) kwargs = { - "x": torch.rand((1, 3, 224, 224)).to("cuda"), + "x": torch.rand((1, 3, 224, 224), dtype=dtype).to("cuda"), "b": torch.tensor(6).to("cuda"), "d": {"value": torch.tensor(8).to("cuda")}, } diff --git a/tests/py/dynamo/models/test_model_refit.py b/tests/py/dynamo/models/test_model_refit.py index 4f165d5278..a8c4d930ce 100644 --- a/tests/py/dynamo/models/test_model_refit.py +++ b/tests/py/dynamo/models/test_model_refit.py @@ -11,7 +11,7 @@ import torch_tensorrt as torch_trt from torch import nn from torch_tensorrt._utils import is_tensorrt_rtx_version_supported -from torch_tensorrt.dynamo import refit_module_weights +from torch_tensorrt.dynamo import partitioning, refit_module_weights from torch_tensorrt.dynamo._refit import ( construct_refit_mapping, get_engine_from_encoded_engine, @@ -68,7 +68,26 @@ def test_mapping(): ) new_gm = exp_program2.module() new_gm = post_lowering(new_gm, settings) - mapping = construct_refit_mapping(new_gm, trt_input, settings) + + # construct_refit_mapping interprets whatever module it is handed, so it has to be + # handed the same subgraph the engine above was built from. refit_module_weights + # always partitions first and maps each accelerated submodule; this test was the + # only caller passing a whole, un-partitioned module. That only worked because the + # model happened to be fully convertible: on Turing (SM 7.5) TensorRT-RTX cannot + # serve resnet18's FP32 fc GEMM, so capability partitioning routes it to PyTorch and + # the engine does not contain it -- while the un-partitioned module still does, with + # no partitioner in construct_refit_mapping to fall back to. Partition here the same + # way compile_module does, so the mapping matches the engine on either architecture. + num_supported_ops, total_ops = partitioning.get_graph_converter_support( + new_gm, settings.torch_executed_ops + ) + partitioned_gm, _ = partitioning.fast_partition( + new_gm, + min_block_size=min_block_size, + torch_executed_ops=settings.torch_executed_ops, + skip_fusion=(num_supported_ops == total_ops), + ) + mapping = construct_refit_mapping(partitioned_gm._run_on_acc_0, trt_input, settings) refitter = trt.Refitter(engine, TRT_LOGGER) weight_list = refitter.get_all_weights() diff --git a/tests/py/dynamo/models/test_turing_capability_guards.py b/tests/py/dynamo/models/test_turing_capability_guards.py index e982a22722..ff0cbb3573 100644 --- a/tests/py/dynamo/models/test_turing_capability_guards.py +++ b/tests/py/dynamo/models/test_turing_capability_guards.py @@ -97,6 +97,35 @@ def forward(self, x): return self.conv(F.pad(x, (0, 1, 0, 1))) +class Linear(nn.Module): + """aten.linear.default: a GEMM the explicit GEMM guards never see. + + Its decomposition into addmm is deliberately disabled, so guarding addmm does + not cover it. + """ + + def __init__(self, dtype=torch.float32): + super().__init__() + self.linear = nn.Linear(16, 32).to(dtype) + + def forward(self, x): + return self.linear(x) + + +class SDPA(nn.Module): + """Fused attention: the GEMMs live inside the converter, not in the graph.""" + + def forward(self, q, k, v): + return torch.ops.aten.scaled_dot_product_attention.default(q, k, v) + + +class EfficientSDPA(nn.Module): + def forward(self, q, k, v): + return torch.ops.aten._scaled_dot_product_efficient_attention.default( + q, k, v, None, False + )[0] + + class Cdist(nn.Module): """A GEMM the explicit GEMM guards never see, emitted inside the cdist converter. @@ -229,6 +258,45 @@ def test_declared_turing_target_keeps_transposed_conv3d_on_trt(self): compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) self.assertGreater(_trt_submodule_count(compiled), 0) + def test_declared_turing_target_falls_back_fp32_linear(self): + mod = Linear() + inputs = (torch.randn(8, 16).cuda(),) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertEqual(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + + def test_declared_turing_target_falls_back_fp32_sdpa(self): + mod = SDPA() + inputs = tuple(torch.randn(4, 8, 32, 16).cuda() for _ in range(3)) + compiled = self._compile( + mod, + inputs, + target_compute_capabilities=[TURING], + decompose_attention=False, + ) + self.assertEqual(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + + def test_declared_turing_target_falls_back_fp32_efficient_sdpa(self): + mod = EfficientSDPA() + inputs = tuple(torch.randn(4, 8, 32, 16).cuda() for _ in range(3)) + compiled = self._compile( + mod, + inputs, + target_compute_capabilities=[TURING], + decompose_attention=False, + ) + self.assertEqual(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + + def test_declared_turing_target_keeps_fp16_linear_on_trt(self): + # The GEMM guards key on FP32 only; FP16 GEMMs run fine on Turing. + mod = Linear(dtype=torch.half) + inputs = (torch.randn(8, 16, dtype=torch.half).cuda(),) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertGreater(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + def test_declared_turing_target_falls_back_cdist_p2_compute_mode_1(self): # compute_mode=1 is "always use the matrix multiply", whatever the row count. mod = Cdist(compute_mode=1) @@ -458,6 +526,61 @@ def test_fp32_gemm_dynamic_is_not_silently_zero(self): ) self.assertGreater(cos.item(), 0.99) + def test_fp32_linear_dynamic_is_not_silently_zero(self): + # Unguarded, this returned an all-zero tensor of the right shape and dtype: + # TensorRT-RTX logged a cuDNN graph-compilation failure from enqueueV3 and + # returned normally, so the caller saw no error at all. + mod = Linear().eval().cuda() + x = torch.randn(8, 16).cuda() + with torch.no_grad(): + ref = mod(x) + dim = torch.export.Dim("batch", min=1, max=64) + ep = torch.export.export(mod, (x,), dynamic_shapes=({0: dim},)) + compiled = torchtrt.compile( + ep, + arg_inputs=[x], + ir="dynamo", + min_block_size=1, + cache_built_engines=False, + reuse_cached_engines=False, + use_python_runtime=True, + ) + with torch.no_grad(): + out = compiled(x) + self.assertFalse(bool(torch.all(out == 0).item())) + cos = F.cosine_similarity( + ref.flatten().unsqueeze(0), out.flatten().unsqueeze(0) + ) + self.assertGreater(cos.item(), 0.99) + + def test_fp32_sdpa_dynamic_is_not_silently_zero(self): + # Same silent failure, reached through a fused attention op instead of a GEMM. + mod = SDPA().eval().cuda() + qkv = tuple(torch.randn(4, 8, 32, 16).cuda() for _ in range(3)) + with torch.no_grad(): + ref = mod(*qkv) + dim = torch.export.Dim("batch", min=1, max=16) + ep = torch.export.export( + mod, qkv, dynamic_shapes=({0: dim}, {0: dim}, {0: dim}) + ) + compiled = torchtrt.compile( + ep, + arg_inputs=list(qkv), + ir="dynamo", + min_block_size=1, + cache_built_engines=False, + reuse_cached_engines=False, + use_python_runtime=True, + decompose_attention=False, + ) + with torch.no_grad(): + out = compiled(*qkv) + self.assertFalse(bool(torch.all(out == 0).item())) + cos = F.cosine_similarity( + ref.flatten().unsqueeze(0), out.flatten().unsqueeze(0) + ) + self.assertGreater(cos.item(), 0.99) + def test_conv3d_falls_back(self): mod = Conv3d() inputs = (torch.randn(1, 4, 8, 8, 8).cuda(),) diff --git a/tests/py/dynamo/models/test_weight_stripped_engine.py b/tests/py/dynamo/models/test_weight_stripped_engine.py index beae4c4ec0..815bceac48 100644 --- a/tests/py/dynamo/models/test_weight_stripped_engine.py +++ b/tests/py/dynamo/models/test_weight_stripped_engine.py @@ -7,6 +7,7 @@ import torch import torch_tensorrt as torch_trt from torch.testing._internal.common_utils import TestCase +from torch_tensorrt._utils import trt_rtx_targets_turing from torch_tensorrt.dynamo import convert_exported_program_to_serialized_trt_engine from torch_tensorrt.dynamo._defaults import TIMING_CACHE_PATH from torch_tensorrt.dynamo._refit import refit_module_weights @@ -20,6 +21,21 @@ import torchvision.models as models +def _turing_safe_gemm_dtype() -> torch.dtype: + """Model dtype for the resnet18 tests in this file. + + On Turing ``gemm_capability_validator`` rejects resnet18's trailing ``fc``, which + breaks these tests two ways: the single-engine entry point has no partitioner to fall + back to and raises, while ``compile`` puts ``fc`` in a PyTorch block whose real bias + is then applied to the stripped engine's zeros, so the output is not the all-zeros + these tests assert. None of them is about FP32 -- they are about engine sizes, + stripping, refit and runtimes, all just as real in FP16 -- so follow the capability + rather than switching them off. Keyed off the same predicate the validator uses so + the two cannot drift. + """ + return torch.float16 if trt_rtx_targets_turing() else torch.float32 + + class TestWeightStrippedEngine(TestCase): @unittest.skipIf( not torch_trt.ENABLED_FEATURES.refit, @@ -79,8 +95,12 @@ def test_three_ways_to_compile(self): "torchvision is not installed", ) def test_compile_weight_stripped_engine(self): - pyt_model = models.resnet18(weights=None).eval().to("cuda") - example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) + # See _turing_safe_gemm_dtype: in FP32 on Turing the fc GEMM falls back to a + # PyTorch block whose real weights survive engine weight-stripping, so the + # all-zeros assertion below cannot hold. + dtype = _turing_safe_gemm_dtype() + pyt_model = models.resnet18(weights=None).eval().to("cuda").to(dtype) + example_inputs = (torch.randn((100, 3, 224, 224), dtype=dtype).to("cuda"),) settings = { "min_block_size": 1, @@ -110,8 +130,12 @@ def test_compile_weight_stripped_engine(self): "torchvision is not installed", ) def test_weight_stripped_engine_sizes(self): - pyt_model = models.resnet18(pretrained=True).eval().to("cuda") - example_inputs = (torch.randn((2, 3, 224, 224)).to("cuda"),) + # See _turing_safe_gemm_dtype: on Turing the FP32 fc GEMM is rejected and this + # entry point has no partitioner to fall back to. The stripped-vs-included size + # comparison below is just as meaningful in FP16. + dtype = _turing_safe_gemm_dtype() + pyt_model = models.resnet18(pretrained=True).eval().to("cuda").to(dtype) + example_inputs = (torch.randn((2, 3, 224, 224), dtype=dtype).to("cuda"),) exp_program = torch.export.export(pyt_model, example_inputs) weight_included_engine = convert_exported_program_to_serialized_trt_engine( exp_program, @@ -155,15 +179,19 @@ def test_weight_stripped_engine_sizes(self): "torchvision is not installed", ) def test_weight_stripped_engine_results(self): - pyt_model = models.resnet18(pretrained=True).eval().to("cuda") - example_inputs = (torch.randn((2, 3, 224, 224)).to("cuda"),) + # See _turing_safe_gemm_dtype: in FP32 on Turing the fc GEMM falls back to a + # PyTorch block whose real weights survive engine weight-stripping, so the + # all-zeros assertion below cannot hold. + dtype = _turing_safe_gemm_dtype() + pyt_model = models.resnet18(pretrained=True).eval().to("cuda").to(dtype) + example_inputs = (torch.randn((2, 3, 224, 224), dtype=dtype).to("cuda"),) # Mark the dim0 of inputs as dynamic batch = torch.export.Dim("batch", min=1, max=200) exp_program = torch.export.export( pyt_model, args=example_inputs, dynamic_shapes={"x": {0: batch}} ) - inputs = [torch.rand((2, 3, 224, 224)).to("cuda")] + inputs = [torch.rand((2, 3, 224, 224), dtype=dtype).to("cuda")] trt_gm = torch_trt.dynamo.compile( exp_program, @@ -530,13 +558,17 @@ def forward(self, x): "torchvision is not installed", ) def test_two_TRTRuntime_in_refitting(self): - pyt_model = models.resnet18(pretrained=True).eval().to("cuda") - example_inputs = (torch.randn((2, 3, 224, 224)).to("cuda"),) + # See _turing_safe_gemm_dtype: in FP32 on Turing the fc GEMM falls back to a + # PyTorch block whose real weights survive engine weight-stripping, so the + # all-zeros assertion below cannot hold. + dtype = _turing_safe_gemm_dtype() + pyt_model = models.resnet18(pretrained=True).eval().to("cuda").to(dtype) + example_inputs = (torch.randn((2, 3, 224, 224), dtype=dtype).to("cuda"),) batch = torch.export.Dim("batch", min=1, max=200) exp_program = torch.export.export( pyt_model, args=example_inputs, dynamic_shapes={"x": {0: batch}} ) - inputs = [torch.rand((2, 3, 224, 224)).to("cuda")] + inputs = [torch.rand((2, 3, 224, 224), dtype=dtype).to("cuda")] pyt_results = pyt_model(*inputs) @@ -620,15 +652,19 @@ def test_refit_identical_engine_weights(self): "Multiple refit requires TensorRT >= 10.14 with INCLUDE_REFIT serialization flag", ) def test_refit_weight_stripped_engine_multiple_times(self): - pyt_model = models.resnet18(pretrained=True).eval().to("cuda") - example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) + # See _turing_safe_gemm_dtype: in FP32 on Turing the fc GEMM falls back to a + # PyTorch block whose real weights survive engine weight-stripping, so the + # all-zeros assertion below cannot hold. + dtype = _turing_safe_gemm_dtype() + pyt_model = models.resnet18(pretrained=True).eval().to("cuda").to(dtype) + example_inputs = (torch.randn((100, 3, 224, 224), dtype=dtype).to("cuda"),) # Mark the dim0 of inputs as dynamic batch = torch.export.Dim("batch", min=1, max=200) exp_program = torch.export.export( pyt_model, args=example_inputs, dynamic_shapes={"x": {0: batch}} ) - inputs = (torch.rand((128, 3, 224, 224)).to("cuda"),) + inputs = (torch.rand((128, 3, 224, 224), dtype=dtype).to("cuda"),) trt_gm = torch_trt.dynamo.compile( exp_program, @@ -654,7 +690,7 @@ def test_refit_weight_stripped_engine_multiple_times(self): msg="refitted engine results should not be all zeros", ) - inputs2 = (torch.rand((64, 3, 224, 224)).to("cuda"),) + inputs2 = (torch.rand((64, 3, 224, 224), dtype=dtype).to("cuda"),) exp_program2 = torch.export.export( pyt_model, args=inputs2, dynamic_shapes={"x": {0: batch}} ) diff --git a/tests/py/dynamo/partitioning/test_001_resource_partitioning.py b/tests/py/dynamo/partitioning/test_001_resource_partitioning.py index 81c34e2bf5..7f4d316367 100644 --- a/tests/py/dynamo/partitioning/test_001_resource_partitioning.py +++ b/tests/py/dynamo/partitioning/test_001_resource_partitioning.py @@ -8,6 +8,7 @@ from torch.fx.passes.splitter_base import Subgraph from torch.ops import aten from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt._utils import trt_rtx_targets_turing from torch_tensorrt.dynamo import partitioning from torch_tensorrt.dynamo.conversion import CompilationSettings from torch_tensorrt.dynamo.lowering import ( @@ -29,6 +30,19 @@ _FIXED_RSS_BYTES = 512 * 1024 * 1024 # 512 MB +def _fp32_gemm_falls_back_to_pytorch() -> bool: + """Whether the trailing FP32 ``nn.Linear`` of these models lands in a PyTorch block. + + TensorRT-RTX cannot serve an FP32 GEMM on Turing (SM 7.5), so + ``gemm_capability_validator`` rejects ``aten.linear.default`` there and capability + partitioning routes the Linear to the fallback, adding one non-accelerated block. + Keyed off the same predicate the validator uses so the expectation cannot drift + from the guard. These tests install no settings into the converter registry, so + the validator resolves the same live-device answer this no-argument call does. + """ + return trt_rtx_targets_turing() + + class TestResourcePartitioning(TestCase): def test_resource_partitioning(self): class net(nn.Module): @@ -183,6 +197,9 @@ def forward(self, x): ) == 4 ), "The graph should have 4 accelerated subgraphs" + # The trailing FP32 Linear becomes an extra PyTorch block wherever TensorRT-RTX + # cannot serve an FP32 GEMM; the accelerated blocks above are unaffected. + expected_gpu_subgraphs = 3 if _fp32_gemm_falls_back_to_pytorch() else 2 assert ( len( [ @@ -191,8 +208,8 @@ def forward(self, x): if "_run_on_gpu" in name ] ) - == 2 - ), "The graph should have 2 non-accelerated subgraphs" + == expected_gpu_subgraphs + ), f"The graph should have {expected_gpu_subgraphs} non-accelerated subgraphs" torch._dynamo.reset() @@ -202,6 +219,8 @@ def test_resource_partitioning_with_capability_partitioning_and_atomic_subgraphs """ After defining the atomic subgraphs, the resource partitioner will not be able to find valid partition in the subgraph. So there should only be 3 accelerated subgraphs and 2 non-accelerated subgraphs. + Where TensorRT-RTX cannot serve the trailing FP32 GEMM the non-accelerated count + is one higher, because the Linear falls back too. """ @register_atomic_subgraph(init_args=(), is_core_aten=True) @@ -232,6 +251,11 @@ def forward( ) return x + # Unregister through addCleanup rather than a trailing statement: an assertion + # that fails below would skip the latter, leaking ReLUConv into the next test + # and making it fail too. + self.addCleanup(ATOMIC_SUBGRAPHS.remove, (ReLUConv, (), True)) + class net(nn.Module): def __init__(self): super().__init__() @@ -315,6 +339,9 @@ def forward(self, x): ) == 3 ), "The graph should have 3 accelerated subgraphs" + # The trailing FP32 Linear becomes an extra PyTorch block wherever TensorRT-RTX + # cannot serve an FP32 GEMM; the accelerated blocks above are unaffected. + expected_gpu_subgraphs = 3 if _fp32_gemm_falls_back_to_pytorch() else 2 assert ( len( [ @@ -323,10 +350,8 @@ def forward(self, x): if "_run_on_gpu" in name ] ) - == 2 - ), "The graph should have 2 non-accelerated subgraphs" - - ATOMIC_SUBGRAPHS.remove((ReLUConv, (), True)) + == expected_gpu_subgraphs + ), f"The graph should have {expected_gpu_subgraphs} non-accelerated subgraphs" torch._dynamo.reset()