refactoring functions for huggingface integration

Author: liangel-02

test/prototype/safetensors/test_safetensors_support.py

@@ -1,7 +1,9 @@
+import json
 import tempfile
 import unittest
 
 import torch
+from safetensors.torch import load_file, save_file
 from torch.testing._internal.common_utils import (
     TestCase,
     run_tests,
@@ -19,6 +21,18 @@
 )
 
 
+def load_data(file_path: str, device: str):
+    loaded_tensors = load_file(file_path, device)
+    with open(file_path, "rb") as f:
+        import struct
+
+        header_size = struct.unpack("<Q", f.read(8))[0]
+        header_bytes = f.read(header_size)
+        header = json.loads(header_bytes)
+        metadata = header.get("__metadata__", {})
+    return loaded_tensors, metadata
+
+
 @unittest.skipIf(not torch.cuda.is_available(), "Need CUDA available")
 @unittest.skipIf(not is_sm_at_least_89(), "Need sm89+")
 class TestSafeTensors(TestCase):
@@ -32,8 +46,10 @@ def test_safetensors(self):
         ref_output = model(*example_inputs)
 
         with tempfile.NamedTemporaryFile() as f:
-            save_tensor_state_dict(model.state_dict(), f.name)
-            reconstructed_dict = load_tensor_state_dict(f.name, device="cuda")
+            tensors_dict, metadata = save_tensor_state_dict(model.state_dict())
+            save_file(tensors_dict, f.name, metadata=metadata)
+            tensor_dict, metadata = load_data(file_path=f.name, device="cuda")
+            reconstructed_dict = load_tensor_state_dict(tensor_dict, metadata)
 
         model = torch.nn.Sequential(
             torch.nn.Linear(32, 256, dtype=torch.bfloat16, device="cuda")

torchao/prototype/safetensors/safetensors_support.py

@@ -1,9 +1,8 @@
 import json
 import logging
-from typing import Dict
+from typing import Any, Dict
 
 import torch
-from safetensors.torch import load_file, save_file
 
 from torchao.prototype.safetensors.safetensors_serialization import (
     Float8TensorAttributeJSONEncoder,
@@ -14,7 +13,10 @@
 logger: logging.Logger = logging.getLogger(__name__)
 
 
-def load_tensor_state_dict(file_path: str, device: str):
+def load_tensor_state_dict(
+    tensor_data: Dict[str, Any],
+    metadata: Dict[str, Any],
+):
     """
     Load a dictionary of tensor subclasses from a safetensors file.
 
@@ -34,20 +36,13 @@ def load_tensor_state_dict(file_path: str, device: str):
             - dtype
 
     Args:
-        file_path: Path to the safetensors file
+        tensor_data: Tensor data,
+        metadata: Tensor attributes,
 
     Returns:
         Dictionary of reconstructed tensor subclasses
     """
-    loaded_tensors = load_file(file_path, device)
-
-    with open(file_path, "rb") as f:
-        import struct
-
-        header_size = struct.unpack("<Q", f.read(8))[0]
-        header_bytes = f.read(header_size)
-        header = json.loads(header_bytes)
-        metadata = header.get("__metadata__", {})
+    combined_data = {**tensor_data, **metadata}
 
     if "tensor_names" not in metadata:
         raise ValueError("No tensors found")
@@ -57,7 +52,7 @@ def load_tensor_state_dict(file_path: str, device: str):
 
     for tensor_name in tensor_names:
         tensor_tensors = {}
-        for key, value in loaded_tensors.items():
+        for key, value in combined_data.items():
             if key.startswith(f"{tensor_name}:"):
                 # Remove the prefix
                 tensor_tensors[key[len(tensor_name) + 1 :]] = value
@@ -73,15 +68,11 @@ def load_tensor_state_dict(file_path: str, device: str):
         else:
             raise ValueError(f"Unsupported tensor type: {tensor_type}")
 
-    logger.info(
-        f"Loaded {len(tensor_names)} tensor subclasses from {file_path} with metadata"
-    )
     return result
 
 
 def save_tensor_state_dict(
     tensor_dict: Dict[str, Dict[str, torch.Tensor]],
-    file_path: str,
 ):
     """
     Save a dictionary of tensor subclasses with appropriate metadata.
@@ -105,7 +96,6 @@ def save_tensor_state_dict(
 
     Args:
         tensor_dict: Dictionary of tensor subclasses to save, with keys as tensor names
-        file_path: Path where to save the tensors
     """
 
     combined_metadata = {}
@@ -136,8 +126,4 @@ def save_tensor_state_dict(
         combined_tensors_dict.update(prefixed_tensors_dict)
 
     combined_metadata["tensor_names"] = json.dumps(list(tensor_dict.keys()))
-
-    save_file(combined_tensors_dict, file_path, metadata=combined_metadata)
-    logger.info(
-        f"Saved {len(tensor_dict)} tensor subclasses to {file_path} with metadata"
-    )
+    return combined_tensors_dict, combined_metadata