Support max_shard_size as string in split_state_dict_into_shards_factory (#2286)

* fix max-shard-size

* for torch also

* add tests

* Fix styling + do not support KiB

---------

Co-authored-by: Lucain Pouget <[email protected]>

Author: SunMarc

src/huggingface_hub/serialization/_base.py

@@ -14,7 +14,7 @@
 """Contains helpers to split tensors into shards."""
 
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, Optional, TypeVar
+from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
 
 from .. import logging
 
@@ -46,7 +46,7 @@ def split_state_dict_into_shards_factory(
     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,
+    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.
@@ -89,6 +89,9 @@ def split_state_dict_into_shards_factory(
     current_shard_size = 0
     total_size = 0
 
+    if isinstance(max_shard_size, str):
+        max_shard_size = parse_size_to_int(max_shard_size)
+
     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
@@ -167,3 +170,44 @@ def split_state_dict_into_shards_factory(
         filename_to_tensors=filename_to_tensors,
         tensor_to_filename=tensor_name_to_filename,
     )
+
+
+SIZE_UNITS = {
+    "TB": 10**12,
+    "GB": 10**9,
+    "MB": 10**6,
+    "KB": 10**3,
+}
+
+
+def parse_size_to_int(size_as_str: str) -> int:
+    """
+    Parse a size expressed as a string with digits and unit (like `"5MB"`) to an integer (in bytes).
+
+    Supported units are "TB", "GB", "MB", "KB".
+
+    Args:
+        size_as_str (`str`): The size to convert. Will be directly returned if an `int`.
+
+    Example:
+
+    ```py
+    >>> parse_size_to_int("5MB")
+    5000000
+    ```
+    """
+    size_as_str = size_as_str.strip()
+
+    # Parse unit
+    unit = size_as_str[-2:].upper()
+    if unit not in SIZE_UNITS:
+        raise ValueError(f"Unit '{unit}' not supported. Supported units are TB, GB, MB, KB. Got '{size_as_str}'.")
+    multiplier = SIZE_UNITS[unit]
+
+    # Parse value
+    try:
+        value = float(size_as_str[:-2].strip())
+    except ValueError as e:
+        raise ValueError(f"Could not parse the size value from '{size_as_str}': {e}") from e
+
+    return int(value * multiplier)

src/huggingface_hub/serialization/_numpy.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Contains numpy-specific helpers."""
 
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Union
 
 from ._base import FILENAME_PATTERN, MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
 
@@ -26,7 +26,7 @@ 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,
+    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.

src/huggingface_hub/serialization/_tensorflow.py

@@ -15,7 +15,7 @@
 
 import math
 import re
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Union
 
 from ._base import MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
 
@@ -28,7 +28,7 @@ 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,
+    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.

src/huggingface_hub/serialization/_torch.py

@@ -15,7 +15,7 @@
 
 import importlib
 from functools import lru_cache
-from typing import TYPE_CHECKING, Dict, Tuple
+from typing import TYPE_CHECKING, Dict, Tuple, Union
 
 from ._base import FILENAME_PATTERN, MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
 
@@ -28,7 +28,7 @@ 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,
+    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.
@@ -67,7 +67,7 @@ def 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():
+    ...     for filename, tensors in state_dict_split.filename_to_tensors.items():
     ...         shard = {tensor: state_dict[tensor] for tensor in tensors}
     ...         safe_save_file(
     ...             shard,

tests/test_serialization.py

@@ -1,4 +1,7 @@
+import pytest
+
 from huggingface_hub.serialization import split_state_dict_into_shards_factory
+from huggingface_hub.serialization._base import parse_size_to_int
 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
@@ -123,3 +126,17 @@ def test_get_tensor_size_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
+
+
+def test_parse_size_to_int():
+    assert parse_size_to_int("1KB") == 1 * 10**3
+    assert parse_size_to_int("2MB") == 2 * 10**6
+    assert parse_size_to_int("3GB") == 3 * 10**9
+    assert parse_size_to_int(" 10 KB ") == 10 * 10**3  # ok with whitespace
+    assert parse_size_to_int("20mb") == 20 * 10**6  # ok with lowercase
+
+    with pytest.raises(ValueError, match="Unit 'IB' not supported"):
+        parse_size_to_int("1KiB")  # not a valid unit
+
+    with pytest.raises(ValueError, match="Could not parse the size value"):
+        parse_size_to_int("1ooKB")  # not a float