Revert "Update export api to be the latest version. (#216)" (#218)

This reverts commit 26909fd.

stack-info: PR: #218, branch: xmfan/stack/15

Author: xmfan

autoparallel/api.py

@@ -8,10 +8,10 @@
 import warnings
 from contextlib import ExitStack, contextmanager
 from types import MethodType
-from typing import Optional, Union
+from typing import Any, Optional, Union
 
 import torch
-from torch._dynamo.functional_export import dynamo_graph_capture_for_export
+from torch._dynamo.functional_export import _dynamo_graph_capture_for_export
 from torch._functorch.aot_autograd import (
     aot_compile_joint_with_descriptors,
     aot_export_joint_with_descriptors,
@@ -23,6 +23,7 @@
 from torch._subclasses import FakeTensorMode
 from torch.distributed.fsdp import MixedPrecisionPolicy
 from torch.distributed.tensor import DeviceMesh
+from torch.export._trace import _restore_state_dict
 from torch.export._unlift import _assign_attr
 from torch.export.unflatten import _AttrKind
 from torch.fx.experimental.symbolic_shapes import ShapeEnv
@@ -165,6 +166,21 @@ def enable_local_map_wrapping():
         yield
 
 
+def _export(model: torch.nn.Module, inputs: tuple[Any]) -> torch.nn.Module:
+    """
+    Thin wrapper around graph capture output that restores the
+    original calling convention and attribute fqn. TODO:
+    1) Use bytecode for calling convention instead of pytree for more
+       seamless UX.
+    2) Attach guards
+    3) Be more careful about tensor constants names.
+    """
+    with torch._dynamo.config.patch(install_free_tensors=True):
+        gm = _dynamo_graph_capture_for_export(model)(*inputs)
+        _restore_state_dict(model, gm)
+        return gm
+
+
 class AutoParallel:
     """
     Args:
@@ -289,7 +305,7 @@ def build_model_graph(self):
         with set_dtype_cast(
             True
         ), enable_local_map_wrapping(), torch._dynamo.utils._disable_saved_tensors_hooks_during_tracing():
-            torch_ir_with_fqn = dynamo_graph_capture_for_export(self.model)(*inputs)
+            torch_ir_with_fqn = _export(self.model, inputs)
             # TODO Cna't use fake mode here because it clashes with the user level
             # fake mode. Ideally dynamo should reuse the user level fake mode.
             self.joint_with_descriptors = aot_export_joint_with_descriptors(