Skip to content
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

Support mark_dynamic #7812

Merged
merged 2 commits into from
Aug 7, 2024
Merged

Support mark_dynamic #7812

merged 2 commits into from
Aug 7, 2024

Conversation

JackCaoG
Copy link
Collaborator

@JackCaoG JackCaoG commented Aug 6, 2024

This pr is to support the use case of

compiled = torch.compile(mode)
torch._dynamo.makr_dynamic(input, 0)
res = compiled(input)

where the 0th dimension of input can be dynamic. I fixed a couple things in this pr

  1. Currently we remove the symints from the args but filter on the tensor type, This is actually problematic because the FX graph actually expects symints to be part of the input. This will likely incur the edge cases where the order/number of args are different from what fx expects

  2. We currently cache each graph by all input shapes. I took a profile and this can be expensive(~1ms for 16 layer decoder). Instead pytorch actually pass us the dynamic dimension as input so we can cache on that instead. If there is no dynamic dimension we will cache on an empty tuple which is cheap.

  3. with dynamic=False and mark_dynami we can no longer easily tell if current graph will be dynamic or not. Since with the 2 above, the caching will be very cheap, so I do the graph_var lookup for every graph instead of dynamic graph only

  4. We can not run the partitioner when there are symints as inputs because partitioner will put the symint and its return into a separate graph which mess up the dynamic dimension look up. For now I just skip partitioner when there is an symint inputs.

@JackCaoG JackCaoG requested a review from wonjoolee95 August 7, 2024 02:30
@JackCaoG JackCaoG marked this pull request as ready for review August 7, 2024 02:31
@JackCaoG
Copy link
Collaborator Author

JackCaoG commented Aug 7, 2024

@ysiraichi FYI, I made the assumption that input to dynamo is either an int or a tensor. Let me know if this is not a correct assumption.

Copy link
Collaborator

@wonjoolee95 wonjoolee95 left a comment

Choose a reason for hiding this comment

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

🚀

@JackCaoG JackCaoG merged commit 08d9595 into master Aug 7, 2024
23 checks passed
yitongh pushed a commit to AlibabaPAI/xla that referenced this pull request Oct 11, 2024
yitongh pushed a commit to AlibabaPAI/xla that referenced this pull request Dec 11, 2024
yitongh pushed a commit to AlibabaPAI/xla that referenced this pull request Dec 11, 2024
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.

2 participants