[FSDP] Option to keep grads in lower prec (#85223)
Reland of https://github.com/pytorch/pytorch/pull/85134, fix is to use fp16 instead of bf16 which is not supported on all platforms.
Differential Revision: [D39565189](https://our.internmc.facebook.com/intern/diff/D39565189/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85223
Approved by: https://github.com/awgu
