[2d] Ensure gradient clear out pending AsyncCollectiveTensor in FSDP Extension #116122
Conversation
This file contains hidden or 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
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/116122
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 9760f30 with merge base 0e63837 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
… extension" cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]
8000
… extension" cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 awgu penguinwu fegin XilunWu fduwjj wz337 tianyu-l wconstab yf225 [ghstack-poisoned]
… extension" As titled, this PR adds gradient hook to the FSDP DTensor extension, given the fact that FlatParam FSDP assumes same sharding layout after backward pass, here we also ensure the DTensor gradients follow the same sharding layout [ghstack-poisoned]
wanchaol
commented
Dec 21, 2023
awgu
approved these changes
Dec 21, 2023
| "mlp.net2": RowwiseParallel(output_layouts=Shard(0)), | ||
| }, | ||
| ) | ||
| distribute_rmsnorm(tp_model.mlp_norm, mesh_1d) |
There was a problem hiding this comment.
IIUC, this means the RMSNorm runs as "sequence parallel"?
There was a problem hiding this comment.
Yep, this is sequence parallel layernorm/rmsnorm, the backward would produce partial dtensor, and the parameter would be replicated dtensor
… extension" As titled, this PR adds gradient hook to the FSDP DTensor extension, given the fact that FlatParam FSDP assumes same sharding layout after backward pass, here we also ensure the DTensor gradients follow the same sharding layout [ghstack-poisoned]
This was referenced Dec 26, 2023
… extension" As titled, this PR adds gradient hook to the FSDP DTensor extension, given the fact that FlatParam FSDP assumes same sharding layout after backward pass, here we also ensure the DTensor gradients follow the same sharding layout [ghstack-poisoned]
… extension" As titled, this PR adds gradient hook to the FSDP DTensor extension, given the fact that FlatParam FSDP assumes same sharding layout after backward pass, here we also ensure the DTensor gradients follow the same sharding layout [ghstack-poisoned]
|
changed the purpose of this PR a bit to make the gradient hook to eagerly wait for the ACT. |
pytorchmergebot
pushed a commit
that referenced
this pull request
Jan 2, 2024
Pull Request resolved: #116244 Approved by: https://github.com/awgu, https://github.com/wz337, https://github.com/fduwjj ghstack dependencies: #116122
awgu
reviewed
Jan 18, 2024
| from torch.distributed._tensor.placement_types import DTensorSpec | ||
|
|
||
|
|
||
| def sync_grad_hook(grad): | ||
| if isinstance(grad, AsyncCollectiveTensor): | ||
| grad.wait() |
There was a problem hiding this comment.
@wanchaol Why do we discard the return value of grad.wait(), which unwraps the ACT? I.e., why not grad = grad.wait()?
There was a problem hiding this comment.
Good catch! I think we should unwrap yes, it's a bug. Could u help submit a fix on this?
There was a problem hiding this comment.
Oh wait, this looks like it was fixed on main already.
Skylion007
pushed a commit
to Skylion007/pytorch
that referenced
this pull request
Feb 12, 2024
…Extension (pytorch#116122) As titled, this PR adds gradient hook to the FSDP DTensor extension, to check if there's gradients that are AsyncCollectiveTensors, if there're some, we eagerly wait there. This is needed because sometimes the parameter's gradient might still pending with AsyncCollectiveTensor, if we directly feed them to FSDP then FSDP would use the ACT's storage to do reduce_scatter, which is wrong. Pull Request resolved: pytorch#116122 Approved by: https://github.com/awgu, https://github.com/fduwjj
atalman
pushed a commit
that referenced
this pull request
Feb 14, 2024
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.
Stack from ghstack (oldest at bottom):
As titled, this PR adds gradient hook to the FSDP DTensor extension, to check if there's gradients that are AsyncCollectiveTensors, if there're some, we eagerly wait there.
This is needed because sometimes the parameter's gradient might still pending with AsyncCollectiveTensor, if we directly feed them to FSDP then FSDP would use the ACT's storage to do reduce_scatter, which is wrong.
cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @fduwjj @wz337 @tianyu-l @wconstab @yf225