checkpoint-engine 0.2.1__tar.gz → 0.2.3__tar.gz

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpoint-engine
-Version: 0.2.1
+Version: 0.2.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
@@ -99,17 +99,15 @@ Use the flexible P2P implementation, notice this will install `mooncake-transfer
 pip install 'checkpoint-engine[p2p]'
 ```
 
-If set `NCCL_IB_HCA` env, checkpoint-engine will use it to auto select net devices for different ranks. Available patterns can be found from [NCCL documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8). If not set, it will read all RDMA devices and try to divide them into each rank.
-
 ## Getting Started
 
-Prepare an H800 or H20 machine with 8 GPUs with latest vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights.
+Prepare an H800 or H20 machine with 8 GPUs with vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights. vLLM version `v0.10.2` is fully tested and recommended.
 
 ```Bash
-cd /opt && git clone https://github.com/vllm-project/vllm && cd vllm
+mkdir -p /opt/vLLM && cd /opt/vLLM
 uv venv --python 3.12 --seed
 source .venv/bin/activate
-VLLM_USE_PRECOMPILED=1 uv pip install --editable .
+uv pip install vllm==0.10.2
 ```
 
 Install checkpoint-engine
@@ -180,6 +178,11 @@ Other unit tests can also be done with pytest. Only test_update.py requires GPUs
 pytest tests/ -m "not gpu"
 ```
 
+### Environment Variables
+- `PS_MAX_BUCKET_SIZE_GB`: An integer is used to set the maximum bucket size for checkpoint-engine. If not set, 8GB is used as default.
+- `PS_P2P_STORE_RDMA_DEVICES`: Comma-separated RDMA devices' names for P2P transfer. If not set, checkpoint-engine will fall back to use `NCCL_IB_HCA` to detect RDMA devices.
+- `NCCL_IB_HCA`: Available patterns can be found from [NCCL documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8). If also not set, all RDMA devices will be used and divided evenly among the ranks.
+
 ## SGLang Integration
 
 Checkpoint Engine provides efficient distributed checkpoint loading for SGLang inference servers, significantly reducing model loading time for large models and multi-node setups.

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/README.md

@@ -75,17 +75,15 @@ Use the flexible P2P implementation, notice this will install `mooncake-transfer
 pip install 'checkpoint-engine[p2p]'
 ```
 
-If set `NCCL_IB_HCA` env, checkpoint-engine will use it to auto select net devices for different ranks. Available patterns can be found from [NCCL documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8). If not set, it will read all RDMA devices and try to divide them into each rank.
-
 ## Getting Started
 
-Prepare an H800 or H20 machine with 8 GPUs with latest vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights.
+Prepare an H800 or H20 machine with 8 GPUs with vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights. vLLM version `v0.10.2` is fully tested and recommended.
 
 ```Bash
-cd /opt && git clone https://github.com/vllm-project/vllm && cd vllm
+mkdir -p /opt/vLLM && cd /opt/vLLM
 uv venv --python 3.12 --seed
 source .venv/bin/activate
-VLLM_USE_PRECOMPILED=1 uv pip install --editable .
+uv pip install vllm==0.10.2
 ```
 
 Install checkpoint-engine
@@ -156,6 +154,11 @@ Other unit tests can also be done with pytest. Only test_update.py requires GPUs
 pytest tests/ -m "not gpu"
 ```
 
+### Environment Variables
+- `PS_MAX_BUCKET_SIZE_GB`: An integer is used to set the maximum bucket size for checkpoint-engine. If not set, 8GB is used as default.
+- `PS_P2P_STORE_RDMA_DEVICES`: Comma-separated RDMA devices' names for P2P transfer. If not set, checkpoint-engine will fall back to use `NCCL_IB_HCA` to detect RDMA devices.
+- `NCCL_IB_HCA`: Available patterns can be found from [NCCL documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8). If also not set, all RDMA devices will be used and divided evenly among the ranks.
+
 ## SGLang Integration
 
 Checkpoint Engine provides efficient distributed checkpoint loading for SGLang inference servers, significantly reducing model loading time for large models and multi-node setups.

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/checkpoint_engine/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.1'
-__version_tuple__ = version_tuple = (0, 2, 1)
+__version__ = version = '0.2.3'
+__version_tuple__ = version_tuple = (0, 2, 3)
 
-__commit_id__ = commit_id = 'g279a908a9'
+__commit_id__ = commit_id = 'g0a6244951'

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/checkpoint_engine/ps.py

@@ -1,6 +1,7 @@
 import argparse
 import concurrent.futures
 import ctypes
+import json
 import os
 import pickle
 import random
