fix: Refactor tensor freezing in Dynamo

Author: gs-olive

py/torch_tensorrt/dynamo/backend/aot_module.py

@@ -1,7 +1,8 @@
 import logging
 from typing import Sequence
 import torch
-from functools import partial
+from torch._dynamo.utils import detect_fake_mode
+import unittest
 import torch._dynamo as td
 from torch._guards import TracingContext
 
@@ -16,15 +17,13 @@
     partition,
     get_submod_inputs,
 )
-from torch_tensorrt.dynamo.lowering._freeze_aot_graph import freeze_autograd_gm
 from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
 from torch_tensorrt.dynamo.conversion import (
     convert_module,
     repair_long_or_double_inputs,
 )
 
-from torch._functorch.aot_autograd import make_boxed_compiler
-from .aot_module import aot_module
+from torch._functorch.aot_autograd import aot_export_joint_simple
 
 
 logger = logging.getLogger(__name__)
@@ -36,8 +35,6 @@ def torch_tensorrt_backend(
 ):
     DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend
 
-    TracingContext.get().fake_mode.allow_non_fake_inputs = True
-
     return DEFAULT_BACKEND(gm, sample_inputs, **kwargs)
 
 
@@ -47,21 +44,25 @@ def aot_torch_tensorrt_aten_backend(
 ):
     settings = parse_dynamo_kwargs(kwargs)
 
-    custom_backend = partial(
-        _pretraced_backend,
-        settings=settings,
-    )
-
     # Perform Pre-AOT Lowering for Module-Level Replacement
     gm = pre_aot_substitutions(gm)
 
-    # Invoke AOTAutograd to translate operators to aten
-    return aot_module(
-        gm,
-        sample_inputs,
-        fw_compiler=make_boxed_compiler(custom_backend),
-        decompositions=get_decompositions(),
-    )
+    fake_mode = detect_fake_mode(sample_inputs)
+
+    # Place backend tracing within FakeTensor context allowing nonfake Tensors
+    with unittest.mock.patch.object(
+        fake_mode, "allow_non_fake_inputs", True
+    ), fake_mode:
+
+        # Invoke AOTAutograd to translate operators to aten
+        graph_module = aot_export_joint_simple(
+            gm,
+            sample_inputs,
+            trace_joint=False,
+            decompositions=get_decompositions(),
+        )
+
+        return _pretraced_backend(graph_module, sample_inputs, settings)
 
 
 def _pretraced_backend(
@@ -81,16 +82,9 @@ def _pretraced_backend(
     try:
         logger.debug("Post-AOT Autograd graph:\n" + str(gm.graph))
 
-        frozen_gm, unfrozen_indices = freeze_autograd_gm(gm, sample_inputs)
-        nonfrozen_inputs = [sample_inputs[idx] for idx in unfrozen_indices]
-
-        frozen_gm.graph.eliminate_dead_code()
-        frozen_gm.graph.lint()
-        frozen_gm.recompile()
-
         trt_compiled = _compile_module(
-            frozen_gm,
-            nonfrozen_inputs,
+            gm,
+            sample_inputs,
             settings=settings,
         )
         return trt_compiled

py/torch_tensorrt/dynamo/backend/backends.py

@@ -8,4 +8,3 @@
 from ._partition import partition, get_submod_inputs, DEFAULT_SINGLE_NODE_PARTITIONS
 from .substitutions import *
 from ._fusers import *
-from ._freeze_aot_graph import *