Improving save_on_cpu's performance by overlapping memory transfers with compute
Labels: oncall: distributed, module: autograd, triaged, enhancement

### 🚀 The feature, motivation and pitch

The current implementation invokes h2d/d2h on the same stream as compute, essentially blocking the GPU from doing other computation. A straightforward approach to improve the performance is to call h2d/d2h on a different stream, allowing GPU's DMA engine to work concurrently with other computation. 

I got a prototype working, but my limited knowledge of pytorch internals prevented me from making this perfect. So I would like to get some help and suggestions, and hopefully making this into pytorch. My code is shown below and the points for discussion are written in the comments.

```python3
from torch.autograd.graph import saved_tensors_hooks
import torch
from typing import Tuple
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ActivationWrapper
from torch.distributed.fsdp._limiter_utils import _FreeEventQueue


class faster_save_on_cpu(saved_tensors_hooks):
    copy = None
    queue = _FreeEventQueue()

    def __init__(self):
        if faster_save_on_cpu.copy is None:
            # make sure we only create the h2d/d2h stream once
            faster_save_on_cpu.copy = torch.cuda.Stream()

        def pack_to_cpu(tensor: torch.Tensor):
            # Here, we may redundantly save full weights tensors to cpu. 
            # For example, by default, a linear layer will save its input, weight 
            # and bias for backward. However, if we already have FSDP, we probably 
            # don't want save the **unsharded** weight and bias to CPU and want to 
            # leave them on GPU to reduce memory transfers. Is there a way to skip 
            # saving weight tensors?

            assert not tensor.is_sparse  # not yet supported
            if not tensor.is_cuda:
                return None, tensor

            self.limit_prefetch()
            faster_save_on_cpu.copy.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(faster_save_on_cpu.copy):
                cpu_tensor = tensor.to("cpu", non_blocking=True)
            self.record_event()
            tensor.record_stream(faster_save_on_cpu.copy)
            return (tensor.device, cpu_tensor)

        def unpack_from_cpu(packed: Tuple[torch.device, torch.Tensor]):
            device, cpu_tensor = packed
            if device is None:
                return cpu_tensor

            self.limit_prefetch()
            with torch.cuda.stream(faster_save_on_cpu.copy):
                # here, we must allocate tensor on the copy stream to allow
                # copy to start early. If pytorch can provide a method to
                # allocate an empty block **immediately** rather than reusing a
                # previous block on the compute stream, the problem below can be solved. 
                tensor = cpu_tensor.to(device, non_blocking=True)
            self.record_event()

            # subsequent gradient computation need to wait for h2d to complete
            torch.cuda.current_stream().wait_stream(faster_save_on_cpu.copy)

            # Now we have a problem. Ideally, we want to avoid this memcpyD2D to save time and memory, 
            # but the tensor is allocated on a different stream as opposed to the compute stream, so
            # to make this work, we need to call .record_stream after this tensor's lifetime ends
            # during backward. However, autogard seems like a blackbox to me, making it not 
            # practical to insert a .record_stream(). As a workaround, we can clone this tensor. 

            # I've tried hacking the destructor of out_tensor by subclassing and overriding Tensor's __del__,
            # but this introduced significant python overhead due to all subsequent tensors having to call
            # .record_stream(). 
            out_tensor = tensor.clone()
            tensor.record_stream(torch.cuda.current_stream())
            return out_tensor

        super().__init__(pack_to_cpu, unpack_from_cpu)

    def limit_prefetch(self):
        # limit the number of h2d/d2h submitted to reduce peak memory usage, like FSDP's limit_all_gathers option
        prev_ev = self.queue.dequeue_if_needed()
        if prev_ev:
            prev_ev.synchronize()

    def record_event(self):
        event = torch.cuda.Event()
        event.record(faster_save_on_cpu.copy)
        self.queue.enqueue(event)

```

### Alternatives

_No response_

### Additional context

_No response_

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @ezyang @albanD @zou3519 @pearu @nikitaved @soulitzer @Lezcano @Varal7