Repository: pytorch
Issue: [Tensor Parallel] Conv2d with replicate inputs and weights raise error in backward
Labels: oncall: distributed,triaged,module: dtensor

### 🐛 Describe the bug

Conv2d with replicate DTensor inputs and weights raise error in backward.

Error Msg:
```
[rank2]: Traceback (most recent call last):
[rank2]:   File "/path/to/tp_conv_bug.py", line 24, in <module>
[rank2]:     res_l.backward(dres)
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/_tensor.py", line 626, in backward
[rank2]:     torch.autograd.backward(
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/autograd/__init__.py", line 347, in backward
[rank2]:     _engine_run_backward(
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/autograd/graph.py", line 823, in _engine_run_backward
[rank2]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/_compile.py", line 32, in inner
[rank2]:     return disable_fn(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 744, in _fn
[rank2]:     return fn(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/distributed/tensor/_api.py", line 343, in __torch_dispatch__
[rank2]:     return DTensor._op_dispatcher.dispatch(
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/distributed/tensor/_dispatch.py", line 163, in dispatch
[rank2]:     return self._custom_op_handlers[op_call](op_call, args, kwargs)  # type: ignore[operator]
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/distributed/tensor/_tp_conv.py", line 267, in convolution_backward_handler
[rank2]:     dtensor.DTensor._op_dispatcher.sharding_propagator.propagate(op_info)
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/distributed/tensor/_sharding_prop.py", line 206, in propagate
[rank2]:     OutputSharding, self.propagate_op_sharding(op_info.schema)
[rank2]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/distributed/tensor/_sharding_prop.py", line 46, in __call__
[rank2]:     return self.cache(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/distributed/tensor/_sharding_prop.py", line 449, in propagate_op_sharding_non_cached
[rank2]:     self._wrap_output_spec_tensor_meta(
[rank2]:   File "/opt/miniconda/envs/torchdev/lib/python3.11/site-packages/torch/distributed/tensor/_sharding_prop.py", line 192, in _wrap_output_spec_tensor_meta
[rank2]:     raise ValueError(
[rank2]: ValueError: ShardingPropagator error: output 0 does not have an associated TensorMeta
```

Code:
```
import os
import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.distributed.tensor import Shard, DTensor, Replicate
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
_world_size = int(os.environ["WORLD_SIZE"])
device_mesh = init_device_mesh(device_type="cuda", mesh_shape=(_world_size,))
conv = nn.Conv2d(64, 64, 3, padding=1).train()
x = torch.randn(1, 64, 32, 32)
x_dt = DTensor.from_local(x, device_mesh, [Replicate()])
w = conv.weight.data
w_dt = torch.nn.Parameter(DTensor.from_local(w, device_mesh, [Replicate()]))

b = conv.bias.data
b_dt = torch.nn.Parameter(DTensor.from_local(b, device_mesh, [Replicate()]))

res = F.conv2d(x_dt, w_dt, b_dt, padding=1)
res_l = res.to_local()
dres = torch.rand_like(res_l)
res_l.backward(dres)

dist.barrier()
dist.destroy_process_group()
```

### Versions

[pip3] numpy==2.1.2
[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+git35c6c7c6
[pip3] torch==2.6.0.dev20241202+cu124
[pip3] torchaudio==2.5.0.dev20241202+cu124
[pip3] torchvision==0.20.0.dev20241202+cu124
[conda] numpy                     2.1.2                    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-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-triton            3.2.0+git35c6c7c6          pypi_0    pypi
[conda] torch                     2.6.0.dev20241202+cu124          pypi_0    pypi
[conda] torchaudio                2.5.0.dev20241202+cu124          pypi_0    pypi
[conda] torchvision               0.20.0.dev20241202+cu124          pypi_0    pypi

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