Introduce FSDPv2 (#6122)
Summary:
This patch introduce a PoC of FSDPv2. The full design doc is here: go/fsdp_v2. A real world use case can be found: https://github.com/pytorch-tpu/transformers/tree/llama2-spmd-fsdp.
Test Plan:
python test/spmd/test_fsdp_v2.py
