[Checkpoint][Test] Add test for optimizer state_dict and resharding to 2d checkpoint test #91092

wz337 · 2022-12-19T09:32:04Z

This PR updates the 2d checkpoint model state test to include:

optimizer state dict test
simple resharding test (pg change)
rename test

pytorch-bot · 2022-12-19T09:32:06Z

🔗 Helpful Links

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

📄 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.

✅ No Failures

As of commit 61cfbdc:
💚 Looks good so far! There are no failures yet. 💚

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

fduwjj

LGTM. Unblock and have a small question here.

fduwjj · 2022-12-21T21:33:19Z

test/distributed/checkpoint/test_2d_fsdp_dt_checkpoint.py

@@ -171,6 +208,29 @@ def test_2d_model_state_checkpoint(self) -> None:
                    else:
                        self.assertEqual(n_p1[1], n_p2[1])

+        def opt_at(opt, idx):
+            return list(iter(opt.state.values()))[idx]


Why do we need "iter" here?

Thanks for catching. iter is unnecessary here.

wz337 · 2022-12-23T01:11:47Z

@pytorchmergebot rebase

pytorchmergebot · 2022-12-23T01:13:40Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2022-12-23T01:13:46Z

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

wz337 · 2022-12-24T06:49:35Z

@pytorchmergebot rebase

pytorchmergebot · 2022-12-24T06:51:31Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2022-12-24T06:51:37Z

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

wz337 · 2022-12-25T03:51:56Z

@pytorchmergebot rebase

pytorchmergebot · 2022-12-25T03:53:54Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2022-12-25T03:53:57Z

Tried to rebase and push PR #91092, but it was already up to date

wz337 · 2022-12-26T02:08:52Z

@pytorchmergebot rebase

pytorchmergebot · 2022-12-26T02:10:43Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2022-12-26T02:10:48Z

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

wz337 · 2022-12-27T04:56:20Z

@pytorchmergebot rebase

pytorchmergebot · 2022-12-27T04:58:29Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2022-12-27T04:58:35Z

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

wz337 · 2022-12-28T01:09:28Z

@pytorchmergebot rebase

pytorchmergebot · 2022-12-28T01:11:13Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2022-12-28T01:11:17Z

Tried to rebase and push PR #91092, but it was already up to date

wz337 · 2022-12-28T03:45:57Z

@pytorchmergebot merge

pytorchmergebot · 2022-12-28T03:47:32Z

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 · 2022-12-28T03:47:36Z

Merge failed

Reason: The following mandatory check(s) failed (Rule Distributed):

pull

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

wz337 · 2022-12-29T07:40:41Z

@pytorchmergebot rebase

pytorchmergebot · 2022-12-29T07:42:23Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2022-12-29T07:42:27Z

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

wz337 · 2023-01-03T16:42:53Z

@pytorchmergebot rebase

pytorchmergebot · 2023-01-03T16:44:43Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2023-01-03T16:44:47Z

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

wz337 · 2023-01-04T19:00:46Z

@pytorchmergebot rebase

pytorchmergebot · 2023-01-04T19:02:40Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2023-01-04T19:02:44Z

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

wz337 · 2023-01-04T23:24:43Z

@pytorchmergebot merge

pytorchmergebot · 2023-01-04T23:26:26Z

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

pytorch-bot bot added the release notes: distributed (fsdp) release notes category label Dec 19, 2022

wz337 marked this pull request as ready for review December 19, 2022 17:48

wz337 requested review from mrshenli, pritamdamania87, zhaojuanmao, rohan-varma, H-Huang, awgu, kwen2501 and wanchaol as code owners December 19, 2022 17:48

wz337 requested a review from fduwjj December 19, 2022 17:48

fduwjj approved these changes Dec 21, 2022

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Dec 28, 2022

wz337 added 7 commits January 4, 2023 19:02

add test_2d_optim_state_checkpoint

28ff5df

test/distributed/checkpoint/test_2d_optim_state_checkpoint.py

c33322e

add model state and optim state with simple resharding

89d303b

update to include resharding test

257fe6c

fix lint'

6281df7

resolve mypy

f337422

fix nit

61cfbdc

pytorchmergebot added the Merged label Jan 4, 2023

pytorchmergebot closed this in 56db21a Jan 4, 2023

wz337 mentioned this pull request Jan 20, 2023

Improve Test for DCP #90960

Closed

[Checkpoint][Test] Add test for optimizer state_dict and resharding to 2d checkpoint test #91092

[Checkpoint][Test] Add test for optimizer state_dict and resharding to 2d checkpoint test #91092

Conversation

wz337 commented Dec 19, 2022 • edited

pytorch-bot bot commented Dec 19, 2022 • edited

🔗 Helpful Links

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

✅ No Failures

fduwjj left a comment

Choose a reason for hiding this comment

fduwjj Dec 21, 2022

Choose a reason for hiding this comment

wz337 Dec 22, 2022

Choose a reason for hiding this comment

wz337 commented Dec 23, 2022

pytorchmergebot commented Dec 23, 2022

pytorchmergebot commented Dec 23, 2022

wz337 commented Dec 24, 2022

pytorchmergebot commented Dec 24, 2022

pytorchmergebot commented Dec 24, 2022

wz337 commented Dec 25, 2022

pytorchmergebot commented Dec 25, 2022

pytorchmergebot commented Dec 25, 2022

wz337 commented Dec 26, 2022

pytorchmergebot commented Dec 26, 2022

pytorchmergebot commented Dec 26, 2022

wz337 commented Dec 27, 2022

pytorchmergebot commented Dec 27, 2022

pytorchmergebot commented Dec 27, 2022

wz337 commented Dec 28, 2022

pytorchmergebot commented Dec 28, 2022

pytorchmergebot commented Dec 28, 2022

wz337 commented Dec 28, 2022

pytorchmergebot commented Dec 28, 2022

Merge started

pytorchmergebot commented Dec 28, 2022

Merge failed

wz337 commented Dec 29, 2022

pytorchmergebot commented Dec 29, 2022

pytorchmergebot commented Dec 29, 2022

wz337 commented Jan 3, 2023

pytorchmergebot commented Jan 3, 2023

pytorchmergebot commented Jan 3, 2023

wz337 commented Jan 4, 2023

pytorchmergebot commented Jan 4, 2023

pytorchmergebot commented Jan 4, 2023

wz337 commented Jan 4, 2023

pytorchmergebot commented Jan 4, 2023

Merge started

wz337 commented Dec 19, 2022 •

edited

pytorch-bot bot commented Dec 19, 2022 •

edited