Use huggingface_hub helper function to split state dict

[SunMarc]

What does this PR do ?

This PR uses the helper function from huggingface-hub to split a state dict into shards instead of shard_checkpoint. This will make maintenance easier for all HF libraries. Similar PR have been created in accelerate and diffusers.

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

[LysandreJik]

This looks good in theory to me; @Wauplin, could you give this a quick look as well?

[SunMarc]

Could you make get_storage_id public/documented since it is used in multiple places (not only when cleaning the state dict) @Wauplin ? Thanks !

[Wauplin]

Looks good to me as well! Thanks for taking care of this @SunMarc :)

If I'm not mistaken there are similar methods for tf / flax in transformers that could benefit from it. Can be done in separate PRs though.

[SunMarc]

If I'm not mistaken there are similar methods for tf / flax in transformers that could benefit from it. Can be done in separate PRs though.

I'll have a look and do it in another PR after this PR gets merged and tested !

Just a quick question @Wauplin , the current helper function in huggingface-hub introduced a regression in transformers (not able to handle kiB, miB). Is is fine ? If so, I can modify the failing test : test_checkpoint_sharding_local_bin

[Wauplin]

Could you make get_storage_id public/documented since it is used in multiple places (not only when cleaning the state dict) @Wauplin ?

@SunMarc Done in huggingface/huggingface_hub#2304

he current helper function in huggingface-hub introduced a regression in transformers (not able to handle kiB, miB). Is is fine ?

I'm not a transformers's maintainer but I would assume it's fine. I'm quite convinced those parameters are never used (expect for the default value) and even less using kib, mib, etc. notation. If others feel strongly about it, we could somehow mitigate it directly in transformers (typically with a max_size.lower().replace("mib", "mb").replace("gib", "gb") + a deprecation warning). But I'd prefer not to update huggingface_hub for those.

[SunMarc]

I'm not a transformers's maintainer but I would assume it's fine. I'm quite convinced those parameters are never used (expect for the default value) and even less using kib, mib, etc. notation. If others feel strongly about it, we could somehow mitigate it directly in transformers (typically with a max_size.lower().replace("mib", "mb").replace("gib", "gb") + a deprecation warning). But I'd prefer not to update huggingface_hub for those.

WDYT @amyeroberts ?

[amyeroberts]

but I would assume it's fine. I'm quite convinced those parameters are never used (expect for the default value) and even less using kib, mib, etc. notation. If others feel strongly about it, we could somehow mitigate it directly in transformers (typically with a max_size.lower().replace("mib", "mb").replace("gib", "gb") + a deprecation warning). But I'd prefer not to update huggingface_hub for those.

@Wauplin @SunMarc Sounds good to me! Agreed, I don't think we need to handle a deprecation cycle here