Add support for FSDP2

[lewtun]

What does this PR do?

This PR adds support for FSDP2 by (a) bumping the lower bound on accelerate and (b) updating prepare_fsdp().

The FSDP2 API is still in beta, so we can hold off on merging this if we feel the accelerate version is too high to bump right now. The alternative would be to check the accelerate version being used and only allow FSDP2 to be run for versions > 1.6.0

In particular, FULL_STATE_DICT is not supported in FSDP2 yet, so it is not possible to save transformers friendly checkpoints. This is currently being worked on in accelerate, so we could hold off until then

A quick test on SFT shows the two FSDP versions are almost identical, modulo grad norms:

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[fabianlim]

@lewtun if its FSDP2, its it gauranteed that summoning named_parameters will give the full params? This looking at the this

[qgallouedec]

if we feel the accelerate version is too high

I think it's ok to bump high since trl is also beta

The alternative would be to check the accelerate version being used and only allow FSDP2 to be run for versions > 1.6.0

So it seems that mergekit requires accelerate>=1.3.0,<1.4.dev0. So the above solution is probably the best.
Maybe just documenting it is enough? Something like # requires accelerate>=1.6 on top of FSDP2 config file?

[qgallouedec]

I tried myself: looks good!

and for the record, FSDP is way faster in my case: