Support for FSDP2

[sjain-profluent]

trafficstars

🚀 Feature Request

Will composer be supporting FSDP2 parallelism available in latest version of pytorch (https://pytorch.org/docs/stable/distributed.fsdp.fully_shard.html)?

Motivation

Pytorch FSDP2 allows better handling of frozen vs non-frozen params which can be difficult in current version of FSDP. The interface is also cleaner and easy to reason about.

[Optional] Implementation

Additional context