huggingface-hub 0.26.3__py3-none-any.whl → 0.26.4__py3-none-any.whl

huggingface_hub/__init__.py

@@ -46,7 +46,7 @@ import sys
 from typing import TYPE_CHECKING
 
 
-__version__ = "0.26.3"
+__version__ = "0.26.4"
 
 # Alphabetical order of definitions is ensured in tests
 # WARNING: any comment added in this dictionary definition will be lost when

huggingface_hub/serialization/_torch.py

@@ -41,6 +41,8 @@ def save_torch_model(
     max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
     metadata: Optional[Dict[str, str]] = None,
     safe_serialization: bool = True,
+    is_main_process: bool = True,
+    shared_tensors_to_discard: Optional[List[str]] = None,
 ):
     """
     Saves a given torch model to disk, handling sharding and shared tensors issues.
@@ -64,6 +66,12 @@ def save_torch_model(
 
     </Tip>
 
+    <Tip warning={true}>
+
+    If your model is a `transformers.PreTrainedModel`, you should pass `model._tied_weights_keys` as `shared_tensors_to_discard` to properly handle shared tensors saving. This ensures the correct duplicate tensors are discarded during saving.
+
+    </Tip>
+
     Args:
         model (`torch.nn.Module`):
             The model to save on disk.
@@ -88,6 +96,13 @@ def save_torch_model(
             Whether to save as safetensors, which is the default behavior. If `False`, the shards are saved as pickle.
             Safe serialization is recommended for security reasons. Saving as pickle is deprecated and will be removed
             in a future version.
+        is_main_process (`bool`, *optional*):
+            Whether the process calling this is the main process or not. Useful when in distributed training like
+            TPUs and need to call this function from all processes. In this case, set `is_main_process=True` only on
+            the main process to avoid race conditions. Defaults to True.
+        shared_tensors_to_discard (`List[str]`, *optional*):
+            List of tensor names to drop when saving shared tensors. If not provided and shared tensors are
+            detected, it will drop the first name alphabetically.
 
     Example:
 
@@ -112,6 +127,8 @@ def save_torch_model(
         metadata=metadata,
         safe_serialization=safe_serialization,
         save_directory=save_directory,
+        is_main_process=is_main_process,
+        shared_tensors_to_discard=shared_tensors_to_discard,
     )
 
 
@@ -124,6 +141,8 @@ def save_torch_state_dict(
     max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
     metadata: Optional[Dict[str, str]] = None,
     safe_serialization: bool = True,
+    is_main_process: bool = True,
+    shared_tensors_to_discard: Optional[List[str]] = None,
 ) -> None:
     """
     Save a model state dictionary to the disk, handling sharding and shared tensors issues.
