add parallel loader

Signed-off-by: yuanyuxing.yyx <[email protected]>

Author: ABNER-1

fastsafetensors/__init__.py

@@ -7,4 +7,5 @@
 
 from .common import SafeTensorsMetadata, SingleGroup, TensorFrame, get_device_numa_node
 from .file_buffer import FilesBufferOnDevice
-from .loader import SafeTensorsFileLoader, fastsafe_open
+from .loader import BaseSafeTensorsFileLoader, SafeTensorsFileLoader, fastsafe_open
+from .parallel_loader import ParallelLoader

fastsafetensors/loader.py

@@ -40,6 +40,7 @@ def __init__(
         disable_cache: bool = True,
         debug_log: bool = False,
         framework="pytorch",
+        **kwargs,
     ):
         self.framework = get_framework_op(framework)
         self.pg = self.framework.get_process_group(pg)
@@ -174,6 +175,7 @@ def __init__(
         disable_cache: bool = True,
         debug_log: bool = False,
         framework="pytorch",
+        **kwargs,
     ):
         self.framework = get_framework_op(framework)
         self.pg = self.framework.get_process_group(pg)
@@ -191,7 +193,14 @@ def __init__(
 
         copier = new_gds_file_copier(self.device, bbuf_size_kb, max_threads, nogds)
         super().__init__(
-            pg, self.device, copier, set_numa, disable_cache, debug_log, framework
+            pg,
+            self.device,
+            copier,
+            set_numa,
+            disable_cache,
+            debug_log,
+            framework,
+            **kwargs,
         )