checkpoint-engine 0.3.1rc0__tar.gz → 0.3.3__tar.gz

{checkpoint_engine-0.3.1rc0 → checkpoint_engine-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpoint-engine
-Version: 0.3.1rc0
+Version: 0.3.3
 Summary: checkpoint-engine is a lightweight, decoupling and efficient weight update middleware
 Project-URL: Homepage, https://github.com/MoonshotAI/checkpoint-engine
 Project-URL: Repository, https://github.com/MoonshotAI/checkpoint-engine

{checkpoint_engine-0.3.1rc0 → checkpoint_engine-0.3.3}/checkpoint_engine/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.1rc0'
-__version_tuple__ = version_tuple = (0, 3, 1, 'rc0')
+__version__ = version = '0.3.3'
+__version_tuple__ = version_tuple = (0, 3, 3)
 
-__commit_id__ = commit_id = 'g09c543af4'
+__commit_id__ = commit_id = 'gf6910d646'

{checkpoint_engine-0.3.1rc0 → checkpoint_engine-0.3.3}/checkpoint_engine/pin_memory.py

@@ -191,6 +191,8 @@ def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
 
 
 def _inplace_pin_memory(files: list[str], rank: int | None = None) -> list[MemoryBuffer]:
+    device_index = torch.cuda.current_device()
+
     def _parse_and_pin_from_safetensors(file_path: str) -> MemoryBuffer:
         """
         safetensors format see https://huggingface.co/docs/safetensors/en/index#format.
@@ -204,9 +206,12 @@ def _inplace_pin_memory(files: list[str], rank: int | None = None) -> list[Memor
             Pin the memory of tensor in-place.
             See: https://github.com/pytorch/pytorch/issues/32167
             """
+            torch.cuda.set_device(device_index)
             cudart = torch.cuda.cudart()
             r = cudart.cudaHostRegister(t.data_ptr(), t.numel() * t.element_size(), 0)
-            assert r == 0, f"pin memory error, error code: {r}"
+            if r != 0:
+                error_msg = cudart.cudaGetErrorString(r)
+                raise RuntimeError(f"pin memory error, error code: {r}, error message: {error_msg}")
 
         # TODO: should only support /dev/shm? but we found files in disk also work?
         size = os.stat(file_path).st_size
@@ -251,6 +256,12 @@ def _inplace_pin_memory(files: list[str], rank: int | None = None) -> list[Memor
         # Remove the file after successfully loading. This will avoid doubling the memory usage.
         # We assume files in /dev/shm/ are temporary files. So it's safe to remove them after loading.
         os.remove(file_path)
+        if not metas:
+            # TODO: should we still return this buffer?
+            assert buffer.nbytes == 0, f"buffer nbytes {buffer.nbytes} should be 0"
+            logger.warning(f"[rank{rank}] no metas found in {file_path}, skip pin memory")
+            return MemoryBuffer(buffer=buffer, size=buffer.nbytes, metas=[], manually_pinned=False)
+
         _pin(buffer)
         logger.info(
             f"[rank{rank}] inplace pin memory for file {file_path} finished, size {buffer.nbytes / 1024 / 1024:.2f}MiB"

{checkpoint_engine-0.3.1rc0 → checkpoint_engine-0.3.3}/checkpoint_engine/ps.py

@@ -391,7 +391,11 @@ class ParameterServer:
                 )
                 cudart = torch.cuda.cudart()
                 r = cudart.cudaHostUnregister(t.data_ptr())
-                assert r == 0, f"unpin memory error, error code: {r}"
+                if r != 0:
+                    error_msg = cudart.cudaGetErrorString(r)
+                    raise RuntimeError(
+                        f"unpin memory error, error code: {r}, error message: {error_msg}"
+                    )
 
             # if the checkpoint is pinned by cudaHostRegister manually, we need to unpin it manually
             try:
@@ -407,7 +411,13 @@ class ParameterServer:
             del self._memory_pool[checkpoint_name]
         # see https://github.com/pytorch/pytorch/blob/31d5c675394705f8a6bc767f80ae14bf4f01246b/torch/csrc/cuda/Module.cpp#L2018
         # this works by using torch>=2.5.0
-        torch._C._host_emptyCache()
+        if self.device_manager.device_type == "cuda":
+            torch._C._host_emptyCache()
+        else:
+            # torch._C._host_emptyCache() is not supported on NPU, so we call gc.collect() to empty host cache.
+            import gc
+
+            gc.collect()
 
     def gather_metas(self, checkpoint_name: str):
         """
@@ -731,6 +741,7 @@ class ParameterServer:
         assert len(self._current_global_parameter_metas) != 0, "parameter metas is empty"
         assert dist.is_initialized(), "process group is not initialized"
 
+        p2p_update = False
         # if both ranks is None or [], it will use fully broadcast to update to all ranks
         if not ranks:
             logger.info(f"[rank{self._rank}] update checkpoint {checkpoint_name}")
@@ -739,6 +750,7 @@ class ParameterServer:
             assert self._p2p_store is not None, "p2p store is not initialized"
             assert ranks, "ranks should be set"
 
+            p2p_update = True
             need_update = self._rank in ranks
             logger.info(
                 f"[rank{self._rank}] update checkpoint {checkpoint_name} p2p, {need_update=} with {ranks=}, "
@@ -764,11 +776,6 @@ class ParameterServer:
             if disable_h2d_buffer
             else torch.empty(bucket_size, dtype=torch.uint8, device=self.device_manager.device_type)
         )
-        # p2p store need to register h2d_buffer to let other ranks read
-        if ranks:
-            h2d_buffer_name = "__h2d_buffer__"
-            if h2d_buffer is not None and self._p2p_store is not None:
-                self._p2p_store.register_named_tensors({h2d_buffer_name: h2d_buffer})
         receiver_rank_buckets: list[tuple[int, H2DBucket]] = []
         for receiver_rank, owner_rank, bucket in buckets:
             if receiver_rank != self._rank:
@@ -778,6 +785,12 @@ class ParameterServer:
         buffer = torch.empty(
             bucket_size * 2, dtype=torch.uint8, device=self.device_manager.device_type
         )
+        if p2p_update:
+            # p2p store need to register buffer to let other ranks read
+            p2p_ipc_buffer_name = "__ipc_buffer__"
+            self._p2p_store.register_named_tensors(
+                {p2p_ipc_buffer_name: buffer if disable_h2d_buffer else h2d_buffer}
+            )
         handle = reduce_tensor(buffer)
 
         buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list)
@@ -823,7 +836,14 @@ class ParameterServer:
                     buffer_b: torch.Tensor = buffer[start : start + bucket.size]
                     if receiver_rank == self._rank:
                         if disable_h2d_buffer:
-                            self._copy_to_buffer(checkpoint_name, bucket, buffer_b)
+                            if p2p_update:
+                                assert bucket == receiver_rank_buckets[i][1]
+                            self._copy_to_buffer(
+                                checkpoint_name,
+                                bucket,
+                                buffer_b,
+                                receiver_rank_buckets[i][0] if p2p_update else None,
+                            )
                         else:
                             buffer_b.data.copy_(h2d_buffer[: bucket.size])
                     dist.broadcast(buffer_b, src=receiver_rank, group=ranks_group)
@@ -850,8 +870,8 @@ class ParameterServer:
             req_thread.join()
             dist.barrier(group=ranks_group)
             socket.close()
-            if ranks and h2d_buffer is not None:
-                self._p2p_store.unregister_named_tensors([h2d_buffer_name])
+            if p2p_update:
+                self._p2p_store.unregister_named_tensors([p2p_ipc_buffer_name])
 
             self.device_manager.device_module.empty_cache()
 

{checkpoint_engine-0.3.1rc0 → checkpoint_engine-0.3.3}/checkpoint_engine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpoint-engine
-Version: 0.3.1rc0
+Version: 0.3.3
 Summary: checkpoint-engine is a lightweight, decoupling and efficient weight update middleware
 Project-URL: Homepage, https://github.com/MoonshotAI/checkpoint-engine
 Project-URL: Repository, https://github.com/MoonshotAI/checkpoint-engine

{checkpoint_engine-0.3.1rc0 → checkpoint_engine-0.3.3}/examples/update.py

@@ -14,7 +14,8 @@ import torch.distributed as dist
 from loguru import logger
 from safetensors import safe_open
 
-from checkpoint_engine.ps import ParameterServer, request_inference_to_update
+from checkpoint_engine import request_inference_to_update
+from checkpoint_engine.ps import ParameterServer
 
 
 @contextmanager