Repository: pytorch
Issue: [DTensor] Error in distribute_module with module._apply
Labels: oncall: distributed,triaged

### 🐛 Describe the bug

In DTensor's README, the use case of distribute_module is given, but I found an error when I tried it. The specific code is as follows:
```python
    mesh = DeviceMesh(device_type="cuda", mesh=[0, 1])

    def shard_params(mod_name, mod, mesh):
        rowwise_placement = [Shard(0)]
        def to_dist_tensor(t):
            return distribute_tensor(t, mesh, rowwise_placement)
        mod._apply(to_dist_tensor)

    sharded_module = distribute_module(MyModule(), mesh, partition_fn=shard_params)
```
The example code in README is:
https://github.com/pytorch/pytorch/blob/683753fb0fc842afa9c43bc2a4576053af529429/torch/distributed/_tensor/README.md?plain=1#L124-L129


And the error is:
```
terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Exception raised from c10_cuda_check_implementation at path/c10/cuda/CUDAException.cpp:44 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0xae (0x7f3d8f9b1e6e in path/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0xf3 (0x7f3d8f967a0b in path/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x3f2 (0x7f3d9e168e72 in path/torch/lib/libc10_cuda.so)
frame #3: <unknown function> + 0x14ef4 (0x7f3d9e131ef4 in path/torch/lib/libc10_cuda.so)
frame #4: <unknown function> + 0x18358 (0x7f3d9e135358 in path/torch/lib/libc10_cuda.so)
frame #5: <unknown function> + 0x1870c (0x7f3d9e13570c in path/torch/lib/libc10_cuda.so)
frame #6: <unknown function> + 0x479024 (0x7f3d6d079024 in path/torch/lib/libtorch_python.so)
frame #7: c10::TensorImpl::~TensorImpl() + 0xd (0x7f3d8f98d79d in path/torch/lib/libc10.so)
frame #8: <unknown function> + 0x72cfd8 (0x7f3d6d32cfd8 in path/torch/lib/libtorch_python.so)
frame #9: THPVariable_subclass_dealloc(_object*) + 0x2b5 (0x7f3d6d32d305 in path/torch/lib/libtorch_python.so)
frame #10: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x4e07e0]
frame #11: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x4e085a]
frame #12: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x4e085a]
frame #13: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x4f1908]
frame #14: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x4f1569]
frame #15: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x4f152d]
frame #16: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x4ce758]
frame #17: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x506c83]
frame #18: _PyEval_EvalFrameDefault + 0x2412 (0x4da402 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #19: _PyEval_EvalCodeWithName + 0x2f1 (0x4d6fb1 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #20: _PyFunction_Vectorcall + 0x19c (0x4e807c in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #21: _PyEval_EvalFrameDefault + 0x6b2 (0x4d86a2 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #22: _PyFunction_Vectorcall + 0x106 (0x4e7fe6 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #23: _PyEval_EvalFrameDefault + 0x399 (0x4d8389 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #24: _PyEval_EvalCodeWithName + 0x2f1 (0x4d6fb1 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #25: _PyFunction_Vectorcall + 0x19c (0x4e807c in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #26: _PyEval_EvalFrameDefault + 0x1150 (0x4d9140 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #27: _PyEval_EvalCodeWithName + 0x2f1 (0x4d6fb1 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #28: PyEval_EvalCodeEx + 0x39 (0x585d79 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #29: PyEval_EvalCode + 0x1b (0x585d3b in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #30: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x5a5a91]
frame #31: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x5a4a9f]
frame #32: PyRun_StringFlags + 0x7b (0x5a24ab in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #33: PyRun_SimpleStringFlags + 0x3b (0x4509c4 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #34: Py_RunMain + 0x278 (0x5a1ad8 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #35: Py_BytesMain + 0x39 (0x579dd9 in /home/shao/miniconda3/envs/torchmaster/bin/python)
frame #36: <unknown function> + 0x29d90 (0x7f3d9f629d90 in /lib/x86_64-linux-gnu/libc.so.6)
frame #37: __libc_start_main + 0x80 (0x7f3d9f629e40 in /lib/x86_64-linux-gnu/libc.so.6)
frame #38: /home/shao/miniconda3/envs/torchmaster/bin/python() [0x579c8d]

Traceback (most recent call last):
  File "dtensor_high.py", line 47, in <module>
    main()
  File "dtensor_high.py", line 44, in main
    mp.spawn(example, args=(world_size,), nprocs=world_size)
  File "path/torch/multiprocessing/spawn.py", line 239, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
  File "path/torch/multiprocessing/spawn.py", line 197, in start_processes
    while not context.join():
  File "path/torch/multiprocessing/spawn.py", line 160, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "path/torch/distributed/c10d_logger.py", line 47, in wrapper
    return func(*args, **kwargs)
  File "path/torch/distributed/distributed_c10d.py", line 3136, in scatter
    work = default_pg.scatter(output_tensors, input_tensors, opts)
RuntimeError: CUDA error: invalid argument
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "path/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/home/test/dtensor/dtensor_high.py", line 32, in example
    sharded_module = distribute_module(MyModule(), mesh, partition_fn=shard_params)
  File "path/torch/distributed/_tensor/api.py", line 519, in distribute_module
    partition_fn(name, submod, device_mesh)
  File "/home/test/dtensor/dtensor_high.py", line 30, in shard_params
    mod._apply(to_dist_tensor)
  File "path/torch/nn/modules/module.py", line 824, in _apply
    param_applied = fn(param)
  File "/home/shao/workspace/test/dtensor/dtensor_high.py", line 29, in to_dist_tensor
    return distribute_tensor(t, mesh, rowwise_placement)
  File "path/torch/distributed/_tensor/api.py", line 433, in distribute_tensor
    output = placement._shard_tensor(local_tensor, device_mesh, idx)
  File "path/torch/distributed/_tensor/placement_types.py", line 168, in _shard_tensor
    mesh.scatter(output, scatter_list, mesh_dim=mesh_dim)
  File "path/torch/distributed/_tensor/device_mesh.py", line 289, in scatter
    fut = scatter(
  File "path/torch/distributed/c10d_logger.py", line 52, in wrapper
    "args": f"{args}, {kwargs}",
  File "path/torch/_tensor.py", line 427, in __repr__
    return torch._tensor_str._str(self, tensor_contents=tensor_contents)
  File "path/torch/_tensor_str.py", line 669, in _str
    return _str_intern(self, tensor_contents=tensor_contents)
  File "path/torch/_tensor_str.py", line 600, in _str_intern
    tensor_str = _tensor_str(self, indent)
  File "path/torch/_tensor_str.py", line 352, in _tensor_str
    formatter = _Formatter(get_summarized_data(self) if summarize else self)
  File "path/torch/_tensor_str.py", line 137, in __init__
    nonzero_finite_vals = torch.masked_select(
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

### Versions

Collecting environment information...
PyTorch version: 2.1.0a0+gitd9c8f9a
Is debug build: False
CUDA used to build PyTorch: 11.5
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04.1) 11.3.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.8.16 (default, Mar  2 2023, 03:21:46)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.19.0-43-generic-x86_64-with-glibc2.17
Is CUDA available: False
CUDA runtime version: 11.5.119
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: False

CPU:
arch：                           x86_64
Versions of relevant libraries:
[pip3] torch==2.1.0a0+gitd9c8f9a
[conda] torch                     2.1.0a0+gitd9c8f9a           dev_0    <develop>
[conda] torchtriton               2.0.0                      py38    pytorch
[conda] torchvision               0.15.2               py38_cu117    pytorch

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu