Repository: pytorch
Issue: test_fsdp_tp_integration fails for non-power-of 2 GPUs
Labels: oncall: distributed

### 🐛 Describe the bug

Running the `test_fsdp_tp_integration` with a number of GPUs that is (likely) not a power of 2 fails with e.g.:
```
torch.testing._internal.common_distributed: [ERROR]   File "/dev/shm//pytorch/test/distributed/fsdp/test_fsdp_tp_integration.py", line 157, in _sync_tp_grads
torch.testing._internal.common_distributed: [ERROR]     per_param_masks = unsharded_zeros.split(splits)
torch.testing._internal.common_distributed: [ERROR]   File "/tmp/easybuild-install/lib/python3.10/site-packages/torch/_tensor.py", line 864, in split
torch.testing._internal.common_distributed: [ERROR]     return torch._VF.split_with_sizes(self, split_size, dim)
torch.testing._internal.common_distributed: [ERROR] RuntimeError: split_with_sizes expects split_sizes to sum exactly to 105 (input tensor's size at dimension 0), but got split_sizes=[20, 4, 16, 4, 48, 12]
```

I investigated a bit: The module is the `SimpleModel` with 
```
        self.net1 = torch.nn.Linear(5, 8)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(8, 4)
        self.net3 = torch.nn.Linear(4, 12)
```

Only `"net1.weight", "net1.bias", "net2.weight"` are sharded leading to a [name-size mapping](https://github.com/pytorch/pytorch/blob/dcf7728fd6c0cee55f060c2afb73c6a9b2176f41/test/distributed/fsdp/test_fsdp_tp_integration.py#L152) of `OrderedDict([('net1.weight', 20), ('net1.bias', 4), ('net2.weight', 16), ('net2.bias', 4), ('net3.weight', 48), ('net3.bias', 12)])` which is the `splits` parameter with a total of 104 elements.

Then we have `unsharded_size = flat_param.numel() * fsdp_world_size = flat_param.numel() * self.world_size // tp_world_size = flat_param.numel() * 6 // 2 = 35 * 3 = 105`:
https://github.com/pytorch/pytorch/blob/dcf7728fd6c0cee55f060c2afb73c6a9b2176f41/test/distributed/fsdp/test_fsdp_tp_integration.py#L154

I imagine there is some rounding going on when distributing the 104 parameters:
- 104 parameters for each of the tensor-parallel shards
- 2 tensor-parallel shards distributed on 6 GPUs -> 1 per 3 GPUs
- 104 / 3 = 34 + 2/3 -> ~35

And the decomposition seems to not take that properly into account.


### Versions

PyTorch version: 2.1.2
Is debug build: False
CUDA used to build PyTorch: 11.7

OS: AlmaLinux release 8.7 (Stone Smilodon) (ppc64le)
GCC version: (GCC) 11.3.0
Clang version: Could not collect
CMake version: version 3.23.1
Libc version: glibc-2.28

Python version: 3.10.4 (main, Feb 11 2023, 06:30:51) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-4.18.0-425.19.2.el8_7.ppc64le-ppc64le-with-glibc2.28
Is CUDA available: True
CUDA runtime version: 11.7.64
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB

Nvidia driver version: 530.30.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False


cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o