Implement FSDP for training large datasets

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/