Use mmap option to load_state_dict (huggingface#28331)

weimingzha0 · Jan 3, 2024 · bffda1f · bffda1f
1 parent 45b1dfa
commit bffda1f
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -515,8 +515,10 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
             map_location = "meta"
         else:
             map_location = "cpu"
-
-        return torch.load(checkpoint_file, map_location=map_location, weights_only=True)
+        extra_args = {}
+        if version.parse(torch.__version__) >= version.parse("2.1.0"):
+            extra_args = {"mmap": True}
+        return torch.load(checkpoint_file, map_location=map_location, weights_only=True, **extra_args)
     except Exception as e:
         try:
             with open(checkpoint_file) as f: