Fix fsdp+pp+te WPS decreasing issue #1139
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
facebook-github-bot
added
the
CLA Signed
jianyuh
requested review from
jspark1105,
jiecaoyu,
vedanuj,
awgu and
vivien-chu
jspark1105
reviewed
Oct 1, 2023
| @@ -1650,6 +1654,9 @@ def _register_post_backward_hooks(self) -> None: | |||
| assert p_tmp.grad_fn is not None | |||
| grad_acc = p_tmp.grad_fn.next_functions[0][0] # Gets its GradAccumulation object. | |||
| handle = grad_acc.register_hook(functools.partial(self._post_backward_hook, p)) | |||
| if not hasattr(p, "_shard_bwd_hooks"): | |||
| p._shard_bwd_hooks = [] | |||
| p._shard_bwd_hooks.append((grad_acc, handle)) | |||
| p._shard_bwd_hook = (grad_acc, handle) | |||
There was a problem hiding this comment.
This should be deleted? See P841842878 CC @awgu
There was a problem hiding this comment.
Commented this line following the style in this file.
Comment on lines
1224
to
1226
| for module_name, module in self.named_modules(): | ||
| if isinstance(module, FullyShardedDataParallel): | ||
| module._module_fqn = module_name |
There was a problem hiding this comment.
These only needed for debugging (when _FSDP_DEBUG is set in P841842878)
Comment on lines
1736
to
1740
| if self.fp32_reduce_scatter: | ||
| # Cast grad to FP32. | ||
| param.grad.data = param.grad.data.float() | ||
| orig_grad_data = param.grad.data.float() | ||
| else: | ||
| orig_grad_data = param.grad.data |
There was a problem hiding this comment.
We should keep param.grad.data = param.grad.data.float() so something like
if self.fp32_reduce_scatter:
# Cast grad to FP32.
param.grad.data = param.grad.data.float()
orig_grad_data = param.grad.data
jspark1105
reviewed
Oct 2, 2023
| @@ -1710,6 +1713,13 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None: | |||
|
|
|||
| # Switch to FP32 shard after backward. | |||
| self._use_fp32_param_shard([param]) | |||
| if self.mixed_precision and self.fp32_reduce_scatter: | |||
There was a problem hiding this comment.
Sorry missed this comment before. Wonder if we can have a separate commit for main_grad related changes from the changes for wps decrease fix (P841842878).
jspark1105
approved these changes
Oct 2, 2023
awgu
approved these changes
Oct 2, 2023
10 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes WPS decreasing with steps, due to the FSDP handle issues.
Credits all go to @awgu @vivien-chu . Merge this into ngoyal_changes_for_pp_fp8 due to the failure observed in test_te.py when reproducing https://github.com/fairinternal/xlformers/pull/1360 w/ @jiecaoyu @jspark1105
Before submitting
PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.