[FSDP][BE] Post-backward hook and `FlatParamHandle` dtype cleanup
Labels: oncall: distributed, triaged, module: fsdp

**Context**
I have been working on moving all gradient tensor allocations in the post-backward hook to the default computation stream out of the post-backward stream to see if that improves memory fragmentation. In this process, I have found that our post-backward hook code and the related mixed precision code is not in the healthiest state: There is non-intuitive branching, non-obvious variable names, and incorrect comments.

Because we may not land the change to move the gradient tensor allocations unless it shows meaningful improvement, this issue is meant to track cleaning up the code regardless.

**`FlatParamHandle` Work Items**
- [x] Rework the storage of parameter and gradient reduction dtypes on the `FlatParamHandle`. Specifically, save the original parameter dtype (`_orig_param_dtype`), the parameter dtype for forward/backward (`_fwd_bwd_param_dtype`), and the gradient reduction dtype (`_reduce_dtype`). If parameter mixed precision is not enabled, then `_fwd_bwd_param_dtype == _orig_param_dtype`, and otherwise, it is the low precision dtype. This separation makes the post-backward reduction logic simpler.
    - The gradient computed by autograd must be of `_fwd_bwd_param_dtype`.
    - The input tensors to the gradient reduction must be of `_reduce_dtype`.
    - The gradient must be cast to `_orig_param_dtype` after the reduction and before the optimizer step.
    - We must perform a tensor allocation for the gradient reduction destination tensor (1) if the gradient needs padding or (2) if the gradient computed by autograd does not have dtype equal to `_reduce_dtype`.
        - (1) can only be true for sharded strategies. `NO_SHARD` never needs padding.
        - (2) is equivalent to checking `_fwd_bwd_param_dtype != _reduce_dtype`.
    - We must perform a tensor allocation for the gradient reduction source tensor if using a sharded strategy. This tensor is padded.
    - We must perform a tensor allocation for the gradient tensor used by the optimizer if `_orig_param_dtype != _reduce_dtype`.
    - For sharded strategies, if `keep_low_precision_grads`, we currently cast the gradient back to `_reduce_dtype` in the post-backward final callback via `prepare_gradient_for_optim()`. IIUC, this is done because we need an existing gradient in the `.grad` field to be able to use the `grad.data = grad.to(dtype)` hack that bypasses the check for matching parameter/gradient dtype. If we want to support applying the optimizer in the backward pass with `keep_low_precision_grads == True`, then we need to move this to the post-backward hook, in which case, this contributes one more tensor allocation.
- [ ] Rename `keep_low_precision_grads` to `keep_grads_in_reduce_dtype` in  the `MixedPrecision` dataclass. This is BC breaking. However, this most directly reflects the semantics (since `reduce_dtype` is allowed to be any dtype, not necessary a low precision one). As a niche example, if a user passes `MixedPrecision(param_dtype=torch.float16, reduce_dtype=torch.float64, keep_low_precision_grads=True)`, in the current implementation, the post-backward hook will cast the gradients to FP32 (to match `_orig_param_dtype`), and the post-backward final callback will cast the gradients to FP64 (to match `_reduce_dtype`).
    https://github.com/pytorch/pytorch/blob/e33f1eeeb73a8d680b8aae7944011389f76faaff/torch/distributed/fsdp/flat_param.py#L1107-L1113

**Post-Backward Hook Work Items**
- [x] Rename `sharded_grad` to `grad_to_offload` since the variable's only purpose is to be offloaded to CPU.
https://github.com/pytorch/pytorch/blob/e33f1eeeb73a8d680b8aae7944011389f76faaff/torch/distributed/fsdp/_runtime_utils.py#L647
https://github.com/pytorch/pytorch/blob/e33f1eeeb73a8d680b8aae7944011389f76faaff/torch/distributed/fsdp/_runtime_utils.py#L654
https://github.com/pytorch/pytorch/blob/e33f1eeeb73a8d680b8aae7944011389f76faaff/torch/distributed/fsdp/_runtime_utils.py#L662-L664
- [x] Fix the comment about why we need to call `record_stream(grad_to_offload, post_backward_stream)` that I added previously. Assume that we do not move any tensor allocations. For sharded strategies, the relevant gradient tensor is allocated in the post-backward stream and consumed in the post-backward stream (offloading to CPU). For `NO_SHARD`, the relevant gradient tensor is allocated in the post-backward stream if it requires a dtype change and is allocated in the default computation stream (by autograd) otherwise. The comment has the directionality wrong and does not mention that the `record_stream` may be a no-op depending on the case.
https://github.com/pytorch/pytorch/blob/e33f1eeeb73a8d680b8aae7944011389f76faaff/torch/distributed/fsdp/_runtime_utils.py#L665-L667

**Communication Hook Work Items**
- I do not have explicit work items here. I only want to mention that the effort to move tensor allocations to the default stream does not apply for communication hook allocations. Given our current design, any allocations in communication hooks will happen in the post-backward stream without choice.



cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501