fix: Add support for passing through build issues

- Add support for `pass_through_build_failures` keyword arg
- Add resnet18 testing to validate feature
- Add minor typo fixes

Author: gs-olive

py/torch_tensorrt/dynamo/backend/__init__.py

@@ -16,6 +16,7 @@
     DEBUG,
     MAX_WORKSPACE_SIZE,
     MIN_BLOCK_SIZE,
+    PASS_THROUGH_BUILD_FAILURES,
 )
 
 
@@ -53,7 +54,8 @@ def compile(
     logger.warn(
         "The Dynamo backend is an experimental feature, for which only the "
         + "following arguments are supported: "
-        + "{enabled_precisions, debug, workspace_size, min_block_size, torch_executed_ops}"
+        + "{enabled_precisions, debug, workspace_size, min_block_size, "
+        + "torch_executed_ops, pass_through_build_failures}"
     )
 
     if not isinstance(inputs, collections.abc.Sequence):
@@ -107,6 +109,7 @@ def create_backend(
     workspace_size: int = MAX_WORKSPACE_SIZE,
     min_block_size: int = MIN_BLOCK_SIZE,
     torch_executed_ops: Sequence[str] = set(),
+    pass_through_build_failures: bool = PASS_THROUGH_BUILD_FAILURES,
     **kwargs,
 ):
     """Create torch.compile backend given specified arguments
@@ -125,6 +128,7 @@ def create_backend(
         workspace_size=workspace_size,
         min_block_size=min_block_size,
         torch_executed_ops=torch_executed_ops,
+        pass_through_build_failures=pass_through_build_failures,
     )
 
     return partial(

py/torch_tensorrt/dynamo/backend/_defaults.py

@@ -5,3 +5,4 @@
 DEBUG = False
 MAX_WORKSPACE_SIZE = 20 << 30
 MIN_BLOCK_SIZE = 5
+PASS_THROUGH_BUILD_FAILURES = False

py/torch_tensorrt/dynamo/backend/_settings.py

@@ -7,6 +7,7 @@
     DEBUG,
     MAX_WORKSPACE_SIZE,
     MIN_BLOCK_SIZE,
+    PASS_THROUGH_BUILD_FAILURES,
 )
 
 
@@ -17,3 +18,4 @@ class CompilationSettings:
     workspace_size: int = MAX_WORKSPACE_SIZE
     min_block_size: int = MIN_BLOCK_SIZE
     torch_executed_ops: Sequence[str] = field(default_factory=set)
+    pass_through_build_failures: bool = PASS_THROUGH_BUILD_FAILURES

py/torch_tensorrt/dynamo/backend/backends.py

@@ -1,7 +1,6 @@
 import logging
 from typing import Sequence
 import torch
-import traceback
 from functools import partial
 import torch._dynamo as td
 
@@ -23,14 +22,10 @@
 from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
 
 
-<<<<<<< HEAD:py/torch_tensorrt/dynamo/backend/backends.py
-@td.register_backend(name="torch_tensorrt")
-=======
 logger = logging.getLogger(__name__)
 
 
-@td.register_backend(name="tensorrt")
->>>>>>> 7e0f4405... feat: Prototype Module-Acceleration in Dynamo:py/torch_tensorrt/dynamo/torch_compile/backends.py
+@td.register_backend(name="torch_tensorrt")
 @fake_tensor_unsupported
 def torch_tensorrt_backend(
     gm: torch.fx.GraphModule,
@@ -85,25 +80,31 @@ def _pretraced_backend(
         Compiled FX GraphModule
     """
     try:
-<<<<<<< HEAD:py/torch_tensorrt/dynamo/backend/backends.py
-        trt_compiled = _compile_module(
-=======
         logger.debug("Post-AOT Autograd graph:\n" + str(gm.graph))
 
-        trt_compiled = compile_module(
->>>>>>> 7e0f4405... feat: Prototype Module-Acceleration in Dynamo:py/torch_tensorrt/dynamo/torch_compile/backends.py
+        trt_compiled = _compile_module(
             gm,
             sample_inputs,
             settings=settings,
         )
         return trt_compiled
     except:
-        traceback.print_exc()
-        print(
+        logger.error(
             "FX2TRT conversion failed on the subgraph. See trace above. "
-            + "Returning GraphModule forward instead."
+            + "Returning GraphModule forward instead.",
+            exc_info=True,
         )
-        return gm.forward
+
+        if not settings.pass_through_build_failures:
+            return gm.forward
+        else:
+            raise AssertionError(
+                "Halting compilation on build failure since "
+                + "pass_through_build_failures was specified as True. "
+                + "To return the default Torch implementation and avoid "
+                + "halting compilation on engine build failures, "
+                + "specify pass_through_build_failures=False."
+            )
 
 
 def _compile_module(

py/torch_tensorrt/dynamo/backend/test/test_resnet.py

@@ -0,0 +1,31 @@
+from torch.testing._internal.common_utils import run_tests, TestCase
+import torch
+import torch_tensorrt
+import torchvision.models as models
+from torch_tensorrt.dynamo.common_utils.test_utils import (
+    COSINE_THRESHOLD,
+    cosine_similarity,
+)
+
+
+class TestResNet18(TestCase):
+    def test_resnet18(ir):
+        model = models.resnet18(pretrained=True).eval().to("cuda")
+        input_ = torch.randn((1, 3, 224, 224)).to("cuda")
+
+        compile_spec = {
+            "inputs": [input_],
+            "enabled_precisions": {torch.float},
+            "pass_through_build_failures": True,
+        }
+
+        trt_mod = torch_tensorrt.dynamo.torch_compile(model, **compile_spec)
+        cos_sim = cosine_similarity(model(input_), trt_mod(input_))
+        assert (
+            cos_sim > COSINE_THRESHOLD,
+            f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
+
+
+if __name__ == "__main__":
+    run_tests()

py/torch_tensorrt/dynamo/common_utils/__init__.py

@@ -7,7 +7,10 @@
 
 from transformers import BertModel
 
-from utils import COSINE_THRESHOLD, cosine_similarity
+from torch_tensorrt.dynamo.common_utils.test_utils import (
+    COSINE_THRESHOLD,
+    cosine_similarity,
+)
 
 
 @pytest.mark.unit
@@ -30,7 +33,7 @@ def test_resnet18(ir):
     cos_sim = cosine_similarity(model(input), trt_mod(input))
     assert (
         cos_sim > COSINE_THRESHOLD,
-        f"Resnet50 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
 
     # Clean up model env
@@ -163,7 +166,7 @@ def test_resnet18_half(ir):
     cos_sim = cosine_similarity(model(input), trt_mod(input))
     assert (
         cos_sim > COSINE_THRESHOLD,
-        f"Resnet50 Half TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        f"Resnet18 Half TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
 
     # Clean up model env