checkpoint-engine 0.1.3__py3-none-any.whl → 0.2.1__py3-none-any.whl

checkpoint_engine/_version.py

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

checkpoint_engine/device_utils.py

@@ -0,0 +1,86 @@
+import os
+import re
+import socket
+import subprocess
+from functools import lru_cache
+
+import torch
+from loguru import logger
+
+
+@lru_cache(maxsize=1)
+def get_ip() -> str:
+    try:
+        # try to get ip from network interface
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+            s.connect(("8.8.8.8", 80))
+            return s.getsockname()[0]
+    except Exception as e:  # noqa: BLE001
+        # fallback to get ip from hostname
+        logger.warning(
+            f"fail to get ip from network interface, fallback to get ip from hostname: {e}"
+        )
+        return socket.gethostbyname(socket.gethostname())
+
+
+def npu_generate_uuid() -> str:
+    str_pid = str(os.getpid())
+    npu_num = 8
+    try:
+        for npu_id in range(npu_num):
+            cmd = ["npu-smi", "info", "-t", "proc-mem", "-i", str(npu_id)]
+            result = subprocess.run(cmd, check=True, capture_output=True, text=True)  # noqa: S603
+            str_result = str(result.stdout)
+            if str_pid in str_result:
+                # In A3 server, one NPU has two chips.
+                match_chip_count = re.search(r"Chip Count[^\d]*(\d+)", str_result)
+                chip_count = int(match_chip_count.group(1))
+                search_after_pid = str_result[str_result.find(str_pid) + len(str_pid) :]
+                match_chip_id = re.search(r"Chip ID[^\d]*(\d+)", search_after_pid)
+                chip_id = int(match_chip_id.group(1))
+                return f"{get_ip()}-{npu_id * chip_count + chip_id}"
+        raise ValueError("The current process is not running on the npu device")
+    except subprocess.CalledProcessError as e:
+        raise ValueError("The current process is not running on the npu device") from e
+
+
+class DeviceManager:
+    def __init__(self):
+        self.device_type = self._detect_device_type()
+        self._setup_device_module()
+
+    def _is_torch_npu_available(self) -> bool:
+        try:
+            if hasattr(torch, "npu") and callable(getattr(torch.npu, "is_available", None)):
+                return torch.npu.is_available()
+            else:
+                return False
+        except ImportError:
+            return False
+
+    def _detect_device_type(self) -> str:
+        if self._is_torch_npu_available():
+            return "npu"
+        elif torch.cuda.is_available():
+            return "cuda"
+        else:
+            raise TypeError("The current device type is not supported")
+
+    def _setup_device_module(self):
+        if self.device_type == "npu":
+            import torch_npu
+
+            self.device_module = torch_npu.npu
+        elif self.device_type == "cuda":
+            self.device_module = torch.cuda
+        else:
+            raise TypeError("The current device type is not supported")
+
+    @property
+    def backend(self) -> str:
+        if self.device_type == "npu":
+            return "hccl"
+        elif self.device_type == "cuda":
+            return "nccl"
+        else:
+            raise TypeError("The current device type is not supported")

checkpoint_engine/ps.py

@@ -4,13 +4,11 @@ import ctypes
 import os
 import pickle
 import random
-import socket
 import threading
 import time
 from collections import defaultdict
 from collections.abc import Callable
 from datetime import timedelta
-from functools import lru_cache
 from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, NamedTuple
 
 import httpx
@@ -23,8 +21,12 @@ from pydantic import BaseModel, PlainSerializer, PlainValidator, WithJsonSchema
 from safetensors.torch import safe_open
 from torch.multiprocessing.reductions import reduce_tensor
 
+from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid
+
 
 if TYPE_CHECKING:
+    from typing import TypeVar
+
     from typing_extensions import TypedDict
 
     class FileMeta(TypedDict):
@@ -34,6 +36,8 @@ if TYPE_CHECKING:
         type: type
         tp_concat_dim: int
 
+    T = TypeVar("T")
+
 
 def _dt_validate(value: Any) -> torch.dtype:
     if isinstance(value, str):
@@ -117,6 +121,7 @@ class MemoryBuffer(BaseModel):
 class MemoryBufferMetaList(BaseModel):
     p2p_store_addr: str | None
     memory_buffer_metas_list: list[MemoryBufferMetas]
+    rdma_device: str
 
 
 class DataToGather(MemoryBufferMetaList):
