Revert "Fix round robin sharding (#121022)"
This reverts commit effdea5fc62c6bf13cb8035f7bfcc205f05a8b6a.
Reverted https://github.com/pytorch/pytorch/pull/121022 on behalf of https://github.com/clee2000 due to made sharding really uneven ([comment](https://github.com/pytorch/pytorch/pull/121022#issuecomment-1986552662))