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

[FSDP] Fix for optim state dict #102901

Closed
wants to merge 3 commits into from

Conversation

rohan-varma
Copy link
Member

@rohan-varma rohan-varma commented Jun 3, 2023

Stack from ghstack (oldest at bottom):

Fix for HSDP + use_orig_params where we need to pass in the PG that
might not be the default.

Differential Revision: D46417327

Fix for HSDP + use_orig_params where we need to pass in the PG that
might not be the default.

Differential Revision: [D46417327](https://our.internmc.facebook.com/intern/diff/D46417327/)

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented Jun 3, 2023

🔗 Helpful Links

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

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

✅ No Failures

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

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

rohan-varma added a commit that referenced this pull request Jun 3, 2023
Fix for HSDP + use_orig_params where we need to pass in the PG that
might not be the default.

Differential Revision: [D46417327](https://our.internmc.facebook.com/intern/diff/D46417327/)

ghstack-source-id: 191112222
Pull Request resolved: #102901
Copy link
Contributor

@wz337 wz337 left a comment

Choose a reason for hiding this comment

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

LGTM!

@@ -1349,6 +1349,7 @@ def _optim_state_dict(
state = (
{} if param_key is None else optim_state_dict["state"][param_key]
)
print(f"RV: calling _Gather_orig_param_state")
Copy link
Contributor

Choose a reason for hiding this comment

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

#nit: Remove print?

Copy link
Contributor

Choose a reason for hiding this comment

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

ditto on the others before landing :)

@@ -1497,7 +1498,9 @@ def _all_gather_optim_state(
object_list: List[StateInfo] = [
processed_state for _ in range(fsdp_state.world_size)
]
dist.all_gather_object(object_list, processed_state)
assert fsdp_state.world_size == fsdp_state.process_group.size()
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: Add an assert message?

@@ -1349,6 +1349,7 @@ def _optim_state_dict(
state = (
{} if param_key is None else optim_state_dict["state"][param_key]
)
print(f"RV: calling _Gather_orig_param_state")
Copy link
Contributor

Choose a reason for hiding this comment

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

ditto on the others before landing :)

Fix for HSDP + use_orig_params where we need to pass in the PG that
might not be the default.

Differential Revision: [D46417327](https://our.internmc.facebook.com/intern/diff/D46417327/)

[ghstack-poisoned]
rohan-varma added a commit that referenced this pull request Jun 5, 2023
Pull Request resolved: #102901

Fix for HSDP + use_orig_params where we need to pass in the PG that
might not be the default.
ghstack-source-id: 191231200

Differential Revision: [D46417327](https://our.internmc.facebook.com/intern/diff/D46417327/)
Fix for HSDP + use_orig_params where we need to pass in the PG that
might not be the default.

Differential Revision: [D46417327](https://our.internmc.facebook.com/intern/diff/D46417327/)

[ghstack-poisoned]
rohan-varma added a commit that referenced this pull request Jun 6, 2023
Pull Request resolved: #102901

Fix for HSDP + use_orig_params where we need to pass in the PG that
might not be the default.
ghstack-source-id: 191337799

Differential Revision: [D46417327](https://our.internmc.facebook.com/intern/diff/D46417327/)
@rohan-varma
Copy link
Member Author

@pytorchbot merge -f "CI done"

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

@wz337 wz337 added suppress-bc-linter Suppresses the failures of API backward-compatibility linter (Lint/bc_linter) and removed suppress-bc-linter Suppresses the failures of API backward-compatibility linter (Lint/bc_linter) labels Jun 6, 2023
@facebook-github-bot facebook-github-bot deleted the gh/rohan-varma/702/head branch June 10, 2023 14:16
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.

4 participants