doc

Author: specture724

checkpoint_engine/ps.py

@@ -969,7 +969,7 @@ def register_checkpoint(
     ) -> 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:
@@ -981,7 +981,7 @@ 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.
+            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:

tests/test_unpin.py

@@ -0,0 +1,74 @@
+import json
+import os
+import random
+import subprocess
+import sys
+import time
+from contextlib import nullcontext
+from types import SimpleNamespace
+
+import torch
+import zmq
+from torch.multiprocessing import Queue, get_context
+
+from checkpoint_engine.device_utils import DeviceManager
+from checkpoint_engine.ps import ParameterServer, _get_physical_gpu_id
+from checkpoint_engine.worker import update_weights_from_ipc
+
+device_manager = DeviceManager()
+
+def gen_test_tensors(rank: int) -> list[tuple[str, torch.Tensor]]:
+    tensors = []
+    for layer in range(random.randint(10, 50)):
+        for num in range(random.randint(50, 100)):
+            r = random.randint(0, 16)
+            if r < 4:
+                dtype = torch.bfloat16
+            elif r < 10:
+                dtype = torch.float16
+            elif r < 14:
+                dtype = torch.float8_e4m3fn
+            else:
+                dtype = torch.float
+            tensors.append(
+                (
+                    f"rank{rank}.layer{layer}.num{num}",
+                    torch.randn([random.randint(100, 500), random.randint(500, 1000)]).to(dtype),
+                )
+            )
+    return tensors
+
+def get_world_size() -> int:
+    return device_manager.device_module.device_count()
+
+def get_files():
+    rank = int(os.getenv("RANK"))
+    named_tensors = dict(gen_test_tensors(rank))
+    import safetensors.torch
+
+    files = []
+    dev_shm_dir = "/dev/shm/checkpoint_engine_tests"  # noqa: S108
+    os.makedirs(dev_shm_dir, exist_ok=True)
+    tensors_in_dev_shm = named_tensors
+    time.sleep(1)
+    dev_shm_files = [
+        os.path.join(dev_shm_dir, f"rank{rank}_checkpoint.safetensors")
+        for _ in range(get_world_size())
+    ]
+    safetensors.torch.save_file(tensors_in_dev_shm, dev_shm_files[rank])
+    time.sleep(1)
+    files.append(dev_shm_files[rank])
+    return files
+
+
+def run_with_files(num):
+    ps = ParameterServer(auto_pg=True)
+    checkpoint_name = "test_with_files"
+    for _ in range(num):
+        files = get_files()
+        ps.register_checkpoint(checkpoint_name, files=files)
+        ps.gather_metas(checkpoint_name)
+        ps.unregister_checkpoint(checkpoint_name)
+
+if __name__ == "__main__":
+    run_with_files(3)