[FSDP] difference between _composable/fsdp or fsdp

I used torch2.3 and found there is torch.distributed.fsdp and torch.distributed._composable.fsdp.

Because I want to understand source codes of fsdp and try to make some contributions, I think fsdp implementations under these two directories are quite different.

Why there are two directories for fsdp and what is the core reason for this difference?

English is not my native language, sorry for my poor expression.

I believe _composable.fsdp refers to the new FSDP2 implementation, but @bdhirsh or others can correct me.

Indeed composable FSDP is being referred to as FSDP2. There are some notes that highlight the differences here torchtitan/docs/fsdp.md at main · pytorch/torchtitan · GitHub

Personally I found FSDP2 exciting because it works when you have heterogeneous data types x.com

1 Like