SIGN IN SIGN UP

Fix AOT export GraphModule serialization (#186418)

GraphModules produced by `aot_export_module` and non-strict export record `_ModuleStackTracer` as their graph tracer. FX pickle loading reconstructs a graph by instantiating a temporary `KeepModules` subclass of that tracer. This fails because `_ModuleStackTracer` requires a scope root, and replaying its normal export behavior would also inline `GraphModule` children and mutate export-only global state.

Add a private `Tracer` factory hook for `GraphModule` deserialization. The default preserves existing custom-tracer behavior, while `_ModuleStackTracer` selects a dedicated replay tracer initialized with the code-only module. That tracer keeps submodules as leaves, preserves distinct paths for shared module aliases, and uses the base tracing path so export-only cleanup and fake-tensor tracking do not run during `torch.load`.

The replay tracer also restores standard FX math autowrapping. `PythonKeyTracer` normally disables it for export tracing, but deserialization must replay generated code containing `math.floor`, `math.ceil`, or `math.trunc` without invoking those functions on `Proxy` objects.

Making `_ModuleStackTracer`'s constructor optional was considered, but that would couple deserialization to its export-only behavior and require a hidden mode flag. The dedicated tracer keeps the two execution modes explicit.

Add regression coverage for `aot_export_module` and non-strict export GraphModules, including graph structure, shared `GraphModule` aliases, retained tracer metadata, preservation of the export fake-tensor map, and dynamic math operations.

Fixes #128554

Generated by my agent

Test Plan:
- `python test/functorch/test_aotdispatch.py TestAOTExport.test_aot_export_module_graph_module_serialization TestAOTExport.test_module_stack_tracer_deserialization_preserves_graph_and_state TestAOTExport.test_module_stack_tracer_deserialization_autowraps_math`
- `python test/test_fx.py TestFX.test_pickle_graphmodule TestFX.test_pickle_custom_import TestFXAPIBackwardCompatibility.test_preserve_unused_attr_after_unpickle`
- `python test/test_fx.py TestFXAPIBackwardCompatibility.test_function_back_compat`
- `python test/test_fx_experimental.py TestFXExperimental.test_meta_tracer`
- `lintrunner -a`
- `git diff --check`

Pull Request resolved: https://github.com/pytorch/pytorch/pull/186418
Approved by: https://github.com/bobrenjc93
J
Jason Ansel committed
c51a5faaa1d700f260ad9dd341b6c365cb708daf
Parent: 988dc63
Committed by PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> on 8/25/2026, 11:17:02 PM