Skip to content

Fullgraph graph capture with dynamo. #159749

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fullgraph graph capture with dynamo. #159749

wants to merge 1 commit into from

Conversation

zhxchen17
Copy link
Contributor

@zhxchen17 zhxchen17 commented Aug 4, 2025

Summary:
Following up on Avik's doc https://docs.google.com/document/d/11RW0Bbkp1QwFbEu8rCNW5d7wUFaEkxbL0uLyqcc2jTk/edit?tab=t.0

We are experimenting with a new API which utilizes torch.compile(fullgraph=True) and intend to use it to replace the old dynamo.export() API.

This PR adds a prototype for the API described in the doc.

Test Plan:
test_misc -- -k test_aot_capture

Rollback Plan:

Differential Revision: D79534608

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela

Copy link

pytorch-bot bot commented Aug 4, 2025

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159749

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 32 Pending

As of commit 321e7c1 with merge base ee9f8ba (image):

NEW FAILURE - The following job has failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79534608

def _backend(gm, example_inputs):
from torch._guards import TracingContext
artifacts.backend_inputs[gm._backend_id] = BackendInput(gm, example_inputs, TracingContext.get().fake_mode)
return gm
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would prefer for us to NOT actually implement it this way (I want to directly call into convert_frame and bypass the eval frame handler entirely)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah I can imagine this may be shadowed by some global context if we dont bypass eval frame things, is this the main issue? how do we get the frame object to feed to convert_frame if we bypass eval_frame, maybe sys.settrace()?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussed offline with @ezyang, it seems essentially we want to make this API only about fullgraph=True which makes sense to me, so I changed the API to fullgraph_capture. eval_frame is not preferred for aesthetic reason and not strictly blocking and can be changed later as impl detail.

facebook-github-bot pushed a commit that referenced this pull request Aug 4, 2025
Summary:

Following up on Avik's doc https://docs.google.com/document/d/11RW0Bbkp1QwFbEu8rCNW5d7wUFaEkxbL0uLyqcc2jTk/edit?tab=t.0

We are experimenting with a new API which utilizes torch.compile(fullgraph=True) and intend to use it to replace the old dynamo.export() API.

This PR adds a prototype for the API described in the doc.

Test Plan:
test_misc -- -k test_aot_capture

Rollback Plan:

Differential Revision: D79534608
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79534608

@zhxchen17 zhxchen17 requested a review from ezyang August 6, 2025 18:32
facebook-github-bot pushed a commit that referenced this pull request Aug 6, 2025
Summary:

Following up on Avik's doc https://docs.google.com/document/d/11RW0Bbkp1QwFbEu8rCNW5d7wUFaEkxbL0uLyqcc2jTk/edit?tab=t.0

We are experimenting with a new API which utilizes torch.compile(fullgraph=True) and intend to use it to replace the old dynamo.export() API.

This PR adds a prototype for the API described in the doc.

Test Plan:
test_misc -- -k test_aot_capture

Rollback Plan:

Differential Revision: D79534608
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79534608

facebook-github-bot pushed a commit that referenced this pull request Aug 12, 2025
Summary:

Following up on Avik's doc https://docs.google.com/document/d/11RW0Bbkp1QwFbEu8rCNW5d7wUFaEkxbL0uLyqcc2jTk/edit?tab=t.0

We are experimenting with a new API which utilizes torch.compile(fullgraph=True) and intend to use it to replace the old dynamo.export() API.

This PR adds a prototype for the API described in the doc.

Test Plan:
test_misc -- -k test_aot_capture

Rollback Plan:

Differential Revision: D79534608
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79534608

@zhxchen17 zhxchen17 changed the title AOT graph capture with dynamo. Fullgraph graph capture with dynamo. Aug 12, 2025
@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 12, 2025
@zhxchen17
Copy link
Contributor Author

To unblock @tugsbayasgalan from export experiments, I will just leave eval_frame there for now since imo our focus is on defining the core data structure for this PR. @ezyang mind reviewing this again?

facebook-github-bot pushed a commit that referenced this pull request Aug 12, 2025
Summary:

Following up on Avik's doc https://docs.google.com/document/d/11RW0Bbkp1QwFbEu8rCNW5d7wUFaEkxbL0uLyqcc2jTk/edit?tab=t.0

We are experimenting with a new API which utilizes torch.compile(fullgraph=True) and intend to use it to replace the old dynamo.export() API.

This PR adds a prototype for the API described in the doc.

Test Plan:
test_misc -- -k test_aot_capture

Rollback Plan:

Differential Revision: D79534608
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79534608

from torch._guards import TracingContext

backend_inputs[gm._backend_id] = BackendInput(
gm, example_inputs, TracingContext.get().fake_mode
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we just do shape env?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tugsbayasgalan I think fake mode can be useful for other consumers. Export as one consumer can easily get shape_env by fake_mode.shape_env

has a standard API to return compilation artifacts.
"""

dynamo_artifacts: DynamoCaptureOutput
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How does dynamo artifacts relate to backend_inputs? Is there common backend_id?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tugsbayasgalan If you read inside DynamoCaptureOutput, you can find common backend_id inside backends field.

Summary:

Following up on Avik's doc https://docs.google.com/document/d/11RW0Bbkp1QwFbEu8rCNW5d7wUFaEkxbL0uLyqcc2jTk/edit?tab=t.0

We are experimenting with a new API which utilizes torch.compile(fullgraph=True) and intend to use it to replace the old dynamo.export() API.

This PR adds a prototype for the API described in the doc.

Test Plan:
test_misc -- -k test_aot_capture

Rollback Plan:

Reviewed By: tugsbayasgalan

Differential Revision: D79534608
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79534608

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants