Bunch of FSDP improvements (#3671)

Commit

200 days ago

Bunch of FSDP improvements (#3671) * Feat: split tests * Feat: finito * Fix * Final, tests pass