Introduce shard-merging util for FSDP

[muellerzr]

What does this PR do?

This PR brings a shard-merging util from PyTorch into accelerate. The reason for doing so is we can specifically save these weights as safetensors rather than .bin.

Applicable to users who use SHARDED_STATE_DICT during FSDP

Note that this is a cpu-bound process. So you just need enough RAM to load the model into memory you want to merge.

New API

Command Line

accelerate merge-weights {checkpoint_dir} {output_dir}

• You can specify if you want to remove the checkpoint_dir via --remove_checkpoint_dir
• By default we save these as safetensors. If you want to disable it, pass in --use_pytorch to save as .bin

Python API

from accelerate.utils import merge_fsdp_weights
output_dir = merge_fsdp_weights(
    checkpoint_dir="my_fsdp_checkpoints",
    output_path="somepath",
    use_safetensors=True,
    remove_checkpoint_dir=True
)

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 or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• 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.

@pacman100 @SunMarc

TODO:

• Finish saving logic
• Write tests
• Write documentation

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[pacman100]

Thank you, @muellerzr, for this nice utility function to merge the sharded FSDP state dicts! It would be great to have an accompanying test for this.

[muellerzr]

@pacman100 hope the tests are to your liking :)

[pacman100]

Thank you @muellerzr for iterating! Very useful feature.

[SunMarc]

Thanks for adding this and creating a great test for that ! Left a few minor comments

[SunMarc]

nice