Repository: pytorch
Issue: [DTensor] use P2P for complicated transformation when redistributing tensor
Labels: oncall: distributed,triaged

### 🚀 The feature, motivation and pitch

# Motivation

For complicated `DTensor` redistribution (e.g. `[S(0), S(1)] -> [S(1), S(0)]`), it's likely that only GPU1 and GPU2 need to communicate (when tensor and mesh are both square) and can be achieved by P2P operations.

The current implementation only applies rule-based redistribution, for the above case, it does the following:
1. `S(1)` -> `R` on mesh dim 1
2. `S(0)` -> `S(1)` on mesh dim 0
3. `R` -> `S(0)` on mesh dim 1

Instead, P2P does:

1. rank1 sends local_tensor to rank2
2. rank2 sends local_tensor to rank1

And they can be done concurrently since there is no data dependency. This helps to optimize both communication volume and intermediate tensor buffer size.

# Experiment

As discussed in #132751, one major concern is that this method cannot utilize comm collectives and might suffer when communicating between 2 nodes. After conducting simple experiments, I believe it still benefits the communication time considering the reduced communication volume.

## Setup
The above case was conducted with a `4*4` mesh (2 nodes, 8 GPUs each, NV8 fully connected and InfiniBand is used), `rule` refers to `main` implementation, and `p2p` refers to the above method. The 2-d square tensor size increases along the x-axis, and execution time is recorded along the y-axis.

## Result

### `[S(0), S(1)] -> [S(1), S(0)]`
![case1_16_gpus_p2p](https://github.com/user-attachments/assets/4923d08b-e557-451f-aefa-c64fe72cd033)

### `[S(0), S(0)] -> [S(1), S(1)]`
![case2_16_gpus_p2p](https://github.com/user-attachments/assets/043d6991-73b7-4e2a-bb85-f5f7218893d2)

# Implementation: Doing P2P and rule-based in a hybrid way

[benchmark file](https://github.com/botbw/pytorch/blob/main/bench1.py)

I do observe P2P suffers in some cases, especially when the redistribution can be done using a single collective. Thus I implemented a draft redistribute function such that it utilizes collectives whenever possible, and uses P2P to handle the rest.

I roughly tested the implementation with different mesh settings: `2*2`, `4*2`, `8*2`, `4*4` (1 node for the first 2 settings and 2 nodes for the last 2), the microbenchmark was done with `(32, 8192, 4096)` tensor redistributing with any placement combination from [R, S(0), S(1), S(2)] (4 ** 4 in total). The performance is as follows (green dots indicate that this implementation reduces communication time):

![torch Size( 2, 2 )_optimization](https://github.com/user-attachments/assets/6a8213fd-0eab-43e4-a690-34f36e1510c5)
![torch Size( 4, 2 )_optimization](https://github.com/user-attachments/assets/d8f8286b-2c3b-46db-912b-df605be8aa43)
![torch Size( 2, 8 )_optimization](https://github.com/user-attachments/assets/ff5207fc-266e-4091-a47b-e83d8f3f6a37)
![torch Size( 4, 4 )_optimization](https://github.com/user-attachments/assets/7deea5ba-9115-4565-a4a7-68ee227a1373)

And the hybrid implementation doesn't hurt the performance we got from *Experiment* section:
### `[S(0), S(1)] -> [S(1), S(0)]`
![case1_16_gpus_hybrid](https://github.com/user-attachments/assets/49133f47-0008-4b36-bf25-5d2c32de3cfd)

### `[S(0), S(0)] -> [S(1), S(1)]`
![case2_16_gpus_hybrid](https://github.com/user-attachments/assets/21b0eca8-9d18-44ef-8b37-685908c533f4)

# Other

I used additional buffers when using P2P for easier implementation so I didn't test on memory optimization. If you guys think P2P makes sense considering the experiment above, do let me know and I'm happy to work on this.

You can find the above experiment code and draft implementation in [this fork](https://github.com/botbw/pytorch/), the p2p implementation passed [test_redistribute_p2p.py](https://github.com/botbw/pytorch/blob/main/test/distributed/_tensor/test_redistribute_p2p.py), which is modified from `test_redistribute.py`

cc: @wanchaol 

### Alternatives

_No response_

### Additional context

_No response_

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