Use size to check same tensor sizes in reduce_scatter and allgather (#84099)
Summary:
Previous code uses tensor.numel() to check if all tensors have the same size in order to switch between reduce_scatter_v v.s. reduce_scatter, same applies to allgather. However, if the user input tensor is zero in the last dimension (e.g., [648632,0]), then numel() returns zero and check_same_numel is always true.
This patch fixes the check to use size rather than numel, to cover the above case.
Differential Revision: D39044439
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84099
Approved by: https://github.com/kwen2501