@@ -249,28 +254,16 @@ def _concat_tp_weights(
     return torch.cat([w for w in tp_weights], dim=tp_concat_dim)
 
 
-def _get_physical_gpu_id(device_index: int | None = None) -> str:
+def _get_physical_gpu_id(device_manager: DeviceManager, device_index: int | None = None) -> str:
     try:
-        return f"GPU-{torch.cuda.get_device_properties(device_index).uuid!s}"
+        if device_manager.device_type == "npu":
+            return f"NPU-{npu_generate_uuid()}"
+        else:
+            return f"GPU-{device_manager.device_module.get_device_properties(device_index).uuid!s}"
     except AssertionError as e:
         raise ValueError(f"fail to get physical gpu id {device_index}") from e
 
 
-@lru_cache(maxsize=1)
-def _get_ip() -> str:
-    try:
-        # try to get ip from network interface
-        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
-            s.connect(("8.8.8.8", 80))
-            return s.getsockname()[0]
-    except Exception as e:  # noqa: BLE001
-        # fallback to get ip from hostname
-        logger.warning(
-            f"fail to get ip from network interface, fallback to get ip from hostname: {e}"
-        )
-        return socket.gethostbyname(socket.gethostname())
-
-
 def _ibv_get_device_list() -> list[str]:
     lib = ctypes.CDLL("libibverbs.so.1")
     lib.ibv_get_device_list.argtypes = [ctypes.POINTER(ctypes.c_int)]  # int *num_devices
@@ -303,14 +296,7 @@ def _get_rdma_devices() -> list[str]:
         return devices_str.split(",")
     # if PS_P2P_STORE_RDMA_DEVICES is not set, try to use NCCL_IB_HCA to get RDMA devices
     hca = os.getenv("NCCL_IB_HCA", None)
-    if hca:
-        hca_list = hca.split(",")
-        if len(hca_list) > 1:
-            # if NCCL_IB_HCA has multiple values, just return
-            return hca_list
-        else:
-            hca = hca_list[0]
-    return [device for device in sorted(_ibv_get_device_list()) if hca is None or hca in device]
+    return _parse_NCCL_IB_HCA(hca or "", _ibv_get_device_list()) or _ibv_get_device_list()
 
 
 def _get_my_rdma_device(local_rank: int, gpu_count: int, devices: list[str]) -> str:
@@ -319,13 +305,90 @@ def _get_my_rdma_device(local_rank: int, gpu_count: int, devices: list[str]) ->
     """
     if not devices:
         raise RuntimeError("no rdma devices found")
-    assert len(devices) <= gpu_count, (
-        f"rdma devices count {len(devices)} should be less than or equal to gpu count {gpu_count}"
-    )
-    assert gpu_count % len(devices) == 0, (
-        f"gpu count {gpu_count} should be divisible by rdma devices count {len(devices)}"
-    )
-    return devices[local_rank // (gpu_count // len(devices))]
+    try:
+        assert len(devices) <= gpu_count, (
+            f"rdma devices count {len(devices)} should be less than or equal to gpu count {gpu_count}"
+        )
+        assert gpu_count % len(devices) == 0, (
+            f"gpu count {gpu_count} should be divisible by rdma devices count {len(devices)}"
+        )
+        return devices[local_rank // (gpu_count // len(devices))]
+    except AssertionError:
+        logger.error(
+            "Please set 'NCCL_IB_HCA' or 'PS_P2P_STORE_RDMA_DEVICES' environment variable to choose proper number of RDMA devices."
+            "The number of RDMA devices should be less than or equal to GPU count, and GPU count should be divisible by the number of RDMA devices."
+            "The acceptable value by NCCL_IB_HCA is documented in 'https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8'."
+        )
+        raise
+
+
+def _parse_NCCL_IB_HCA(value: str, available_devices: list[str]) -> list[str]:
+    """
+    The acceptable value by NCCL_IB_HCA is documented in https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8.
+    The Python version parser is referred to the CPP parser in NCCL: https://github.com/NVIDIA/nccl/blob/v2.28.3-1/src/transport/net_ib.cc#L658-L662.
+
+    The list is comma-separated; port numbers are NOT supported yet.
+    An optional prefix '^' indicates the list is an exclude list.
+    A second optional prefix '=' indicates that the tokens are exact names, otherwise by default NCCL would treat each token as a prefix.
+    Please note that when '^' and '=' appear together, only '^=' is allowed, '=^' is not supported.
+
+    Examples:
+    - `NCCL_IB_HCA="mlx5"`: Use all cards starting with `mlx5`.
+    - `NCCL_IB_HCA="=mlx5_0,mlx5_1"`: Use specific cards `mlx5_0` and `mlx5_1`.
+    - `NCCL_IB_HCA="^mlx5"`: Use all cards except those starting with `mlx5`.
+    - `NCCL_IB_HCA="^=mlx5_0,mlx5_1"`: Use all cards except `mlx5_0` and `mlx5_1`.
+    """
+    max_hcas = 32
+    if not value or value.strip() == "":
+        return available_devices[:max_hcas]
+
+    value = value.strip()
+    result = []
+    is_exclude = value.startswith("^")
+    if is_exclude:
+        value = value.removeprefix("^")
+    is_exact_match = value.startswith("=")
+    if is_exact_match:
+        value = value.removeprefix("=")
+
+    device_specs = [spec.strip() for spec in value.split(",") if spec.strip()]
+
+    result = _resolve_device_specs(device_specs, is_exact_match, available_devices)
+    if is_exclude:
+        result = [dev for dev in available_devices if dev not in result]
+    if len(result) > max_hcas:
+        result = result[:max_hcas]
+
+    logger.info(f"RDMA Devices from 'NCCL_IB_HCA': {result}")
+
+    return result
+
+
+def _resolve_device_specs(
+    device_specs: list[str], is_exact_match: bool, available_devices: list[str]
+) -> list[str]:
+    devices = set()
+    for spec in device_specs:
+        parts = spec.split(":", 1)
+        device_name = parts[0].strip()
+        # HACK: mooncake transfer engine does not support port specification yet, so we ignore it
+        # port = parts[1].strip() if len(parts) > 1 else None
+        base_devices = (
+            [device_name]
+            if device_name in available_devices
+            else []
+            if is_exact_match
+            else [dev for dev in available_devices if dev.startswith(device_name)]
+        )
+
+        if not base_devices:
+            logger.warning(f"No RDMA device match {device_name=} where {is_exact_match=}.")
+            continue
+
+        for base_dev in base_devices:
+            devices.add(base_dev)
+
+    return sorted(devices)
 
 
 def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
@@ -490,8 +553,12 @@ def request_inference_to_update(
 
 
 def _gen_h2d_buckets(
-    global_metas: dict[int, MemoryBufferMetaList], bucket_size: int
-) -> list[tuple[int, H2DBucket]]:
+    global_metas: dict[int, MemoryBufferMetaList],
+    bucket_size: int,
+    local_topo: dict[str, set[int]],
+    remote_topo: dict[str, set[int]],
+    ranks: list[int] | None = None,
+) -> list[tuple[int, int, H2DBucket]]:
     buckets: list[tuple[int, H2DBucket]] = []
 
     for owner_rank, items in global_metas.items():
@@ -514,7 +581,73 @@ def _gen_h2d_buckets(
         assert buckets[-1][1].size > 0, (
             f"buckets[-1][1].size {buckets[-1][1].size} should be greater than 0"
         )
-    return buckets
+    ranks_set = set(ranks) if ranks else set()
+    actual_local_topo = (
+        {k: v & ranks_set for k, v in local_topo.items() if v & ranks_set} if ranks else local_topo
+    )
+    # if ranks is empty, assign the owner_rank as receiver_rank, this is used for colocate architecture
+    if not ranks:
+        return [(owner_rank, owner_rank, bucket) for owner_rank, bucket in buckets]
+    else:
+        return _assign_receiver_ranks(buckets, actual_local_topo, remote_topo)
+
+
+def _assign_receiver_ranks(
+    buckets: list[tuple[int, "T"]],
+    local_topo: dict[str, set[int]],
+    remote_topo: dict[str, set[int]],
+) -> list[tuple[int, int, "T"]]:
+    """
+    (owner_rank, bucket) -> (receiver_rank, owner_rank, bucket)
+
+    Assign receiver ranks to buckets. If ranks is empty, assign the owner_rank as receiver_rank.
+    GPU-rdma_device topology will be considered to make full use of the bandwidth.
+    """
+    if not buckets:
+        logger.warning("bucket list is empty, no need to assign receiver ranks")
+        return []
+    rank_to_rdma_device = {
+        rank: rdma_device for rdma_device, ranks in remote_topo.items() for rank in ranks
+    }
+
+    # group buckets by owner RDMA devices
+    buckets_by_rdma_device = defaultdict(list)
+    for owner_rank, bucket in buckets:
+        owner_rdma_device = rank_to_rdma_device[owner_rank]
+        buckets_by_rdma_device[owner_rdma_device].append((owner_rank, bucket))
+
+    buckets_matrix = list(buckets_by_rdma_device.values())
+    assert buckets_matrix, "buckets_matrix should not be empty"
+
+    # Select receiver ranks. We use the minimum rank in each local RDMA device group as receiver rank
+    num_receivers = min(len(local_topo), len(buckets_by_rdma_device))
+    receiver_list = [min(ranks) for ranks in list(local_topo.values())[:num_receivers]]
+
+    flattened_buckets = [
+        buckets_matrix[row][col]
+        for col in range(
+            max(len(matrix_row) for matrix_row in buckets_matrix) if buckets_matrix else 0
+        )
+        for row in range(len(buckets_matrix))
+        if col < len(buckets_matrix[row])
+    ]
+
+    buckets_with_receiver = []
+    assigned_cnt = 0
+    while assigned_cnt < len(flattened_buckets):
+        occupied_devices = set()
+        for receiver_rank in receiver_list:
+            if assigned_cnt >= len(flattened_buckets):
+                break
+            owner_rank, bucket = flattened_buckets[assigned_cnt]
+            rdma_device = rank_to_rdma_device[owner_rank]
+            if rdma_device in occupied_devices:
+                break
+            buckets_with_receiver.append((receiver_rank, owner_rank, bucket))
+            occupied_devices.add(rdma_device)
+            assigned_cnt += 1
+
+    return buckets_with_receiver
 
 
 def _get_master_port(master_port: int | None = None) -> int:
@@ -525,21 +658,44 @@ def _get_master_port(master_port: int | None = None) -> int:
     return master_port
 
 
+def _get_bcast_rank_map(world_size: int, ranks: list[int] | None) -> dict[int, int]:
+    """
+    map the real ranks (receiver_rank) to the bcast ranks (0 ~ len(ranks) - 1),
+    which are generated in self.init_process_group_for_ranks
+    """
+    bcast_rank_map: dict[int, int] = {}
+    if not ranks:
+        bcast_rank_map = {r: r for r in range(world_size)}
+    else:
+        for i, r in enumerate(ranks):
+            bcast_rank_map[r] = i
+    return bcast_rank_map
+
+
 class P2PStore:
-    def __init__(self):
+    def __init__(self, device_manager: DeviceManager):
         from mooncake.engine import TransferEngine
 
         self.rank = int(os.getenv("RANK"))
-        gpu_count = torch.cuda.device_count()
+        gpu_count = device_manager.device_module.device_count()
         local_rank = self.rank % gpu_count
-        device = _get_my_rdma_device(local_rank, gpu_count, _get_rdma_devices())
-        self.ip = _get_ip()
+        device_type = device_manager.device_type
+        if device_type == "npu" and os.getenv("PS_P2P_STORE_RDMA_DEVICES") is None:
+            self.device = ""
+        else:
+            self.device = _get_my_rdma_device(local_rank, gpu_count, _get_rdma_devices())
+        self.ip = get_ip()
 
         # we will start at most 8 ps processes, so we use 8 retries to avoid port conflicts in extreme cases
         retry_count = 8
         for i in range(retry_count):
             self.engine = TransferEngine()
-            ret = self.engine.initialize(self.ip, "P2PHANDSHAKE", "rdma", device)
+            ret = self.engine.initialize(
+                self.ip,
+                "P2PHANDSHAKE",
+                "ascend_direct" if device_type == "npu" else "rdma",
+                self.device,
+            )
             if ret == 0:
                 break
             # sleep 0.5 ~ 2.0s, to avoid port conflicts when two processes retry at the same time
@@ -553,7 +709,7 @@ class P2PStore:
         self.port = self.engine.get_rpc_port()
         self.named_tensors: dict[str, torch.Tensor] = {}
         logger.info(
-            f"[rank{self.rank}] p2p store initialized, addr is {self.addr}, rdma device is {device}"
+            f"[rank{self.rank}] p2p store initialized, addr is {self.addr}, rdma device is {self.device}"
         )
 
     @property
@@ -606,15 +762,18 @@ class ParameterServer:
         Args:
             auto_pg: Whether to automatically initialize the process group.
                 Notice that if auto_pg is True, will destroy the process group after update.
-            mem_fraction: The proportion (as a fraction) of the current free CUDA memory for allocation.
+            mem_fraction: The proportion (as a fraction) of the current free device memory for allocation.
         """
         self._rank = rank or int(os.environ.get("RANK", None))
         self._world_size = world_size or int(os.environ.get("WORLD_SIZE", None))
-        self._gpu_count = gpu_count or torch.cuda.device_count()
+        self.device_manager = DeviceManager()
+        self._gpu_count = gpu_count or self.device_manager.device_module.device_count()
         self._local_rank = self._rank % self._gpu_count
         self._auto_pg = auto_pg
         self._all_hosts = []
         self._global_device_uuids: list[str] = []
+        self._local_rdma_devices: dict[str, set[int]] = defaultdict(set)
+        self._remote_rdma_devices: dict[str, set[int]] = defaultdict(set)
         self._mem_fraction = mem_fraction or 0.9
 
         assert self._rank is not None and self._rank >= 0, self._rank
@@ -622,7 +781,7 @@ class ParameterServer:
         assert (
             self._gpu_count is not None
             and self._gpu_count > 0
-            and self._gpu_count <= torch.cuda.device_count()
+            and self._gpu_count <= self.device_manager.device_module.device_count()
         ), self._gpu_count
         assert (
             self._mem_fraction is not None and self._mem_fraction > 0 and self._mem_fraction <= 1
@@ -634,15 +793,17 @@ class ParameterServer:
         self._memory_pool: dict[str, list[MemoryBuffer]] = {}
         # 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.
+        device_index = self._local_rank
+        self.device_manager.device_module.set_device(device_index)
         try:
-            self._p2p_store = P2PStore()
+            self._p2p_store = P2PStore(self.device_manager)
         except ImportError as e:
             logger.warning(f"[rank{self._rank}] fail to initialize p2p store due to {e}")
             self._p2p_store = None
 
-        device_index = self._local_rank
-        torch.cuda.set_device(device_index)
-        self._device_uuid = _get_physical_gpu_id(device_index)
+        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 _logger_rank0(self, msg: str):
         if self._local_rank == 0:
@@ -653,6 +814,13 @@ class ParameterServer:
 
     def load_metas(self, metas: dict[int, MemoryBufferMetaList]):
         self._current_global_parameter_metas = metas
+        self._remote_rdma_devices = defaultdict(set)
+        for i, meta in self._current_global_parameter_metas.items():
+            assert meta.rdma_device is not None, "meta.rdma_device should not be None"
+            assert meta.p2p_store_addr is not None, "meta.p2p_store_addr should not be None"
+            self._remote_rdma_devices[
+                meta.rdma_device + "@" + meta.p2p_store_addr.split(":")[0]
+            ].add(i)
 
     def register_checkpoint(
         self,
@@ -724,13 +892,15 @@ class ParameterServer:
                 for x in self._memory_pool.get(checkpoint_name, [])
             ],
             p2p_store_addr=None if self._p2p_store is None else self._p2p_store.addr,
-            host_ip=_get_ip(),
+            host_ip=get_ip(),
             device_uuid=self._device_uuid,
+            rdma_device=self._rdma_device or "",
         )
 
         dist.all_gather_object(metas_lst, metas)
 
         self._current_global_parameter_metas = {}
+
         num_parameters = 0
         all_hosts: list[str] = []
         global_device_uuids: list[str] = []
@@ -741,12 +911,24 @@ class ParameterServer:
             if not self._global_device_uuids:
                 global_device_uuids.append(metas_buckets.device_uuid)
             if metas_buckets.memory_buffer_metas_list:
-                self._current_global_parameter_metas[i] = metas_buckets
+                self._current_global_parameter_metas[i] = MemoryBufferMetaList(
+                    memory_buffer_metas_list=metas_buckets.memory_buffer_metas_list,
+                    p2p_store_addr=metas_buckets.p2p_store_addr,
+                    rdma_device=metas_buckets.rdma_device,
+                )
                 num_parameters += sum(len(x.metas) for x in metas_buckets.memory_buffer_metas_list)
+            self._local_rdma_devices[
+                metas_buckets.rdma_device + "@" + metas_buckets.p2p_store_addr.split(":")[0]
+                if metas_buckets.p2p_store_addr
+                else metas_buckets.host_ip
+            ].add(i)
         if not self._all_hosts:
             self._all_hosts = all_hosts
         if not self._global_device_uuids:
             self._global_device_uuids = global_device_uuids
+        # Sender node and Receiver node have the same GPU-rdma_device topology is considered as default.
+        # Rewrite the sender's topology (_remote_rdma_devices) by calling load_metas.
+        self._remote_rdma_devices = self._local_rdma_devices.copy()
         logger.info(
             f"[rank{self._rank}] gather parameter metas finished, num_parameters: {num_parameters}"
         )
@@ -775,7 +957,7 @@ class ParameterServer:
             is_master=self._rank == 0,
         )
         dist.init_process_group(
-            backend="nccl",
+            backend=self.device_manager.backend,
             world_size=self._world_size,
             rank=self._rank,
             timeout=timeout,
@@ -801,6 +983,7 @@ class ParameterServer:
                 If set, will use p2p to update to the ranks, this is flexible to update to a group of ranks,
                 which is useful in disaggregated architecture.
         """
+        assert req_func is not None, "req_func is required"
         try:
             # if both ranks is None or [], it will use fully broadcast to update to all ranks
             if not ranks:
@@ -808,32 +991,31 @@ class ParameterServer:
                     self.init_process_group()
                 self._update_per_bucket(checkpoint_name, req_func)
             else:
-                if not self._auto_pg and self._rank not in ranks:
-                    return
                 if self._auto_pg:
                     if dist.is_initialized():
                         dist.destroy_process_group()
                         # HACK: wait 2s to ensure destroy is finished
                         time.sleep(2)
-                    if self._rank not in ranks:
-                        return
                     self.init_process_group_for_ranks(ranks)
-                self._update_per_bucket_p2p(checkpoint_name, req_func, ranks)
-            if self._auto_pg:
-                dist.destroy_process_group()
-
-            torch.cuda.empty_cache()
+                if self._rank not in ranks:
+                    return
+                self._update_per_bucket(checkpoint_name, req_func, ranks)
 
-            logger.info(
-                f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} done. "
-                f"Current CUDA allocated {torch.cuda.memory_allocated() / 1024 / 1024} MB, "
-                f"reserved {torch.cuda.memory_reserved() / 1024 / 1024} MB."
-            )
         except Exception as e:
             logger.exception(
                 f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} error {e}"
             )
             raise
+        finally:
+            if self._auto_pg and (not ranks or self._rank in ranks):
+                dist.destroy_process_group()
+
+            self.device_manager.device_module.empty_cache()
+            logger.info(
+                f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} done. "
+                f"Current device allocated {self.device_manager.device_module.memory_allocated() / 1024 / 1024} MB, "
+                f"reserved {self.device_manager.device_module.memory_reserved() / 1024 / 1024} MB."
+            )
 
     def _bind_zmq_socket(self) -> tuple[zmq.Socket, list[tuple[str, str]]]:
         def zmq_handle(device_uuid: str) -> str:
@@ -850,14 +1032,16 @@ class ParameterServer:
         # auto detect bucket size
         tensor = torch.tensor(
             [
-                # proportion of current cuda free memory bytes
-                int(float(torch.cuda.mem_get_info()[0]) * self._mem_fraction),
+                # proportion of current device free memory bytes
+                int(
+                    float(self.device_manager.device_module.mem_get_info()[0]) * self._mem_fraction
+                ),
                 # we use negative value to reuse allreduce min operation
                 # for getting the max value of zmq_addr_counter in all ranks
                 -self._zmq_addr_counter,
             ],
             dtype=torch.int64,
-            device="cuda",
+            device=self.device_manager.device_type,
         )
         dist.all_reduce(tensor, op=dist.ReduceOp.MIN)
         tensor = tensor.cpu()
@@ -920,7 +1104,7 @@ class ParameterServer:
         assert offset == bucket.size, f"offset {offset} != bucket_size {bucket.size}"
         if owner_rank is not None:
             self._p2p_store.batch_transfer_sync_read(target_addr, buf_ptrs, remote_ptrs, lens)
-        torch.cuda.synchronize()
+        self.device_manager.device_module.synchronize()
 
     def init_process_group_for_ranks(
         self,
@@ -960,73 +1144,12 @@ class ParameterServer:
             master_addr, master_port, len(ranks), is_master=rank == 0, timeout=timeout
         )
         dist.init_process_group(
-            backend="nccl", world_size=len(ranks), rank=rank, timeout=timeout, store=store
-        )
-
-    def _update_per_bucket_p2p(
-        self,
-        checkpoint_name: str,
-        req_func: Callable[[list[tuple[str, str]]], None],
-        ranks: list[int],
-    ):
-        assert self._p2p_store is not None, "p2p store is not initialized"
-        assert ranks, "ranks should be set"
-        if len(self._current_global_parameter_metas) == 0:
-            raise ValueError("parameter metas is empty")
-        assert dist.is_initialized(), (
-            "process group is not initialized when update model per bucket p2p"
-        )
-
-        need_update = self._rank in ranks
-        logger.info(
-            f"[rank{self._rank}] update checkpoint {checkpoint_name} p2p, {need_update=} with {ranks=}, "
-            f"gpu_count {self._gpu_count}, world_size {self._world_size}"
-        )
-
-        if not need_update:
-            return
-
-        # first execute a barrier to avoid subsequent cuda oom
-        dist.barrier()
-
-        bucket_size, _ = self._detect_bucket_size(disable_h2d_buffer=True)
-        buffer = torch.empty(bucket_size * 2, dtype=torch.uint8, device="cuda")
-        ipc_buffer_name = "__ipc_buffer___"
-        self._p2p_store.register_named_tensors({ipc_buffer_name: buffer})
-        logger.info(
-            f"[rank{self._rank}] register buffer, shape={buffer.shape}, dtype={buffer.dtype}, data_ptr={buffer.data_ptr()}, nbytes={buffer.nbytes}"
+            backend=self.device_manager.backend,
+            world_size=len(ranks),
+            rank=rank,
+            timeout=timeout,
+            store=store,
         )
-        handle = reduce_tensor(buffer)
-
-        buckets = _gen_h2d_buckets(self._current_global_parameter_metas, bucket_size)
-        socket, socket_paths = self._bind_zmq_socket()
-        req_thread = threading.Thread(
-            target=req_func,
-            args=(socket_paths,),
-        )
-        req_thread.start()
-        socket.send_pyobj(handle)
-        for gidx, (owner_rank, bucket) in enumerate(buckets):
-            self._logger_rank0(
-                f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} owner_rank {owner_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. "
-            )
-            _buffer = buffer[gidx % 2 * bucket_size : gidx % 2 * bucket_size + bucket.size]
-            if dist.get_rank() == 0:
-                self._copy_to_buffer(checkpoint_name, bucket, _buffer, owner_rank)
-            # broadcast the collected data to all ranks
-            dist.broadcast(_buffer, src=0)
-            socket.recv()
-            dist.barrier()
-            socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size))
-
-        socket.recv()
-        socket.send_pyobj(None)
-        socket.recv()
-        req_thread.join()
-        dist.barrier()
-        socket.close()
-        self._p2p_store.unregister_named_tensors([ipc_buffer_name])
-        torch.cuda.empty_cache()
 
     def _get_addr_ptrs(self, owner_rank: int) -> tuple[str, list[tuple[int, int]]]:
         addr = self._current_global_parameter_metas[owner_rank].p2p_store_addr
@@ -1057,38 +1180,65 @@ class ParameterServer:
         self,
         checkpoint_name: str,
         req_func: Callable[[list[tuple[str, str]]], None],
+        ranks: list[int] | None = None,
     ):
-        if len(self._current_global_parameter_metas) == 0:
-            raise ValueError("parameter metas is empty")
-
+        assert len(self._current_global_parameter_metas) != 0, "parameter metas is empty"
         assert dist.is_initialized(), "process group is not initialized"
+        # 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}")
+        # if ranks is set, it will use p2p to update to the ranks
+        else:
+            assert self._p2p_store is not None, "p2p store is not initialized"
+            assert ranks, "ranks should be set"
 
-        logger.info(f"[rank{self._rank}] update checkpoint {checkpoint_name}")
+            need_update = self._rank in ranks
+            logger.info(
+                f"[rank{self._rank}] update checkpoint {checkpoint_name} p2p, {need_update=} with {ranks=}, "
+                f"gpu_count {self._gpu_count}, world_size {self._world_size}"
+            )
+
+            if not need_update:
+                return
+            # first execute a barrier to avoid subsequent device oom
+            dist.barrier()
 
         bucket_size, disable_h2d_buffer = self._detect_bucket_size()
-        buckets = _gen_h2d_buckets(self._current_global_parameter_metas, bucket_size)
+        buckets = _gen_h2d_buckets(
+            self._current_global_parameter_metas,
+            bucket_size,
+            self._local_rdma_devices,
+            self._remote_rdma_devices,
+            ranks,
+        )
 
         h2d_buffer: torch.Tensor | None = (
             None
             if disable_h2d_buffer
-            else torch.empty(bucket_size, dtype=torch.uint8, device="cuda")
+            else torch.empty(bucket_size, dtype=torch.uint8, device=self.device_manager.device_type)
         )
-
-        owner_rank_buckets: list[H2DBucket] = []
-        for owner_rank, bucket in buckets:
-            if owner_rank != self._rank:
+        # 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:
                 continue
-            owner_rank_buckets.append(bucket)
+            receiver_rank_buckets.append((owner_rank, bucket))
 
-        buffer = torch.empty(bucket_size * 2, dtype=torch.uint8, device="cuda")
+        buffer = torch.empty(
+            bucket_size * 2, dtype=torch.uint8, device=self.device_manager.device_type
+        )
         handle = reduce_tensor(buffer)
 
-        buckets_by_owner_rank: dict[int, list[H2DBucket]] = defaultdict(list)
+        buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list)
         max_len = 0
-        for owner_rank, bucket in buckets:
-            buckets_by_owner_rank[owner_rank].append(bucket)
-            if len(buckets_by_owner_rank[owner_rank]) > max_len:
-                max_len = len(buckets_by_owner_rank[owner_rank])
+        for receiver_rank, _, bucket in buckets:
+            buckets_by_receiver_rank[receiver_rank].append(bucket)
+            if len(buckets_by_receiver_rank[receiver_rank]) > max_len:
+                max_len = len(buckets_by_receiver_rank[receiver_rank])
 
         socket, socket_paths = self._bind_zmq_socket()
         req_thread = threading.Thread(
@@ -1099,43 +1249,66 @@ class ParameterServer:
         socket.send_pyobj(handle)
 
         gidx = 0
-        for i in range(max_len):
-            if i < len(owner_rank_buckets) and not disable_h2d_buffer:
-                self._copy_to_buffer(checkpoint_name, owner_rank_buckets[i], h2d_buffer)
-
-            for owner_rank, _buckets in buckets_by_owner_rank.items():
-                if i >= len(_buckets):
-                    continue
-                bucket = _buckets[i]
-                alloc, reserved = (
-                    torch.cuda.memory_allocated() / 1024 / 1024,
-                    torch.cuda.memory_reserved() / 1024 / 1024,
-                )
-                self._logger_rank0(
-                    f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} owner_rank {owner_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. "
-                    f"Current CUDA allocated {alloc:.2f} MB, "
-                    f"reserved {reserved:.2f} MB."
-                )
-                start = gidx % 2 * bucket_size
-                buffer_b: torch.Tensor = buffer[start : start + bucket.size]
-                if owner_rank == self._rank:
-                    if disable_h2d_buffer:
-                        self._copy_to_buffer(checkpoint_name, bucket, buffer_b)
-                    else:
-                        buffer_b.data.copy_(h2d_buffer[: bucket.size])
-                dist.broadcast(buffer_b, src=owner_rank)
-                socket.recv()
-                dist.barrier()
-                socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size))
-                gidx += 1
-
-        socket.recv()
-        socket.send_pyobj(None)
-        socket.recv()
-        req_thread.join()
-        dist.barrier()
-        socket.close()
-        torch.cuda.empty_cache()
+        ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64)
+        bcast_rank_map = _get_bcast_rank_map(self._world_size, ranks)
+        try:
+            for i in range(max_len):
+                if i < len(receiver_rank_buckets) and not disable_h2d_buffer:
+                    self._copy_to_buffer(
+                        checkpoint_name,
+                        receiver_rank_buckets[i][1],
+                        h2d_buffer,
+                        receiver_rank_buckets[i][0] if ranks else None,
+                    )
+                for receiver_rank, _buckets in buckets_by_receiver_rank.items():
+                    if i >= len(_buckets):
+                        continue
+                    bucket = _buckets[i]
+                    alloc, reserved = (
+                        self.device_manager.device_module.memory_allocated() / 1024 / 1024,
+                        self.device_manager.device_module.memory_reserved() / 1024 / 1024,
+                    )
+                    self._logger_rank0(
+                        f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} receiver_rank {receiver_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. "
+                        f"Current device allocated {alloc:.2f} MB, "
+                        f"reserved {reserved:.2f} MB."
+                    )
+                    start = gidx % 2 * bucket_size
+                    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)
+                        else:
+                            buffer_b.data.copy_(h2d_buffer[: bucket.size])
+                    brank = bcast_rank_map[receiver_rank]
+                    dist.broadcast(buffer_b, src=brank)
+                    resp = socket.recv()
+                    if resp != b"":
+                        exception_obj = pickle.loads(resp)
+                        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}"
+                        )
+                        ret_code.fill_(1)
+                    dist.all_reduce(ret_code, op=dist.ReduceOp.SUM)
+                    self.device_manager.device_module.synchronize()
+                    if ret_code.item() != 0:
+                        # quit early if any rank failed
+                        socket.send_pyobj(RuntimeError("Some workers failed to update weights"))
+                        raise RuntimeError("Failed to update weights due to remote errors")
+                    socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size))
+                    gidx += 1
+
+            socket.recv()
+            socket.send_pyobj(None)
+            socket.recv()
+        finally:
+            req_thread.join()
+            dist.barrier()
+            socket.close()
+            if ranks and h2d_buffer is not None:
+                self._p2p_store.unregister_named_tensors([h2d_buffer_name])
+
+            self.device_manager.device_module.empty_cache()
 
 
 def _init_api(ps: ParameterServer) -> Any:
@@ -1153,6 +1326,7 @@ def _init_api(ps: ParameterServer) -> Any:
         update_url: str | None = None
         inference_group_ranks: list[int] = []
         timeout: float = 300.0
+        uds: str | None = None
 
     def wrap_exception(func: Callable[[], None]) -> Response:
         try:
@@ -1185,7 +1359,9 @@ def _init_api(ps: ParameterServer) -> Any:
                 return
             if req.inference_group_ranks:
                 socket_paths = [socket_paths[i] for i in req.inference_group_ranks]
-            request_inference_to_update(req.update_url, dict(socket_paths), timeout=req.timeout)
+            request_inference_to_update(
+                req.update_url, dict(socket_paths), timeout=req.timeout, uds=req.uds
+            )
 
         return wrap_exception(lambda: ps.update(checkpoint_name, update_func, ranks=req.ranks))
 

checkpoint_engine/worker.py

@@ -5,6 +5,8 @@ from typing import TypedDict
 import torch
 import zmq
 
+from checkpoint_engine.device_utils import DeviceManager, npu_generate_uuid
+
 
 def _rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor:
     func, args = handle
@@ -53,51 +55,105 @@ def update_weights_from_ipc(
     socket = zmq_ctx.socket(zmq.REP)
     socket.connect(zmq_handle)
     buffer: torch.Tensor | None = None
-    while True:
-        payload: tuple[Callable, tuple] | list[FlattenedTensorMetadata] | None = socket.recv_pyobj()
-        if payload is None:
-            # means the update is done
-            if post_hook is not None:
-                post_hook()
-            torch.cuda.synchronize()
-            socket.send(b"")
-            break
-        if isinstance(payload, tuple):
-            # an ipc handle that vLLM can use `func, args = handle`
-            # and `func(*args)` to rebuild GPU tensor.
-            buffer = _rebuild_ipc(payload, device_id)
-            assert buffer.dtype == torch.uint8
-            socket.send(b"")
-            continue
-        assert isinstance(payload, list)
-        run(_extract_weights(payload, buffer))
-        torch.cuda.synchronize()
+    device_manager = DeviceManager()
+    try:
+        ipc_handle: tuple[Callable, tuple] = socket.recv_pyobj()
+        assert isinstance(ipc_handle, tuple)
+        buffer = _rebuild_ipc(ipc_handle, device_id)
+        assert buffer.dtype == torch.uint8
         socket.send(b"")
+    except Exception as e:
+        socket.send_pyobj(e)
+        socket.recv()  # wait for ack
+        raise
+    try:
+        while True:
+            payload: list[FlattenedTensorMetadata] | Exception | None = socket.recv_pyobj()
+            if payload is None:  # done signal
+                if post_hook is not None:
+                    post_hook()
+                device_manager.device_module.synchronize()
+                socket.send(b"")
+                break
+            if isinstance(payload, list):  # still updating weights
+                try:
+                    run(_extract_weights(payload, buffer))
+                    device_manager.device_module.synchronize()
+                    socket.send(b"")
+                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)
+            elif isinstance(
+                payload, Exception
+            ):  # error occurred, got force quit signal from Parameter Server
+                raise payload
+            else:
+                raise TypeError(f"Unexpected payload type: {type(payload)}")
 
-    socket.close()
-    del buffer
-    gc.collect()
-    torch.cuda.empty_cache()
+    finally:
+        socket.close()
+        del buffer
+        gc.collect()
+        device_manager.device_module.empty_cache()
 
 
 class VllmColocateWorkerExtension:
     """
-    The class for vLLM's worker to inherit from, in the colocate setting.
-    By defining an extension class, the code can work no matter what is
-    the underlying worker class. This way, the code can be compatible
-    with both vLLM V0 and V1.
-    NOTE: we define this class in a separate module, and the main module
-    should pass the full qualified name as `worker_extension_cls` argument.
+    Worker extension for vLLM to update weights from checkpoint-engine.
+
+    This class provides a worker extension mechanism that allows vLLM workers to receive
+    and apply weight updates from the checkpoint-engine via IPC (Inter-Process Communication).
+    The methods in this worker extension will be injected into the vLLM worker class and
+    are callable from the `collective_rpc` API, enabling seamless weight updates for both
+    vLLM V0 and V1 versions.
+
+    Note:
+        This class is defined in a separate module. The fully qualified name
+        `checkpoint_engine.worker.VllmColocateWorkerExtension` should be passed as the
+        `worker_extension_cls` argument when initializing the vLLM worker.
     """
 
     def update_weights_from_ipc(self, zmq_handles: dict[str, str]):
+        """
+        Update model weights from checkpoint-engine via IPC communication.
+
+        This method establishes a ZMQ connection to the checkpoint-engine and receives
+        weight updates through a shared memory buffer. The update process includes:
+        1. Receiving IPC handles to reconstruct shared memory tensors
+        2. Extracting flattened metadata describing tensor weights in the shared memory tensor
+        3. Loading weights into the model
+        4. Post-processing weights after loading
+
+        Args:
+            zmq_handles: A dictionary mapping device UUIDs to ZMQ socket handles.
+                        The device UUID is platform-specific:
+                        - For CUDA: UUID from `current_platform.get_device_uuid()`
+                        - For NPU: Format "NPU-{generated_uuid}"
+
+        Raises:
+            ValueError: If the device type is not supported (not CUDA or NPU).
+            AssertionError: If the device is not properly initialized.
+
+        Note:
+            This method is called by vLLM's collective RPC mechanism. The ZMQ context
+            is lazily initialized on first call and reused for subsequent updates.
+        """
         from vllm.model_executor.model_loader.utils import process_weights_after_loading
         from vllm.platforms import current_platform
 
+        # vllm-ascend not init device
+        if current_platform.device_type == "npu" and self.device is None:
+            self.device = torch.device(f"npu:{self.local_rank}")
         assert self.device is not None
         if not hasattr(self, "_zmq_ctx") or self._zmq_ctx is None:
             self._zmq_ctx = zmq.Context()
-        device_uuid = current_platform.get_device_uuid(self.device.index)
+        if current_platform.device_type == "cuda":
+            device_uuid = current_platform.get_device_uuid(self.device.index)
+        elif current_platform.device_type == "npu":
+            device_uuid = f"NPU-{npu_generate_uuid()}"
+        else:
+            raise ValueError(f"Unsupported device type: {current_platform.device_type}")
         update_weights_from_ipc(
             self._zmq_ctx,
             zmq_handles[device_uuid],

{checkpoint_engine-0.1.3.dist-info → checkpoint_engine-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: checkpoint-engine
-Version: 0.1.3
+Version: 0.2.1
 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
@@ -38,8 +38,8 @@ updating our [Kimi-K2](https://github.com/MoonshotAI/Kimi-K2) model (1 Trillion
 
 The core weight update logic is in `ParameterServer` class, a service colocated with inference engines. It provides two implementations of weight update: Broadcast and P2P.
 
-- **Broadcast**: Used when a large number of inference instances need to update weights in synchronous. This is the fastest implementation and should be used as the default update method. See `_update_per_bucket`.
-- **P2P**: Used when new inference instances are dynamically added (due to restarts or dynamic availability) while the existing instances are already serving requests. Under this scenario, to avoid affecting the workloads on existing instances, we use the [`mooncake-transfer-engine`](https://github.com/kvcache-ai/Mooncake?tab=readme-ov-file#use-python-package) to P2P send weights from CPUs in existing instances to GPUs in new instances. See `_update_per_bucket_p2p`.
+- **Broadcast**: Used when a large number of inference instances need to update weights in synchronous. This is the fastest implementation and should be used as the default update method. See `_update_per_bucket` with `ranks == None or []`.
+- **P2P**: Used when new inference instances are dynamically added (due to restarts or dynamic availability) while the existing instances are already serving requests. Under this scenario, to avoid affecting the workloads on existing instances, we use the [`mooncake-transfer-engine`](https://github.com/kvcache-ai/Mooncake?tab=readme-ov-file#use-python-package) to P2P send weights from CPUs in existing instances to GPUs in new instances. See `_update_per_bucket` with `ranks` specified.
 
 ### Optimized Weight Broadcast
 In the *Broadcast* implementation, the checkpoint-engine holds references to sharded weights in CPU memory, and need to efficiently broadcast them to a cluster of inference instances, often under a different sharding pattern.
@@ -60,16 +60,22 @@ It then executes the transfer, where it controls the inference engine through a
 
 Pipelining naturally requires more GPU memory. When memory is not enough, checkpoint-engine will fallback to serial execution.
 
+### Optimized P2P Bucket Assignment
+In the *P2P* implementation, checkpoint-engine needs to send weights from existing instances to new instances.
+To minimize the overall transfer time, checkpoint-engine optimizes the bucket assignment for each sender-receiver pair.
+The optimization goal is to make full use of the available network bandwidth for each sender and receiver.
+See [issue #25](https://github.com/MoonshotAI/checkpoint-engine/issues/25)
+
 ## Benchmark
 
 | Model                                | Device Info  | GatherMetas | Update (Broadcast) | Update (P2P)            |
 | :----------------------------------- | :----------- | :---------- |:-------------------| :---------------------- |
-| GLM-4.5-Air (BF16)                   | 8xH800 TP8  | 0.17s       | 3.94s (1.42GiB)    | 8.83s (4.77GiB)         |
-| Qwen3-235B-A22B-Instruct-2507 (BF16) | 8xH800 TP8  | 0.46s       | 6.75s (2.69GiB)    | 16.47s (4.05GiB)        |
-| DeepSeek-V3.1 (FP8)                  | 16xH20 TP16  | 1.44s       | 12.22s (2.38GiB)   | 25.77s (3.61GiB)        |
-| Kimi-K2-Instruct (FP8)               | 16xH20 TP16  | 1.81s       | 15.45s (2.93GiB)   | 36.24s (4.46GiB)        |
-| DeepSeek-V3.1 (FP8)                  | 256xH20 TP16 | 1.40s       | 13.88s (2.54GiB)   | 33.30s (3.86 GiB) |
-| Kimi-K2-Instruct (FP8)               | 256xH20 TP16 | 1.88s       | 21.50s (2.99GiB)   | 34.49s (4.57 GiB) |
+| GLM-4.5-Air (BF16)                   | 8xH800 TP8   | 0.12s       | 3.47s (3.02GiB)    | 4.12s (3.02GiB)         |
+| Qwen3-235B-A22B-Instruct-2507 (BF16) | 8xH800 TP8   | 0.33s       | 6.22s (2.67GiB)    | 7.10s (2.68GiB)         |
+| DeepSeek-V3.1 (FP8)                  | 16xH20 TP16  | 1.17s       | 10.19s (5.39GiB)   | 11.80s (5.41GiB)        |
+| Kimi-K2-Instruct (FP8)               | 16xH20 TP16  | 1.33s       | 14.36s (5.89GiB)   | 17.49s (5.91GiB)        |
+| DeepSeek-V3.1 (FP8)                  | 256xH20 TP16 | 0.80s       | 11.33s (8.00GiB)   | 11.81s (8.00GiB)        |
+| Kimi-K2-Instruct (FP8)               | 256xH20 TP16 | 1.22s       | 16.04s (8.00GiB)   | 16.75s (8.00GiB)        |
 
 All results above are tested by [`examples/update.py`](./examples/update.py) and use [vLLM v0.10.2rc1](https://github.com/vllm-project/vllm/tree/v0.10.2rc1) as inference engine. Some notes:
 
@@ -77,6 +83,7 @@ All results above are tested by [`examples/update.py`](./examples/update.py) and
 * Device Info: we tested various combination of devices and parallelism setups. For example, a 256-GPU TP16 setup means that we deploy 16 vLLM instances, each with 16-way tensor parallelism.
 * Since update duration is related to IPC bucket size, we provide the bucket size in the table.
 * The P2P time were tested for updating no more than two nodes (16 GPUs) (`ParameterServer.update(ranks=range(0, 16))`) out of the entire cluster.
+* We bind each GPU to its corresponding NUMA node to ensure stable H2D transfer speeds.
 
 ## Installation
 
@@ -92,7 +99,7 @@ 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. If not set, it will read all RDMA devices and try to divide them into each rank.
+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
 
@@ -162,14 +169,64 @@ A [PR](https://github.com/vllm-project/vllm/pull/24488) is opened to the vLLM pr
 Run a simple correctness test for checkpoint_engine
 
 ```bash
-torchrun --nproc-per-node 8 tests/test_update.py
+pytest tests/test_update.py
+```
+
+`test_update.py` are only designed to run with `pytest`. Please don't run it directly with `torchrun`.
+
+Other unit tests can also be done with pytest. Only test_update.py requires GPUs, other tests can be run on CPUs. Only to run CPU tests, use:
+
+```bash
+pytest tests/ -m "not gpu"
+```
+
+## 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.
+
+### Quick Start
+
+**1. Install checkpoint-engine:**
+```bash
+pip install 'checkpoint-engine[p2p]'
+```
+
+**2. Launch SGLang server:**
+```bash
+python -m sglang.launch_server \
+    --model-path $MODEL_PATH \
+    --tp 8 \
+    --load-format dummy \
+    --wait-for-initial-weights
+```
+
+**3. Run checkpoint engine:**
+```bash
+python -m sglang.srt.checkpoint_engine.update \
+    --update-method broadcast \
+    --checkpoint-path $MODEL_PATH \
+    --inference-parallel-size 8
 ```
 
+### Multi-Node Setup
+
+For 2-node setup, run the same commands on both nodes with appropriate `--host` and distributed training parameters.
+
+### Key Options
+
+**SGLang Server:**
+- `--wait-for-initial-weights`: Wait for checkpoint engine before becoming ready
+- `--load-format dummy`: Enable overlapping initialization tasks
+
+**Checkpoint Engine:**
+- `--update-method`: Choose `broadcast`, `p2p`, or `all`
+- `--inference-parallel-size`: Number of parallel processes
+- `--checkpoint-path`: Model checkpoint directory
+
 ## Limitations and Future Work
 
-- This project is currently only tested with vLLM. But it is easy to integrate with other frameworks like SGLang.
+- This project is currently tested with vLLM and SGLang. Integration with other frameworks is planned for future releases.
 - The perfect three-stage pipeline mentioned in our paper is currently not implemented. This could be useful for architectures where H2D and broadcast do not conflict in PCIE.
-- The P2P update method is currently not the optimal implementation since it will receive data only in rank 0 and broadcast to others synchronizely. This is a potential optimization in the future.
 
 ## Acknowledgments
 

checkpoint_engine-0.2.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+checkpoint_engine/__init__.py,sha256=Zj4I008kn9R6fYr0lVBzcQSnvckLpX2s1ljCOOqV1c8,87
+checkpoint_engine/_version.py,sha256=vYqoJTG51NOUmYyL0xt8asRK8vUT4lGAdal_EZ59mvw,704
+checkpoint_engine/device_utils.py,sha256=iKrof60j3CY3fStRTq3DRTt_kE1vYoEWHhAeyh0lByA,3020
+checkpoint_engine/ps.py,sha256=Mgfr_MXYYZ_6JKqD5kIGDBWCYNWAtIPfAppP_cFu604,57781
+checkpoint_engine/worker.py,sha256=5TzDgTPew6Ts9sMOzecalLCR1p_ZwfeKPdzzr68kAQg,6564
+checkpoint_engine-0.2.1.dist-info/licenses/LICENCE,sha256=D3gPmHKpGtF1yxYNhqjtBtZY_brZjDotJTzpnmClzlY,1067
+checkpoint_engine-0.2.1.dist-info/METADATA,sha256=7E7NhehWHpS6QVkLup-oCm350wdbiZX8CY3jmmJP0bU,11315
+checkpoint_engine-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+checkpoint_engine-0.2.1.dist-info/top_level.txt,sha256=66sik_1eLakLYmcllOEJzFaNbSfjsueuP0tHYEzhMSs,18
+checkpoint_engine-0.2.1.dist-info/RECORD,,

checkpoint_engine-0.1.3.dist-info/RECORD

@@ -1,9 +0,0 @@
-checkpoint_engine/__init__.py,sha256=Zj4I008kn9R6fYr0lVBzcQSnvckLpX2s1ljCOOqV1c8,87
-checkpoint_engine/_version.py,sha256=q5nF98G8SoVeJqaknL0xdyxtv0egsqb0fK06_84Izu8,704
-checkpoint_engine/ps.py,sha256=9dXRXi0QDPoRYrgGKAYvdmDFBXejgusjR0ltbii5_B0,49134
-checkpoint_engine/worker.py,sha256=ZmJTHeNPbnE8sPInfrghj9jeRDkMUSQO906o1UoJv-E,3748
-checkpoint_engine-0.1.3.dist-info/licenses/LICENCE,sha256=D3gPmHKpGtF1yxYNhqjtBtZY_brZjDotJTzpnmClzlY,1067
-checkpoint_engine-0.1.3.dist-info/METADATA,sha256=y96dMjEOKWaO_PA0h5BX8G3Ku7Tt1jCU09uIf8iYgic,9322
-checkpoint_engine-0.1.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-checkpoint_engine-0.1.3.dist-info/top_level.txt,sha256=66sik_1eLakLYmcllOEJzFaNbSfjsueuP0tHYEzhMSs,18
-checkpoint_engine-0.1.3.dist-info/RECORD,,