[FSDP2][ez] Replaced `groupby` with `all` for same-dtype check (#119825)
The `groupby` logic to check if all all-gather inputs have the same dtype is not so readable. Let us use `all` instead.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119825
Approved by: https://github.com/Skylion007
ghstack dependencies: #119550, #118136, #118223, #118755
