Repository: pytorch
Issue: [DTensor]`Linear` fails on 3D DTensor with `batch size > 1` and `Replicate` input redistributed from `Shard`
Labels: oncall: distributed,module: dtensor

### 🐛 Describe the bug

`Linear` fails on 3D DTensor with `batch size > 1` and Replicate input from shard (not divisible by TP size).

* Error Message

```
[rank3]: Traceback (most recent call last):
[rank3]:   File "/path/to/pytorch_bug/linear_bug.py", line 34, in <module>
[rank3]:     mod(x_dt)
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank3]:     return self._call_impl(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank3]:     return forward_call(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/path/to/pytorch_bug/linear_bug.py", line 20, in forward
[rank3]:     return self.layer(x)
[rank3]:            ^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank3]:     return self._call_impl(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl
[rank3]:     return inner()
[rank3]:            ^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1793, in inner
[rank3]:     result = forward_call(*args, **kwargs)
[rank3]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/nn/modules/linear.py", line 125, in forward
[rank3]:     return F.linear(input, self.weight, self.bias)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/_compile.py", line 32, in inner
[rank3]:     return disable_fn(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 751, in _fn
[rank3]:     return fn(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/distributed/tensor/_api.py", line 343, in __torch_dispatch__
[rank3]:     return DTensor._op_dispatcher.dispatch(
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/distributed/tensor/_dispatch.py", line 216, in dispatch
[rank3]:     local_results = op_call(*local_tensor_args, **op_info.local_kwargs)
[rank3]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]:   File "/opt/miniconda/envs/torchtitan/lib/python3.11/site-packages/torch/_ops.py", line 722, in __call__
[rank3]:     return self._op(*args, **kwargs)
[rank3]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
```

* Code

`torchrun --nnodes=1 --nproc-per-node=4 --standalone /path/to/pytorch_bug/linear_bug.py`

```Python


from torch.distributed.tensor import Shard, DTensor, Replicate
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
)
_world_size = int(os.environ["WORLD_SIZE"])
device_mesh = init_device_mesh(device_type="cuda", mesh_shape=(_world_size,))
class Mod(nn.Module):
    def __init__(self):
        super(Mod, self).__init__()
        self.layer = nn.Linear(64, 64)
    def forward(self, x):
        return self.layer(x)
mod_ref = Mod().cuda()
mod = Mod().cuda()
parallelize_module(mod, device_mesh, {
    "layer": ColwiseParallel()
})
x = torch.randn(4, 111, 64).cuda()
x_dt = DTensor.from_local(x, device_mesh, [Replicate()])
x_shard = x_dt.redistribute(device_mesh, [Shard(1)])
x_dt = x_shard.redistribute(device_mesh, [Replicate()])
# works
mod_ref(x)
mod_ref(x_dt._local_tensor)
x_dt_2 = DTensor.from_local(x_dt._local_tensor, device_mesh, [Replicate()])
mod(x_dt_2)
# error
mod(x_dt)
dist.barrier()
dist.destroy_process_group()
```

### Versions

```
PyTorch version: 2.6.0.dev20241222+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.11.10 (main, Oct  3 2024, 07:29:13) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.210-4-velinux1-amd64-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
Nvidia driver version: 535.86.10
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==2.1.3
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] pytorch-triton==3.2.0+git0d4682f0
[pip3] torch==2.6.0.dev20241222+cu124
[pip3] torchaudio==2.6.0.dev20241222+cu124
[pip3] torchdata==0.9.0
[pip3] torchpippy==0.2.0+1bcb2bf
[pip3] torchtitan==0.0.2
[pip3] torchvision==0.22.0.dev20241222+cu124
[pip3] triton==3.1.0
[conda] blas                      1.0                         mkl  
[conda] cuda-cudart               12.1.105                      0    nvidia
[conda] cuda-cupti                12.1.105                      0    nvidia
[conda] cuda-libraries            12.1.0                        0    nvidia
[conda] cuda-nvrtc                12.1.105                      0    nvidia
[conda] cuda-nvtx                 12.1.105                      0    nvidia
[conda] cuda-opencl               12.4.127                      0    nvidia
[conda] cuda-runtime              12.1.0                        0    nvidia
[conda] libcublas                 12.1.0.26                     0    nvidia
[conda] libcufft                  11.0.2.4                      0    nvidia
[conda] libcurand                 10.3.5.147                    0    nvidia
[conda] libcusolver               11.4.4.55                     0    nvidia
[conda] libcusparse               12.0.2.55                     0    nvidia
[conda] libnvjitlink              12.1.105                      0    nvidia
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-service               2.4.0           py311h5eee18b_1  
[conda] mkl_fft                   1.3.11          py311h5eee18b_0  
[conda] mkl_random                1.2.8           py311ha02d727_0  
[conda] numpy                     2.1.3           py311h08b1b3b_0  
[conda] numpy-base                2.1.3           py311hf175353_0  
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-cusparselt-cu12    0.6.2                    pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] pytorch-cuda              12.1                 ha16c6d3_6    pytorch
[conda] pytorch-triton            3.2.0+git0d4682f0          pypi_0    pypi
[conda] torch                     2.6.0.dev20241222+cu124          pypi_0    pypi
[conda] torchaudio                2.6.0.dev20241222+cu124          pypi_0    pypi
[conda] torchdata                 0.9.0                    pypi_0    pypi
[conda] torchpippy                0.2.0+1bcb2bf            pypi_0    pypi
[conda] torchtitan                0.0.2                    pypi_0    pypi
[conda] torchvision               0.22.0.dev20241222+cu124          pypi_0    pypi
[conda] triton                    3.1.0                    pypi_0    pypi
```

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