94057c3d3e
PR Test (NPU) / check-changes (push) Has been cancelled
PR Test (NPU) / pr-gate (push) Has been cancelled
PR Test (NPU) / set-image-config (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
PR Test (Arm64) / pr-gate (push) Has been cancelled
PR Test (Arm64) / check-changes (push) Has been cancelled
PR Test (Arm64) / build-test (push) Has been cancelled
PR Test (sgl-router) / gate (push) Has been cancelled
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
PR Test (sgl-router) / finish (push) Has been cancelled
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
PR Test (Xeon) / pr-gate (push) Has been cancelled
PR Test (Xeon) / check-changes (push) Has been cancelled
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
PR Test (XPU) / check-changes (push) Has been cancelled
PR Test (XPU) / pr-gate (push) Has been cancelled
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / finish (push) Has been cancelled
CI Model Inventory / build-inventory (push) Has been cancelled
Lint / lint (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
PR Test (SMG) / build-wheel (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
PR Test (SMG) / python-unit-tests (push) Has been cancelled
PR Test (SMG) / unit-tests (push) Has been cancelled
PR Test (SMG) / benchmarks (push) Has been cancelled
PR Test (SMG) / chat-completions (push) Has been cancelled
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
PR Test (SMG) / e2e (push) Has been cancelled
PR Test (SMG) / docker-build-test (push) Has been cancelled
PR Test (SMG) / k8s-integration (push) Has been cancelled
PR Test (SMG) / finish (push) Has been cancelled
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Release SGLang Kernels / release-musa43 (push) Has been cancelled
258 lines
7.5 KiB
Plaintext
258 lines
7.5 KiB
Plaintext
---
|
|
title: "Checkpoint Engine Integration"
|
|
metatags:
|
|
description: "SGLang checkpoint engine: distributed model weight loading, parallel multi-node setup, broadcast and P2P modes. Reduces loading time for large models."
|
|
---
|
|
The SGLang checkpoint engine integration provides an efficient way to load model weights using a distributed checkpoint loading system. This feature significantly reduces model loading time, especially for large models and multi-node setups, by parallelizing the weight loading process across multiple processes and nodes.
|
|
|
|
## Overview
|
|
|
|
The checkpoint engine integration allows SGLang to:
|
|
- Load model weights in parallel using multiple processes
|
|
- Distribute weight loading across multiple nodes to increase effective disk bandwidth
|
|
- Overlap weight loading with other initialization tasks like CUDA graph capture
|
|
- Support both single-node and multi-node deployments
|
|
|
|
## Installation
|
|
|
|
First, install the checkpoint engine package:
|
|
|
|
```bash Command
|
|
pip install 'checkpoint-engine[p2p]'
|
|
```
|
|
|
|
## Architecture
|
|
|
|
The system consists of two main components:
|
|
|
|
1. **SGLang Server**: Runs with `--wait-for-initial-weights` flag to wait for weights before becoming ready
|
|
2. **Checkpoint Engine Workers**: Separate processes (managed by torchrun) that load and distribute model weights
|
|
|
|
The checkpoint engine uses a parameter server architecture with support for:
|
|
- **Broadcast mode**: Weights are broadcast from loading processes to inference processes
|
|
- **P2P mode**: Direct peer-to-peer weight transfer between processes
|
|
- **All mode**: Combination of both broadcast and P2P methods
|
|
|
|
## Usage Examples
|
|
|
|
### Single Node Setup
|
|
|
|
**Terminal 1 - Launch SGLang Server:**
|
|
```bash Command
|
|
python -m sglang.launch_server \
|
|
--model-path Qwen/Qwen3-8B \
|
|
--tp 8 \
|
|
--load-format dummy \
|
|
--wait-for-initial-weights
|
|
```
|
|
|
|
**Terminal 2 - Run Checkpoint Engine:**
|
|
|
|
Using sglang entrypoint:
|
|
```bash Command
|
|
python -m sglang.srt.checkpoint_engine.update \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 8
|
|
```
|
|
|
|
Using torchrun directly:
|
|
```bash Command
|
|
torchrun --nproc-per-node 8 \
|
|
examples/checkpoint_engine/update.py \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 8
|
|
```
|
|
|
|
### Multi-Node Setup (2 Nodes)
|
|
|
|
**Node 0:**
|
|
|
|
Launch SGLang server:
|
|
```bash Command
|
|
python -m sglang.launch_server \
|
|
--model-path Qwen/Qwen3-8B \
|
|
--tp 8 \
|
|
--load-format dummy \
|
|
--wait-for-initial-weights \
|
|
--host [IP]
|
|
```
|
|
|
|
Run checkpoint engine:
|
|
|
|
Using sglang entrypoint (recommended):
|
|
```bash Command
|
|
python -m sglang.srt.checkpoint_engine.update \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 8
|
|
```
|
|
|
|
Using torchrun directly:
|
|
```bash Command
|
|
torchrun --nproc-per-node 8 \
|
|
--nnodes 2 \
|
|
--node-rank 0 \
|
|
--master-addr [IP] \
|
|
--master-port 29500 \
|
|
examples/checkpoint_engine/update.py \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 8
|
|
```
|
|
|
|
**Node 1:**
|
|
|
|
Launch SGLang server:
|
|
```bash Command
|
|
python -m sglang.launch_server \
|
|
--model-path Qwen/Qwen3-8B \
|
|
--tp 8 \
|
|
--load-format dummy \
|
|
--wait-for-initial-weights \
|
|
--host [IP]
|
|
```
|
|
|
|
Run checkpoint engine:
|
|
|
|
Using sglang entrypoint (recommended):
|
|
```bash Command
|
|
python -m sglang.srt.checkpoint_engine.update \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 8
|
|
```
|
|
|
|
Using torchrun directly:
|
|
```bash Command
|
|
torchrun --nproc-per-node 8 \
|
|
--nnodes 2 \
|
|
--node-rank 1 \
|
|
--master-addr [IP] \
|
|
--master-port 29500 \
|
|
examples/checkpoint_engine/update.py \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 8
|
|
```
|
|
|
|
### Multi-Node Setup with Tensor Parallelism (TP=16)
|
|
|
|
**Node 0:**
|
|
|
|
Launch SGLang server:
|
|
```bash Command
|
|
python -m sglang.launch_server \
|
|
--model-path Qwen/Qwen3-8B \
|
|
--tp 8 \
|
|
--load-format dummy \
|
|
--wait-for-initial-weights \
|
|
--host [IP] \
|
|
--dist-init-addr [IP]:9120 \
|
|
--nnodes 2 \
|
|
--node-rank 0
|
|
```
|
|
|
|
Run checkpoint engine:
|
|
|
|
Using sglang entrypoint (recommended):
|
|
```bash Command
|
|
python -m sglang.srt.checkpoint_engine.update \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 16
|
|
```
|
|
|
|
Using torchrun directly:
|
|
```bash Command
|
|
torchrun --nproc-per-node 8 \
|
|
--nnodes 2 \
|
|
--node-rank 0 \
|
|
--master-addr [IP] \
|
|
--master-port 29500 \
|
|
examples/checkpoint_engine/update.py \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 16
|
|
```
|
|
|
|
**Node 1:**
|
|
|
|
Launch SGLang server:
|
|
```bash Command
|
|
python -m sglang.launch_server \
|
|
--model-path Qwen/Qwen3-8B \
|
|
--tp 8 \
|
|
--load-format dummy \
|
|
--wait-for-initial-weights \
|
|
--host [IP] \
|
|
--dist-init-addr [IP]:9120 \
|
|
--nnodes 2 \
|
|
--node-rank 1
|
|
```
|
|
|
|
Run checkpoint engine:
|
|
|
|
Using sglang entrypoint (recommended):
|
|
```bash Command
|
|
python -m sglang.srt.checkpoint_engine.update \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 16
|
|
```
|
|
|
|
Using torchrun directly:
|
|
```bash Command
|
|
torchrun --nproc-per-node 8 \
|
|
--nnodes 2 \
|
|
--node-rank 1 \
|
|
--master-addr [IP] \
|
|
--master-port 29500 \
|
|
examples/checkpoint_engine/update.py \
|
|
--update-method broadcast \
|
|
--checkpoint-path /path/to/Qwen/Qwen3-8B/ \
|
|
--inference-parallel-size 16
|
|
```
|
|
|
|
## Configuration Options
|
|
|
|
### SGLang Server Options
|
|
|
|
- `--load-format dummy`: Use dummy format for initial loading (allows overlapping with other tasks)
|
|
- `--wait-for-initial-weights`: Wait for checkpoint engine to provide weights before becoming ready
|
|
- `--host`: Host address for multi-node setups
|
|
- `--dist-init-addr`: Distributed initialization address for tensor parallelism
|
|
|
|
### Checkpoint Engine Options
|
|
|
|
- `--update-method`: Weight update method (`broadcast`, `p2p`, or `all`)
|
|
- `--checkpoint-path`: Path to model checkpoint directory
|
|
- `--inference-parallel-size`: Number of inference parallel processes
|
|
- `--endpoint`: SGLang server endpoint (default: `http://localhost:19730`)
|
|
- `--checkpoint-name`: Name for the checkpoint (default: `my-checkpoint-iter-0`)
|
|
- `--save-metas-file`: File to save checkpoint metadata
|
|
- `--load-metas-file`: File to load checkpoint metadata from
|
|
- `--uds`: Unix domain socket path for communication
|
|
- `--weight-version`: Version identifier for weights
|
|
|
|
## Performance Benefits
|
|
|
|
The checkpoint engine provides significant time savings in two main aspects:
|
|
|
|
1. **Multi-node Loading**: Each node only loads a portion of weights from disk, effectively increasing disk bandwidth. More participating nodes provide greater acceleration. Preliminary tests show 20-second acceleration when loading DeepSeek-R1 on H20-3e with two nodes.
|
|
|
|
2. **Single Process Optimization**: Using dummy format allows overlapping disk-to-CPU transfer with CUDA graph capture and other initialization tasks, providing additional time savings.
|
|
|
|
## Troubleshooting
|
|
|
|
- Ensure checkpoint engine package is installed: `pip install 'checkpoint-engine[p2p]'`
|
|
- Verify network connectivity between nodes in multi-node setups
|
|
- Check that the checkpoint path contains valid model files
|
|
- Monitor logs for connection errors between SGLang server and checkpoint engine
|
|
- Use `--sleep-time` parameter to add delays if needed for debugging
|
|
|
|
## References
|
|
|
|
- [Checkpoint Engine Repository](https://github.com/MoonshotAI/checkpoint-engine)
|