Repository: pytorch
Issue: Incorrect Results with Tensor Parallelism
Labels: oncall: distributed

### 🐛 Describe the bug

I am trying a basic Tensor Parallel implementation on a 2 layer MLP using `ColwiseParallel` followed by a `RowwiseParallel`. I would expect the final output of the MLP to be the same in the Tensor Parallel version compared to the non-parallelized version. However, the output tensors are different.

```python
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel
from torch.distributed.tensor.placement_types import Replicate, Shard

class MLP(nn.Module):
    def __init__(
        self,
        dim: int,
        expand_ratio: int,
        mp_mesh,
        _parallelize=True
    ):
        super().__init__()
        self.mp_mesh = mp_mesh

        self.proj_in = nn.Linear(dim, dim * expand_ratio)
        self.act = nn.GELU("tanh")
        self.proj_out = nn.Linear(dim * expand_ratio, dim)

    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
        x = self.proj_in(x)
        x = self.act(x)
        x = self.proj_out(x)
        return x
    

if __name__ == "__main__":
    import os
    from torch.distributed.device_mesh import init_device_mesh
    import torch.distributed.tensor as dtensor

    torch.manual_seed(0)

    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f'cuda:{local_rank}')
    mesh = init_device_mesh("cuda", (8,))

    head_dim = 80
    num_heads = 24
    d_model = head_dim * num_heads
    text_seq_len = 10

    model = MLP(d_model, expand_ratio=4, mp_mesh=mesh, _parallelize=parallelize).to(device).to(torch.bfloat16)

    dtext = dtensor.randn((text_seq_len, d_model), dtype=torch.bfloat16, device_mesh=mesh, placements=[Replicate()])

    text = dtext.full_tensor()

    text_output = model(text)

    model = parallelize_module(model, device_mesh=mesh, parallelize_plan={
        "proj_in": ColwiseParallel(use_local_output=True), 
        "proj_out": RowwiseParallel(use_local_output=True)})

    parallel_text_out = model(dtext)

    if local_rank == 0:
        print("Text output: ", text_output)
        print("Parallel text output: ", parallel_text_out)
        assert text_output.size() == parallel_text_out.size()
        assert torch.allclose(text_output, parallel_text_out) # This assertion fails
```

I run this on a single node with 8 GPUs via `torchrun --nproc_per_node=8 torch_tp_test.py`.

But the assertion fails with 

```
Text output:  tensor([[-0.1299, -0.1758, -0.0344,  ...,  0.1128, -0.2178, -0.0466],
        [-0.0226,  0.1167,  0.1768,  ..., -0.0160, -0.0405, -0.2656],
        [-0.1641, -0.0554,  0.2715,  ...,  0.1475,  0.0967,  0.1309],
        ...,
        [-0.0820, -0.0391,  0.2480,  ..., -0.0525, -0.0962,  0.0903],
        [-0.0179, -0.0850, -0.1641,  ..., -0.2451,  0.0364, -0.0962],
        [-0.2676,  0.0332, -0.2500,  ..., -0.0410, -0.2412,  0.2930]],
       device='cuda:0', dtype=torch.bfloat16, grad_fn=<AddmmBackward0>)
Parallel text output:  AsyncCollectiveTensor(tensor([[-0.1309, -0.1758, -0.0334,  ...,  0.1108, -0.2188, -0.0471],
        [-0.0234,  0.1162,  0.1758,  ..., -0.0176, -0.0381, -0.2676],
        [-0.1621, -0.0549,  0.2695,  ...,  0.1455,  0.0967,  0.1318],
        ...,
        [-0.0825, -0.0366,  0.2480,  ..., -0.0537, -0.0977,  0.0898],
        [-0.0181, -0.0830, -0.1621,  ..., -0.2451,  0.0361, -0.0977],
        [-0.2676,  0.0325, -0.2490,  ..., -0.0410, -0.2402,  0.2930]],
       device='cuda:0', dtype=torch.bfloat16))
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/amogkamsetty/torch_tp_test.py", line 88, in <module>
[rank0]:     assert torch.allclose(text_output, parallel_text_out)
[rank0]: AssertionError
```

### Versions

```
Collecting environment information...
PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (x86_64)
GCC version: (conda-forge gcc 12.1.0-17) 12.1.0
Clang version: Could not collect
CMake version: version 3.30.0
Libc version: glibc-2.35

Python version: 3.10.14 (main, May  6 2024, 19:42:50) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-6.1.85+-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA H100 80GB HBM3
GPU 1: NVIDIA H100 80GB HBM3
GPU 2: NVIDIA H100 80GB HBM3
GPU 3: NVIDIA H100 80GB HBM3
GPU 4: NVIDIA H100 80GB HBM3
GPU 5: NVIDIA H100 80GB HBM3
GPU 6: NVIDIA H100 80GB HBM3
GPU 7: NVIDIA H100 80GB HBM3

Nvidia driver version: 550.90.07
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.7
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.7
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        52 bits physical, 57 bits virtual
Byte Order:                           Little Endian
CPU(s):                               208
On-line CPU(s) list:                  0-207
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Xeon(R) Platinum 8481C CPU @ 2.70GHz
CPU family:                           6
Model:                                143
Thread(s) per core:                   2
Core(s) per socket:                   52
Socket(s):                            2
Stepping:                             8
BogoMIPS:                             5399.99
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rtm avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 arat avx512vbmi umip avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid cldemote movdiri movdir64b fsrm md_clear serialize amx_bf16 avx512_fp16 amx_tile amx_int8 arch_capabilities
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            4.9 MiB (104 instances)
L1i cache:                            3.3 MiB (104 instances)
L2 cache:                             208 MiB (104 instances)
L3 cache:                             210 MiB (2 instances)
NUMA node(s):                         2
NUMA node0 CPU(s):                    0-51,104-155
NUMA node1 CPU(s):                    52-103,156-207
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:             Mitigation; Enhanced / Automatic IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[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-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] torch==2.5.1
[pip3] torch-tb-profiler==0.3.1
[pip3] torchaudio==2.0.1+3b40834
[pip3] torchmetrics==1.4.0.post0
[pip3] torchtyping==0.1.4
[pip3] triton==3.1.0
[pip3] vllm_nccl_cu12==2.18.1.0.4.0
[conda] numpy                     1.26.4                   pypi_0    pypi
[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-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] torch                     2.5.1                    pypi_0    pypi
[conda] torch-tb-profiler         0.3.1                    pypi_0    pypi
[conda] torchaudio                2.0.1+3b40834            pypi_0    pypi
[conda] torchmetrics              1.4.0.post0              pypi_0    pypi
[conda] torchtyping               0.1.4                    pypi_0    pypi
[conda] triton                    3.1.0                    pypi_0    pypi
[conda] vllm-nccl-cu12            2.18.1.0.4.0             pypi_0    pypi
```

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