Implementation:Sgl project Sglang Checkpoint Update

Knowledge Sources	Sgl_project_Sglang
Domains	Model Serving, Distributed Computing
Last Updated	2026-02-10 00:00 GMT

Overview

Demonstrates how to hot-swap model weights on a running SGLang server using the checkpoint engine with distributed execution via torchrun.

Description

update.py is a distributed script designed to be launched via torchrun. It reads model weights from safetensors checkpoint files, splits them across ranks using either file-level (split_checkpoint_files) or tensor-level (split_tensors) partitioning, registers them with a ParameterServer, and pushes updates to a running SGLang server via the /update_weights_from_ipc HTTP endpoint.

The script supports two update methods: broadcast (default) where weights are pushed from a single source, and p2p (peer-to-peer) with explicit rank assignment. It also supports a "join" mode that loads previously saved metadata to add new workers to an existing parameter server group.

Key functions include check_sglang_ready() for polling server readiness, req_inference() for creating HTTP request closures, update_weights() for the main update pipeline, and join() for adding workers to an existing group. The timer context manager provides performance instrumentation.

Usage

Use this script for online weight update workflows such as RLHF training loops, where model weights need to be hot-swapped on a running inference server without restart. It demonstrates the checkpoint engine's IPC-based weight transfer capability for distributed tensor-parallel inference setups.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: examples/checkpoint_engine/update.py
• Lines: 1-241

Signature

@contextmanager
def timer(msg: str): ...

def check_sglang_ready(endpoint: str, inference_parallel_size: int, uds: str | None = None): ...
def split_checkpoint_files(checkpoint_path: str, rank: int, world_size: int) -> list[str]: ...
def split_tensors(checkpoint_path: str, rank: int, world_size: int) -> dict[str, torch.Tensor]: ...
def req_inference(endpoint: str, inference_parallel_size: int, timeout: float = 300.0,
                  uds: str | None = None, weight_version: str | None = None) -> Callable[[list[tuple[str, str]]], None]: ...
def update_weights(ps: ParameterServer, checkpoint_name: str, checkpoint_files: list[str],
                   named_tensors: dict[str, torch.Tensor], req_func: Callable,
                   inference_parallel_size: int, endpoint: str,
                   save_metas_file: str | None = None,
                   update_method: Literal["broadcast", "p2p", "all"] = "broadcast",
                   uds: str | None = None): ...
def join(ps: ParameterServer, checkpoint_name: str, load_metas_file: str,
         req_func: Callable, inference_parallel_size: int, endpoint: str,
         uds: str | None = None): ...

Import

import argparse
import json
import os
import pickle
import time
from collections import defaultdict
from collections.abc import Callable
from contextlib import contextmanager
from typing import Literal

import httpx
import torch
import torch.distributed as dist
from checkpoint_engine.ps import ParameterServer
from loguru import logger
from safetensors import safe_open

I/O Contract

Inputs

Name	Type	Required	Description
--checkpoint-path	str	Yes (unless join mode)	Path to the safetensors checkpoint directory
--endpoint	str	No (default: "http://localhost:19730")	SGLang server endpoint URL
--inference-parallel-size	int	No (default: 8)	Number of inference parallel workers
--checkpoint-name	str	No (default: "my-checkpoint-iter-0")	Name identifier for the checkpoint
--update-method	str	No (default: "broadcast")	Update method: broadcast, p2p, or all
--save-metas-file	str	No	Path to save metadata for later join operations
--load-metas-file	str	No	Path to load metadata for join mode
--uds	str	No	Unix domain socket path for communication
--weight-version	str	No	Weight version identifier
--sleep-time	int	No (default: 0)	Time to sleep after update completes

Outputs

Name	Type	Description
Weight update	HTTP POST	Pushes weight updates to SGLang server via /update_weights_from_ipc endpoint
Metadata file	pickle file	Optionally saves parameter server metadata for join operations

Usage Examples

# Launch server with wait-for-initial-weights:
# python -m sglang.launch_server --model-path /workspace/Qwen/Qwen3-4B/ \
#   --tensor-parallel-size 2 --port 19730 --load-format dummy \
#   --checkpoint-engine-wait-weights-before-ready --mem-fraction-static 0.7

# Run weight update via torchrun:
# torchrun --nproc-per-node 2 update.py \
#   --update-method broadcast \
#   --checkpoint-path /workspace/Qwen/Qwen3-4B/ \
#   --inference-parallel-size 2

Related Pages

• Environment:Sgl_project_Sglang_Distributed