Repository: huggingface/huggingface_hub
Pull request #1938 (2023-12-22 14:21:29)

Framework-agnostic `split_state_dict_into_shards` helper
**(feedback welcome :pray:)**

PR based on an idea from @LysandreJik. Goal is to have a framework-agnostic helper to split a state dict into shards that can be reused in `transformers`, `peft`, `accelerate`, etc. The scope of this method is yet to be defined. At the moment, it takes as input a `state_dict`, a "max_size_per_shard" and a filename and returns an index + a "filename => tensor" mapping.


```py
>>> import json
>>> import os
>>> from safetensors.torch import save_file as safe_save_file
>>> from huggingface_hub import split_torch_state_dict_into_shards

>>> def save_state_dict(state_dict: Dict[str, torch.Tensor], save_directory: str):
...     state_dict_split = split_torch_state_dict_into_shards(state_dict)
...     for filename, tensors in state_dict_split.filename_to_tensors.values():
...         shard = {tensor: state_dict[tensor] for tensor in tensors}
...         safe_save_file(
...             shard,
...             os.path.join(save_directory, filename),
...             metadata={"format": "pt"},
...         )
...     if state_dict_split.is_sharded:
...         index = {
...             "metadata": state_dict_split.metadata,
...             "weight_map": state_dict_split.tensor_to_filename,
...         }
...         with open(os.path.join(save_directory, "model.safetensors.index.json"), "w") as f:
...             f.write(json.dumps(index, indent=2))
```

**Currently in PR:**
- take state_dict + threshold as input
- take filename as input as a pattern (e.g. `"model{suffix}.safetensors"`, `"tf_model{suffix}.h5"`, `"pytorch_model{suffix}.bin"`)
- group tensors in shards
- respect storage id if tensors have to be saved together (still to be done)
- build index with metadata (total size) + weights_map
- return shards (a list of state_dict) + index (a jsonable dict)
- support for torch, tensorflow, numpy

**Currently not in PR:**
- add framework to index (e.g. `"pt"`)
- provide filename for the index =>how?
- save tensors to files =>do we want to provide a helper for that? Especially to save index in correct file + weights in a consistent way)
- deserialize/load sharded model (will most probably never be done in `huggingface_hub`)

