Refactor make_fx to better support hop subgraph tracing #125267
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125267
Note: Links to docs will display an error until the docs builds have been completed. ❌ 8 New Failures, 2 Unrelated FailuresAs of commit 2f6e5d7 with merge base b522e65 ( NEW FAILURES - The following jobs have failed:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
ydwu4
added a commit
that referenced
this pull request
Apr 30, 2024
ghstack-source-id: 31c1a4b4f8edad489e5d8e17862d194e1a05e951 Pull Request resolved: #125267
[ghstack-poisoned]
ydwu4
added a commit
that referenced
this pull request
Apr 30, 2024
ghstack-source-id: 29cdd64041ecdd97056b75d49ccf625dd45d30ea Pull Request resolved: #125267
[ghstack-poisoned]
ydwu4
added a commit
that referenced
this pull request
Apr 30, 2024
ghstack-source-id: d6244d2808300faa00e4c74d9795c829d473d6b8 Pull Request resolved: #125267
[ghstack-poisoned]
ydwu4
added a commit
that referenced
this pull request
May 1, 2024
ghstack-source-id: d96625ccfd8cf522b0ca10319dc83e70a31f81fc Pull Request resolved: #125267
[ghstack-poisoned]
ydwu4
added a commit
that referenced
this pull request
May 1, 2024
ghstack-source-id: f6d4b1d1304bb4d314793fc5bfacc66e9e8abb4d Pull Request resolved: #125267
This is to replace #122972. We extract the states of make_fx out into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs. The entry point for tracings subgraphas is 'reenter_make_fx(fn)'. Test Plan: Existing tests and remove a related task on torch function. [ghstack-poisoned]
ydwu4
commented
May 2, 2024
Code movement + minor rewrites. We extract the states of make_fx out into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs. The entry point for tracings subgraphas is 'reenter_make_fx(fn)'. Test Plan: Existing tests and remove a related task on torch function. [ghstack-poisoned]
Code movement + minor rewrites. We extract the states of make_fx out into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs. The entry point for tracings subgraphas is 'reenter_make_fx(fn)'. Test Plan: Existing tests and remove a related task on torch function. [ghstack-poisoned]
ydwu4
commented
May 7, 2024
|
|
||
| # All context managers and their states should be initialized before tracing based on the inputs | ||
| # and configurations. | ||
| # After tracing, thier states should be cleaned except for shape_env. |
There was a problem hiding this comment.
shape_env is leaked out because of t.shape_env = self.fake_tensor_mode.shape_env
Chillee
reviewed
May 7, 2024
|
yeah, this pr is supposed to be a code movement only change. Existing tests should suffice. Any suspicious places you want to test? |
Chillee
approved these changes
May 7, 2024
|
Ok looks fine to me. |
ezyang
reviewed
May 8, 2024
|
|
||
| # All context managers and their states should be initialized before tracing based on the inputs | ||
| # and configurations. | ||
| # After tracing, thier states should be cleaned except for shape_env. |
Code movement + minor rewrites. We extract the states of make_fx out and encapsulate them into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs, the actual logic for tracing subgraph is in the next diff. Test Plan: Existing tests. [ghstack-poisoned]
Code movement + minor rewrites. We extract the states of make_fx out and encapsulate them into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs, the actual logic for tracing subgraph is in the next diff. Test Plan: Existing tests. [ghstack-poisoned]
Code movement + minor rewrites. We extract the states of make_fx out and encapsulate them into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs, the actual logic for tracing subgraph is in the next diff. Test Plan: Existing tests. [ghstack-poisoned]
Code movement + minor rewrites. We extract the states of make_fx out and encapsulate them into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs, the actual logic for tracing subgraph is in the next diff. Test Plan: Existing tests. [ghstack-poisoned]
pytorchmergebot
pushed a commit
that referenced
this pull request
May 15, 2024
Adds trace_subgraph to _MakefxTracer, the motivation is in #122972. Also migrate all existing usage of reenter_make_fx to the new sub-tracer. Previously, the torch function mode for creating torch_fn metadata won't be re-enetered when we're in ProxyTensorMode (since it's inside of __torch_function__). This PR reconstruct the torch function mode based on parent tracer's config and reentered the torch function mode so the metadata is shown in the graph. **Test Plan:** Existing tests. We have a bunch of make_fx tests for cond, map and while_loop. Also remove expected failure for torch_fn since reenter_make_fx is able to re-construct torch function modes. Also fixes #124643 Pull Request resolved: #125363 Approved by: https://github.com/Chillee ghstack dependencies: #125267
ZelboK
pushed a commit
to ZelboK/pytorch
that referenced
this pull request
May 19, 2024
Code movement + minor rewrites. We extract the states of make_fx out and encapsulate them into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs, the actual logic for tracing subgraph is in the next diff. Test Plan: Existing tests. Pull Request resolved: pytorch#125267 Approved by: https://github.com/Chillee
ZelboK
pushed a commit
to ZelboK/pytorch
that referenced
this pull request
May 19, 2024
Adds trace_subgraph to _MakefxTracer, the motivation is in pytorch#122972. Also migrate all existing usage of reenter_make_fx to the new sub-tracer. Previously, the torch function mode for creating torch_fn metadata won't be re-enetered when we're in ProxyTensorMode (since it's inside of __torch_function__). This PR reconstruct the torch function mode based on parent tracer's config and reentered the torch function mode so the metadata is shown in the graph. **Test Plan:** Existing tests. We have a bunch of make_fx tests for cond, map and while_loop. Also remove expected failure for torch_fn since reenter_make_fx is able to re-construct torch function modes. Also fixes pytorch#124643 Pull Request resolved: pytorch#125363 Approved by: https://github.com/Chillee ghstack dependencies: pytorch#125267
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
Code movement + minor rewrites. We extract the states of make_fx out and encapsulate them into a _MakefxTracer class. This allows us to create a new make_fx_tracer when tracing subgraphs, the actual logic for tracing subgraph is in the next diff.
Test Plan:
Existing tests.