@@ -147,6 +166,12 @@ def save_torch_state_dict(
 
     </Tip>
 
+    <Tip warning={true}>
+
+    If your model is a `transformers.PreTrainedModel`, you should pass `model._tied_weights_keys` as `shared_tensors_to_discard` to properly handle shared tensors saving. This ensures the correct duplicate tensors are discarded during saving.
+
+    </Tip>
+
     Args:
         state_dict (`Dict[str, torch.Tensor]`):
             The state dictionary to save.
@@ -171,6 +196,13 @@ def save_torch_state_dict(
             Whether to save as safetensors, which is the default behavior. If `False`, the shards are saved as pickle.
             Safe serialization is recommended for security reasons. Saving as pickle is deprecated and will be removed
             in a future version.
+        is_main_process (`bool`, *optional*):
+            Whether the process calling this is the main process or not. Useful when in distributed training like
+            TPUs and need to call this function from all processes. In this case, set `is_main_process=True` only on
+            the main process to avoid race conditions. Defaults to True.
+        shared_tensors_to_discard (`List[str]`, *optional*):
+            List of tensor names to drop when saving shared tensors. If not provided and shared tensors are
+            detected, it will drop the first name alphabetically.
 
     Example:
 
@@ -192,7 +224,8 @@ def save_torch_state_dict(
             else constants.PYTORCH_WEIGHTS_FILE_PATTERN
         )
 
-    # Imports correct library
+    if metadata is None:
+        metadata = {}
     if safe_serialization:
         try:
             from safetensors.torch import save_file as save_file_fn
@@ -201,7 +234,13 @@ def save_torch_state_dict(
                 "Please install `safetensors` to use safe serialization. "
                 "You can install it with `pip install safetensors`."
             ) from e
-
+        # Clean state dict for safetensors
+        state_dict = _clean_state_dict_for_safetensors(
+            state_dict,
+            metadata,
+            force_contiguous=force_contiguous,
+            shared_tensors_to_discard=shared_tensors_to_discard,
+        )
     else:
         from torch import save as save_file_fn  # type: ignore[assignment]
 
@@ -210,13 +249,6 @@ def save_torch_state_dict(
             "pickled models from untrusted sources. If you intend to share your model, we strongly recommend "
             "using safe serialization by installing `safetensors` with `pip install safetensors`."
         )
-
-    # Clean state dict for safetensors
-    if metadata is None:
-        metadata = {}
-    if safe_serialization:
-        state_dict = _clean_state_dict_for_safetensors(state_dict, metadata, force_contiguous=force_contiguous)
-
     # Split dict
     state_dict_split = split_torch_state_dict_into_shards(
         state_dict, filename_pattern=filename_pattern, max_shard_size=max_shard_size
@@ -459,7 +491,10 @@ def storage_ptr(tensor: "torch.Tensor") -> Union[int, Tuple[Any, ...]]:
 
 
 def _clean_state_dict_for_safetensors(
-    state_dict: Dict[str, "torch.Tensor"], metadata: Dict[str, str], force_contiguous: bool = True
+    state_dict: Dict[str, "torch.Tensor"],
+    metadata: Dict[str, str],
+    force_contiguous: bool = True,
+    shared_tensors_to_discard: Optional[List[str]] = None,
 ):
     """Remove shared tensors from state_dict and update metadata accordingly (for reloading).
 
@@ -467,7 +502,7 @@ def _clean_state_dict_for_safetensors(
 
     Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L155.
     """
-    to_removes = _remove_duplicate_names(state_dict)
+    to_removes = _remove_duplicate_names(state_dict, discard_names=shared_tensors_to_discard)
     for kept_name, to_remove_group in to_removes.items():
         for to_remove in to_remove_group:
             if metadata is None:

{huggingface_hub-0.26.3.dist-info → huggingface_hub-0.26.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: huggingface-hub
-Version: 0.26.3
+Version: 0.26.4
 Summary: Client library to download and publish models, datasets and other repos on the huggingface.co hub
 Home-page: https://github.com/huggingface/huggingface_hub
 Author: Hugging Face, Inc.

{huggingface_hub-0.26.3.dist-info → huggingface_hub-0.26.4.dist-info}/RECORD

@@ -1,4 +1,4 @@
-huggingface_hub/__init__.py,sha256=TvADZbi2Jw5oeg_Mu-yZ7xo66rN_x7RTs1F6c9gt2aI,35993
+huggingface_hub/__init__.py,sha256=kXk5sqwCRAGkjc6nnpI_Mu-LGaIoinBcbXytSzhca3s,35993
 huggingface_hub/_commit_api.py,sha256=Y9eTaW4bYzxtrZsSniVtfeAuFafqx8x1ofMI5es8hvM,31057
 huggingface_hub/_commit_scheduler.py,sha256=nlJS_vnLb8i92NLrRwJX8Mg9QZ7f3kfLbLlQuEd5YjU,13647
 huggingface_hub/_inference_endpoints.py,sha256=wzjD8P68VpUDHzIDbXzFXsM2Y-aNVSAap7BXsZFuthk,16750
@@ -79,7 +79,7 @@ huggingface_hub/inference/_generated/types/zero_shot_object_detection.py,sha256=
 huggingface_hub/serialization/__init__.py,sha256=z5MLxMqz0Y2qST-3Lj0PZHUONL-SGRlc0g4Z6MdL6rw,988
 huggingface_hub/serialization/_base.py,sha256=JZneES-HgcRH9C2SQehIGRDtT7nS7emu-RRV4ZjB6xo,8124
 huggingface_hub/serialization/_tensorflow.py,sha256=zHOvEMg-JHC55Fm4roDT3LUCDO5zB9qtXZffG065RAM,3625
-huggingface_hub/serialization/_torch.py,sha256=i6UFAHk1MDx_RONaXYolsISVa0V3a_YH-bdQtCYnmtg,26498
+huggingface_hub/serialization/_torch.py,sha256=KlCRgLarzegkbfUmb73h82p2vDvgKrWw03ltQ7klI2Q,28685
 huggingface_hub/templates/datasetcard_template.md,sha256=W-EMqR6wndbrnZorkVv56URWPG49l7MATGeI015kTvs,5503
 huggingface_hub/templates/modelcard_template.md,sha256=4AqArS3cqdtbit5Bo-DhjcnDFR-pza5hErLLTPM4Yuc,6870
 huggingface_hub/utils/__init__.py,sha256=aMEsiXGi93z-dXz1W7FFma71tAMeKw0SoKVZSQUeE_4,3525
@@ -109,9 +109,9 @@ huggingface_hub/utils/insecure_hashlib.py,sha256=OjxlvtSQHpbLp9PWSrXBDJ0wHjxCBU-
 huggingface_hub/utils/logging.py,sha256=Cp03s0uEl3kDM9XHQW9a8GAoExODQ-e7kEtgMt-_To8,4728
 huggingface_hub/utils/sha.py,sha256=OFnNGCba0sNcT2gUwaVCJnldxlltrHHe0DS_PCpV3C4,2134
 huggingface_hub/utils/tqdm.py,sha256=jQiVYwRG78HK4_54u0vTtz6Kt9IMGiHy3ixbIn3h2TU,9368
-huggingface_hub-0.26.3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-huggingface_hub-0.26.3.dist-info/METADATA,sha256=xeabjkgCLGRol8NXGS0ftagVPKUfTaq2iaKfd_k5P9M,13091
-huggingface_hub-0.26.3.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-huggingface_hub-0.26.3.dist-info/entry_points.txt,sha256=Y3Z2L02rBG7va_iE6RPXolIgwOdwUFONyRN3kXMxZ0g,131
-huggingface_hub-0.26.3.dist-info/top_level.txt,sha256=8KzlQJAY4miUvjAssOAJodqKOw3harNzuiwGQ9qLSSk,16
-huggingface_hub-0.26.3.dist-info/RECORD,,
+huggingface_hub-0.26.4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+huggingface_hub-0.26.4.dist-info/METADATA,sha256=PqF3HduxB5xJmjEI0-bMVfg2z72wgsioKo-ZHyH226g,13091
+huggingface_hub-0.26.4.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+huggingface_hub-0.26.4.dist-info/entry_points.txt,sha256=Y3Z2L02rBG7va_iE6RPXolIgwOdwUFONyRN3kXMxZ0g,131
+huggingface_hub-0.26.4.dist-info/top_level.txt,sha256=8KzlQJAY4miUvjAssOAJodqKOw3harNzuiwGQ9qLSSk,16
+huggingface_hub-0.26.4.dist-info/RECORD,,