Skip to content
This repository was archived by the owner on Aug 5, 2025. It is now read-only.

Trace local graph and then replace nodes with DT's dispatch graph#460

Merged
mrshenli merged 22 commits intomainfrom
gh/mrshenli/3/head
Oct 3, 2022
Merged

Trace local graph and then replace nodes with DT's dispatch graph#460
mrshenli merged 22 commits intomainfrom
gh/mrshenli/3/head

Conversation

@mrshenli
Copy link
Copy Markdown
Contributor

@mrshenli mrshenli commented Sep 15, 2022

mrshenli added a commit that referenced this pull request Sep 15, 2022
@mrshenli
Copy link
Copy Markdown
Contributor Author

land after pytorch/pytorch#84980 shows up in PT nightly

args = tree_map(lambda n: node_to_obj[n], node.args)
kwargs = tree_map(lambda n: node_to_obj[n], node.kwargs)

out = operator_dispatch(
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

So this is doing "concrete" tracing, where we're tracing based on a concrete DTensor. Is this just for testing purposes or do we want to actually rely on concrete tracing for the final solution?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This is not yet tracing. It is getting the intermediate DTensor outputs, so that we know what DTensor specs to pass to the next operator. Tracing is on line 318, which only applies to local tensors.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Got it, yeah what I mean is that this corresponds to doing concrete tracing -- where you need the concrete tensors in order to propagate sharding.

mrshenli added a commit that referenced this pull request Sep 19, 2022
mrshenli added a commit that referenced this pull request Sep 19, 2022
mrshenli added a commit that referenced this pull request Sep 19, 2022
mrshenli added a commit that referenced this pull request Sep 19, 2022
mrshenli added a commit that referenced this pull request Sep 19, 2022
mrshenli added a commit that referenced this pull request Sep 19, 2022
@mrshenli mrshenli mentioned this pull request Oct 3, 2022
mrshenli added a commit that referenced this pull request Oct 3, 2022
@mrshenli mrshenli requested a review from wanchaol October 3, 2022 15:47
mrshenli added a commit that referenced this pull request Oct 3, 2022
@mrshenli mrshenli changed the base branch from gh/mrshenli/3/base to main October 3, 2022 16:39
@mrshenli mrshenli changed the base branch from main to gh/mrshenli/3/base October 3, 2022 16:39
Copy link
Copy Markdown
Contributor

@aazzolini aazzolini left a comment

Choose a reason for hiding this comment

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

Looks good, thanks!

@mrshenli mrshenli changed the base branch from gh/mrshenli/3/base to main October 3, 2022 17:29
@mrshenli mrshenli merged commit c72e4ae into main Oct 3, 2022
@facebook-github-bot facebook-github-bot deleted the gh/mrshenli/3/head branch November 3, 2022 14:19
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants