hotfix: add a switch to disable inplace pinning of tensors (#68)

* feat: add a switch to disable inplace pinning of tensors

* Update checkpoint_engine/ps.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Anjie Hou <149605198+specture724@users.noreply.github.com>

* doc

---------

Signed-off-by: Anjie Hou <149605198+specture724@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 2 people

checkpoint_engine/ps.py

@@ -622,19 +622,25 @@ def _register_checkpoint(
     named_tensors: dict[str, torch.Tensor],
     rank: int | None = None,
     shared_pin_memory: list[MemoryBuffer] | None = None,
+    inplace_pin: bool = False,
 ) -> list[MemoryBuffer]:
     logger.info(
         f"[rank{rank}] start to register checkpoint with {len(files)} files and {len(named_tensors)} named_tensors"
     )
     if not files and not named_tensors:
         return []
     memory_buffers: list[MemoryBuffer] = []
-    files_to_inplace_pin = [
-        file
-        for file in files
-        if file.startswith("/dev/shm/") and file.endswith(".safetensors")  # noqa: S108
-    ]
-    files_to_normal_pin = [file for file in files if file not in files_to_inplace_pin]
+    if inplace_pin:
+        logger.info(f"[rank{rank}] allow inplace pin memory for /dev/shm/ safetensors files")
+        files_to_inplace_pin = [
+            file
+            for file in files
+            if file.startswith("/dev/shm/") and file.endswith(".safetensors")  # noqa: S108
+        ]
+        files_to_normal_pin = [file for file in files if file not in files_to_inplace_pin]
+    else:
+        files_to_normal_pin = files
+        files_to_inplace_pin = []
     if files_to_normal_pin or named_tensors:
         memory_buffers.extend(
             _normal_pin_memory(
@@ -973,10 +979,11 @@ def register_checkpoint(
         files: list[str] | None = None,
         named_tensors: dict[str, torch.Tensor] | None = None,
         use_shared_memory_pool: bool = False,
+        use_inplace_pin_memory: bool = False,
     ) -> None:
         """
         Register a checkpoint to the parameter server. Both files and named_tensors will be registered together.
-        Warning: .safetensors files in /dev/shm/ will be pinned in-place, and the files will be REMOVED after pinning.
+        Warning: if `use_inplace_pin_memory` is True, .safetensors files in /dev/shm/ will be pinned in-place, and the files will be REMOVED after pinning.
         Please make sure to copy the files to disks if you need to keep them.
 
         Args:
@@ -988,6 +995,8 @@ def register_checkpoint(
                 cannot accommodate checkpoints with different memory requirements.
                 To free the actual memory of the shared pool or to modify its shape,
                 please unregister the current user of the shared memory pool using `unregister_checkpoint` with `force=True`.
+            use_inplace_pin_memory: If True, allows inplace pin memory for /dev/shm/ safetensors files. This option is ignored when ``use_shared_memory_pool`` is True.
+                Currently, this feature is experimental and may crash.
         """
         try:
             if use_shared_memory_pool:
@@ -1016,7 +1025,10 @@ def register_checkpoint(
                     f"checkpoint {checkpoint_name} already registered"
                 )
                 self._memory_pool[checkpoint_name] = _register_checkpoint(
-                    files=files or [], named_tensors=named_tensors or {}, rank=self._rank
+                    files=files or [],
+                    named_tensors=named_tensors or {},
+                    rank=self._rank,
+                    inplace_pin=use_inplace_pin_memory,
                 )
                 if self._p2p_store is not None:
                     self._register_parameters_to_p2p_store(checkpoint_name)

tests/test_update.py

@@ -218,7 +218,8 @@ def run_with_files(
     if rank == 0:
         import shutil
 
-        os.removedirs(dev_shm_dir)
+        # this test should be run under use_inplace_pin_memory=False. Otherwise, the files in /dev/shm/ will be deleted.
+        shutil.rmtree(dev_shm_dir)
         shutil.rmtree(disk_dir)
     assert proc.exitcode == 0