检查点引擎集成

Checkpoint Engine 集成#

SGLang 的 checkpoint engine 集成提供了一种使用分布式检查点加载系统高效加载模型权重的方法。该功能通过在多个进程和节点之间并行化权重加载过程，显著缩短了模型加载时间，特别是对于大型模型和多节点设置。

概述#

checkpoint engine 集成使 SGLang 能够：

  • 使用多个进程并行加载模型权重

  • 在多个节点之间分布权重加载，以提高有效磁盘带宽

  • 使权重加载与 CUDA graph 捕获等其他初始化任务重叠进行

  • 支持单节点和多节点部署

安装#

首先，安装 checkpoint engine 软件包

pip install 'checkpoint-engine[p2p]'

架构#

该系统由两个主要组件组成：

  1. SGLang 服务端：带有 --wait-for-initial-weights 标志运行，在就绪前等待权重

  2. Checkpoint Engine 工作进程：由 torchrun 管理的独立进程，负责加载和分发模型权重

checkpoint engine 采用参数服务器架构，支持以下模式：

  • Broadcast 模式：权重从加载进程广播到推理进程

  • P2P 模式：进程间直接进行点对点权重传输

  • All 模式：结合了广播和点对点两种方法

使用示例#

单节点设置#

终端 1 - 启动 SGLang 服务端

python -m sglang.launch_server \
    --model-path Qwen/Qwen3-8B \
    --tp 8 \
    --load-format dummy \
    --wait-for-initial-weights

终端 2 - 运行 Checkpoint Engine

使用 sglang 入口点

python -m sglang.srt.checkpoint_engine.update \
    --update-method broadcast \
    --checkpoint-path /path/to/Qwen/Qwen3-8B/ \
    --inference-parallel-size 8

直接使用 torchrun

torchrun --nproc-per-node 8 \
    examples/checkpoint_engine/update.py \
    --update-method broadcast \
    --checkpoint-path /path/to/Qwen/Qwen3-8B/ \
    --inference-parallel-size 8

多节点设置（2 个节点）#

节点 0

启动 SGLang 服务端

python -m sglang.launch_server \
    --model-path Qwen/Qwen3-8B \
    --tp 8 \
    --load-format dummy \
    --wait-for-initial-weights \
    --host [IP]

运行 checkpoint engine

使用 sglang 入口点（推荐）

python -m sglang.srt.checkpoint_engine.update \
    --update-method broadcast \
    --checkpoint-path /path/to/Qwen/Qwen3-8B/ \
    --inference-parallel-size 8

直接使用 torchrun

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

节点 1

启动 SGLang 服务端

python -m sglang.launch_server \
    --model-path Qwen/Qwen3-8B \
    --tp 8 \
    --load-format dummy \
    --wait-for-initial-weights \
    --host [IP]

运行 checkpoint engine

使用 sglang 入口点（推荐）

python -m sglang.srt.checkpoint_engine.update \
    --update-method broadcast \
    --checkpoint-path /path/to/Qwen/Qwen3-8B/ \
    --inference-parallel-size 8

直接使用 torchrun

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

带有张量并行（TP=16）的多节点设置#

节点 0

启动 SGLang 服务端

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

运行 checkpoint engine

使用 sglang 入口点（推荐）

python -m sglang.srt.checkpoint_engine.update \
    --update-method broadcast \
    --checkpoint-path /path/to/Qwen/Qwen3-8B/ \
    --inference-parallel-size 16

直接使用 torchrun

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

节点 1

启动 SGLang 服务端

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

运行 checkpoint engine

使用 sglang 入口点（推荐）

python -m sglang.srt.checkpoint_engine.update \
    --update-method broadcast \
    --checkpoint-path /path/to/Qwen/Qwen3-8B/ \
    --inference-parallel-size 16

直接使用 torchrun

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

配置选项#

SGLang 服务端选项#

  • --load-format dummy：使用虚拟格式进行初始加载（允许与其他任务重叠）

  • --wait-for-initial-weights：在就绪前等待 checkpoint engine 提供权重

  • --host：多节点设置的主机地址

  • --dist-init-addr：用于张量并行的分布式初始化地址

Checkpoint Engine 选项#

  • --update-method：权重更新方法（broadcast、p2p 或 all）

  • --checkpoint-path：模型检查点目录路径

  • --inference-parallel-size：推理并行进程数

  • --endpoint：SGLang 服务端端点（默认：https://:19730）

  • --checkpoint-name：检查点名称（默认：my-checkpoint-iter-0）

  • --save-metas-file：保存检查点元数据的文件

  • --load-metas-file：加载检查点元数据的文件

  • --uds：用于通信的 Unix 域套接字路径

  • --weight-version：权重的版本标识符

性能优势#

checkpoint engine 在两个主要方面提供了显著的时间节省：

  1. 多节点加载：每个节点只从磁盘加载部分权重，有效增加了磁盘带宽。参与节点越多，加速效果越明显。初步测试显示，在两台 H20-3e 节点上加载 DeepSeek-R1 时，可实现 20 秒的加速。

  2. 单进程优化：使用 dummy 格式允许磁盘到 CPU 的传输与 CUDA graph 捕获及其他初始化任务并行，从而进一步节省时间。

故障排除#

  • 确保已安装 checkpoint engine 包：pip install 'checkpoint-engine[p2p]'

  • 在多节点设置中验证节点间的网络连通性

  • 检查检查点路径是否包含有效的模型文件

  • 监控日志中 SGLang 服务端与 checkpoint engine 之间的连接错误

  • 如果调试需要，使用 --sleep-time 参数添加延迟

参考资料#

目录