Add a param for save format in Storage Writer (pytorch#150025)

Summary: add a param to specify to the storage writer how to save tensors. Write now the only options are safetensors and torch.save.

Test Plan:
(lintrunner) [[email protected] /data/users/ankitageorge/fbsource/fbcode/caffe2 (1d57cb27b)]$ buck2 test 'fbcode//mode/opt' fbcode//caffe2/test/distributed/checkpoint:test_hf_storage
File changed: fbcode//caffe2/torch/distributed/checkpoint/filesystem.py
Buck UI: https://www.internalfb.com/buck2/e80cc963-e34a-4876-b6f4-7ce2794e48dd
Test UI: https://www.internalfb.com/intern/testinfra/testrun/3659174965882569
Network: Up: 32KiB  Down: 1.9KiB  (reSessionID-ef9fa764-a40a-451b-ab58-08eabe7a9422)
Executing actions. Remaining     0/4                                                                                             3.4s exec time total
Command: test.     Finished 2 local
Time elapsed: 19.6s
Tests finished: Pass 4. Fail 0. Fatal 0. Skip 0. Build failure 0

Reviewed By: saumishr

Differential Revision: D70271943

Pull Request resolved: pytorch#150025
Approved by: https://github.com/saumishr

Author: ankitageorge

Diff for: torch/distributed/checkpoint/_fsspec_filesystem.py

@@ -15,6 +15,7 @@
     FileSystemBase,
     FileSystemReader,
     FileSystemWriter,
+    SerializationFormat,
 )
 
 
