Support HSDP + Monolith Checkpointing by mvpatel2000 · Pull Request #128446 · pytorch/pytorch

mvpatel2000 · 2024-06-11T20:30:19Z

Fixes #128444. Rank 0 check should be in the same group as the broadcast

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @LucasLLC @MeetVadakkanchery @mhorowitz

pytorch-bot · 2024-06-11T20:30:22Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

inductor_torchbench_perf jobs are broken due to numpy 2.0 update

❌ 1 New Failure, 1 Unrelated Failure

As of commit 7422e78 with merge base d0c0892 ():

NEW FAILURE - The following job has failed:

periodic / win-vs2019-cuda11.8-py3 / test (default, 1, 4, windows.g5.4xlarge.nvidia.gpu) (gh)
test_decomp 25/25 failed!

FLAKY - The following job failed but was likely due to flakiness present on trunk:

inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_timm, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (similar failure)
gluon_inception_v3

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

Skylion007 · 2024-06-12T04:56:05Z

@pytorchbot rebase

pytorchmergebot · 2024-06-12T04:57:32Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2024-06-12T04:57:34Z

Tried to rebase and push PR #128446, but it was already up to date. Try rebasing against main by issuing:
@pytorchbot rebase -b main

Skylion007 · 2024-06-13T14:16:49Z

@pytorchbot rebase

pytorchmergebot · 2024-06-13T14:18:14Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

pytorchmergebot · 2024-06-13T14:18:22Z

Successfully rebased patch-6 onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout patch-6 && git pull --rebase)

fegin · 2024-06-14T07:41:15Z

torch/distributed/fsdp/_optim_utils.py

 ) -> Dict[str, Any]:
    objects: List[Any] = [None]
-    if fsdp_state.rank == 0:
+    if dist.get_rank(group) == 0:


This looks legit since we use dist.broadcast_object_list(objects, src=0, group=group) in the next few lines. If fsdp_state.rank mismatches dist.broadcast_object_list(objects, src=0, group=group), the broadcast won't work.

fegin · 2024-06-14T07:42:16Z

This PR looks legit, it seems that the HSDP1 does not support broadcast optimizer state_dict correctly. cc., @awgu I'll review it again after all the tests pass.

mvpatel2000 · 2024-06-17T16:50:19Z

@fegin is the test failure related to regressions in pytorch main or this PR? They look unrelated so not sure what action item is here

…cast_from_rank0" HSDP1 + broadcast_from_rank0 actually behaves differently from FSDP1 + broadcast_from_rank0. So we need an unittest to cover this use case. This test relies on the fix from #128446. Differential Revision: [D58621436](https://our.internmc.facebook.com/intern/diff/D58621436/) cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k LucasLLC MeetVadakkanchery mhorowitz [ghstack-poisoned]

HSDP1 + broadcast_from_rank0 actually behaves differently from FSDP1 + broadcast_from_rank0. So we need an unittest to cover this use case. This test relies on the fix from #128446. Differential Revision: [D58621436](https://our.internmc.facebook.com/intern/diff/D58621436/) cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k LucasLLC MeetVadakkanchery mhorowitz [ghstack-poisoned]

fegin

#128755 is the unittest to verify this PR.

fegin · 2024-06-17T16:58:03Z

@pytorchbot merge -f "The test failure is unrelated."

pytorchmergebot · 2024-06-17T16:59:32Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

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

…cast_from_rank0" HSDP1 + broadcast_from_rank0 actually behaves differently from FSDP1 + broadcast_from_rank0. So we need an unittest to cover this use case. This test relies on the fix from #128446. Differential Revision: [D58621436](https://our.internmc.facebook.com/intern/diff/D58621436/) cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k LucasLLC MeetVadakkanchery mhorowitz [ghstack-poisoned]

HSDP1 + broadcast_from_rank0 actually behaves differently from FSDP1 + broadcast_from_rank0. So we need an unittest to cover this use case. This test relies on the fix from #128446. Differential Revision: [D58621436](https://our.internmc.facebook.com/intern/diff/D58621436/) cc mrshenli pritamdamania87 zhaojuanmao satgera gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k LucasLLC MeetVadakkanchery mhorowitz [ghstack-poisoned]

Fixes #128444. Rank 0 check should be in the same group as the broadcast Pull Request resolved: #128446 Approved by: https://github.com/fegin

HSDP1 + broadcast_from_rank0 actually behaves differently from FSDP1 + broadcast_from_rank0. So we need an unittest to cover this use case. This test relies on the fix from #128446. Differential Revision: [D58621436](https://our.internmc.facebook.com/intern/diff/D58621436/) Pull Request resolved: #128755 Approved by: https://github.com/Skylion007, https://github.com/wz337 ghstack dependencies: #128685

