Handle shared layers in `save_torch_state_dict` + add `save_torch_model`

Wauplin commented 1 week ago

Partially resolve https://github.com/huggingface/huggingface_hub/issues/2065. Follow-up PR after https://github.com/huggingface/huggingface_hub/pull/2314.

In https://github.com/huggingface/huggingface_hub/pull/2314, we introduce save_torch_state_dict. This new PR:

adds logic to deduplicate shared layers in safetensors. This is mostly taken from safetensors's torch helpers (see here). See slack thread (private) for discussions around this. See also https://huggingface.co/docs/safetensors/torch_shared_tensors for more details.
adds save_torch_model to directly save a torch nn.Module
renames internal methods get_tf_storage_size / get_torch_storage_size and make them public + documented
tests and documentation have also been updated.

A last follow-up PR should had load_torch_state_dict / load_torch_model helpers as well to correctly reload those files, including the shared layers.

I'm pinging transformers/accelerate/diffusers cores maintainers for visibility as well. Feel free to comment if someone should be done differently.

HuggingFaceDocBuilderDev commented 1 week ago

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 commented 23 hours ago

Thanks for the reviews! Let's ship it yes :smile:

huggingface / huggingface_hub

Handle shared layers in `save_torch_state_dict` + add `save_torch_model` #2373