Skip to content

[Traceable FSDP2] Add partial-graph (graph-break) unit tests #131747

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.

Sign up for GitHub

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

Jump to bottom
Closed
wants to merge 5 commits into from
Closed

[Traceable FSDP2] Add partial-graph (graph-break) unit tests #131747

wants to merge 5 commits into from

Conversation

yf225
Copy link
Contributor

@yf225 yf225 commented Jul 25, 2024 •
edited
Loading

@yf225
Update
31392c7 
[ghstack-poisoned]
@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jul 25, 2024 •
edited
Loading

🔗 Helpful Links

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

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

✅ No Failures

As of commit 6a35e66 with merge base 89bdd9c (image):
💚 Looks good so far! There are no failures yet. 💚

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

@pytorch-bot pytorch-bot bot added oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: distributed (fsdp) release notes category labels Jul 25, 2024
@yf225 yf225 requested a review from bdhirsh July 25, 2024 06:41
@yf225 yf225 mentioned this pull request Jul 25, 2024
Closed
bdhirsh
bdhirsh reviewed Jul 25, 2024
View reviewed changes
test/distributed/_composable/fsdp/test_fully_shard_compile.py
*self._create_transformer_factory_fns(), "aot_eager", fullgraph=True
)
def test_transformer_backend_aot_eager(self):
for fullgraph in [True, False]:
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: @parametrize("graph_break", [True, False]) might make debugging the test easier if e.g. the graph break path fails at some point (example: https://github.com/pytorch/pytorch/blob/main/test/autograd/test_functional.py#L684)

Copy link
Collaborator

Choose a reason for hiding this comment

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

The issue with distributed is that @parametrize directly would initialize NCCL PG for every parametrized test, which is slow. If you have what @yf225 has, it would reuse the same NCCL PG for all subtests.

I think @kwen2501 added MultiProcContinousTest to maybe address this. I am not sure exactly the limitations though.

Copy link
Contributor

Choose a reason for hiding this comment

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

ahh

bdhirsh
bdhirsh reviewed Jul 25, 2024
View reviewed changes
test/distributed/_composable/fsdp/test_fully_shard_compile.py Outdated
torch._dynamo.graph_break()
return orig_fn(*args, **kwargs)

def _mock_sdpa(self, fullgraph):
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: maybe the name should imply that you're just conditionally adding a graph break to sdpa? with _maybe_add_graph_break_to_sdpa(...)

bdhirsh
bdhirsh reviewed Jul 25, 2024
View reviewed changes
test/distributed/_composable/fsdp/test_fully_shard_compile.py Outdated
).run(code)
else:
self.assertTrue(
len(triton_codes) >= 3,
Copy link
Contributor

Choose a reason for hiding this comment

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

do you think it would be useful to be a bit stricter about what we assert? If I understand properly, my understanding is something like:

3 graphs total (2 fw's due to graph break, 1 unified bw due to compiled autograd)

fw 1: only comm is a all_gather_out, and one set_() (performs the weight AG, but does not free the weight)
fw2: no comms, only contains a set_() (just finishes the fw compute and frees the weight)
bw: contains all_gather_out, reduce_scatter_out, and 2set_() ops (fully gathers weight, does bw compute and fress)

(or at least, why is the assert >=3 and not just == 3?)

Copy link
Contributor Author

Choose a reason for hiding this comment

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

yes I think we want to be more strict here - there is some recompile happening when there is graph break, I'll look into it with @anijain2305 and then add more strict checks

bdhirsh
bdhirsh approved these changes Jul 25, 2024
View reviewed changes
Copy link
Contributor

@bdhirsh bdhirsh left a comment

Choose a reason for hiding this comment

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

very cool to see graph breaks don't error :)

@yf225
Update
14383ca 
[ghstack-poisoned]
@yf225 yf225 added topic: not user facing topic category oncall: pt2 and removed release notes: distributed (fsdp) release notes category labels Jul 25, 2024
@yf225
Copy link
Contributor Author

yf225 commented Jul 25, 2024

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jul 25, 2024
@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@yf225
Update
7a96fbf 
[ghstack-poisoned]
@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: New commits were pushed while merging. Please rerun the merge command.

Details for Dev Infra team Raised by workflow job

@yf225
Copy link
Contributor Author

yf225 commented Jul 25, 2024

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@yf225
Update
72b6d7c 
[ghstack-poisoned]
@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: New commits were pushed while merging. Please rerun the merge command.

Details for Dev Infra team Raised by workflow job

@yf225
Update
6a35e66 
[ghstack-poisoned]
@yf225
Copy link
Contributor Author

yf225 commented Jul 25, 2024

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: 68 mandatory check(s) failed. The first few are:

Dig deeper by viewing the failures on hud

Details for Dev Infra team Raised by workflow job

Failing merge rule: Core Maintainers

@yf225 yf225 closed this Jul 25, 2024
@yf225 yf225 reopened this Jul 25, 2024
@yf225 yf225 added the keep-going Don't stop on first failure, keep running tests until the end label Jul 25, 2024
@yf225
Copy link
Contributor Author

yf225 commented Jul 25, 2024

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: 32 jobs have failed, first few of them are: trunk / linux-focal-cuda12.4-py3.10-gcc9-sm86 / build, trunk / pytorch-linux-focal-py3-clang9-android-ndk-r21e-build / build (default, 1, 1, linux.2xlarge), trunk / linux-focal-cuda11.8-py3.10-gcc9-experimental-split-build-test / test (distributed, 1, 3, linux.8xlarge.nvidia.gpu), trunk / linux-focal-cuda11.8-py3.10-gcc9-experimental-split-build-test / test (distributed, 2, 3, linux.8xlarge.nvidia.gpu), trunk / linux-focal-cuda11.8-py3.10-gcc9-experimental-split-build-test / test (distributed, 3, 3, linux.8xlarge.nvidia.gpu)

Details for Dev Infra team Raised by workflow job

@yf225
Copy link
Contributor Author

yf225 commented Jul 26, 2024

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

bigfootjon pushed a commit that referenced this pull request Jul 31, 2024
@yf225 @bigfootjon
[Traceable FSDP2] Add partial-graph (graph-break) unit tests (#131747)
c4f53ff 
Pull Request resolved: #131747
Approved by: https://github.com/bdhirsh

(cherry picked from commit 236d055)
@henrylhtsang henrylhtsang mentioned this pull request Jul 31, 2024
Closed
@github-actions github-actions bot deleted the gh/yf225/89/head branch August 26, 2024 02:00
enter-ctrl9 pushed a commit to enter-ctrl9/pytorch11 that referenced this pull request Sep 15, 2024
@yf225
[Traceable FSDP2] Add partial-graph (graph-break) unit tests
a779add 
ghstack-source-id: ee7a75f
Pull Request resolved: pytorch/pytorch#131747
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@bdhirsh bdhirsh bdhirsh approved these changes

@awgu awgu awgu left review comments

Assignees

No one assigned

Labels

ciflow/trunk Trigger trunk jobs on your pull request keep-going Don't stop on first failure, keep running tests until the end Merged oncall: distributed Add this issue/PR to distributed oncall triage queue oncall: pt2 topic: not user facing topic category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@yf225 @pytorchmergebot @bdhirsh @awgu