Use shard saving from huggingface_hub

[SunMarc]

What does this PR do?

This PR replaces the shard saving utility function to use the one from huggingface hub. The goal is to make it easier to maintain this piece of code that it used across many libraries such as transformers, peft and diffusers.

TODO:

• Decide if we should also use id_tensor_storage from huggingface_hub or not (it is called get_storage_id there). We use it in clean_state_dict_for_safetensors to preprocess the state_dict before sharding it.
  Edit: after some thoughts, it is better to just use the get_storage_id from huggingface_hub. I will use it when it becomes public.
• Also, we need to wait for the release on huggingface_hub, after this PR gets merged.

This is POC for now. I need to test it a bit and also see if it is doable to perform the change in other libraries.

cc @Wauplin @sayakpaul

[sayakpaul]

Could this be a done in a threaded manner or not really so as to preserve readability?

[SunMarc]

I'll have a look !

[muellerzr]

Not sure why threading may matter here, this only happens on the main accelerate process on purpose

[muellerzr]

Thanks for updating us to use the 🤗 Hub!

This looks great! Made some comments (in others) on deprecation and some general advice. Otherwise looks great (once tests pass)

[Wauplin]

We use it in clean_state_dict_for_safetensors to preprocess the state_dict before sharding it.

Haven't tried locally but is clean_state_dict_for_safetensors even needed now that you use huggingface_hub's implementation? If yes, we might think of a way to clean this in split_torch_state_dict_into_shards directly? (or later, if we implement a real save_state_dict method in huggingface_hub.

Otherwise, if accelerate.id_tensor_storage and huggingface_hub.get_storage_id does the same job, better to use the huggingface_hub one. I can make it public/documented if needed.

Also, we need to wait for the release on huggingface_hub, after this huggingface/huggingface_hub#2286 gets merged.

It's merged. I can make a patch release on Monday :)

[Wauplin]

@SunMarc Patch release 0.23.2 has been shipped: https://github.com/huggingface/huggingface_hub/releases/tag/v0.23.2. Feel free to use it :)

[SunMarc]

Haven't tried locally but is clean_state_dict_for_safetensors even needed now that you use huggingface_hub's implementation. If yes, we might think of a way to clean this in split_torch_state_dict_into_shards directly? (or later, if we implement a real save_state_dict method in huggingface_hub.

Yes, we still need that. I think it will be better to include this is a real save_state_dict method in huggingface_hub (have it as a separate function). I will check if we can do something that also works for transformers library.

Otherwise, if accelerate.id_tensor_storage and huggingface_hub.get_storage_id does the same job, better to use the huggingface_hub one. I can make it public/documented if needed.

Sounds good ! However, it think that if you have a function that clean the state_dict just like here, we don't even need to make it public since we only use that function when cleaning the state dict + splitting

[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.

[Wauplin]

Looks good to me! Thanks for integrating this :)

[muellerzr]

Merging now that release has been cut (so this is not included while we test!)