Repository: pytorch
Issue: Support backward hook optimizers in FSDP
Labels: oncall: distributed,triaged,module: fsdp

### 🚀 The feature, motivation and pitch

I'm currently optimizing the [Lightning reference implementation of LLaMA](https://github.com/Lightning-AI/lit-llama) (7B), although the following will be generally applicable to any LLM with high memory pressure. The default configuration (24GB model sharded across 4x40GB A100s) is just on the cusp of being able to run (Between weights, 2 AdamW states, and gradients for the shard **logical** GPU memory caps out around 27GB, although I don't think that captures comms buffers) and the profile shows clear signs of allocator thrashing.

After a bit of hacking I came up with this monstrosity:
```
def add_optimizer_hooks(
    model,
    optimizers: Dict[torch.nn.Parameter, torch.optim.Optimizer],  # Per-parameter optimizers
):
    """Ugly FSDP analog to torch.distributed.optim._apply_optimizer_in_backward
    
    FSDP changes acc_grad every step, so we need to apply this before *each* `backward()`
    call, unlike the normal recipe where we only apply it once.
    """

    param_handles = torch.distributed.fsdp._traversal_utils._get_fsdp_handles(model)
    assert set(model.parameters()) == {i.flat_param for i in param_handles} == set(optimizers.keys())

    # We need to use the post backward stream so updates apply gradients are accumulated
    stream = torch.distributed.fsdp._common_utils._get_module_fsdp_state(model)._streams["post_backward"]

    for h in param_handles:
        # We're going to call this early, so if we don't override to a no-op FSDP proper will call it again and assert fail.
        h.prepare_gradient_for_optim = lambda: None

        p = h.flat_param
        assert hasattr(p, "_post_backward_hook_state")
        fsdp_acc_grad, _ = p._post_backward_hook_state

        def _opt_hook(optimizer, p, h, *_unused):
            assert p._post_backward_called

            with torch.cuda.stream(stream):
                # Use the class to get at `prepare_gradient_for_optim`
                h.__class__.prepare_gradient_for_optim(h)
                assert p.grad is not None
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)  # Cool that this is now the default
                assert p.grad is None

        fsdp_acc_grad.register_hook(functools.partial(_opt_hook, optimizers[p], p, h))
```


<img width="350" alt="Screenshot 2023-04-05 at 8 43 23 AM" src="https://user-images.githubusercontent.com/13089297/230133300-0865ec63-45e5-416c-aa66-091c16c3ef3e.png"> <img width="339" alt="Screenshot 2023-04-05 at 8 43 47 AM" src="https://user-images.githubusercontent.com/13089297/230133388-ab6303bb-a2fa-4c22-baff-a1f075bf5a68.png">

More importantly, that's enough to get out of the high contention regime and decreases the step time close to an order of magnitude. But given how much FSDP internal state I had to crack open to get things running (and I'm sure I missed plenty...) it's really only suitable as a PoC.

CC @rohan-varma @albanD @zdevito

### Alternatives

I know there's been more general discussion of creating optimizers on the fly so there might be a better alternative to the big list of single Tensor optimizers.

### Additional context

_No response_

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu