[FSDP] Add gradient predivide factor to avoid overflow/underflow with large world size #565
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
myleott
approved these changes
Apr 1, 2021
| @@ -1071,7 +1078,7 @@ def _post_backward_hook(self, param: Parameter, *unused: Any) -> None: | |||
|
|
|||
| if self.world_size > 1: | |||
There was a problem hiding this comment.
change to if self.gradient_predivide_factor > 1
| @@ -1098,7 +1105,7 @@ def _post_reduction_hook(self, param: Parameter, reduced_grad: torch.Tensor) -> | |||
| assert torch.cuda.current_stream() == self._streams["post_backward"] | |||
| assert param.grad is not None | |||
| self.assert_state(TrainingState.BACKWARD_POST) | |||
| param.grad.data = reduced_grad | |||
| param.grad.data = reduced_grad.div_(self.world_size/self.gradient_predivide_factor) | |||
There was a problem hiding this comment.
change to:
param.grad.data = reduced_grad
if self.gradient_postdivide_factor > 1:
param.grad.data.div_(self.gradient_postdivide_factor)
sshleifer
approved these changes
Apr 1, 2021
|
@shruti-bh I pushed the lint fixes and merged this 😄 |
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.
This PR tries to optimally scale gradients when the world_size is big.
For example,
reduce_scatter(gradient) / world_sizemight overflow andreduce_scatter(gradient / world_size)might underflow with world_size=1024. We might prefer insteadreduce_scatter(gradient / 32) / 32.See Nvidia’s note about this
Testing:
After 10 updates, I see ~ consistent loss/gnorm with the main branch and this branch.
Main:
2021-03-31 16:55:49 | INFO | train | epoch 001 | loss 16.361 | nll_loss 16.324 | moe_gate_loss 2.56751 | overflow_expert1 24.327 | overflow_expert2 53.304 | entropy_gating 1.994 | expert1_balance_top 68.159 | expert1_balance_bottom 2.701 | unused_expert1_count 0 | expert2_balance_top 54.38 | expert2_balance_bottom 4.655 | unused_expert2_count 0 | ppl 82053.2 | wps 212239 | ups 6.47 | wpb 32768 | bsz 32 | num_updates 10 | lr 0.002 | gnorm 4.115 | loss_scale 32 | train_wall 4 | gb_free 29.2 | wall 24This Branch:
2021-03-31 16:57:48 | INFO | train | epoch 001 | loss 16.361 | nll_loss 16.324 | moe_gate_loss 2.56767 | overflow_expert1 24.324 | overflow_expert2 53.299 | entropy_gating 1.994 | expert1_balance_top 68.156 | expert1_balance_bottom 2.7 | unused_expert1_count 0 | expert2_balance_top 54.396 | expert2_balance_bottom 4.65 | unused_expert2_count 0 | ppl 82051.3 | wps 216940 | ups 6.61 | wpb 32768 | bsz 32 | num_updates 10 | lr 0.002 | gnorm 4.116 | loss_scale 32 | train_wall 6 | gb_free 29.2 | wall 24After 100 updates, the loss/gnorm numbers are slightly different after the 2nd decimal place between the main branch and this branch
Main:
2021-03-31 17:01:50 | INFO | train | epoch 001 | loss 11.751 | nll_loss 11.713 | moe_gate_loss 2.63956 | overflow_expert1 43.949 | overflow_expert2 64.359 | entropy_gating 2.022 | expert1_balance_top 84.331 | expert1_balance_bottom 1.088 | unused_expert1_count 0 | expert2_balance_top 75.253 | expert2_balance_bottom 1.84 | unused_expert2_count 0 | ppl 3356.47 | wps 229459 | ups 7 | wpb 32768 | bsz 32 | num_updates 100 | lr 0 | gnorm 0.931 | loss_scale 32 | train_wall 16 | gb_free 29.2 | wall 38This Branch:
2021-03-31 16:59:44 | INFO | train | epoch 001 | loss 11.747 | nll_loss 11.706 | moe_gate_loss 2.79818 | overflow_expert1 44.95 | overflow_expert2 65.851 | entropy_gating 2.003 | expert1_balance_top 84.568 | expert1_balance_bottom 0.411 | unused_expert1_count 0 | expert2_balance_top 76.632 | expert2_balance_bottom 0.981 | unused_expert2_count 0 | ppl 3341.38 | wps 227973 | ups 6.96 | wpb 32768 | bsz 32 | num_updates 100 | lr 0 | gnorm 0.937 | loss_scale 32 | train_wall 18 | gb_free 29.2 | wall 37Would love advice on how to proceed given these differences which would tend to emerge since we are scaling the gradients a bit differently now.