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

checkpoint_engine-0.2.1/.github/workflows/cpu-tests.yml

@@ -0,0 +1,30 @@
+name: CPU Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    types: [opened, synchronize, reopened]
+
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install pytest
+        pip install .[p2p]
+    - name: Do CPU tests with pytest
+      run: |
+        pytest -v -m "not gpu" tests/

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

@@ -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.1.3 → checkpoint_engine-0.2.1}/README.md

@@ -14,8 +14,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.
@@ -36,16 +36,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:
 
@@ -53,6 +59,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
 
@@ -68,7 +75,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
 
@@ -138,14 +145,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.1.3 → checkpoint_engine-0.2.1}/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 = 'g8a60e65ba'
+__commit_id__ = commit_id = 'g279a908a9'

checkpoint_engine-0.2.1/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")