Implement FSDP for training large datasets

[TheRealVish]

Definition of done: Implement training large models using FSDP to accelerate training on large datasets.

Reference: https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/

[lucylililiwang]

Hi, Can I please work on this issue? Thank you!