pytorch/torch
zhxchen17 757975ad50 [export] Unified graph capture with fullgraph_capture. (#165562)
Summary:
_dynamo_graph_capture_for_export in the current form has the compability issue
with the main torch.compile() path despite we reuse fullgraph_capture as the
bytecode tracer. The reason is that we flip on many export specific flags
and even trace with a wrapped function which will cause divergence with
torch.compile() again.

This PR instead creates a new implementation of dynamo_graph_capture_for_export
which 100% relies on fullgraph capture and post-processing on CaptureOutput so
that we can avoid the inversion of phases in PT2 compiler stack.

This also benefits precompile workflow since we want to have a feature that
only accepts pytree inputs and ship portable python wrappers in package. In
other words, I think the code here is sharable between export and precompile
for exporting portable graph.

Test Plan:
===================================================================== test session starts =====================================================================
platform linux -- Python 3.12.11, pytest-7.3.2, pluggy-1.6.0
rootdir: /data/users/zhxchen17/pytorch
configfile: pytest.ini
plugins: xdoctest-1.1.0, hypothesis-5.35.1, xdist-3.3.1, subtests-0.13.1, rerunfailures-14.0, flakefinder-1.1.0, cpp-2.3.0, anyio-4.10.0
collected 9 items
Running 9 items in this shard

test/distributed/tensor/test_dtensor_export.py ........x                                                                                                [100%]

================================================================ 8 passed, 1 xfailed in 11.42s ================================================================

Reviewers:

Subscribers:

Tasks:

Tags:

Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/165562
Approved by: https://github.com/tugsbayasgalan
2025-10-22 20:44:55 +00:00
..
_awaits
_C Warn if AccumulateGrad stream does not match producer node stream (#165065) 2025-10-22 17:33:27 +00:00
_C_flatbuffer
_custom_op [2/N] Fix ruff warnings (#164460) 2025-10-04 03:40:32 +00:00
_decomp Back out "Do not decompose in functionalization/proxy tensor if autograd wouldn't have decomposed (#164939)" (#165910) 2025-10-21 16:36:38 +00:00
_dispatch Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
_dynamo [export] Unified graph capture with fullgraph_capture. (#165562) 2025-10-22 20:44:55 +00:00
_export [torch.export] Rmoving unused constants - add support for corner case (#165205) 2025-10-14 20:26:28 +00:00
_functorch Allow GraphPickler to pickle graph modules containing AOTCompiled subgraphs (#165844) 2025-10-22 17:03:49 +00:00
_higher_order_ops [scan x vmap] support scan in vmap (#165580) 2025-10-22 09:46:00 +00:00
_inductor inductor_provenance: Correctly handle null provenance (#166019) 2025-10-22 18:21:57 +00:00
_lazy Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
_library Enable all flake8-logging-format rules (#164655) 2025-10-19 00:59:28 +00:00
_logging [annotation] add logging for debugging annotation (#165797) 2025-10-20 21:27:38 +00:00
_numpy Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
_prims Fix self assignment (#165816) 2025-10-18 18:51:52 +00:00
_prims_common [2/N] Use "is" in python type comparison (#165142) 2025-10-10 15:36:44 +00:00
_refs Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
_strobelight Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
_subclasses Enable PLC0414 on ruff (#165828) 2025-10-22 04:56:52 +00:00
_vendor
accelerator
amp Revert "[AMP][Refactor] Autocast dtype handling to simplify device-specific c… (#165221)" 2025-10-22 00:26:57 +00:00
ao Revert "[Code Clean] Clean asserts in torch/ao/quantization (root, quantizer, backend_config) (#165433)" 2025-10-21 22:10:19 +00:00
autograd Warn if AccumulateGrad stream does not match producer node stream (#165065) 2025-10-22 17:33:27 +00:00
backends Add type suppressions to _inductor/runtime (#165918) 2025-10-21 02:54:22 +00:00
compiler Megacache integration (#163533) 2025-10-15 22:49:15 +00:00
contrib
cpu Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
csrc Warn if AccumulateGrad stream does not match producer node stream (#165065) 2025-10-22 17:33:27 +00:00
cuda Add type suppressions to _inductor/runtime (#165918) 2025-10-21 02:54:22 +00:00
distributed Revert "shrink_group implementation to expose ncclCommShrink API (#164518)" 2025-10-21 20:24:14 +00:00
distributions [1/N] Use "is" in python type comparison (#165037) 2025-10-10 12:36:50 +00:00
export Back out "Do not decompose in functionalization/proxy tensor if autograd wouldn't have decomposed (#164939)" (#165910) 2025-10-21 16:36:38 +00:00
fft
func
futures [5/N] Apply ruff UP035 rule (#164423) 2025-10-02 07:31:11 +00:00
fx Allow GraphPickler to pickle graph modules containing AOTCompiled subgraphs (#165844) 2025-10-22 17:03:49 +00:00
headeronly Move toUnderlying to headeronly (#165694) 2025-10-22 05:31:16 +00:00
jit Fix missing brackets (#165138) 2025-10-10 17:23:31 +00:00
legacy
lib [2/N] Mark unused parameters in C++ code (#165121) 2025-10-15 03:04:39 +00:00
linalg Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
masked [1/N] Use "is" in python type comparison (#165037) 2025-10-10 12:36:50 +00:00
monitor
mps
mtia Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
multiprocessing Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
nativert [2/N] Mark unused parameters in C++ code (#165121) 2025-10-15 03:04:39 +00:00
nested Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
nn Enable PLW0127 in ruff (#165851) 2025-10-21 03:30:57 +00:00
numa Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
onnx Enable all flake8-logging-format rules (#164655) 2025-10-19 00:59:28 +00:00
optim [RFC] Add pyrefly to lintrunner (#165179) 2025-10-16 20:07:09 +00:00
package [1/N] Use "is" in python type comparison (#165037) 2025-10-10 12:36:50 +00:00
profiler Enable PLC0414 on ruff (#165828) 2025-10-22 04:56:52 +00:00
quantization [RFC] Add pyrefly to lintrunner (#165179) 2025-10-16 20:07:09 +00:00
signal Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
sparse [RFC] Add pyrefly to lintrunner (#165179) 2025-10-16 20:07:09 +00:00
special
testing [AMD][gfx1100] test_decompose_mem_bound_mm.py tolerance increase (#165625) 2025-10-22 01:38:48 +00:00
utils [DebugMode] refactor logs into _DebugCalls (#165376) 2025-10-22 19:01:56 +00:00
xpu Add a new API torch.xpu.is_tf32_supported for Intel GPU (#163141) 2025-10-12 12:11:57 +00:00
__config__.py
__future__.py
__init__.py [BE][Ez]: Update torch.is_tensor documentation (#165841) 2025-10-19 09:24:11 +00:00
_appdirs.py
_classes.py
_compile.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_custom_ops.py
_environment.py
_guards.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_jit_internal.py [2/N] More ruff SIM fixes (#165031) 2025-10-14 14:22:54 +00:00
_linalg_utils.py
_lobpcg.py Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
_lowrank.py
_meta_registrations.py Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
_namedtensor_internals.py
_ops.py [2/N] More ruff SIM fixes (#165031) 2025-10-14 14:22:54 +00:00
_python_dispatcher.py
_size_docs.py
_sources.py
_storage_docs.py
_streambase.py
_tensor_docs.py [reland] Allow setting grad_dtype on leaf tensors (#164751) 2025-10-08 20:23:13 +00:00
_tensor_str.py Enable all PIE rules on ruff (#165814) 2025-10-18 07:36:18 +00:00
_tensor.py Pyrefly suppressions 6/n (#164877) 2025-10-08 02:30:57 +00:00
_thread_safe_fork.py
_torch_docs.py Clarrifying input output angle unit in the docs for trigonometric fun… (#161248) 2025-10-18 11:53:48 +00:00
_utils_internal.py Revert "Call internal log_compilation_event if it exists (#164855)" 2025-10-09 22:38:45 +00:00
_utils.py Enable ruff rule E721 (#165162) 2025-10-13 01:48:55 +00:00
_VF.py
_vmap_internals.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
_weights_only_unpickler.py [4/N] Apply ruff UP035 rule to python code (#164206) 2025-10-01 19:05:53 +00:00
CMakeLists.txt
custom_class_detail.h Mark unused parameters in C++ code (#164912) 2025-10-09 06:23:25 +00:00
custom_class.h Mark unused parameters in C++ code (#164912) 2025-10-09 06:23:25 +00:00
extension.h
functional.py [2/N] Fix ruff warnings (#164460) 2025-10-04 03:40:32 +00:00
header_only_apis.txt Move toUnderlying to headeronly (#165694) 2025-10-22 05:31:16 +00:00
hub.py Enable PLC1802 on ruff (#165813) 2025-10-18 05:44:14 +00:00
library.h Mark unused parameters in C++ code (#164912) 2025-10-09 06:23:25 +00:00
library.py Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
overrides.py Add scaled_grouped_mm_v2 and python API (#165154) 2025-10-15 17:47:23 +00:00
py.typed
quasirandom.py
random.py Revert "Add device argument to torch.random.get_rng_state (#163034)" 2025-10-04 15:25:45 +00:00
return_types.py
script.h
serialization.py Add initial suppressions for pyrefly (#164177) 2025-10-02 20:57:41 +00:00
storage.py [2/N] Use "is" in python type comparison (#165142) 2025-10-10 15:36:44 +00:00
torch_version.py
types.py Enable PLC0414 on ruff (#165828) 2025-10-22 04:56:52 +00:00
version.py.tpl