mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
757975ad50
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 |
||
|---|---|---|
| .. | ||
| backends | ||
| polyfills | ||
| repro | ||
| variables | ||
| __init__.py | ||
| _trace_wrapped_higher_order_op.py | ||
| aot_compile_types.py | ||
| aot_compile.py | ||
| bytecode_analysis.py | ||
| bytecode_transformation.py | ||
| cache_size.py | ||
| callback.py | ||
| code_context.py | ||
| codegen.py | ||
| compiled_autograd.py | ||
| comptime.py | ||
| config.py | ||
| convert_frame.py | ||
| create_parameter_op.py | ||
| current_scope_id.py | ||
| debug_utils.py | ||
| decorators.py | ||
| device_interface.py | ||
| distributed.py | ||
| eval_frame.py | ||
| exc.py | ||
| external_utils.py | ||
| funcname_cache.py | ||
| functional_export.py | ||
| graph_break_hints.py | ||
| graph_break_registry.json | ||
| graph_bytecode_inputs.py | ||
| graph_deduplication.py | ||
| graph_region_tracker.py | ||
| graph_utils.py | ||
| guards.py | ||
| hooks.py | ||
| logging.py | ||
| metrics_context.py | ||
| mutation_guard.py | ||
| output_graph.py | ||
| package.py | ||
| pgo.py | ||
| precompile_context.py | ||
| profiler.py | ||
| replay_record.py | ||
| resume_execution.py | ||
| side_effects.py | ||
| source.py | ||
| symbolic_convert.py | ||
| tensor_version_op.py | ||
| test_case.py | ||
| test_dont_skip_tracing_functions.py | ||
| test_minifier_common.py | ||
| testing.py | ||
| trace_rules.py | ||
| types.py | ||
| utils.py | ||