@@ -18,7 +19,7 @@ import torch.distributed as dist
 import zmq
 from loguru import logger
 from pydantic import BaseModel, PlainSerializer, PlainValidator, WithJsonSchema
-from safetensors.torch import safe_open
+from safetensors.torch import _getdtype, safe_open
 from torch.multiprocessing.reductions import reduce_tensor
 
 from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid
@@ -92,6 +93,7 @@ class ParameterMeta(BaseModel):
     name: str
     dtype: _TorchDtype
     shape: _TorchSize
+    aligned_size: int
 
 
 class BucketRange(NamedTuple):
@@ -140,7 +142,7 @@ def _align_size(dtype: torch.dtype, shape: torch.Size) -> int:
 def _to_named_tensor(metas: list[ParameterMeta], offset: int = 0) -> list[dict]:
     ret = []
     for meta in metas:
-        size = _align_size(meta.dtype, meta.shape)
+        size = meta.aligned_size
         ret.append(
             {
                 "name": meta.name,
@@ -422,6 +424,7 @@ def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
                     name=parameter_name,
                     shape=meta["shape"],
                     dtype=meta["dtype"],
+                    aligned_size=_align_size(meta["dtype"], meta["shape"]),
                 )
             tp_meta = tp_metas[parameter_name]
             if tp_meta.concat_dim != -1:
@@ -431,7 +434,10 @@ def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
             shape = list(parameter_metas[name].shape)
             shape[tp_meta.concat_dim] = shape[tp_meta.concat_dim] * tp_meta.size
             parameter_metas[name] = ParameterMeta(
-                name=name, shape=torch.Size(shape), dtype=parameter_metas[name].dtype
+                name=name,
+                shape=torch.Size(shape),
+                dtype=parameter_metas[name].dtype,
+                aligned_size=_align_size(parameter_metas[name].dtype, torch.Size(shape)),
             )
         weights_in_cpu = [parameters_with_tp[name][key] for key in sorted(parameters_with_tp[name])]
         # TODO: here concat is serial, which may be slow
@@ -449,17 +455,85 @@ def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
     return parameters
 
 
-def _register_checkpoint(
-    *,
+def _inplace_pin_memory(files: list[str], rank: int | None = None) -> list[MemoryBuffer]:
+    def _parse_and_pin_from_safetensors(file_path: str) -> MemoryBuffer:
+        """
+        safetensors format see https://huggingface.co/docs/safetensors/en/index#format.
+        We load the safetensors file as bytes, then parse the header manually to get parameter metas.
+        The actual tensor data is in the remaining bytes and is naturally aligned.
+        We pin the remaining bytes as the buffer, making pinning faster.
+        """
+
+        def _pin(t: torch.Tensor):
+            """
+            Pin the memory of tensor in-place.
+            See: https://github.com/pytorch/pytorch/issues/32167
+            """
+            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}"
+
+        # TODO: should only support /dev/shm? but we found files in disk also work?
+        size = os.stat(file_path).st_size
+        flag_size = 8
+        t = torch.from_file(file_path, True, size, dtype=torch.uint8)
+        assert t.nbytes > flag_size, (
+            f"tensor nbytes {t.nbytes} should be greater than flag_size {flag_size}"
+        )
+        start_pos = (
+            int.from_bytes(t[0:flag_size].numpy().tobytes(), byteorder="little", signed=False)
+            + flag_size
+        )
+        header_tensor = t[flag_size:start_pos]
+        header = json.loads(header_tensor.numpy().tobytes())
+        if "__metadata__" in header:
+            header.pop("__metadata__")
+
+        metas: list[ParameterMeta] = []
+        offset = 0
+        try:
+            for name, meta in sorted(header.items(), key=lambda x: x[1]["data_offsets"]):
+                start, end = meta["data_offsets"]
+                # safetensors format ensures offsets are aligned
+                assert offset == start, f"offset {offset} should be equal to start {start}"
+                metas.append(
+                    ParameterMeta(
+                        name=name,
+                        dtype=_getdtype(meta["dtype"]),
+                        shape=torch.Size(meta["shape"]),
+                        aligned_size=end - start,
+                    )
+                )
+                offset = end
+        except Exception as e:
+            logger.error(f"fail to parse safetensors header from {file_path}: {e}")
+            raise
+
+        buffer = t[start_pos:]
+        assert offset == buffer.nbytes, (
+            f"offset {offset} should be equal to buffer.nbytes {buffer.nbytes}"
+        )
+        # 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)
+        _pin(buffer)
+        logger.info(
+            f"[rank{rank}] inplace pin memory for file {file_path} finished, size {buffer.nbytes / 1024 / 1024:.2f}MiB"
+        )
+        return MemoryBuffer(buffer=buffer, size=buffer.nbytes, metas=metas)
+
+    memory_buffers: list[MemoryBuffer] = []
+    with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
+        memory_buffers = list(executor.map(_parse_and_pin_from_safetensors, files))
+    return memory_buffers
+
+
+def _normal_pin_memory(
     files: list[str],
     named_tensors: dict[str, torch.Tensor],
     rank: int | None = None,
+    shared_pin_memory: list[MemoryBuffer] | None = None,
 ) -> 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 []
     parameters = _load_checkpoint(files)
     if named_tensors:
         parameters.update(named_tensors)
@@ -469,13 +543,16 @@ def _register_checkpoint(
         size: int
         metas: list[ParameterMeta]
 
-    buckets: list[MemoryBucket] = [MemoryBucket(size=0, metas=[])]
+    buckets: list[MemoryBucket] = []
+    buckets.append(MemoryBucket(size=0, metas=[]))
     for name, tensor in sorted(parameters.items()):
         size = _align_size(tensor.dtype, tensor.shape)
         if buckets[-1].size + size > bucket_size:
             assert buckets[-1], f"buckets[{len(buckets) - 1}] should not be empty"
             buckets.append(MemoryBucket(size=0, metas=[]))
-        buckets[-1].metas.append(ParameterMeta(name=name, shape=tensor.shape, dtype=tensor.dtype))
+        buckets[-1].metas.append(
+            ParameterMeta(name=name, shape=tensor.shape, dtype=tensor.dtype, aligned_size=size)
+        )
         buckets[-1].size += size
 
     memory_buffers = [
@@ -483,16 +560,34 @@ def _register_checkpoint(
         for bucket in buckets
     ]
 
-    def register_pin_memory(idx: int, size: int) -> tuple[int, torch.Tensor]:
-        buffer = torch.empty(size, dtype=torch.uint8, pin_memory=True)
-        return idx, buffer
+    def register_pin_memory(
+        idx: int, size: int, shared_pin_memory: list[MemoryBuffer] | None = None
+    ) -> tuple[int, torch.Tensor]:
+        if shared_pin_memory:
+            # If shared_pin_memory is provided, reuse the pin memory buffer, do not allocate new one
+            # Reusing pin memory only support fixed shape of checkpoints, which is registered the first time
+            assert idx < len(shared_pin_memory), (
+                f"idx {idx} should be less than shared_pin_memory length {len(shared_pin_memory)}"
+            )
+            assert shared_pin_memory[idx].size == size, (
+                f"shared_pin_memory[{idx}].size {shared_pin_memory[idx].size} should be equal to {size}"
+            )
+            return idx, shared_pin_memory[idx].buffer
+        else:
+            buffer = torch.empty(size, dtype=torch.uint8, pin_memory=True)
+            return idx, buffer
 
     def register_tensor(buffer: torch.Tensor, offset: int, tensor: torch.Tensor):
         buffer[offset : offset + tensor.nbytes] = tensor.view(-1).view(dtype=torch.uint8)
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
         futures = [
-            executor.submit(register_pin_memory, idx, bucket.size)
+            executor.submit(
+                register_pin_memory,
+                idx,
+                bucket.size,
+                shared_pin_memory,
+            )
             for idx, bucket in enumerate(buckets)
         ]
         new_futures = []
@@ -518,6 +613,45 @@ def _register_checkpoint(
                 offset += size
         for future in concurrent.futures.as_completed(new_futures):
             future.result()
+        return memory_buffers
+
+
+def _register_checkpoint(
+    *,
+    files: list[str],
+    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] = []
+    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(
+                files=files_to_normal_pin,
+                named_tensors=named_tensors,
+                rank=rank,
+                shared_pin_memory=shared_pin_memory,
+            )
+        )
+    if files_to_inplace_pin:
+        memory_buffers.extend(_inplace_pin_memory(files_to_inplace_pin, rank=rank))
     return memory_buffers
 
 
@@ -566,7 +700,7 @@ def _gen_h2d_buckets(
         for idx, metas in enumerate(items.memory_buffer_metas_list):
             start_offset, offset = 0, 0
             for meta in metas.metas:
-                s = _align_size(meta.dtype, meta.shape)
+                s = meta.aligned_size
                 if buckets[-1][1].size + s > bucket_size:
                     if offset - start_offset > 0:
                         buckets[-1][1].ranges.append(
@@ -747,6 +881,8 @@ class P2PStore:
 
 
 class ParameterServer:
+    shared_memory_pool_name = "__shared_memory_pool__"
+
     def __init__(
         self,
         *,
@@ -790,7 +926,10 @@ class ParameterServer:
         self._zmq_ctx = zmq.Context()
         self._zmq_addr_counter = 0
 
+        # stores the name of the checkpoint currently using the shared memory pool, or empty string if none
+        self._current_shared_memory_pool_user: str = ""
         self._memory_pool: dict[str, list[MemoryBuffer]] = {}
+        self._memory_pool[self.shared_memory_pool_name] = []
         # dict key is owner_rank, value is a bucket metas list in owner_rank
         self._current_global_parameter_metas: dict[int, MemoryBufferMetaList] = {}
         # NPU transfer engine initialization requires prior set_device.
@@ -805,6 +944,17 @@ class ParameterServer:
         self._device_uuid = _get_physical_gpu_id(self.device_manager, device_index)
         self._rdma_device = None if self._p2p_store is None else self._p2p_store.device
 
+    def _get_memory_pool(self, checkpoint_name: str) -> list[MemoryBuffer]:
+        if checkpoint_name == self._current_shared_memory_pool_user:
+            assert self._memory_pool[self.shared_memory_pool_name], (
+                f"shared memory pool is not initialized, but checkpoint {checkpoint_name} is using it"
+            )
+            return self._memory_pool[self.shared_memory_pool_name]
+        elif checkpoint_name in self._memory_pool:
+            return self._memory_pool[checkpoint_name]
+        else:
+            raise RuntimeError(f"checkpoint {checkpoint_name} is not registered")
+
     def _logger_rank0(self, msg: str):
         if self._local_rank == 0:
             logger.info(msg)
@@ -828,46 +978,103 @@ class ParameterServer:
         *,
         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: 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:
             checkpoint_name: The name of the checkpoint.
             files: The safetensors files to register.
             named_tensors: The named tensors to register.
+            use_shared_memory_pool: If True, uses a reusable shared pin memory pool instead of allocating new memory.
+                Only one checkpoint can use the shared pool at a time. The pool's shape is fixed on first use and
+                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:
-            assert checkpoint_name not in self._memory_pool, (
-                f"checkpoint {checkpoint_name} already registered"
-            )
-            self._memory_pool[checkpoint_name] = _register_checkpoint(
-                files=files or [], named_tensors=named_tensors or {}, rank=self._rank
-            )
-            if self._p2p_store is not None:
-                self._register_parameters_to_p2p_store(checkpoint_name)
+            if use_shared_memory_pool:
+                logger.info(
+                    f"[rank{self._rank}] checkpoint {checkpoint_name} use shared memory pool"
+                )
+                assert self._current_shared_memory_pool_user == "", (
+                    f"cannot register checkpoint {checkpoint_name} to shared memory pool, "
+                    f"since checkpoint {self._current_shared_memory_pool_user} is already using shared memory pool. "
+                    f"This registration may cause unexpected conflicts."
+                )
+                # Since we set the uninitialized shared memory pool to empty list,
+                # we can check whether this is the first time to use shared memory pool
+                _is_first_time = not self._memory_pool[self.shared_memory_pool_name]
+                self._memory_pool[self.shared_memory_pool_name] = _register_checkpoint(
+                    files=files or [],
+                    named_tensors=named_tensors or {},
+                    rank=self._rank,
+                    shared_pin_memory=self._memory_pool[self.shared_memory_pool_name],
+                )
+                self._current_shared_memory_pool_user = checkpoint_name
+                if self._p2p_store is not None and _is_first_time:
+                    self._register_parameters_to_p2p_store(checkpoint_name)
+            else:
+                assert checkpoint_name not in self._memory_pool, (
+                    f"checkpoint {checkpoint_name} already registered"
+                )
+                self._memory_pool[checkpoint_name] = _register_checkpoint(
+                    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)
         except Exception:
             logger.exception(
                 f"[rank{self._rank}] fail to register checkpoint {checkpoint_name} with files {files}"
             )
-            if self._p2p_store is not None:
+            if self._p2p_store is not None and not use_shared_memory_pool:
                 self._unregister_parameters_from_p2p_store(checkpoint_name)
             self.unregister_checkpoint(checkpoint_name)
             raise
 
-    def unregister_checkpoint(self, checkpoint_name: str):
+    def unregister_checkpoint(self, checkpoint_name: str, force: bool = False) -> None:
         """
         Unregister a checkpoint from the parameter server. This function will also unregister the checkpoint
         from p2p store if p2p store is initialized.
+        Args:
+            checkpoint_name: The name of the checkpoint.
+            force: This flag is designed for shared memory pool user. If True, the memory for shared memory pool itself will be freed.
+                    If False, only the checkpoint name will be unregistered, and the shared memory pool will be kept for future use.
         """
-        if checkpoint_name not in self._memory_pool:
+        if (
+            checkpoint_name not in self._memory_pool
+            and checkpoint_name != self._current_shared_memory_pool_user
+        ):
+            logger.warning(
+                f"[rank{self._rank}] unregister checkpoint name {checkpoint_name} not found"
+            )
+            return
+
+        if checkpoint_name == self._current_shared_memory_pool_user and not force:
+            self._current_shared_memory_pool_user = ""
             return
+
         if self._p2p_store is not None:
             num_unregistered = self._unregister_parameters_from_p2p_store(checkpoint_name)
             logger.info(
                 f"[rank{self._rank}] unregister {num_unregistered} parameters from p2p store for checkpoint {checkpoint_name}"
             )
-        del self._memory_pool[checkpoint_name]
+
+        if checkpoint_name == self._current_shared_memory_pool_user:
+            self._current_shared_memory_pool_user = ""
+            del self._memory_pool[self.shared_memory_pool_name]
+            self._memory_pool[self.shared_memory_pool_name] = []
+        else:
+            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()
@@ -882,6 +1089,10 @@ class ParameterServer:
             self.init_process_group()
         assert dist.is_initialized(), "process group is not initialized"
         metas_lst: list[DataToGather | None] = [None for _ in range(self._world_size)]  # type: ignore
+        try:
+            memory_pool = self._get_memory_pool(checkpoint_name)
+        except RuntimeError:
+            memory_pool = []
         metas = DataToGather(
             memory_buffer_metas_list=[
                 MemoryBufferMetas(
@@ -889,7 +1100,7 @@ class ParameterServer:
                     ptr=x.buffer.data_ptr(),
                     size=x.size,
                 )
-                for x in self._memory_pool.get(checkpoint_name, [])
+                for x in memory_pool
             ],
             p2p_store_addr=None if self._p2p_store is None else self._p2p_store.addr,
             host_ip=get_ip(),
@@ -1050,7 +1261,7 @@ class ParameterServer:
         for items in self._current_global_parameter_metas.values():
             for metas_list in items.memory_buffer_metas_list:
                 for meta in metas_list.metas:
-                    max_tensor_bytes = max(max_tensor_bytes, _align_size(meta.dtype, meta.shape))
+                    max_tensor_bytes = max(max_tensor_bytes, meta.aligned_size)
         free_bytes_divided_3 = free_bytes // (3 * _ALIGN_SIZE) * _ALIGN_SIZE
         if max_tensor_bytes <= free_bytes_divided_3 and not disable_h2d_buffer:
             self._logger_rank0(f"[rank{self._rank}] use h2d buffer")
@@ -1095,7 +1306,7 @@ class ParameterServer:
                 remote_ptrs.append(ptrs[b.idx][0] + b.offset)
                 lens.append(b.size)
             else:
-                pool = self._memory_pool[checkpoint_name][b.idx]
+                pool = self._get_memory_pool(checkpoint_name)[b.idx]
                 buffer[offset : offset + b.size].data.copy_(
                     pool.buffer[b.offset : b.offset + b.size],
                     non_blocking=True,
@@ -1158,22 +1369,32 @@ class ParameterServer:
 
     def _register_parameters_to_p2p_store(self, checkpoint_name: str):
         assert self._p2p_store is not None, "p2p store is not initialized"
-        pool = self._memory_pool[checkpoint_name]
+        pool = self._get_memory_pool(checkpoint_name)
         if len(pool) == 0:
             return
         named_tensors, tensor_ptrs = {}, []
+        register_name = (
+            checkpoint_name
+            if checkpoint_name != self._current_shared_memory_pool_user
+            else self.shared_memory_pool_name
+        )
         for idx, memory_buffer in enumerate(pool):
-            named_tensors[f"memory_pool_{checkpoint_name}_{idx}"] = memory_buffer.buffer
+            named_tensors[f"memory_pool_{register_name}_{idx}"] = memory_buffer.buffer
             tensor_ptrs.append((memory_buffer.buffer.data_ptr(), memory_buffer.size))
         self._p2p_store.register_named_tensors(named_tensors)
 
     def _unregister_parameters_from_p2p_store(self, checkpoint_name: str) -> int:
         assert self._p2p_store is not None, "p2p store is not initialized"
-        pool = self._memory_pool[checkpoint_name]
+        pool = self._get_memory_pool(checkpoint_name)
         if len(pool) == 0:
             return 0
+        unregister_name = (
+            checkpoint_name
+            if checkpoint_name != self._current_shared_memory_pool_user
+            else self.shared_memory_pool_name
+        )
         return self._p2p_store.unregister_named_tensors(
-            [f"memory_pool_{checkpoint_name}_{idx}" for idx, _ in enumerate(pool)]
+            [f"memory_pool_{unregister_name}_{idx}" for idx, _ in enumerate(pool)]
         )
 
     def _update_per_bucket(
@@ -1284,9 +1505,9 @@ class ParameterServer:
                     dist.broadcast(buffer_b, src=brank)
                     resp = socket.recv()
                     if resp != b"":
-                        exception_obj = pickle.loads(resp)
+                        msg = resp.decode("utf-8")
                         logger.error(
-                            f"[rank{self._rank}] receive error response '{type(exception_obj).__name__}: {exception_obj}' from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}"
+                            f"[rank{self._rank}] receive error response from rank {receiver_rank} for bucket {gidx} in checkpoint {checkpoint_name}: {msg}"
                         )
                         ret_code.fill_(1)
                     dist.all_reduce(ret_code, op=dist.ReduceOp.SUM)

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/checkpoint_engine/worker.py

@@ -1,4 +1,5 @@
 import gc
+import traceback
 from collections.abc import Callable
 from typing import TypedDict
 
@@ -63,7 +64,8 @@ def update_weights_from_ipc(
         assert buffer.dtype == torch.uint8
         socket.send(b"")
     except Exception as e:
-        socket.send_pyobj(e)
+        msg = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+        socket.send_string(msg)
         socket.recv()  # wait for ack
         raise
     try:
@@ -83,7 +85,8 @@ def update_weights_from_ipc(
                 except Exception as e:  # noqa: BLE001
                     # Send exception back to Parameter Server.
                     # Don't raise here. Because all workers should quit in the same way by receiving the exception from PS
-                    socket.send_pyobj(e)
+                    msg = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+                    socket.send_string(msg)
             elif isinstance(
                 payload, Exception
             ):  # error occurred, got force quit signal from Parameter Server

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/checkpoint_engine.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpoint-engine
-Version: 0.2.1
+Version: 0.2.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
@@ -99,17 +99,15 @@ Use the flexible P2P implementation, notice this will install `mooncake-transfer
 pip install 'checkpoint-engine[p2p]'
 ```
 
-If set `NCCL_IB_HCA` env, checkpoint-engine will use it to auto select net devices for different ranks. Available patterns can be found from [NCCL documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8). If not set, it will read all RDMA devices and try to divide them into each rank.
-
 ## Getting Started
 
-Prepare an H800 or H20 machine with 8 GPUs with latest vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights.
+Prepare an H800 or H20 machine with 8 GPUs with vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights. vLLM version `v0.10.2` is fully tested and recommended.
 
 ```Bash
-cd /opt && git clone https://github.com/vllm-project/vllm && cd vllm
+mkdir -p /opt/vLLM && cd /opt/vLLM
 uv venv --python 3.12 --seed
 source .venv/bin/activate
-VLLM_USE_PRECOMPILED=1 uv pip install --editable .
+uv pip install vllm==0.10.2
 ```
 
 Install checkpoint-engine
@@ -180,6 +178,11 @@ Other unit tests can also be done with pytest. Only test_update.py requires GPUs
 pytest tests/ -m "not gpu"
 ```
 
+### Environment Variables
+- `PS_MAX_BUCKET_SIZE_GB`: An integer is used to set the maximum bucket size for checkpoint-engine. If not set, 8GB is used as default.
+- `PS_P2P_STORE_RDMA_DEVICES`: Comma-separated RDMA devices' names for P2P transfer. If not set, checkpoint-engine will fall back to use `NCCL_IB_HCA` to detect RDMA devices.
+- `NCCL_IB_HCA`: Available patterns can be found from [NCCL documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8). If also not set, all RDMA devices will be used and divided evenly among the ranks.
+
 ## SGLang Integration
 
 Checkpoint Engine provides efficient distributed checkpoint loading for SGLang inference servers, significantly reducing model loading time for large models and multi-node setups.

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/checkpoint_engine.egg-info/SOURCES.txt

@@ -23,5 +23,6 @@ figures/overlap-update-and-copy.png
 figures/pipeline.png
 patches/vllm_fp8.patch
 tests/test_assign_receiver_ranks.py
+tests/test_pin_memory.py
 tests/test_rdma_parser.py
 tests/test_update.py

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/examples/update.py

@@ -100,8 +100,9 @@ def update_weights(
     update_method: Literal["broadcast", "p2p", "all"] = "broadcast",
     uds: str | None = None,
 ):
-    ps.register_checkpoint(checkpoint_name, files=checkpoint_files, named_tensors=named_tensors)
     ps.init_process_group()
+    dist.barrier()
+    ps.register_checkpoint(checkpoint_name, files=checkpoint_files, named_tensors=named_tensors)
     check_vllm_ready(endpoint, inference_parallel_size, uds)
     dist.barrier()
     with timer("Gather metas"):
@@ -173,7 +174,9 @@ if __name__ == "__main__":
             args.uds,
         )
     else:
-        if os.path.exists(os.path.join(args.checkpoint_path, "model.safetensors.index.json")):
+        if os.path.exists(
+            os.path.join(args.checkpoint_path, "model.safetensors.index.json")
+        ) and not args.checkpoint_path.startswith("/dev/shm/"):  # noqa: S108
             named_tensors = split_tensors(args.checkpoint_path, rank, world_size)
             checkpoint_files = []
         else:

checkpoint_engine-0.2.3/tests/test_pin_memory.py

@@ -0,0 +1,77 @@
+import os
+
+import pytest
+import torch
+
+from checkpoint_engine.ps import ParameterServer
+
+
+def generate_dummy_checkpoint() -> dict[str, torch.Tensor]:
+    """
+    Generate dummy checkpoint data
+    """
+    named_tensors = {
+        "layer1.weight": torch.randn(1024, 1024),
+        "layer1.bias": torch.randn(1024),
+        "layer2.weight": torch.randn(2048, 1024),
+        "layer2.bias": torch.randn(2048),
+    }
+    return named_tensors
+
+
+@pytest.mark.gpu
+def test_register_pin_memory():
+    os.environ["RANK"] = "0"
+    os.environ["WORLD_SIZE"] = "1"
+    ps = ParameterServer()
+    checkpoint1 = generate_dummy_checkpoint()
+    checkpoint_shared1 = generate_dummy_checkpoint()
+    checkpoint2 = generate_dummy_checkpoint()
+    checkpoint_shared2 = generate_dummy_checkpoint()
+    checkpoint_shared3 = generate_dummy_checkpoint()
+    checkpoint_shared3["layer3.weight"] = torch.randn(4096, 2048)
+    checkpoint_shared3["layer3.bias"] = torch.randn(4096)
+    ps.register_checkpoint("test_checkpoint1", named_tensors=checkpoint1)
+    ps.unregister_checkpoint("test_checkpoint1")
+    assert "test_checkpoint1" not in ps._memory_pool
+    ps.register_checkpoint(
+        "test_checkpoint_shared1", named_tensors=checkpoint_shared1, use_shared_memory_pool=True
+    )
+    ps.register_checkpoint("test_checkpoint2", named_tensors=checkpoint2)
+    assert "test_checkpoint_shared1" not in ps._memory_pool
+    assert "__shared_memory_pool__" in ps._memory_pool
+    assert ps._current_shared_memory_pool_user == "test_checkpoint_shared1"
+    assert "test_checkpoint2" in ps._memory_pool
+    try:
+        ps.register_checkpoint(
+            "test_checkpoint_shared2", named_tensors=checkpoint_shared2, use_shared_memory_pool=True
+        )  # this will fail
+    except AssertionError:
+        print("Caught expected AssertionError when registering second shared memory pool user")
+    assert "test_checkpoint_shared2" not in ps._memory_pool
+    assert ps._current_shared_memory_pool_user == "test_checkpoint_shared1"
+    ps.unregister_checkpoint("test_checkpoint_shared1")
+    assert ps._current_shared_memory_pool_user == ""
+    assert "__shared_memory_pool__" in ps._memory_pool
+    ps.register_checkpoint(
+        "test_checkpoint_shared2", named_tensors=checkpoint_shared2, use_shared_memory_pool=True
+    )
+    assert "test_checkpoint_shared2" not in ps._memory_pool
+    assert "__shared_memory_pool__" in ps._memory_pool
+    assert ps._current_shared_memory_pool_user == "test_checkpoint_shared2"
+    ps.unregister_checkpoint("test_checkpoint1")  # this will trigger an warning
+    assert "test_checkpoint1" not in ps._memory_pool
+    ps.unregister_checkpoint("test_checkpoint2")
+    assert "test_checkpoint2" not in ps._memory_pool
+    ps.unregister_checkpoint("test_checkpoint_shared2", force=True)
+    assert ps._current_shared_memory_pool_user == ""
+    assert "__shared_memory_pool__" in ps._memory_pool
+    ps.register_checkpoint(
+        "test_checkpoint_shared3", named_tensors=checkpoint_shared3, use_shared_memory_pool=True
+    )
+    assert "test_checkpoint_shared3" not in ps._memory_pool
+    assert "__shared_memory_pool__" in ps._memory_pool
+    assert ps._current_shared_memory_pool_user == "test_checkpoint_shared3"
+    ps.unregister_checkpoint("test_checkpoint_shared3")
+    assert ps._current_shared_memory_pool_user == ""
+    assert "__shared_memory_pool__" in ps._memory_pool

{checkpoint_engine-0.2.1 → checkpoint_engine-0.2.3}/tests/test_update.py

@@ -82,7 +82,7 @@ def checker_proc_with_error(
         try:
             trigger_error(socket_paths)
         except RuntimeError as e:
-            assert str(e) == "Failed to update weights due to remote errors"
+            assert str(e) == "Some workers failed to update weights"
 
 
 def checker_proc(rank: int, device_uuid: str, named_tensors: dict[str, torch.Tensor], queue: Queue):
@@ -96,7 +96,7 @@ def checker_proc(rank: int, device_uuid: str, named_tensors: dict[str, torch.Ten
         for name, weight in weights:
             if name not in named_tensors:
                 continue
-            assert (weight == named_tensors[name]).all()
+            assert (weight == named_tensors[name]).all(), f"Tensor {name} does not match!"
             names_to_check[name] = True
 
     def check_weights(names_to_check: dict[str, bool], socket_paths: list[tuple[str, str]]):
@@ -163,6 +163,67 @@ def run(
     assert proc.exitcode == 0
 
 
+def run_with_files(
+    checker_func: callable,
+):
+    rank = int(os.getenv("RANK"))
+    ctx = get_context("spawn")
+    queue = ctx.Queue()
+    _device_uuid = _get_physical_gpu_id(device_manager, rank)
+    ps = ParameterServer(auto_pg=True)
+    _device_uuid = _get_physical_gpu_id(ps.device_manager, rank)
+    named_tensors = dict(gen_test_tensors(rank))
+
+    # Save 1/3 tensors to /dev/shm/ as .safetensors files
+    # Save 1/3 tensors to ./tmp (disk) as .safetensors files
+    # Keep 1/3 tensors in memory
+    import safetensors.torch
+
+    files = []
+    dev_shm_dir = "/dev/shm/checkpoint_engine_tests"  # noqa: S108
+    disk_dir = "/tmp/checkpoint_engine_tests"  # noqa: S108
+    os.makedirs(dev_shm_dir, exist_ok=True)
+    os.makedirs(disk_dir, exist_ok=True)
+    tensors_items = list(named_tensors.items())
+    tensors_in_dev_shm = named_tensors
+    tensors_in_dev_shm = dict(tensors_items[: len(tensors_items) // 2])
+    tensors_in_disk = dict(tensors_items[len(tensors_items) // 3 : 2 * len(tensors_items) // 3])
+    tensors_in_memory = dict(tensors_items[1 * len(tensors_items) // 2 :])
+    disk_files = [
+        os.path.join(disk_dir, f"rank{_rank}_checkpoint.safetensors")
+        for _rank in range(get_world_size())
+    ]
+    safetensors.torch.save_file(tensors_in_disk, disk_files[rank])
+    time.sleep(1)
+    files.append(disk_files[rank])
+    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])
+
+    checkpoint_name = "test_with_files"
+    proc = ctx.Process(target=checker_func, args=(rank, _device_uuid, named_tensors, queue))
+    proc.start()
+    ps.register_checkpoint(checkpoint_name, named_tensors=tensors_in_memory, files=files)
+    ps.gather_metas(checkpoint_name)
+    ps.update(checkpoint_name, queue.put, ranks=[])
+    # sleep 3s to wait process group is destroyed
+    time.sleep(3)
+    ps.unregister_checkpoint(checkpoint_name)
+    queue.put(None)
+    proc.join()
+    if rank == 0:
+        import shutil
+
+        # 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
+
+
 @pytest.mark.gpu
 @pytest.mark.parametrize(
     "test_name,rank_list",
@@ -211,6 +272,37 @@ def test_update(test_name: str, rank_list: list[list[int]] | None):
     assert result.returncode == 0
 
 
+@pytest.mark.gpu
+def test_update_with_files(test_name: str = "test_with_files"):
+    world_size = device_manager.device_module.device_count()
+    assert world_size >= 2, "This test requires at least 2 GPUs."
+    master_addr = "localhost"
+    master_port = 25400
+    cmd = [
+        "torchrun",
+        "--nproc_per_node",
+        str(world_size),
+        "--master_addr",
+        master_addr,
+        "--master_port",
+        str(master_port),
+        __file__,
+        test_name,
+        "[]",
+    ]
+
+    result = subprocess.run(  # noqa: S603
+        cmd,
+        capture_output=False,
+        text=True,
+        cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        shell=False,
+        check=False,
+    )
+
+    assert result.returncode == 0
+
+
 if __name__ == "__main__":
     run_with_pytest = "PYTEST_CURRENT_TEST" in os.environ
     if not run_with_pytest:
@@ -230,5 +322,7 @@ if __name__ == "__main__":
             expected_exception=RuntimeError,
             exception_msg="Failed to update weights due to remote errors",
         )
+    elif test_type == "test_with_files":
+        run_with_files(checker_proc)
     else:
         raise ValueError(f"Unknown TEST_TYPE: {test_type}")