@@ -115,6 +116,7 @@ def __init__(
         per_thread_copy_ahead: int = 10_000_000,
         overwrite: bool = True,
         _extensions: Optional[Sequence[StreamTransformExtension]] = None,
+        serialization_format: SerializationFormat = SerializationFormat.TORCH_SAVE,
         **kwargs,
     ) -> None:
         """
@@ -139,6 +141,7 @@ def __init__(
             per_thread_copy_ahead,
             overwrite=overwrite,
             _extensions=_extensions,
+            serialization_format=serialization_format,
         )
         self.fs = FileSystem()
         self.path = self.fs.init_path(path, **kwargs)

Diff for: torch/distributed/checkpoint/_hf_storage.py

@@ -11,6 +11,7 @@
     _FqnToFileMapping,
     _HuggingFaceLoadPlanner,
 )
+from torch.distributed.checkpoint.filesystem import SerializationFormat
 from torch.distributed.checkpoint.metadata import (
     BytesStorageMetadata,
     Metadata,
@@ -64,7 +65,11 @@ def __init__(
         if HfFileSystem.protocol not in fsspec.available_protocols():
             fsspec.register_implementation(HfFileSystem.protocol, HfFileSystem)
 
-        super().__init__(path=path, token=token)
+        super().__init__(
+            path=path,
+            token=token,
+            serialization_format=SerializationFormat.SAFETENSORS,
+        )
         self._fqn_to_index_mapping: dict[str, int] = fqn_to_index_mapping
 
     def prepare_local_plan(self, plan: SavePlan) -> SavePlan:
@@ -99,7 +104,7 @@ def write_data(
                 (self.fs.concat_path(self.path, file_name), file_name, write_items)
             )
 
-        return super()._write_data(planner, file_queue, safe_tensors=True)
+        return super()._write_data(planner, file_queue)
 
     def finish(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
         metadata_to_write = {}

Diff for: torch/distributed/checkpoint/filesystem.py

@@ -13,6 +13,7 @@
 from collections.abc import Generator, Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
+from enum import Enum
 from io import UnsupportedOperation
 from pathlib import Path
 from typing import Any, Callable, cast, IO, Optional, Union
@@ -49,7 +50,13 @@
 from torch.futures import Future
 
 
-__all__ = ["FileSystemWriter", "FileSystemReader", "FileSystem", "FileSystemBase"]
+__all__ = [
+    "FileSystemWriter",
+    "FileSystemReader",
+    "FileSystem",
+    "FileSystemBase",
+    "SerializationFormat",
+]
 
 _metadata_fn: str = ".metadata"
 
@@ -72,6 +79,11 @@ class _StoragePrefix:
     prefix: str
 
 
+class SerializationFormat(Enum):
+    TORCH_SAVE = "torch_save"
+    SAFETENSORS = "safetensors"
+
+
 DEFAULT_SUFFIX = ".distcp"
 
 
@@ -298,7 +310,7 @@ def _write_item(
     data: Union[io.BytesIO, torch.Tensor],
     write_item: WriteItem,
     storage_key: str,
-    safe_tensors: bool = False,
+    serialization_format: SerializationFormat,
 ) -> WriteResult:
     offset = stream.tell()
 
@@ -312,12 +324,14 @@ def _write_item(
     else:
         assert isinstance(data, torch.Tensor)
         assert data.device == torch.device("cpu")
-        if not safe_tensors:
+        if serialization_format == SerializationFormat.TORCH_SAVE:
             torch.save(data, transform_to)
 
     transform_to.close()
 
-    if not safe_tensors or isinstance(data, io.BytesIO):
+    if serialization_format == SerializationFormat.TORCH_SAVE or isinstance(
+        data, io.BytesIO
+    ):
         length = stream.tell() - offset
     else:
         length = data.numel() * data.element_size()
@@ -349,7 +363,7 @@ def _write_files_from_queue(
     inflight_threshhold: int,
     use_fsync: bool,
     thread_count: int,
-    safe_tensors: bool,
+    serialization_format: SerializationFormat,
 ) -> None:
     try:
         while True:
@@ -397,7 +411,7 @@ def _write_files_from_queue(
                             data,
                             write_item,
                             storage_key,
-                            safe_tensors,
+                            serialization_format,
                         )
                     )
 
@@ -411,12 +425,12 @@ def _write_files_from_queue(
                             tensor,
                             write_item,
                             storage_key,
-                            safe_tensors,
+                            serialization_format,
                         )
                     )
                     tensor_dict[write_item.index.fqn] = tensor
 
-                if safe_tensors:
+                if serialization_format == SerializationFormat.SAFETENSORS:
                     from safetensors.torch import save  # type: ignore[import-not-found]
 
                     stream.write(save(tensor_dict))
@@ -549,6 +563,7 @@ def __init__(
         per_thread_copy_ahead: int = 10_000_000,
         overwrite: bool = True,
         _extensions: Optional[Sequence[StreamTransformExtension]] = None,
+        serialization_format: SerializationFormat = SerializationFormat.TORCH_SAVE,
         *args: Any,
         **kwargs: Any,
     ) -> None:
@@ -576,6 +591,7 @@ def __init__(
         self.save_id = _generate_uuid()
         self.overwrite = overwrite
         self.transforms = _StorageWriterTransforms(_extensions)
+        self.serialization_format = serialization_format
 
     def reset(self, checkpoint_id: Union[str, os.PathLike, None] = None) -> None:
         if checkpoint_id:
@@ -638,7 +654,6 @@ def _write_data(
         self,
         planner: SavePlanner,
         file_queue: queue.Queue,
-        safe_tensors: bool = False,
     ) -> Future[list[WriteResult]]:
         result_queue: queue.Queue = queue.Queue()
 
@@ -655,7 +670,7 @@ def _write_data(
                     self.per_thread_copy_ahead,
                     self.sync_files,
                     self.thread_count,
-                    safe_tensors,
+                    self.serialization_format,
                 ),
             )
             t.start()
@@ -670,7 +685,7 @@ def _write_data(
             inflight_threshhold=self.per_thread_copy_ahead,
             use_fsync=self.sync_files,
             thread_count=self.thread_count,
-            safe_tensors=safe_tensors,
+            serialization_format=self.serialization_format,
         )
 
         for t in threads:
@@ -892,6 +907,7 @@ def __init__(
         cache_staged_state_dict: bool = False,
         overwrite: bool = True,
         _extensions: Optional[Sequence[StreamTransformExtension]] = None,
+        serialization_format: SerializationFormat = SerializationFormat.TORCH_SAVE,
     ) -> None:
         """
         Initialize the writer pointing to `path`.
@@ -919,6 +935,7 @@ def __init__(
             per_thread_copy_ahead=per_thread_copy_ahead,
             overwrite=overwrite,
             _extensions=_extensions,
+            serialization_format=serialization_format,
         )
         BlockingAsyncStager.__init__(
             self,