[Safetensors] Fix mmap for Windows system (#8734)

* fix mmap for Windows system

Author: DrownFish19

paddlenlp/trainer/plugins/unified_checkpoint.py

@@ -17,6 +17,7 @@
 import json
 import multiprocessing
 import os
+import sys
 
 import numpy as np
 import paddle
@@ -67,10 +68,12 @@
 from paddlenlp.utils.tools import get_env_device
 
 if is_safetensors_available():
-    # from safetensors import safe_open
     from safetensors.numpy import save_file as safe_save_file
 
-    from paddlenlp.utils.safetensors import fast_safe_open as safe_open
+    if sys.platform.startswith("win"):
+        from safetensors import safe_open
+    else:
+        from paddlenlp.utils.safetensors import fast_safe_open as safe_open
 
 FP32_MASTER = "fp32_master_0"
 optimizer_scalar_name = [

paddlenlp/transformers/model_utils.py

@@ -20,6 +20,7 @@
 import json
 import os
 import re
+import sys
 import tempfile
 import warnings
 from contextlib import contextmanager
@@ -108,12 +109,13 @@ def unwrap_optimizer(optimizer, optimizer_instances=()):
 
 
 if is_safetensors_available():
-
-    # from safetensors import safe_open
     from safetensors.numpy import load_file as safe_load_file
     from safetensors.numpy import save_file as safe_save_file
 
-    from paddlenlp.utils.safetensors import fast_safe_open as safe_open
+    if sys.platform.startswith("win"):
+        from safetensors import safe_open
+    else:
+        from paddlenlp.utils.safetensors import fast_safe_open as safe_open
 
 
 def prune_linear_layer(layer: nn.Linear, index: paddle.Tensor, dim: int = 0) -> nn.Linear:

tests/testing_utils.py

@@ -296,6 +296,22 @@ def decorator(func):
     return decorator
 
 
+def skip_platform(*platform):
+    """decorator which can detect that it will skip the specific platform
+
+    Args:
+        platform (str): the name of platform, including win32, cygwin, linux, and darwin
+    """
+
+    def decorator(func):
+        for plat in platform:
+            if sys.platform.startswith(plat):
+                return unittest.skip(f"platform<{plat}> matched, so to skip this test")(func)
+        return func
+
+    return decorator
+
+
 def is_slow_test() -> bool:
     """check whether is the slow test
 

tests/transformers/test_safetensors.py

@@ -17,12 +17,12 @@
 import unittest
 
 import numpy as np
-
-# from safetensors import safe_open
 from safetensors.numpy import load_file, save_file
 
 from paddlenlp.utils.safetensors import fast_load_file, fast_safe_open
 
+from ..testing_utils import skip_platform
+
 
 class FastSafetensors(unittest.TestCase):
     def setUp(self):
@@ -42,6 +42,7 @@ def setUp(self):
             count += 1
         print(self.weigth_map)
 
+    @skip_platform("win32", "cygwin")
     def test_load_file(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
             path = os.path.join(tmpdirname, "test.safetensors")
@@ -52,6 +53,7 @@ def test_load_file(self):
                 np.testing.assert_equal(v, sf_load[k])
                 np.testing.assert_equal(v, fs_sf_load[k])
 
+    @skip_platform("win32", "cygwin")
     def test_safe_open(self):
         with tempfile.TemporaryDirectory() as tmpdirname:
             path = os.path.join(tmpdirname, "test.safetensors")