Repository: pytorch
Issue: Cannot create and distribute array in torch.func.grad
Labels: oncall: distributed,triaged,module: functorch,module: dtensor

While experimenting with PyTorch's DTensor functionality in my single node "cpu" environment, I've encountered an unexpected limitation. Specifically, in the provided code example, I'm attempting to use torch.func.grad to compute gradients of a forward function that involves distributed tensors (DTensors) created with distribute_tensor. 

It raises :
"RuntimeError: You are attempting to call Tensor.requires_grad_() (or perhaps using torch.autograd.functional.* APIs) inside of a function being transformed by a functorch transform. This is unsupported, please attempt to use the functorch transforms (e.g. grad, vjp, jacrev, jacfwd, hessian) or call requires_grad_() outside of a function being transformed instead."

Repro:

```python

from functools import partial

import os
import torch
import torch.distributed as dist
from torch.func import grad
from torch.distributed._tensor import DeviceMesh, init_device_mesh, distribute_tensor, Shard, Replicate


def forward(weight, input, device_mesh: DeviceMesh):
    res = input @ weight
    rand = distribute_tensor(torch.rand(res.shape), device_mesh, [Replicate()])
    res += rand
    return res.mean()

def run():
    rank = int(os.environ['RANK'])
    world_size = int(os.environ['WORLD_SIZE'])

    # Initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    device_mesh = init_device_mesh("cpu", (world_size,))

    input = torch.ones(256, 128)
    weight = torch.ones(128, 256)

    input_d = distribute_tensor(input, device_mesh, [Shard(0)])
    weight_d = distribute_tensor(weight, device_mesh, [Replicate()])

    _forward = partial(forward, input=input_d, device_mesh=device_mesh)
    grads = grad(_forward)(weight_d)
    
    if rank == 0:
        print(grads)

    # Clean up the process group
    dist.destroy_process_group()

if __name__ == "__main__":
    run()
```

Terminal script: 
`torchrun --nproc_per_node=4 test.py`

### Versions

PyTorch version: 2.4.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.6.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.1.0.2.5)
CMake version: version 3.28.1
Libc version: N/A

Python version: 3.11.0 (main, Mar  1 2023, 12:33:14) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.6.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2 Pro

Versions of relevant libraries:
[pip3] numpy==1.26.0
[pip3] torch==2.4.0
[conda] numpy                     1.26.0                   pypi_0    pypi
[conda] torch                     2.4.0                    pypi_0    pypi

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @zou3519 @Chillee @samdow @kshitij12345 @tianyu-l @janeyx99