Make inline tests to use new exporter and fix some issues around it (pytorch#162682)

inline_and_install_module export variant is our long term state so it is better to use the new tracer for this. It also uncovered bunch of minor bugs because with inline_and_install_module, the nn_module_stack generation is changed a bit.

Differential Revision: [D82478648](https://our.internmc.facebook.com/intern/diff/D82478648)
Pull Request resolved: pytorch#162682
Approved by: https://github.com/zhxchen17
ghstack dependencies: pytorch#162557, pytorch#162558, pytorch#162559

Author: tugsbayasgalan

test/export/test_export_with_inline_and_install.py

@@ -3,8 +3,9 @@
 
 import unittest
 
-from torch._dynamo import config
+from torch._dynamo import config as dynamo_config
 from torch._dynamo.testing import make_test_cls_with_patches
+from torch._export import config as export_config
 
 
 try:
@@ -44,8 +45,9 @@ def make_dynamic_cls(cls):
         cls_a,
         cls_prefix,
         "",
-        (config, "install_free_tensors", True),
-        (config, "inline_inbuilt_nn_modules", True),
+        (export_config, "use_new_tracer_experimental", True),
+        (dynamo_config, "install_free_tensors", True),
+        (dynamo_config, "inline_inbuilt_nn_modules", True),
         xfail_prop="_expected_failure_inline_and_install",
     )
 
@@ -71,7 +73,30 @@ def make_dynamic_cls(cls):
 unittest.expectedFailure(
     InlineAndInstallStrictExportTestExport.test_buffer_util_inline_and_install_strict  # noqa: F821
 )
-
+# this is because we can't preserve stacktrace
+unittest.expectedFailure(
+    InlineAndInstallStrictExportTestExport.test_stack_trace_make_fx_inline_and_install_strict  # noqa: F821
+)
+# this is because we marked unlift hooks to be dynamo skip traced
+unittest.expectedFailure(
+    InlineAndInstallStrictExportTestExport.test_custom_tag_metadata_re_export_inline_and_install_strict  # noqa: F821
+)
+unittest.expectedFailure(
+    InlineAndInstallStrictExportTestExport.test_from_node_metadata_export_inline_and_install_strict  # noqa: F821
+)
+unittest.expectedFailure(
+    InlineAndInstallStrictExportTestExport.test_module_inline_and_install_strict  # noqa: F821
+)
+unittest.expectedFailure(
+    InlineAndInstallStrictExportTestExport.test_module_with_dict_container_inp_out_inline_and_install_strict  # noqa: F821
+)
+unittest.expectedFailure(
+    InlineAndInstallStrictExportTestExport.test_retrace_pre_autograd_inline_and_install_strict  # noqa: F821
+)
+# this is because detect leak test has export root
+unittest.expectedFailure(
+    InlineAndInstallStrictExportTestExport.test_detect_leak_strict_inline_and_install_strict  # noqa: F821
+)
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests

torch/_dynamo/functional_export.py

@@ -49,20 +49,63 @@ def post_process_error_msg(
     return constraint_violation_error
 
 
-def clean_nn_module_stack(graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
+def clean_nn_module_stack(
+    graph_module: torch.fx.GraphModule, is_inline_builtin=False
+) -> torch.fx.GraphModule:
+    """
+    Clean up nn_module_stack metadata by removing export_root references.
+
+    Removes the _export_root module references from nn_module_stack metadata
+    in graph nodes, which are artifacts from the export process. Fixes two patterns:
+
+    1. Keys: Removes "__export_root_" and "__modules['_export_root']_" prefixes
+       - Normal case: "L__self____export_root_child" -> "L__self__child"
+       - inline_builtin case: Uses numeric ID strings like "140468831433840"
+
+    2. Values: Removes "._export_root" and "._modules['_export_root']" from child names
+       e.g., "L['self']._export_root.child" -> "L['self'].child"
+       e.g., "L['self']._modules['_export_root'].child" -> "L['self'].child"
+
+    Also removes the root export entry "L__self____export_root" entirely.
+
+    Args:
+        graph_module: The GraphModule to clean up
+        is_inline_builtin: If True, keys are numeric ID strings and self references
+                          (L['self']) are filtered out
+
+    Returns:
+        The cleaned GraphModule (modified in-place)
+    """
     for node in graph_module.graph.nodes:
-        if "nn_module_stack" in node.meta:
-            nn_module_stack = node.meta["nn_module_stack"].copy()
-            first_key = next(iter(nn_module_stack.keys()))
-            if "export_root" in first_key:
-                del nn_module_stack[first_key]
-            nn_module_stack_corrected = {}
-            for k, v in nn_module_stack.items():
-                k_new = "".join(k.split("__export_root"))
-                child_name, child_class = v
-                child_name = child_name.replace("._export_root", "")
-                nn_module_stack_corrected[k_new] = (child_name, child_class)
-            node.meta["nn_module_stack"] = nn_module_stack_corrected
+        if "nn_module_stack" not in node.meta:
+            continue
+
+        nn_module_stack = node.meta["nn_module_stack"].copy()
+
+        if "L__self____export_root" in nn_module_stack:
+            del nn_module_stack["L__self____export_root"]
+
+        # Clean up remaining entries
+        cleaned_stack = {}
+        for key, (child_name, child_class) in nn_module_stack.items():
+            # Clean key by removing export_root patterns
+            clean_key = key.replace("__modules['_export_root']_", "").replace(
+                "__export_root_", ""
+            )
+
+            # Clean child_name by removing export_root patterns
+            clean_name = child_name.replace("._modules['_export_root']", "").replace(
+                "._export_root", ""
+            )
+
+            # Skip self reference for inline builtin case
+            if is_inline_builtin and clean_name == "L['self']":
+                continue
+
+            cleaned_stack[clean_key] = (clean_name, child_class)
+
+        node.meta["nn_module_stack"] = cleaned_stack
+
     return graph_module
 
 
@@ -71,7 +114,11 @@ def clean_export_root(graph_module: torch.fx.GraphModule) -> None:
 
     # Clean parameter names: L__self____export_root_param -> L__self___param
     def clean_name(name) -> str:
-        return name.replace("__export_root_", "_") if "__export_root_" in name else name
+        if "____modules___export_root_" in name:
+            return name.replace("____modules___export_root_", "_")
+        if "__export_root_" in name:
+            return name.replace("__export_root_", "_")
+        return name
 
     # Update get_attr nodes in-place
     for node in graph_module.graph.nodes:
@@ -409,7 +456,9 @@ def inner(*args: Any, **kwargs: Any) -> torch.fx.GraphModule:
             )
             transformed_graph.recompile()
 
-            clean_nn_module_stack(transformed_graph)
+            clean_nn_module_stack(
+                transformed_graph, torch._dynamo.config.inline_inbuilt_nn_modules
+            )
             clean_export_root(transformed_graph)
 
             transformed_graph.meta["module_call_specs"] = module_call_spec