Fixes #128444. Rank 0 check should be in the same group as the broadcast Pull Request resolved: #128446 Approved by: https://github.com/fegin (cherry picked from commit 153362f)

HSDP1 + broadcast_from_rank0 actually behaves differently from FSDP1 + broadcast_from_rank0. So we need an unittest to cover this use case. This test relies on the fix from #128446. Differential Revision: [D58621436](https://our.internmc.facebook.com/intern/diff/D58621436/) Pull Request resolved: #128755 Approved by: https://github.com/Skylion007, https://github.com/wz337 ghstack dependencies: #128685 (cherry picked from commit fe8558b)

…129255) HSDP1 + broadcast_from_rank0 actually behaves differently from FSDP1 + broadcast_from_rank0. So we need an unittest to cover this use case. This test relies on the fix from #128446. Differential Revision: [D58621436](https://our.internmc.facebook.com/intern/diff/D58621436/) Pull Request resolved: #128755 Approved by: https://github.com/Skylion007, https://github.com/wz337 ghstack dependencies: #128685 (cherry picked from commit fe8558b)

Fixes #128444. Rank 0 check should be in the same group as the broadcast Pull Request resolved: #128446 Approved by: https://github.com/fegin (cherry picked from commit 153362f) Co-authored-by: Mihir Patel <mihir.v.patel7@gmail.com>

pytorch-bot bot added oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: distributed (fsdp) release notes category labels Jun 11, 2024

mvpatel2000 mentioned this pull request Jun 11, 2024

HSDP + set_optimizer_state_dict errors with monolithic checkpointing #128444

Closed

pytorchbot added the open source label Jun 11, 2024

colesbury requested a review from fegin June 12, 2024 01:36

colesbury added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jun 12, 2024

weifengpy added release notes: distributed (checkpoint) and removed open source release notes: distributed (fsdp) release notes category labels Jun 13, 2024

pytorchbot added the open source label Jun 13, 2024

Update _optim_utils.py

7422e78

pytorchmergebot force-pushed the patch-6 branch from b3f800d to 7422e78 Compare June 13, 2024 14:18

Skylion007 changed the title ~~Support HSDP + Monolith Checkpoiting~~ Support HSDP + Monolith Checkpointing Jun 13, 2024

pytorch-bot bot added the ciflow/inductor label Jun 13, 2024

fegin added ciflow/periodic Trigger jobs ran periodically on master (periodic.yml) on the PR ciflow/trunk Trigger trunk jobs on your pull request labels Jun 14, 2024

fegin reviewed Jun 14, 2024

View reviewed changes

fegin added the module: distributed_checkpoint label Jun 14, 2024

fegin mentioned this pull request Jun 14, 2024

[DSD] Add unittest to verify HSDP1 + broadcast_from_rank0 #128755

Closed

fegin approved these changes Jun 17, 2024

View reviewed changes

pytorchmergebot added the merging label Jun 17, 2024

pytorchmergebot closed this in 153362f Jun 17, 2024

pytorchmergebot added Merged and removed merging labels Jun 17, 2024

PaliC pushed a commit that referenced this pull request Jun 17, 2024

Support HSDP + Monolith Checkpointing (#128446)

8f88f71

Fixes #128444. Rank 0 check should be in the same group as the broadcast Pull Request resolved: #128446 Approved by: https://github.com/fegin

fegin mentioned this pull request Jun 21, 2024

Support HSDP + Monolith Checkpointing (#128446) #129254

Merged

This was referenced Jun 21, 2024

[DSD] Add unittest to verify HSDP1 + broadcast_from_rank0 (#128755) #129255

Merged

[v.2.4.0] Release Tracker #128436

Closed

Conversation

mvpatel2000 commented Jun 11, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jun 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❗ 1 Active SEVs

❌ 1 New Failure, 1 Unrelated Failure

Uh oh!

Skylion007 commented Jun 12, 2024

Uh oh!

pytorchmergebot commented Jun 12, 2024

Uh oh!

pytorchmergebot commented Jun 12, 2024

Uh oh!

Skylion007 commented Jun 13, 2024

Uh oh!

pytorchmergebot commented Jun 13, 2024

Uh oh!

pytorchmergebot commented Jun 13, 2024

Uh oh!

fegin Jun 14, 2024

Choose a reason for hiding this comment

Uh oh!

fegin commented Jun 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mvpatel2000 commented Jun 17, 2024

Uh oh!

fegin left a comment

Choose a reason for hiding this comment

Uh oh!

fegin commented Jun 17, 2024

Uh oh!

pytorchmergebot commented Jun 17, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

mvpatel2000 commented Jun 11, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Jun 11, 2024 •

edited

Loading

fegin commented Jun 14, 2024 •

edited

Loading