From c1dd5085b686eaeec5e2c5195a0a8689620d48cb Mon Sep 17 00:00:00 2001 From: Nan Date: Tue, 7 Jul 2026 08:13:19 +0000 Subject: [PATCH] [RL] Add /pull_weights: engine-side pull of published weights into a host-local checkpoint Add the disaggregated receiver used by Stitch: POST /pull_weights fans out across the deployment, walks the published weight_v{N} chain from the nearest full anchor, and replays zstd-compressed XOR deltas with per-tensor checksums into a host-local checkpoint for update_weights_from_disk. Harden eventual-consistency and recovery semantics by reading and size-verifying each source blob in memory before mutation, distinguishing incomplete sources from corrupt local state, syncing dirty pages before advancing the version marker, deduplicating work under a per-host lock, and reseeding from the latched pristine boot checkpoint after a torn apply. Seed shards copy in parallel with 16 MiB streaming reads and progress logging. Ported from the Stitch production fork. Upstreaming keeps the engine-side receiver independent of trainer framework and omits the separate partial-reload optimization tier. --- python/pyproject.toml | 1 + python/sglang/srt/entrypoints/http_server.py | 13 + python/sglang/srt/managers/io_struct.py | 13 + python/sglang/srt/managers/scheduler.py | 2 + .../scheduler_components/weight_updater.py | 37 ++ .../srt/managers/tokenizer_control_mixin.py | 12 + python/sglang/srt/server_args.py | 4 + .../srt/weight_sync/local_checkpoint.py | 458 ++++++++++++++++++ .../weight_sync/test_local_checkpoint.py | 317 ++++++++++++ 9 files changed, 857 insertions(+) create mode 100644 python/sglang/srt/weight_sync/local_checkpoint.py create mode 100644 test/registered/weight_sync/test_local_checkpoint.py diff --git a/python/pyproject.toml b/python/pyproject.toml index f0146cebae60..62369cb1f481 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -87,6 +87,7 @@ dependencies = [ "uvloop", "watchfiles", "xgrammar==0.2.1", + "xxhash", "zstandard", ] diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 9a1aa72c560d..4824ac651043 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -130,6 +130,7 @@ ParseFunctionCallReq, PauseGenerationReqInput, ProfileReq, + PullWeightsReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, SendWeightsToRemoteInstanceReqInput, @@ -1187,6 +1188,18 @@ async def update_weights_from_disk( ) +@app.post("/pull_weights") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def pull_weights(obj: Annotated[PullWeightsReqInput, Body()], request: Request): + """Materialize published weights on every engine host.""" + success, message = await _global_state.tokenizer_manager.pull_weights(obj, request) + + content = {"success": success, "message": message} + return ORJSONResponse( + content, status_code=HTTPStatus.OK if success else HTTPStatus.BAD_REQUEST + ) + + @app.post("/init_weights_send_group_for_remote_instance") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def init_weights_send_group_for_remote_instance( diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index fa359c548306..271e53dcddae 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1562,6 +1562,19 @@ class UpdateWeightFromDiskReqOutput(BaseReq, kw_only=True): num_paused_requests: int = 0 +class PullWeightsReqInput(BaseReq, kw_only=True): + """Request to materialize a published weight version.""" + + local_checkpoint_dir: str + source_dir: str + target_version: int + + +class PullWeightsReqOutput(BaseReq, kw_only=True): + success: bool + message: str + + class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True): names: List[str] dtypes: List[str] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c46a8c596812..3b832b240cec 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -121,6 +121,7 @@ OpenSessionReqInput, PauseGenerationReqInput, ProfileReq, + PullWeightsReqInput, ReleaseMemoryOccupationReqInput, RemoveExternalCorpusReqInput, RemoveExternalCorpusReqOutput, @@ -1392,6 +1393,7 @@ def init_request_dispatcher(self): CheckWeightsReqInput, self.weight_updater.check_weights, ), + (PullWeightsReqInput, self.weight_updater.pull_weights), (SlowDownReqInput, self.slow_down), ( ProfileReq, diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py index 7013ce04fe8a..e83e7129e3d5 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py @@ -28,6 +28,8 @@ GetWeightsByNameReqOutput, InitWeightsUpdateGroupReqInput, InitWeightsUpdateGroupReqOutput, + PullWeightsReqInput, + PullWeightsReqOutput, ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, @@ -84,6 +86,7 @@ class SchedulerWeightUpdaterManager: metrics_collector: Optional[Any] = None offload_tags: set = field(default_factory=set) stashed_model_static_state: Any = None + _pull_weights_base_dir: Optional[str] = None @contextmanager def _observe_weight_load(self, source: str) -> Iterator[None]: @@ -122,6 +125,40 @@ def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): success=success, message=message, num_paused_requests=0 ) + def pull_weights(self, recv_req: PullWeightsReqInput): + """Materialize a published weight version on every host.""" + from sglang.srt.weight_sync import local_checkpoint + + server_args = self.tp_worker.model_runner.server_args + if self._pull_weights_base_dir is None: + self._pull_weights_base_dir = server_args.model_path + try: + local_checkpoint.pull( + local_checkpoint_dir=recv_req.local_checkpoint_dir, + base_dir=self._pull_weights_base_dir, + source_dir=recv_req.source_dir, + target_version=recv_req.target_version, + pre_read_hook=server_args.custom_pull_weights_pre_read_hook, + ) + success, message = True, "Success." + except Exception: + success, message = False, traceback.format_exc() + logger.error(message) + + world_size = ( + torch.distributed.get_world_size(group=self.tp_cpu_group) + if torch.distributed.is_initialized() + else 1 + ) + if world_size > 1: + results = [None] * world_size + torch.distributed.all_gather_object( + results, (success, message), group=self.tp_cpu_group + ) + success = all(ok for ok, _ in results) + message = "; ".join(msg for ok, msg in results if not ok) or message + return PullWeightsReqOutput(success=success, message=message) + def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput): """Initialize the online model parameter update group.""" success, message = self.tp_worker.init_weights_update_group(recv_req) diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index 86b7b378f896..1fa8f39f0f85 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -51,6 +51,8 @@ ProfileReq, ProfileReqOutput, ProfileReqType, + PullWeightsReqInput, + PullWeightsReqOutput, ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, RemoveExternalCorpusReqInput, @@ -105,6 +107,7 @@ ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), ("check_weights", CheckWeightsReqOutput), + ("pull_weights", PullWeightsReqOutput), ("slow_down", SlowDownReqOutput), ("flush_cache", FlushCacheReqOutput), ("add_external_corpus", AddExternalCorpusReqOutput), @@ -770,6 +773,15 @@ async def resume_memory_occupation( self.auto_create_handle_loop() await self.resume_memory_occupation_communicator(obj) + async def pull_weights( + self: TokenizerManager, + obj: PullWeightsReqInput, + request: Optional[fastapi.Request] = None, + ) -> Tuple[bool, str]: + self.auto_create_handle_loop() + results = await self.pull_weights_communicator(obj) + return FanOutCommunicator.merge_results(results) + async def check_weights( self: TokenizerManager, obj: CheckWeightsReqInput, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2f93fcd4b566..ece73bc32313 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2596,6 +2596,10 @@ class ServerArgs: nargs="*", ), ] = None + custom_pull_weights_pre_read_hook: A[ + Optional[str], + "Import path of hook(source_dir, target_version) called before /pull_weights reads shared storage.", + ] = None weight_loader_disable_mmap: A[ bool, "Disable mmap while loading weight using safetensors.", diff --git a/python/sglang/srt/weight_sync/local_checkpoint.py b/python/sglang/srt/weight_sync/local_checkpoint.py new file mode 100644 index 000000000000..08f64618bb2d --- /dev/null +++ b/python/sglang/srt/weight_sync/local_checkpoint.py @@ -0,0 +1,458 @@ +"""Materialize published full and delta weights in a host-local checkpoint.""" + +from __future__ import annotations + +import fcntl +import glob +import importlib +import json +import logging +import mmap +import os +import struct +import threading +import zlib +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from typing import NamedTuple, Optional + +import numpy as np +import zstandard + +logger = logging.getLogger(__name__) + +_DELTA_WORKERS = min(32, (os.cpu_count() or 8)) +_SEED_COPY_WORKERS = int(os.environ.get("SGLANG_SEED_COPY_WORKERS", "8")) +_SEED_COPY_CHUNK_SIZE = 16 << 20 +_XOR_CHUNK_SIZE = 2 << 20 +_SYNC_DIR = ".weight_sync" + + +class _DeltaItem(NamedTuple): + name: str + compressed_data: memoryview + path: str + offset: int + num_bytes: int + expected_checksum: Optional[str] + + +def pull( + local_checkpoint_dir: str, + base_dir: str, + source_dir: str, + target_version: int, + pre_read_hook: Optional[str] = None, +) -> None: + """Apply published versions through ``target_version``. + + Missing or incomplete source data raises ``FileNotFoundError``. Other apply + failures trigger one clean reseed and replay before being raised. + """ + with _pull_lock(local_checkpoint_dir): + applied_version = _read_applied_version(local_checkpoint_dir) + if applied_version is not None and applied_version >= target_version: + return + if target_version > 0 and pre_read_hook: + _load_hook(pre_read_hook)(source_dir, target_version) + try: + _pull_locked( + local_checkpoint_dir, + base_dir, + source_dir, + target_version, + force_reseed=False, + ) + except FileNotFoundError: + _log_pull_not_found(source_dir, target_version) + raise + except Exception: + logger.exception( + "Pull to version %d failed; reseeding and replaying", + target_version, + ) + _pull_locked( + local_checkpoint_dir, + base_dir, + source_dir, + target_version, + force_reseed=True, + ) + + +def _pull_locked( + local_checkpoint_dir: str, + base_dir: str, + source_dir: str, + target_version: int, + force_reseed: bool, +) -> None: + applied_version = ( + None if force_reseed else _read_applied_version(local_checkpoint_dir) + ) + search_floor = applied_version if applied_version is not None else 0 + base_version = target_version + while base_version > search_floor and _is_delta( + _version_dir(source_dir, base_version) + ): + base_version -= 1 + if applied_version is None or base_version > applied_version: + seed_dir = ( + base_dir if base_version == 0 else _version_dir(source_dir, base_version) + ) + _reset_checkpoint(seed_dir, local_checkpoint_dir, base_version) + else: + base_version = applied_version + for version in range(base_version + 1, target_version + 1): + _apply_delta(local_checkpoint_dir, _version_dir(source_dir, version)) + + +def _load_hook(path: str): + module_path, _, name = path.rpartition(".") + return getattr(importlib.import_module(module_path), name) + + +def _log_pull_not_found(source_dir: str, target_version: int) -> None: + """Log source state after a missing or incomplete version.""" + target_dir = _version_dir(source_dir, target_version) + try: + versions = sorted(n for n in os.listdir(source_dir) if n.startswith("weight_v")) + except OSError as exc: + versions = [f""] + target_contents = None + if os.path.isdir(target_dir): + try: + target_contents = sorted(os.listdir(target_dir)) + except OSError as exc: + target_contents = [f""] + logger.error( + "Missing weight version %d: versions=%s, target=%s, target_exists=%s, " + "target_contents=%s, latest=%s", + target_version, + versions, + target_dir, + os.path.isdir(target_dir), + target_contents, + _read_latest_pointer_for_log(source_dir), + ) + + +def _read_latest_pointer_for_log(source_dir: str) -> str: + for path in ( + os.path.join(source_dir, "latest"), + os.path.join(os.path.dirname(source_dir.rstrip("/")), "latest"), + ): + try: + with open(path) as f: + return f"{path}={f.read().strip()!r}" + except OSError: + continue + return "" + + +def _version_dir(source_dir: str, version: int) -> str: + return os.path.join(source_dir, f"weight_v{version:06d}") + + +def _is_delta(version_dir: str) -> bool: + """Return whether a version index declares a delta encoding.""" + if not os.path.isdir(version_dir): + raise FileNotFoundError(f"published weight version missing: {version_dir}") + try: + with open(os.path.join(version_dir, "model.safetensors.index.json")) as f: + return "delta_encoding" in json.load(f).get("metadata", {}) + except FileNotFoundError: + return False + + +class _Adler32: + """Expose Adler-32 through the incremental hash interface.""" + + def __init__(self): + self._value = 1 + + def update(self, data) -> None: + self._value = zlib.adler32(data, self._value) + + def hexdigest(self) -> str: + return f"{self._value:08x}" + + +def _new_hasher(algorithm: str): + if algorithm == "xxh3-128": + import xxhash + + return xxhash.xxh3_128() + if algorithm == "blake3": + import blake3 + + return blake3.blake3() + if algorithm == "adler32": + return _Adler32() + raise KeyError(f"unknown checksum algorithm {algorithm!r}") + + +def _checksum(algorithm: str, buffer) -> str: + hasher = _new_hasher(algorithm) + hasher.update(buffer) + return hasher.hexdigest() + + +@contextmanager +def _pull_lock(local_checkpoint_dir: str): + sync_dir = os.path.join(local_checkpoint_dir, _SYNC_DIR) + os.makedirs(sync_dir, exist_ok=True) + with open(os.path.join(sync_dir, "lock"), "w") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(lock_file, fcntl.LOCK_UN) + + +def _read_applied_version(local_checkpoint_dir: str) -> Optional[int]: + try: + with open(os.path.join(local_checkpoint_dir, _SYNC_DIR, "state.json")) as f: + return int(json.load(f)["version"]) + except FileNotFoundError: + return None + + +def _write_applied_version(local_checkpoint_dir: str, version: int) -> None: + path = os.path.join(local_checkpoint_dir, _SYNC_DIR, "state.json") + temp_path = path + ".tmp" + with open(temp_path, "w") as f: + json.dump({"version": f"{version:06d}"}, f) + f.flush() + os.fsync(f.fileno()) + os.replace(temp_path, path) + + +def _drop_page_cache(path: str) -> None: + """Evict a file from the page cache when supported.""" + if not hasattr(os, "posix_fadvise"): + return + try: + fd = os.open(path, os.O_RDONLY) + try: + os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED) + finally: + os.close(fd) + except OSError: + pass + + +def _reset_checkpoint(source_dir: str, local_checkpoint_dir: str, version: int) -> None: + """Replace the local checkpoint with a full checkpoint.""" + logger.info( + "Pulling full checkpoint v%d %s -> %s", + version, + source_dir, + local_checkpoint_dir, + ) + os.makedirs(local_checkpoint_dir, exist_ok=True) + source_files = [entry for entry in os.scandir(source_dir) if entry.is_file()] + total_gb = sum(entry.stat().st_size for entry in source_files) / 1e9 + copied_gb = 0.0 + progress_lock = threading.Lock() + + def copy_file(entry) -> None: + nonlocal copied_gb + destination = os.path.join(local_checkpoint_dir, entry.name) + if not ( + os.path.exists(destination) and os.path.samefile(entry.path, destination) + ): + with open(entry.path, "rb") as source, open(destination, "wb") as output: + while chunk := source.read(_SEED_COPY_CHUNK_SIZE): + output.write(chunk) + _drop_page_cache(entry.path) + with progress_lock: + copied_gb += entry.stat().st_size / 1e9 + logger.info("Seeding base v%d: %.0f/%.0f GB", version, copied_gb, total_gb) + + worker_count = min(_SEED_COPY_WORKERS, len(source_files) or 1) + with ThreadPoolExecutor(max_workers=worker_count) as pool: + list(pool.map(copy_file, source_files)) + source_names = {entry.name for entry in source_files} + for entry in os.scandir(local_checkpoint_dir): + if entry.is_file() and entry.name not in source_names: + os.remove(entry.path) + for entry in source_files: + copied_size = os.path.getsize(os.path.join(local_checkpoint_dir, entry.name)) + source_size = entry.stat().st_size + if copied_size != source_size: + raise RuntimeError( + f"size mismatch copying {entry.name}: " + f"source={source_size}, local={copied_size}" + ) + _write_applied_version(local_checkpoint_dir, version) + + +def _tensor_locations(checkpoint_dir: str) -> dict: + """Map tensor names to file paths, byte offsets, and sizes.""" + locations = {} + for path in glob.glob(os.path.join(checkpoint_dir, "*.safetensors")): + with open(path, "rb") as f: + (header_len,) = struct.unpack(" Optional[int]: + """Return the size declared by a safetensors payload.""" + if len(blob) < 8: + return None + header_len = struct.unpack(" None: + """Apply and verify one delta version in place.""" + with open(os.path.join(version_dir, "model.safetensors.index.json")) as f: + index = json.load(f) + metadata = index["metadata"] + applied_version = _read_applied_version(local_checkpoint_dir) + if applied_version == int(metadata["version"]): + return + for blob_name in sorted(set(index.get("weight_map", {}).values())): + if not os.path.exists(os.path.join(version_dir, blob_name)): + raise FileNotFoundError( + f"incomplete source version {version_dir}: missing blob {blob_name}" + ) + if applied_version != int(metadata["base_version"]): + raise RuntimeError( + f"out-of-order delta: local at {applied_version}, " + f"delta builds on {metadata['base_version']}" + ) + if metadata["compression_format"] != "zstd": + raise NotImplementedError( + f"compression {metadata['compression_format']!r} not supported" + ) + encoding = metadata["delta_encoding"] + checksum_algorithm = metadata["checksum_format"] + tensor_locations = _tensor_locations(local_checkpoint_dir) + mapped_files = {} + checksum_mismatches = [] + mismatch_lock = threading.Lock() + delta_blobs = [] # Keep blobs alive while memoryviews reference them. + delta_items = [] + try: + for delta_file in sorted(glob.glob(os.path.join(version_dir, "*.safetensors"))): + with open(delta_file, "rb") as f: + blob = f.read() + expected_size = _expected_safetensors_size(blob) + if expected_size is None or len(blob) != expected_size: + raise FileNotFoundError( + f"incomplete source blob {delta_file}: {len(blob)}B, header " + f"declares {expected_size}B" + ) + delta_blobs.append(blob) + (header_len,) = struct.unpack(" None: + region = np.ndarray( + (item.num_bytes,), + dtype=np.uint8, + buffer=mapped_files[item.path][1], + offset=item.offset, + ) + hasher = _new_hasher(checksum_algorithm) + reader = zstandard.ZstdDecompressor().stream_reader(item.compressed_data) + position = 0 + while position < item.num_bytes: + block = reader.read(min(_XOR_CHUNK_SIZE, item.num_bytes - position)) + if not block: + break + chunk = np.frombuffer(block, dtype=np.uint8) + region[position : position + chunk.size] ^= chunk + hasher.update(region[position : position + chunk.size]) + position += chunk.size + if hasher.hexdigest() != item.expected_checksum: + with mismatch_lock: + checksum_mismatches.append(item.name) + + def apply_overwrite(item: _DeltaItem) -> None: + delta = np.frombuffer( + zstandard.ZstdDecompressor().decompress(item.compressed_data), + dtype=np.uint8, + ) + region = np.ndarray( + (item.num_bytes,), + dtype=np.uint8, + buffer=mapped_files[item.path][1], + offset=item.offset, + ) + count = int.from_bytes(delta[:4], "little") + positions = np.frombuffer(delta[4 : 4 + 4 * count], dtype=" str: + return f"{zlib.adler32(bytes(data), 1):08x}" + + +class _Publisher: + SHARD = "model-00001-of-00001.safetensors" + + def __init__(self, root): + self.base_dir = os.path.join(root, "base") + self.source_dir = os.path.join(root, "published") + os.makedirs(self.base_dir) + os.makedirs(self.source_dir) + rng = np.random.default_rng(7) + self.state = { + "layer.a": rng.integers(0, 256, 4096, dtype=np.uint8).tobytes(), + "layer.b": rng.integers(0, 256, 2048, dtype=np.uint8).tobytes(), + } + self.versions = {0: dict(self.state)} + _write_safetensors(os.path.join(self.base_dir, self.SHARD), self.state) + + def publish_delta(self, version, changed): + version_dir = os.path.join(self.source_dir, f"weight_v{version:06d}") + os.makedirs(version_dir) + payloads = {} + checksums = {} + for name, new in changed.items(): + old = self.state[name] + diff = ( + np.frombuffer(new, dtype=np.uint8) ^ np.frombuffer(old, dtype=np.uint8) + ).tobytes() + payloads[name] = zstandard.ZstdCompressor().compress(diff) + checksums[name] = _adler32_hex(new) + self.state[name] = new + self.versions[version] = dict(self.state) + _write_safetensors( + os.path.join(version_dir, self.SHARD), payloads, metadata=checksums + ) + with open(os.path.join(version_dir, "model.safetensors.index.json"), "w") as f: + json.dump( + { + "metadata": { + "version": f"{version:06d}", + "base_version": f"{version - 1:06d}", + "delta_encoding": "xor", + "compression_format": "zstd", + "checksum_format": "adler32", + }, + "weight_map": {name: self.SHARD for name in payloads}, + }, + f, + ) + + +def _read_local(local_dir): + path = os.path.join(local_dir, _Publisher.SHARD) + with open(path, "rb") as f: + (header_len,) = struct.unpack("