[FSDP2] Fix issue with set_reduce_scatter_divide_factor errors and MixedPrecisionPolicy #155964
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/155964
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit ce6d607 with merge base 9620994 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/inductor
oncall: distributed
| # uses NCCL's PreMulSum, which only allows data type half, float, | ||
| # or double. Set reduce_dtype as orig_dtype so that if won't be cast by | ||
| # mp_policy. | ||
| self._reduce_dtype = self._orig_dtype |
There was a problem hiding this comment.
this breaks mixed precision contract? is it silently ignoring reduce_dtype from mp policy?
There was a problem hiding this comment.
Yeah, the casted dtype by mp_policy cannot run under PreMulSum is they are not float.
| @skipIfRocm # regressed in ROCm 6.4, but ROCm 6.5 fixes it | ||
| @skip_if_lt_x_gpu(2) | ||
| @requires_nccl_version((2, 10), "Need NCCL 2.10+ for bf16 collectives") | ||
| def test_set_reduce_scatter_divide_factor(self): |
There was a problem hiding this comment.
modify existing unit test instead?
| return unpackPreMulSum<at::Half, ncclHalf>(reduceOp, comm); | ||
| case ncclFloat: | ||
| return unpackPreMulSum<float, ncclFloat>(reduceOp, comm); | ||
| case ncclBfloat16: |
There was a problem hiding this comment.
add ProcessGroupNCCL unit test here?
There was a problem hiding this comment.
@kwen2501 seems nccl can take all the datatype they support. that's why we added ncclBfloat16 https://github.com/NVIDIA/nccl/blob/72d2432094d6ae36abd6e511c3a16a2d052dbf94/src/enqueue.cc#L2467-L2480
There was a problem hiding this comment.
rebase on top of #155915
mori360
changed the title
|
@pytorchmergebot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
3 participants
fix #155223
cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k