-
Notifications
You must be signed in to change notification settings - Fork 22.6k
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
[FSDP] Fix for optim state dict #102901
Conversation
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]
🔗 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 FailuresAs of commit 84be448: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
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
There was a problem hiding this 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") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
#nit: Remove print?
There was a problem hiding this comment.
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() |
There was a problem hiding this comment.
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") |
There was a problem hiding this comment.
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]
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]
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/)
@pytorchbot merge -f "CI done" |
Merge startedYour 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 |
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