[FSDP] New rate limiter and memory reuse system
Labels: release notes: distributed (fsdp)

### 1. Motivation
Today's FSDP uses CPU-GPU synchronization to limit issuance rate of all-gather. This stalls the CPU thread from issuance new tasks when waiting for GPU computation to finish. See "CPU Thread Activities".
<img width="910" alt="Screenshot 2023-09-28 at 1 08 16 AM" src="https://github.com/pytorch/pytorch/assets/6676466/1d07138f-b8fc-4e06-8bfb-69b8575c2df4">

And when the CPU thread resumes, there is an 235 us overhead between the time the synchronization is met and the time a new all-gather is launched. See "Overhead (zoomed in)". This overhead would turn into GPU utilization gap.

<img width="1011" alt="Screenshot 2023-09-28 at 1 12 24 AM" src="https://github.com/pytorch/pytorch/assets/6676466/d381ea12-c5b3-4c2b-aa96-deb33475f409">

### 2. Proposal
Replace CPU synchronization with stream wait in limitation of all-gather issuance. 

We ask the unshard stream to wait for the resharding event of the "-2" all-gather, which marks the compute completion of the "-2" FSDP instance.

In order for the CUDA caching allocator to work with the stream wait, we ask it to hold the all-gather buffer till the "+2" all-gather. We do so by stashing the all-gather buffer into the existing "free event queue", and free it only when we are about to issue the "+2" all-gather (after "+2" all-gather has established wait relationship).

CPU timeline as follows:
![fsdp_rate_limiter](https://github.com/pytorch/pytorch/assets/6676466/a5089bc0-d9ab-4345-b12e-7ae3c4a62b43)

### 2.1 Additional Changes
**2.1.1 Moved `_free_storage` from reshard to +2 unshard.**
![free_storage](https://github.com/pytorch/pytorch/assets/6676466/7b912f65-95bf-4494-bbaf-fd7f2945e5b4)

**2.1.2 Consolidated `BACKWARD_PRE` and `BACKWARD_POST` options for backward prefetch.**
i.e. we would just use BACKWARD_POST, for example, the completion of layer 12 backward would trigger all-gather of layer 10 (with a stream wait, of course)

### 2.2 Backward wait system
The backward wait is implemented as follows:
![backward](https://github.com/pytorch/pytorch/assets/6676466/47e9e5fc-ff2b-4494-968f-4f0d75086dbc)

Location 1:
We need to allocate two buffers to achieve compute-communication overlap. This is auto-satisfied by the state of the free-event queue at the end of forward -- which still has two last layers stashed. In fact, by such stashing, we save two all-gathers (for the last two layers) compared to today's FSDP. 

Location 2:
What happen here are:
Compute on unshard buffer B done -> Unshard stream wait for compute stream -> Free unshard buffer B -> Switch to unshard stream -> Allocate unshard buffer B+1 -> Perform all-gather on B+1

Location 3:
Switch back to default stream, call `split_and_view` to get individual parameter -> Default stream wait for unshard stream. We switch back to default stream so that the backward op for `split_and_view` would be also on default stream. Otherwise, autograd would use a separate stream for this backward op and call `record_stream` on the param tensors.

### 3. Evaluation
![380210060_682184966891666_4956717331232594368_n](https://github.com/pytorch/pytorch/assets/6676466/7b67e75e-068b-49d7-bf18-e0ff7206b6ea)

Cc: @awgu @albanD @janeyx99