The current implementation is inspired by the torch implementation (see [here](https://github.com/huggingface/transformers/blob/74d9d0cebb0263a3f8ab9c280569170cc74651d0/src/transformers/modeling_utils.py#L330)). It support torch, tensorflow and numpy. This PR is still in draft so nothing is set in stone. In particular, depending on the scope we want, inputs and outputs can be adapted to be as user-friendly as possible (and still flexible).

Ping @amyeroberts @ArthurZucker @muellerzr on this (and please ping others if relevant). 

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_1938). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
Wauplin: Linking internal [slack thread](https://huggingface.slack.com/archives/C02EMARJ65P/p1704362118984569?thread_ts=1704359866.316609&cid=C02EMARJ65P) discussing it  (cc @pcuenca).
LysandreJik: Pretty excited by this PR! The current skeleton looks good to me. I'm wondering if it wouldn't make sense to make the addition of `TensorT` custom classes simpler by defining it as a class with the necessary overrideable methods, namely:
- getting the storage IDs
- getting the tensor sizes
- anything else that might be worthwhile

and then we'd start with the definition of torch/tf/numpy methods but adding extras will therefore be super simple as long as those three methods are implemented.
Wauplin: Thanks for the review and the idea @LysandreJik! Will have a look on how I could make the framework-specific stuff simpler :+1: 
Wauplin: @LysandreJik I switched to a more functional programming design (which should be easier to test and maintain as you suggested). I realized that we don't need to test if the tensor if from tensorflow, numpy or torch each time we handle a new one. Instead I'm defining 1 method per framework and it's at the user discretion to use the correct one. I don't see a situation where a user doesn't know which type of tensor they are using. WDYT of the current design? 

If that's fine, I'll clean this up, add some tests and document it a bit.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) Report
Attention: `128 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`0c272d5`)](https://app.codecov.io/gh/huggingface/huggingface_hub/commit/0c272d506e390f2d7b9dac68159595845c7f8e3b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) 82.37% compared to head [(`ec50ac9`)](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) 81.08%.

| [Files](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) | Patch % | Lines |
|---|---|---|
| [src/huggingface\_hub/serialization/\_base.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9zZXJpYWxpemF0aW9uL19iYXNlLnB5) | 0.00% | [48 Missing :warning: ](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) |
| [src/huggingface\_hub/serialization/\_torch.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9zZXJpYWxpemF0aW9uL190b3JjaC5weQ==) | 0.00% | [48 Missing :warning: ](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) |
| [src/huggingface\_hub/serialization/\_tensorflow.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9zZXJpYWxpemF0aW9uL190ZW5zb3JmbG93LnB5) | 0.00% | [19 Missing :warning: ](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) |
| [src/huggingface\_hub/serialization/\_numpy.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9zZXJpYWxpemF0aW9uL19udW1weS5weQ==) | 0.00% | [8 Missing :warning: ](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) |
| [src/huggingface\_hub/serialization/\_\_init\_\_.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9zZXJpYWxpemF0aW9uL19faW5pdF9fLnB5) | 0.00% | [5 Missing :warning: ](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1938      +/-   ##
==========================================
- Coverage   82.37%   81.08%   -1.30%     
==========================================
  Files          66       71       +5     
  Lines        8240     8368     +128     
==========================================
- Hits         6788     6785       -3     
- Misses       1452     1583     +131     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1938?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface).

LysandreJik: Yes, the current design looks great to me!
Wauplin: PR is now ready to be reviewed. I have added an example on how to use it with torch. We could add a "save" method for each framework that saves the state dict to files but that will be done in a follow-up PR.
Wauplin: Thanks for the thorough review @pcuenca! I have addressed or replied to all of your comments :)

Patch:
diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml
index 7d778183..50326714 100644
--- a/docs/source/en/_toctree.yml
+++ b/docs/source/en/_toctree.yml
@@ -82,3 +82,5 @@
       title: TensorBoard logger
     - local: package_reference/webhooks_server
       title: Webhooks server
+    - local: package_reference/serialization
+      title: Serialization
diff --git a/docs/source/en/package_reference/serialization.md b/docs/source/en/package_reference/serialization.md
new file mode 100644
index 00000000..545b486c
--- /dev/null
+++ b/docs/source/en/package_reference/serialization.md
@@ -0,0 +1,29 @@
+<!--⚠️ Note that this file is in Markdown but contain specific syntax for our doc-builder (similar to MDX) that may not be
+rendered properly in your Markdown viewer.
+-->
+
+# Serialization
+
+`huggingface_hub` contains helpers to help ML libraries to serialize models weights in a standardized way. This part of the lib is still under development and will be improved in future releases. The goal is to harmonize how weights are serialized on the Hub, both to remove code duplication across libraries and to foster conventions on the Hub.
+
+## Split state dict into shards
+
+At the moment, this module contains a single helper that takes a state dictionary (e.g. a mapping between layer names and related tensors) and split it into several shards, while creating a proper index in the process. This helper is available for `torch`, `tensorflow` and `numpy` tensors and is designed to be easily extended to any other ML frameworks.
+
+### split_numpy_state_dict_into_shards
+
+[[autodoc]] huggingface_hub.split_numpy_state_dict_into_shards
+
+### split_tf_state_dict_into_shards
+
+[[autodoc]] huggingface_hub.split_tf_state_dict_into_shards
+
+### split_torch_state_dict_into_shards
+
+[[autodoc]] huggingface_hub.split_torch_state_dict_into_shards
+
+### split_state_dict_into_shards_factory
+
+This is the underlying factory from which each framework-specific helper is derived. In practice, you are not expected to use this factory directly except if you need to adapt it to a framework that is not yet supported. If that is the case, please let us know by [opening a new issue](https://github.com/huggingface/huggingface_hub/issues/new) on the `huggingface_hub` repo.
+
+[[autodoc]] huggingface_hub.split_state_dict_into_shards_factory
diff --git a/src/huggingface_hub/__init__.py b/src/huggingface_hub/__init__.py
index f4aac31a..990e048f 100644
--- a/src/huggingface_hub/__init__.py
+++ b/src/huggingface_hub/__init__.py
@@ -287,6 +287,13 @@ _SUBMOD_ATTRS = {
     "repository": [
         "Repository",
     ],
+    "serialization": [
+        "StateDictSplit",
+        "split_numpy_state_dict_into_shards",
+        "split_state_dict_into_shards_factory",
+        "split_tf_state_dict_into_shards",
+        "split_torch_state_dict_into_shards",
+    ],
     "utils": [
         "CacheNotFound",
         "CachedFileInfo",
@@ -629,6 +636,13 @@ if TYPE_CHECKING:  # pragma: no cover
         SpaceCardData,  # noqa: F401
     )
     from .repository import Repository  # noqa: F401
+    from .serialization import (
+        StateDictSplit,  # noqa: F401
+        split_numpy_state_dict_into_shards,  # noqa: F401
+        split_state_dict_into_shards_factory,  # noqa: F401
+        split_tf_state_dict_into_shards,  # noqa: F401
+        split_torch_state_dict_into_shards,  # noqa: F401
+    )
     from .utils import (
         CachedFileInfo,  # noqa: F401
         CachedRepoInfo,  # noqa: F401
diff --git a/src/huggingface_hub/serialization/__init__.py b/src/huggingface_hub/serialization/__init__.py
new file mode 100644
index 00000000..feede2d0
--- /dev/null
+++ b/src/huggingface_hub/serialization/__init__.py
@@ -0,0 +1,19 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ruff: noqa: F401
+"""Contains helpers to serialize tensors."""
+from ._base import StateDictSplit, split_state_dict_into_shards_factory
+from ._numpy import split_numpy_state_dict_into_shards
+from ._tensorflow import split_tf_state_dict_into_shards
+from ._torch import split_torch_state_dict_into_shards
diff --git a/src/huggingface_hub/serialization/_base.py b/src/huggingface_hub/serialization/_base.py
new file mode 100644
index 00000000..48712da8
--- /dev/null
+++ b/src/huggingface_hub/serialization/_base.py
@@ -0,0 +1,168 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Contains helpers to split tensors into shards."""
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional, TypeVar
+
+from .. import logging
+
+
+TensorT = TypeVar("TensorT")
+TensorSizeFn_T = Callable[[TensorT], int]
+StorageIDFn_T = Callable[[TensorT], Optional[Any]]
+
+MAX_SHARD_SIZE = 5_000_000_000  # 5GB
+FILENAME_PATTERN = "model{suffix}.safetensors"
+
+logger = logging.get_logger(__file__)
+
+
+@dataclass
+class StateDictSplit:
+    is_sharded: bool = field(init=False)
+    metadata: Dict[str, Any]
+    filename_to_tensors: Dict[str, List[str]]
+    tensor_to_filename: Dict[str, str]
+
+    def __post_init__(self):
+        self.is_sharded = len(self.filename_to_tensors) > 1
+
+
+def split_state_dict_into_shards_factory(
+    state_dict: Dict[str, TensorT],
+    *,
+    get_tensor_size: TensorSizeFn_T,
+    get_storage_id: StorageIDFn_T = lambda tensor: None,
+    filename_pattern: str = FILENAME_PATTERN,
+    max_shard_size: int = MAX_SHARD_SIZE,
+) -> StateDictSplit:
+    """
+    Split a model state dictionary in shards so that each shard is smaller than a given size.
+
+    The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization
+    made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we
+    have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not
+    [6+2+2GB], [6+2GB], [6GB].
+
+    <Tip warning={true}>
+
+    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    size greater than `max_shard_size`.
+
+    </Tip>
+
+    Args:
+        state_dict (`Dict[str, Tensor]`):
+            The state dictionary to save.
+        get_tensor_size (`Callable[[Tensor], int]`):
+            A function that returns the size of a tensor in bytes.
+        get_storage_id (`Callable[[Tensor], Optional[Any]]`, *optional*):
+            A function that returns a unique identifier to a tensor storage. Multiple different tensors can share the
+            same underlying storage. This identifier is guaranteed to be unique and constant for this tensor's storage
+            during its lifetime. Two tensor storages with non-overlapping lifetimes may have the same id.
+        filename_pattern (`str`, *optional*):
+            The pattern to generate the files names in which the model will be saved. Pattern must be a string that
+            can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix`
+            Defaults to `"model{suffix}.safetensors"`.
+        max_shard_size (`int` or `str`, *optional*):
+            The maximum size of each shard, in bytes. Defaults to 5GB.
+
+    Returns:
+        [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them.
+    """
+    storage_id_to_tensors: Dict[Any, List[str]] = {}
+
+    shard_list: List[Dict[str, TensorT]] = []
+    current_shard: Dict[str, TensorT] = {}
+    current_shard_size = 0
+    total_size = 0
+
+    for key, tensor in state_dict.items():
+        # when bnb serialization is used the weights in the state dict can be strings
+        # check: https://github.com/huggingface/transformers/pull/24416 for more details
+        if isinstance(tensor, str):
+            logger.info("Skipping tensor %s as it is a string (bnb serialization)", key)
+            continue
+
+        # If a `tensor` shares the same underlying storage as another tensor, we put `tensor` in the same `block`
+        storage_id = get_storage_id(tensor)
+        if storage_id is not None:
+            if storage_id in storage_id_to_tensors:
+                # We skip this tensor for now and will reassign to correct shard later
+                storage_id_to_tensors[storage_id].append(key)
+                continue
+            else:
+                # This is the first tensor with this storage_id, we create a new entry
+                # in the storage_id_to_tensors dict => we will assign the shard id later
+                storage_id_to_tensors[storage_id] = [key]
+
+        # Compute tensor size
+        tensor_size = get_tensor_size(tensor)
+
+        # If this tensor is bigger than the maximal size, we put it in its own shard
+        if tensor_size > max_shard_size:
+            total_size += tensor_size
+            shard_list.append({key: tensor})
+            continue
+
+        # If this tensor is going to tip up over the maximal size, we split.
+        # Current shard already has some tensors, we add it to the list of shards and create a new one.
+        if current_shard_size + tensor_size > max_shard_size:
+            shard_list.append(current_shard)
+            current_shard = {}
+            current_shard_size = 0
+
+        # Add the tensor to the current shard
+        current_shard[key] = tensor
+        current_shard_size += tensor_size
+        total_size += tensor_size
+
+    # Add the last shard
+    if len(current_shard) > 0:
+        shard_list.append(current_shard)
+    nb_shards = len(shard_list)
+
+    # Loop over the tensors that share the same storage and assign them together
+    for storage_id, keys in storage_id_to_tensors.items():
+        # Let's try to find the shard where the first tensor of this storage is and put all tensors in the same shard
+        for shard in shard_list:
+            if keys[0] in shard:
+                for key in keys:
+                    shard[key] = state_dict[key]
+                break
+
+    # If we only have one shard, we return it => no need to build the index
+    if nb_shards == 1:
+        filename = filename_pattern.format(suffix="")
+        return StateDictSplit(
+            metadata={"total_size": total_size},
+            filename_to_tensors={filename: list(state_dict.keys())},
+            tensor_to_filename={key: filename for key in state_dict.keys()},
+        )
+
+    # Now that each tensor is assigned to a shard, let's assign a filename to each shard
+    tensor_name_to_filename = {}
+    filename_to_tensors = {}
+    for idx, shard in enumerate(shard_list):
+        filename = filename_pattern.format(suffix=f"-{idx+1:05d}-of-{nb_shards:05d}")
+        for key in shard:
+            tensor_name_to_filename[key] = filename
+        filename_to_tensors[filename] = list(shard.keys())
+
+    # Build the index and return
+    return StateDictSplit(
+        metadata={"total_size": total_size},
+        filename_to_tensors=filename_to_tensors,
+        tensor_to_filename=tensor_name_to_filename,
+    )
diff --git a/src/huggingface_hub/serialization/_numpy.py b/src/huggingface_hub/serialization/_numpy.py
new file mode 100644
index 00000000..b334a095
--- /dev/null
+++ b/src/huggingface_hub/serialization/_numpy.py
@@ -0,0 +1,67 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Contains numpy-specific helpers."""
+from typing import TYPE_CHECKING, Dict
+
+from ._base import FILENAME_PATTERN, MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
+
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+def split_numpy_state_dict_into_shards(
+    state_dict: Dict[str, "np.ndarray"],
+    *,
+    filename_pattern: str = FILENAME_PATTERN,
+    max_shard_size: int = MAX_SHARD_SIZE,
+) -> StateDictSplit:
+    """
+    Split a model state dictionary in shards so that each shard is smaller than a given size.
+
+    The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization
+    made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we
+    have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not
+    [6+2+2GB], [6+2GB], [6GB].
+
+    <Tip warning={true}>
+
+    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    size greater than `max_shard_size`.
+
+    </Tip>
+
+    Args:
+        state_dict (`Dict[str, np.ndarray]`):
+            The state dictionary to save.
+        filename_pattern (`str`, *optional*):
+            The pattern to generate the files names in which the model will be saved. Pattern must be a string that
+            can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix`
+            Defaults to `"model{suffix}.safetensors"`.
+        max_shard_size (`int` or `str`, *optional*):
+            The maximum size of each shard, in bytes. Defaults to 5GB.
+
+    Returns:
+        [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them.
+    """
+    return split_state_dict_into_shards_factory(
+        state_dict,
+        max_shard_size=max_shard_size,
+        filename_pattern=filename_pattern,
+        get_tensor_size=get_tensor_size,
+    )
+
+
+def get_tensor_size(tensor: "np.ndarray") -> int:
+    return tensor.nbytes
diff --git a/src/huggingface_hub/serialization/_tensorflow.py b/src/huggingface_hub/serialization/_tensorflow.py
new file mode 100644
index 00000000..865d9db3
--- /dev/null
+++ b/src/huggingface_hub/serialization/_tensorflow.py
@@ -0,0 +1,93 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Contains tensorflow-specific helpers."""
+import math
+import re
+from typing import TYPE_CHECKING, Dict
+
+from ._base import MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
+
+
+if TYPE_CHECKING:
+    import tensorflow as tf
+
+
+def split_tf_state_dict_into_shards(
+    state_dict: Dict[str, "tf.Tensor"],
+    *,
+    filename_pattern: str = "tf_model{suffix}.h5",
+    max_shard_size: int = MAX_SHARD_SIZE,
+) -> StateDictSplit:
+    """
+    Split a model state dictionary in shards so that each shard is smaller than a given size.
+
+    The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization
+    made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we
+    have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not
+    [6+2+2GB], [6+2GB], [6GB].
+
+    <Tip warning={true}>
+
+    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    size greater than `max_shard_size`.
+
+    </Tip>
+
+    Args:
+        state_dict (`Dict[str, Tensor]`):
+            The state dictionary to save.
+        filename_pattern (`str`, *optional*):
+            The pattern to generate the files names in which the model will be saved. Pattern must be a string that
+            can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix`
+            Defaults to `"tf_model{suffix}.h5"`.
+        max_shard_size (`int` or `str`, *optional*):
+            The maximum size of each shard, in bytes. Defaults to 5GB.
+
+    Returns:
+        [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them.
+    """
+    return split_state_dict_into_shards_factory(
+        state_dict,
+        max_shard_size=max_shard_size,
+        filename_pattern=filename_pattern,
+        get_tensor_size=get_tensor_size,
+    )
+
+
+def get_tensor_size(tensor: "tf.Tensor") -> int:
+    # Return `math.ceil` since dtype byte size can be a float (e.g., 0.125 for tf.bool).
+    # Better to overestimate than underestimate.
+    return math.ceil(tensor.numpy().size * _dtype_byte_size_tf(tensor.dtype))
+
+
+def _dtype_byte_size_tf(dtype) -> float:
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+    Taken from https://github.com/huggingface/transformers/blob/74d9d0cebb0263a3f8ab9c280569170cc74651d0/src/transformers/modeling_tf_utils.py#L608.
+    NOTE: why not `tensor.numpy().nbytes`?
+    Example:
+    ```py
+    >>> _dtype_byte_size(tf.float32)
+    4
+    ```
+    """
+    import tensorflow as tf
+
+    if dtype == tf.bool:
+        return 1 / 8
+    bit_search = re.search(r"[^\d](\d+)$", dtype.name)
+    if bit_search is None:
+        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
+    bit_size = int(bit_search.groups()[0])
+    return bit_size // 8
diff --git a/src/huggingface_hub/serialization/_torch.py b/src/huggingface_hub/serialization/_torch.py
new file mode 100644
index 00000000..e479e780
--- /dev/null
+++ b/src/huggingface_hub/serialization/_torch.py
@@ -0,0 +1,199 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Contains pytorch-specific helpers."""
+import importlib
+from functools import lru_cache
+from typing import TYPE_CHECKING, Dict, Tuple
+
+from ._base import FILENAME_PATTERN, MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
+
+
+if TYPE_CHECKING:
+    import torch
+
+
+def split_torch_state_dict_into_shards(
+    state_dict: Dict[str, "torch.Tensor"],
+    *,
+    filename_pattern: str = FILENAME_PATTERN,
+    max_shard_size: int = MAX_SHARD_SIZE,
+) -> StateDictSplit:
+    """
+    Split a model state dictionary in shards so that each shard is smaller than a given size.
+
+    The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization
+    made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we
+    have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not
+    [6+2+2GB], [6+2GB], [6GB].
+
+    <Tip warning={true}>
+
+    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    size greater than `max_shard_size`.
+
+    </Tip>
+
+    Args:
+        state_dict (`Dict[str, torch.Tensor]`):
+            The state dictionary to save.
+        filename_pattern (`str`, *optional*):
+            The pattern to generate the files names in which the model will be saved. Pattern must be a string that
+            can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix`
+            Defaults to `"model{suffix}.safetensors"`.
+        max_shard_size (`int` or `str`, *optional*):
+            The maximum size of each shard, in bytes. Defaults to 5GB.
+
+    Returns:
+        [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them.
+
+    Example:
+    ```py
+    >>> import json
+    >>> import os
+    >>> from safetensors.torch import save_file as safe_save_file
+    >>> from huggingface_hub import split_torch_state_dict_into_shards
+
+    >>> def save_state_dict(state_dict: Dict[str, torch.Tensor], save_directory: str):
+    ...     state_dict_split = split_torch_state_dict_into_shards(state_dict)
+    ...     for filename, tensors in state_dict_split.filename_to_tensors.values():
+    ...         shard = {tensor: state_dict[tensor] for tensor in tensors}
+    ...         safe_save_file(
+    ...             shard,
+    ...             os.path.join(save_directory, filename),
+    ...             metadata={"format": "pt"},
+    ...         )
+    ...     if state_dict_split.is_sharded:
+    ...         index = {
+    ...             "metadata": state_dict_split.metadata,
+    ...             "weight_map": state_dict_split.tensor_to_filename,
+    ...         }
+    ...         with open(os.path.join(save_directory, "model.safetensors.index.json"), "w") as f:
+    ...             f.write(json.dumps(index, indent=2))
+    ```
+    """
+    return split_state_dict_into_shards_factory(
+        state_dict,
+        max_shard_size=max_shard_size,
+        filename_pattern=filename_pattern,
+        get_tensor_size=get_tensor_size,
+        get_storage_id=get_storage_id,
+    )
+
+
+def get_storage_id(tensor: "torch.Tensor") -> Tuple["torch.device", int, int]:
+    """
+    Return unique identifier to a tensor storage.
+
+    Multiple different tensors can share the same underlying storage. For
+    example, "meta" tensors all share the same storage, and thus their identifier will all be equal. This identifier is
+    guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with
+    non-overlapping lifetimes may have the same id.
+
+    Taken from https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/pytorch_utils.py#L278.
+    """
+    if tensor.device.type == "xla" and is_torch_tpu_available():
+        # NOTE: xla tensors dont have storage
+        # use some other unique id to distinguish.
+        # this is a XLA tensor, it must be created using torch_xla's
+        # device. So the following import is safe:
+        import torch_xla
+
+        unique_id = torch_xla._XLAC._xla_get_tensor_id(tensor)
+    else:
+        unique_id = storage_ptr(tensor)
+
+    return tensor.device, unique_id, get_storage_size(tensor)
+
+
+def get_tensor_size(tensor: "torch.Tensor") -> int:
+    return tensor.numel() * tensor.element_size()
+
+
+@lru_cache()
+def is_torch_tpu_available(check_device=True):
+    """
+    Checks if `torch_xla` is installed and potentially if a TPU is in the environment
+
+    Taken from https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/utils/import_utils.py#L463.
+    """
+    if importlib.util.find_spec("torch_xla") is not None:
+        if check_device:
+            # We need to check if `xla_device` can be found, will raise a RuntimeError if not
+            try:
+                import torch_xla.core.xla_model as xm
+
+                _ = xm.xla_device()
+                return True
+            except RuntimeError:
+                return False
+        return True
+    return False
+
+
+def storage_ptr(tensor: "torch.Tensor") -> int:
+    """
+    Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L11C1-L20C21.
+    """
+    try:
+        return tensor.untyped_storage().data_ptr()
+    except Exception:
+        # Fallback for torch==1.10
+        try:
+            return tensor.storage().data_ptr()
+        except NotImplementedError:
+            # Fallback for meta storage
+            return 0
+
+
+def get_storage_size(tensor: "torch.Tensor") -> int:
+    """
+    Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L31C1-L41C59
+    """
+    try:
+        return tensor.untyped_storage().nbytes()
+    except AttributeError:
+        # Fallback for torch==1.10
+        try:
+            return tensor.storage().size() * _get_dtype_size(tensor.dtype)
+        except NotImplementedError:
+            # Fallback for meta storage
+            # On torch >=2.0 this is the tensor size
+            return tensor.nelement() * _get_dtype_size(tensor.dtype)
+
+
+@lru_cache()
+def _get_dtype_size(dtype: "torch.dtype") -> int:
+    """
+    Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L344
+    """
+    import torch
+
+    # torch.float8 formats require 2.1; we do not support these dtypes on earlier versions
+    _float8_e4m3fn = getattr(torch, "float8_e4m3fn", None)
+    _float8_e5m2 = getattr(torch, "float8_e5m2", None)
+    _SIZE = {
+        torch.int64: 8,
+        torch.float32: 4,
+        torch.int32: 4,
+        torch.bfloat16: 2,
+        torch.float16: 2,
+        torch.int16: 2,
+        torch.uint8: 1,
+        torch.int8: 1,
+        torch.bool: 1,
+        torch.float64: 8,
+        _float8_e4m3fn: 1,
+        _float8_e5m2: 1,
+    }
+    return _SIZE[dtype]
diff --git a/src/huggingface_hub/utils/__init__.py b/src/huggingface_hub/utils/__init__.py
index dfc2faab..7a0a9e63 100644
--- a/src/huggingface_hub/utils/__init__.py
+++ b/src/huggingface_hub/utils/__init__.py
@@ -78,6 +78,7 @@ from ._runtime import (
     is_jinja_available,
     is_notebook,
     is_numpy_available,
+    is_package_available,
     is_pillow_available,
     is_pydantic_available,
     is_pydot_available,
diff --git a/src/huggingface_hub/utils/_runtime.py b/src/huggingface_hub/utils/_runtime.py
index 82aa41ce..7915d154 100644
--- a/src/huggingface_hub/utils/_runtime.py
+++ b/src/huggingface_hub/utils/_runtime.py
@@ -69,7 +69,7 @@ def _get_version(package_name: str) -> str:
     return _package_versions.get(package_name, "N/A")
 
 
-def _is_available(package_name: str) -> bool:
+def is_package_available(package_name: str) -> bool:
     return _get_version(package_name) != "N/A"
 
 
@@ -85,7 +85,7 @@ def get_hf_hub_version() -> str:
 
 # aiohttp
 def is_aiohttp_available() -> bool:
-    return _is_available("aiohttp")
+    return is_package_available("aiohttp")
 
 
 def get_aiohttp_version() -> str:
@@ -94,7 +94,7 @@ def get_aiohttp_version() -> str:
 
 # FastAI
 def is_fastai_available() -> bool:
-    return _is_available("fastai")
+    return is_package_available("fastai")
 
 
 def get_fastai_version() -> str:
@@ -103,7 +103,7 @@ def get_fastai_version() -> str:
 
 # Fastcore
 def is_fastcore_available() -> bool:
-    return _is_available("fastcore")
+    return is_package_available("fastcore")
 
 
 def get_fastcore_version() -> str:
@@ -112,7 +112,7 @@ def get_fastcore_version() -> str:
 
 # FastAI
 def is_gradio_available() -> bool:
-    return _is_available("gradio")
+    return is_package_available("gradio")
 
 
 def get_gradio_version() -> str:
@@ -121,7 +121,7 @@ def get_gradio_version() -> str:
 
 # Graphviz
 def is_graphviz_available() -> bool:
-    return _is_available("graphviz")
+    return is_package_available("graphviz")
 
 
 def get_graphviz_version() -> str:
@@ -130,7 +130,7 @@ def get_graphviz_version() -> str:
 
 # hf_transfer
 def is_hf_transfer_available() -> bool:
-    return _is_available("hf_transfer")
+    return is_package_available("hf_transfer")
 
 
 def get_hf_transfer_version() -> str:
@@ -139,7 +139,7 @@ def get_hf_transfer_version() -> str:
 
 # Numpy
 def is_numpy_available() -> bool:
-    return _is_available("numpy")
+    return is_package_available("numpy")
 
 
 def get_numpy_version() -> str:
@@ -148,7 +148,7 @@ def get_numpy_version() -> str:
 
 # Jinja
 def is_jinja_available() -> bool:
-    return _is_available("jinja")
+    return is_package_available("jinja")
 
 
 def get_jinja_version() -> str:
@@ -157,7 +157,7 @@ def get_jinja_version() -> str:
 
 # Pillow
 def is_pillow_available() -> bool:
-    return _is_available("pillow")
+    return is_package_available("pillow")
 
 
 def get_pillow_version() -> str:
@@ -166,7 +166,7 @@ def get_pillow_version() -> str:
 
 # Pydantic
 def is_pydantic_available() -> bool:
-    if not _is_available("pydantic"):
+    if not is_package_available("pydantic"):
         return False
     # For Pydantic, we add an extra check to test whether it is correctly installed or not. If both pydantic 2.x and
     # typing_extensions<=4.5.0 are installed, then pydantic will fail at import time. This should not happen when
@@ -195,7 +195,7 @@ def get_pydantic_version() -> str:
 
 # Pydot
 def is_pydot_available() -> bool:
-    return _is_available("pydot")
+    return is_package_available("pydot")
 
 
 def get_pydot_version() -> str:
@@ -204,7 +204,7 @@ def get_pydot_version() -> str:
 
 # Tensorboard
 def is_tensorboard_available() -> bool:
-    return _is_available("tensorboard")
+    return is_package_available("tensorboard")
 
 
 def get_tensorboard_version() -> str:
@@ -213,7 +213,7 @@ def get_tensorboard_version() -> str:
 
 # Tensorflow
 def is_tf_available() -> bool:
-    return _is_available("tensorflow")
+    return is_package_available("tensorflow")
 
 
 def get_tf_version() -> str:
@@ -222,7 +222,7 @@ def get_tf_version() -> str:
 
 # Torch
 def is_torch_available() -> bool:
-    return _is_available("torch")
+    return is_package_available("torch")
 
 
 def get_torch_version() -> str:


Test patch:
diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index d676e9f0..c4d9375e 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -115,10 +115,12 @@ jobs:
               # Cannot be on same line since '_tf*' checks if tensorflow is NOT imported by default
               eval "$PYTEST ../tests/test_tf*"
               eval "$PYTEST ../tests/test_keras*"
+              eval "$PYTEST ../tests/test_serialization.py"
             ;;
 
             torch)
-              eval "$PYTEST ../tests/test_hub_mixin*"
+            eval "$PYTEST ../tests/test_hub_mixin*"
+            eval "$PYTEST ../tests/test_serialization.py"
             ;;
 
           esac
diff --git a/tests/test_hub_mixin_pytorch.py b/tests/test_hub_mixin_pytorch.py
index 9e131e1f..6aeb11a1 100644
--- a/tests/test_hub_mixin_pytorch.py
+++ b/tests/test_hub_mixin_pytorch.py
@@ -12,25 +12,12 @@ from huggingface_hub.hub_mixin import PyTorchModelHubMixin
 from huggingface_hub.utils import SoftTemporaryDirectory, is_torch_available
 
 from .testing_constants import ENDPOINT_STAGING, TOKEN, USER
-from .testing_utils import repo_name
+from .testing_utils import repo_name, requires
 
 
 if is_torch_available():
     import torch.nn as nn
 
-
-def require_torch(test_case):
-    """
-    Decorator marking a test that requires PyTorch.
-    These tests are skipped when PyTorch isn't installed.
-    """
-    if not is_torch_available():
-        return unittest.skip("test requires PyTorch")(test_case)
-    else:
-        return test_case
-
-
-if is_torch_available():
     CONFIG = {"num": 10, "act": "gelu_fast"}
 
     class DummyModel(nn.Module, PyTorchModelHubMixin):
@@ -46,7 +33,7 @@ else:
     DummyModel = None
 
 
-@require_torch
+@requires("torch")
 @pytest.mark.usefixtures("fx_cache_dir")
 class PytorchHubMixinTest(unittest.TestCase):
     cache_dir: Path
diff --git a/tests/test_serialization.py b/tests/test_serialization.py
new file mode 100644
index 00000000..004f3d3e
--- /dev/null
+++ b/tests/test_serialization.py
@@ -0,0 +1,125 @@
+from huggingface_hub.serialization import split_state_dict_into_shards_factory
+from huggingface_hub.serialization._numpy import get_tensor_size as get_tensor_size_numpy
+from huggingface_hub.serialization._tensorflow import get_tensor_size as get_tensor_size_tensorflow
+from huggingface_hub.serialization._torch import get_tensor_size as get_tensor_size_torch
+
+from .testing_utils import requires
+
+
+DUMMY_STATE_DICT = {
+    "layer_1": [6],
+    "layer_2": [10],
+    "layer_3": [30],
+    "layer_4": [2],
+    "layer_5": [2],
+}
+
+
+def _dummy_get_storage_id(item):
+    return None
+
+
+def _dummy_get_tensor_size(item):
+    return sum(item)
+
+
+def test_single_shard():
+    state_dict_split = split_state_dict_into_shards_factory(
+        DUMMY_STATE_DICT,
+        get_storage_id=_dummy_get_storage_id,
+        get_tensor_size=_dummy_get_tensor_size,
+        max_shard_size=100,  # large shard size => only one shard
+        filename_pattern="file{suffix}.dummy",
+    )
+    assert not state_dict_split.is_sharded
+    assert state_dict_split.filename_to_tensors == {
+        # All layers fit in one shard => no suffix in filename
+        "file.dummy": ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5"],
+    }
+    assert state_dict_split.tensor_to_filename == {
+        "layer_1": "file.dummy",
+        "layer_2": "file.dummy",
+        "layer_3": "file.dummy",
+        "layer_4": "file.dummy",
+        "layer_5": "file.dummy",
+    }
+    assert state_dict_split.metadata == {"total_size": 50}
+
+
+def test_multiple_shards():
+    state_dict_split = split_state_dict_into_shards_factory(
+        DUMMY_STATE_DICT,
+        get_storage_id=_dummy_get_storage_id,
+        get_tensor_size=_dummy_get_tensor_size,
+        max_shard_size=10,  # small shard size => multiple shards
+        filename_pattern="file{suffix}.dummy",
+    )
+
+    assert state_dict_split.is_sharded
+    assert state_dict_split.filename_to_tensors == {
+        # layer 4 and 5 could go in this one but assignment is not optimal, and it's fine
+        "file-00001-of-00004.dummy": ["layer_1"],
+        "file-00002-of-00004.dummy": ["layer_3"],
+        "file-00003-of-00004.dummy": ["layer_2"],
+        "file-00004-of-00004.dummy": ["layer_4", "layer_5"],
+    }
+    assert state_dict_split.tensor_to_filename == {
+        "layer_1": "file-00001-of-00004.dummy",
+        "layer_3": "file-00002-of-00004.dummy",
+        "layer_2": "file-00003-of-00004.dummy",
+        "layer_4": "file-00004-of-00004.dummy",
+        "layer_5": "file-00004-of-00004.dummy",
+    }
+    assert state_dict_split.metadata == {"total_size": 50}
+
+
+def test_tensor_same_storage():
+    state_dict_split = split_state_dict_into_shards_factory(
+        {
+            "layer_1": [1],
+            "layer_2": [2],
+            "layer_3": [1],
+            "layer_4": [2],
+            "layer_5": [1],
+        },
+        get_storage_id=lambda x: (x[0]),  # dummy for test: storage id based on first element
+        get_tensor_size=_dummy_get_tensor_size,
+        max_shard_size=1,
+    )
+    assert state_dict_split.is_sharded
+    assert state_dict_split.filename_to_tensors == {
+        "model-00001-of-00002.safetensors": ["layer_2", "layer_4"],
+        "model-00002-of-00002.safetensors": ["layer_1", "layer_3", "layer_5"],
+    }
+    assert state_dict_split.tensor_to_filename == {
+        "layer_1": "model-00002-of-00002.safetensors",
+        "layer_2": "model-00001-of-00002.safetensors",
+        "layer_3": "model-00002-of-00002.safetensors",
+        "layer_4": "model-00001-of-00002.safetensors",
+        "layer_5": "model-00002-of-00002.safetensors",
+    }
+    assert state_dict_split.metadata == {"total_size": 3}  # count them once
+
+
+@requires("numpy")
+def test_get_tensor_size_numpy():
+    import numpy as np
+
+    assert get_tensor_size_numpy(np.array([1, 2, 3, 4, 5], dtype=np.float64)) == 5 * 8
+    assert get_tensor_size_numpy(np.array([1, 2, 3, 4, 5], dtype=np.float16)) == 5 * 2
+
+
+@requires("tensorflow")
+def test_get_tensor_size_tensorflow():
+    import tensorflow as tf
+
+    assert get_tensor_size_tensorflow(tf.constant([1, 2, 3, 4, 5], dtype=tf.float64)) == 5 * 8
+    assert get_tensor_size_tensorflow(tf.constant([1, 2, 3, 4, 5], dtype=tf.float16)) == 5 * 2
+
+
+@requires("torch")
+def test_get_tensor_size_torch():
+    import torch
+
+    assert get_tensor_size_torch(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)) == 5 * 8
+    assert get_tensor_size_torch(torch.tensor([1, 2, 3, 4, 5], dtype=torch.float16)) == 5 * 2
diff --git a/tests/test_webhooks_server.py b/tests/test_webhooks_server.py
index 18b6e07d..5d18af1f 100644
--- a/tests/test_webhooks_server.py
+++ b/tests/test_webhooks_server.py
@@ -5,7 +5,7 @@ from fastapi import Request
 
 from huggingface_hub.utils import capture_output, is_gradio_available
 
-from .testing_utils import require_webhooks
+from .testing_utils import requires
 
 
 if is_gradio_available():
@@ -95,7 +95,7 @@ def test_deserialize_payload_example_without_comment() -> None:
     assert payload.comment is None
 
 
-@require_webhooks
+@requires("gradio")
 class TestWebhooksServerDontRun(unittest.TestCase):
     def test_add_webhook_implicit_path(self):
         # Test adding a webhook
@@ -133,7 +133,7 @@ class TestWebhooksServerDontRun(unittest.TestCase):
                 pass
 
 
-@require_webhooks
+@requires("gradio")
 class TestWebhooksServerRun(unittest.TestCase):
     HEADERS_VALID_SECRET = {"x-webhook-secret": "my_webhook_secret"}
     HEADERS_WRONG_SECRET = {"x-webhook-secret": "wrong_webhook_secret"}
diff --git a/tests/testing_utils.py b/tests/testing_utils.py
index 2fb6d314..8307e1fd 100644
--- a/tests/testing_utils.py
+++ b/tests/testing_utils.py
@@ -15,7 +15,11 @@ from unittest.mock import Mock, patch
 import pytest
 import requests
 
-from huggingface_hub.utils import is_gradio_available, logging, reset_sessions
+from huggingface_hub.utils import (
+    is_package_available,
+    logging,
+    reset_sessions,
+)
 from tests.testing_constants import ENDPOINT_PRODUCTION, ENDPOINT_PRODUCTION_URL_SCHEME
 
 
@@ -114,17 +118,19 @@ def require_git_lfs(test_case):
         return test_case
 
 
-def require_webhooks(test_case):
+def requires(package_name: str):
     """
-    Decorator to mark tests that requires `webhooks` extra (i.e. gradio, fastapi, pydantic).
-
-    git-lfs requires additional dependencies, and tests are skipped by default. Set the RUN_GIT_LFS_TESTS environment
-    variable to a truthy value to run them.
+    Decorator marking a test that requires PyTorch.
+    These tests are skipped when PyTorch isn't installed.
     """
-    if not is_gradio_available():
-        return unittest.skip("Skip webhook test")(test_case)
-    else:
-        return test_case
+
+    def _inner(test_case):
+        if not is_package_available(package_name):
+            return unittest.skip(f"Test requires '{package_name}'")(test_case)
+        else:
+            return test_case
+
+    return _inner
 
 
 class RequestWouldHangIndefinitelyError(Exception):