fix: Refactor tensor freezing in Dynamo

Author: gs-olive

py/torch_tensorrt/dynamo/backend/aot_module.py

@@ -1,22 +1,19 @@
 from __future__ import annotations
 
 import logging
-from functools import partial
+import unittest
 from typing import Any, Callable, Sequence
 
 import torch
 import torch._dynamo as td
-from torch._functorch.aot_autograd import make_boxed_compiler
-from torch._guards import TracingContext
+from torch._dynamo.utils import detect_fake_mode
+from torch._functorch.aot_autograd import aot_export_joint_simple
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.compile import compile_module
 from torch_tensorrt.dynamo.lowering._decompositions import get_decompositions
-from torch_tensorrt.dynamo.lowering._freeze_aot_graph import freeze_autograd_gm
 from torch_tensorrt.dynamo.lowering._pre_aot_lowering import pre_aot_substitutions
 from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
 
-from .aot_module import aot_module
-
 logger = logging.getLogger(__name__)
 
 
@@ -37,8 +34,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)
 
 
@@ -48,21 +43,26 @@ def aot_torch_tensorrt_aten_backend(
 ) -> torch.nn.Module:
     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(settings.enable_experimental_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(
+                settings.enable_experimental_decompositions
+            ),
+        )
+
+        return _pretraced_backend(graph_module, sample_inputs, settings)
 
 
 def _pretraced_backend(
@@ -82,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