检查点引擎集成

检查点引擎集成#

SGLang 检查点引擎集成提供了一种使用分布式检查点加载系统高效加载模型权重的途径。通过跨多个进程和节点并行化权重加载过程，这一功能显著减少了模型加载时间，特别是对于大型模型和多节点设置。

概述#

检查点引擎集成允许 SGLang：

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

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

  • 将权重加载与其他初始化任务（如 CUDA 图捕获）重叠

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

安装#

首先安装检查点引擎包：

pip install 'checkpoint-engine[p2p]'

架构#

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

  1. SGLang 服务器：使用 --wait-for-initial-weights 标志运行，等待权重加载完成后再变为就绪状态

  2. 检查点引擎工作进程：独立进程（由 torchrun 管理），负责加载和分发模型权重

检查点引擎使用参数服务器架构，支持：

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

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

  • All 模式：广播和 P2P 方法的组合

使用示例#

单节点设置#

终端 1 - 启动 SGLang 服务器：

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

终端 2 - 运行检查点引擎：

使用 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]

运行检查点引擎：

使用 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]

运行检查点引擎：

使用 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

运行检查点引擎：

使用 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

运行检查点引擎：

使用 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：等待检查点引擎提供权重后再变为就绪状态

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

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

检查点引擎选项#

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

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

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

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

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

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

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

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

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

性能优势#

检查点引擎在以下两个方面提供了显著的时间节省：

  1. 多节点加载：每个节点仅从磁盘加载部分权重，有效增加了磁盘带宽。更多的参与节点提供了更大的加速。初步测试显示，在使用两个节点的 H20-3e 上加载 DeepSeek-R1 时，可节省 20 秒。

  2. 单进程优化：使用虚拟格式允许磁盘到 CPU 的传输与 CUDA 图捕获和其他初始化任务重叠，提供额外的时间节省。

故障排除#

  • 确保已安装检查点引擎包：pip install 'checkpoint-engine[p2p]'

  • 验证多节点设置中节点间的网络连接

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

  • 监控 SGLang 服务器和检查点引擎之间的连接错误日志

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

参考资料#

目录