Repository: pytorch
Issue: [DTensor] enforcing foreach for optimizer causes jittering
Labels: triaged,module: dtensor

### 🐛 Describe the bug

The optimizer performance of DTensor is not as good as that of plain tensor. Enforcing `foreach` kernels can improve performance, but it introduces unexpected jittering between multiple GPUs.

`register_hook + foreach=True`

```bash
rank 4, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    634.32 ms, forward     231.13 ms, backward     365.47 ms
rank 1, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    586.95 ms, forward     183.30 ms, backward     365.26 ms
rank 0, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    585.46 ms, forward     181.09 ms, backward     364.70 ms
rank 2, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    558.44 ms, forward     157.19 ms, backward     365.53 ms
rank 5, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    574.07 ms, forward     171.96 ms, backward     365.14 ms
rank 7, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    567.33 ms, forward     163.65 ms, backward     365.11 ms
rank 3, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    619.27 ms, forward     208.34 ms, backward     363.71 ms
rank 6, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    564.20 ms, forward     159.40 ms, backward     364.88 ms
```


`register_hook + foreach=False`
```bash
rank 3, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    690.94 ms, forward     151.91 ms, backward     362.24 ms
rank 1, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    690.35 ms, forward     151.00 ms, backward     361.53 ms
rank 4, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    691.06 ms, forward     150.77 ms, backward     361.45 ms
rank 7, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    691.15 ms, forward     151.25 ms, backward     362.51 ms
rank 6, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    691.31 ms, forward     151.99 ms, backward     362.77 ms
rank 5, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    690.95 ms, forward     151.43 ms, backward     362.29 ms
rank 0, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    691.17 ms, forward     150.08 ms, backward     361.26 ms
rank 2, global input shape: [16, 3, 512, 2048], 20 iterations, latency per iteration    691.04 ms, forward     150.06 ms, backward     362.68 ms
```

reproducing script: [convnext_example.py](https://github.com/pytorch/pytorch/blob/main/torch/distributed/_tensor/examples/convnext_example.py)
```bash
python3 convnext_example.py
```

We can see jittering by passing `foreach=True` to [L201](https://github.com/pytorch/pytorch/blob/main/torch/distributed/_tensor/examples/convnext_example.py#L201)

### Versions

docker image: ghcr.io/pytorch/pytorch-nightly:latest

@wanchaol 

cc @wanchaol @XilunWu @tianyu-l