From 4694fedf9b5d133a91a3b9d06856aea9461cd96a Mon Sep 17 00:00:00 2001 From: Nan Date: Mon, 6 Jul 2026 18:36:56 +0000 Subject: [PATCH] Disk weight sync: engines pull published weights on every host (/pull_weights) Follow-up to #2089. That PR fanned the host-local delta apply out from slime via Ray (all_engine_actors, one actor per host of each multi-node engine) because only node 0 of a multi-node engine has an HTTP server. That leaks engine topology into slime and cannot work for external rollout engines, where slime only has an endpoint. A new /pull_weights endpoint (shipped as the standalone docker/patch/latest/sglang-pull_weights.patch) has each engine pull the published weights onto every host it spans, riding the existing control-request broadcast; a per-host flock collapses co-located ranks to one pull, and the reply is all-gathered across the TP group so success means every host holds a verified checkpoint. The pull is artifact-driven: a published version is either a full HF checkpoint (copied as-is, resetting the chain) or a delta against its predecessor (patched in place, per-tensor checksums), detected from the index metadata; a fresh host seeds from the newest full version at or below the target, or from the engine's own model path for a pure-delta stream. The shared-filesystem refresh hook follows sglang's custom-weight-loader convention (--custom-pull-weights-pre-read-hook), forwarded through the existing --sglang-* passthrough. slime now only talks to one endpoint per engine: all_engine_actors and SGLangEngine.sync_local_checkpoint are removed, the delta updater calls pull_weights(v) before the reload, and pull_weights(0) during baseline capture replaces the init-thread base materialization (still overlapped with the snapshot gather). Full-mode disk sync pulls to local disk too when --update-weight-local-checkpoint-dir is set. The post-write hook is renamed --custom-update-weight-post-write-path since it now serves full and delta alike, and two pre-existing full-disk bugs on non-POSIX shared filesystems are fixed: every writing rank creates the version dir and renames its own shards, and the post-write hook runs on every rank. Validated end to end on GLM-4.7-Flash non-colocated (2 trainer nodes + one 2-node tp16 rollout engine, object-store-backed shared filesystem): delta mode 4/4 pulls with 0 checksum mismatches (~23s steady-state sync); full mode 4/4 pulls, one pull per host per version, engines reloading from the pulled local checkpoint. --- docker/Dockerfile | 5 +- docker/patch/latest/sglang-pull_weights.patch | 555 ++++++++++++++++++ docs/en/advanced/delta-weight-sync.md | 50 +- docs/en/advanced/external-rollout-engines.md | 4 +- docs/en/get_started/customization.md | 22 + docs/zh/advanced/delta-weight-sync.md | 23 +- docs/zh/advanced/external-rollout-engines.md | 4 +- docs/zh/get_started/customization.md | 19 + examples/delta_weight_sync/README.md | 13 +- .../run-glm4.7-30B-A3B-delta.sh | 6 +- slime/backends/megatron_utils/actor.py | 7 +- .../megatron_utils/hf_checkpoint_saver.py | 36 +- .../update_weight/update_weight_from_disk.py | 18 +- .../update_weight_from_disk_delta.py | 42 +- .../update_weight_from_distributed.py | 1 - .../update_weight_from_tensor.py | 1 - slime/backends/sglang_utils/sglang_engine.py | 64 +- slime/ray/actor_group.py | 10 +- slime/ray/rollout.py | 3 +- slime/utils/arguments.py | 33 +- slime/utils/disk_delta.py | 187 +----- tests/utils/test_hf_checkpoint_saver.py | 11 +- 22 files changed, 784 insertions(+), 330 deletions(-) create mode 100644 docker/patch/latest/sglang-pull_weights.patch diff --git a/docker/Dockerfile b/docker/Dockerfile index f71f04f742..eb4ba7b183 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -105,18 +105,19 @@ ARG ENABLE_SGLANG_PATCH=1 COPY docker/patch/${PATCH_VERSION}/sglang.patch \ docker/patch/${PATCH_VERSION}/sglang-top_p.patch \ docker/patch/${PATCH_VERSION}/sglang-release_hicache.patch \ + docker/patch/${PATCH_VERSION}/sglang-pull_weights.patch \ /sgl-workspace/sglang/ RUN if [ "$ENABLE_SGLANG_PATCH" = "1" ]; then \ cd /sgl-workspace/sglang && \ git update-index --refresh && \ - for patch in sglang.patch sglang-top_p.patch sglang-release_hicache.patch; do \ + for patch in sglang.patch sglang-top_p.patch sglang-release_hicache.patch sglang-pull_weights.patch; do \ git apply --3way "$patch" || exit 1; \ if git grep -n '^<<<<<<< ' -- .; then \ echo "Patch failed to apply cleanly. Please resolve conflicts." && \ exit 1; \ fi; \ done && \ - rm sglang.patch sglang-top_p.patch sglang-release_hicache.patch; \ + rm sglang.patch sglang-top_p.patch sglang-release_hicache.patch sglang-pull_weights.patch; \ fi # ====================================== Install main package ============================================ diff --git a/docker/patch/latest/sglang-pull_weights.patch b/docker/patch/latest/sglang-pull_weights.patch new file mode 100644 index 0000000000..c0bde131f2 --- /dev/null +++ b/docker/patch/latest/sglang-pull_weights.patch @@ -0,0 +1,555 @@ +diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py +index 2c881d9..89c653a 100644 +--- a/python/sglang/srt/entrypoints/http_server.py ++++ b/python/sglang/srt/entrypoints/http_server.py +@@ -129,6 +129,7 @@ from sglang.srt.managers.io_struct import ( + PauseGenerationReqInput, + PostProcessWeightsReqInput, + ProfileReqInput, ++ PullWeightsReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, + SendWeightsToRemoteInstanceReqInput, +@@ -1271,6 +1272,19 @@ async def post_process_weights(req: PostProcessWeightsReqInput, request: Request + ) + + ++@app.post("/pull_weights") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def pull_weights(obj: PullWeightsReqInput, request: Request): ++ """Have every host of this deployment pull published weight deltas into its ++ local checkpoint (materialized from the model path on first use).""" ++ success, message = await _global_state.tokenizer_manager.pull_weights(obj, request) ++ ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ + @app.post("/update_weight_version") + @auth_level(AuthLevel.ADMIN_OPTIONAL) + async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Request): +diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py +index 2d544ca..21465f4 100644 +--- a/python/sglang/srt/managers/io_struct.py ++++ b/python/sglang/srt/managers/io_struct.py +@@ -1699,6 +1699,24 @@ class CheckWeightsReqOutput(BaseReq): + payload: Optional[Dict] = None + + ++@dataclass ++class PullWeightsReqInput(BaseReq): ++ # Host-local checkpoint dir the pulled weights land in; seeded from the ++ # server's model path when the published stream has no full version. ++ local_checkpoint_dir: str ++ # Shared dir the publisher writes weight_v{N:06d}/ version dirs under; each ++ # version is a full HF checkpoint or a delta against the previous version. ++ source_dir: str ++ # The version to bring the local checkpoint up to. ++ target_version: int ++ ++ ++@dataclass ++class PullWeightsReqOutput(BaseReq): ++ success: bool ++ message: str ++ ++ + @dataclass + class SlowDownReqInput(BaseReq): + forward_sleep_time: Optional[float] +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index bd154c1..d2cd752 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -126,6 +126,7 @@ from sglang.srt.managers.io_struct import ( + PauseGenerationReqInput, + PostProcessWeightsReqInput, + ProfileReq, ++ PullWeightsReqInput, + ReleaseMemoryOccupationReqInput, + RemoveExternalCorpusReqInput, + RemoveExternalCorpusReqOutput, +@@ -1339,6 +1340,10 @@ class Scheduler( + 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 9ab3abe..c2ac697 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 @@ from sglang.srt.managers.io_struct import ( + InitWeightsUpdateGroupReqOutput, + PostProcessWeightsReqInput, + PostProcessWeightsReqOutput, ++ PullWeightsReqInput, ++ PullWeightsReqOutput, + ReleaseMemoryOccupationReqInput, + ReleaseMemoryOccupationReqOutput, + ResumeMemoryOccupationReqInput, +@@ -276,6 +278,40 @@ class SchedulerWeightUpdaterManager: + + return ResumeMemoryOccupationReqOutput() + ++ def pull_weights(self, recv_req: PullWeightsReqInput): ++ """Sync this host's local checkpoint up to recv_req.target_version. ++ ++ Every rank runs the pull; a per-host file lock collapses co-located ++ ranks to one pull. Success is gathered across the TP group (all nodes), ++ so the reply only reports success once every host holds a verified ++ checkpoint. ++ """ ++ from sglang.srt.weight_sync import local_checkpoint ++ ++ server_args = self.tp_worker.model_runner.server_args ++ try: ++ local_checkpoint.pull( ++ local_checkpoint_dir=recv_req.local_checkpoint_dir, ++ base_dir=server_args.model_path, ++ 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) ++ ++ tp_size = torch.distributed.get_world_size(group=self.tp_cpu_group) ++ if tp_size > 1: ++ results = [None] * tp_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 check_weights(self, recv_req: CheckWeightsReqInput): + try: + payload = self.tp_worker.model_runner.check_weights(action=recv_req.action) +diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py +index ee25e5e..4988bd3 100644 +--- a/python/sglang/srt/managers/tokenizer_control_mixin.py ++++ b/python/sglang/srt/managers/tokenizer_control_mixin.py +@@ -53,6 +53,8 @@ from sglang.srt.managers.io_struct import ( + ProfileReq, + ProfileReqOutput, + ProfileReqType, ++ PullWeightsReqInput, ++ PullWeightsReqOutput, + ReleaseMemoryOccupationReqInput, + ReleaseMemoryOccupationReqOutput, + RemoveExternalCorpusReqInput, +@@ -99,6 +101,7 @@ _COMMUNICATOR_SPECS = [ + ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput), + ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput), + ("post_process_weights", PostProcessWeightsReqOutput), ++ ("pull_weights", PullWeightsReqOutput), + ("get_weights_by_name", GetWeightsByNameReqOutput), + ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), + ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), +@@ -767,6 +770,15 @@ class TokenizerControlMixin: + results = await self.post_process_weights_communicator(obj) + return FanOutCommunicator.merge_results(results) + ++ 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 6c77ff6..12f31e8 100644 +--- a/python/sglang/srt/server_args.py ++++ b/python/sglang/srt/server_args.py +@@ -850,6 +850,7 @@ class ServerArgs: + + # For model weight update and weight loading + custom_weight_loader: Optional[List[str]] = None ++ custom_pull_weights_pre_read_hook: Optional[str] = None + weight_loader_disable_mmap: bool = False + weight_loader_prefetch_checkpoints: bool = False + weight_loader_prefetch_num_threads: int = 4 +@@ -7095,6 +7096,12 @@ class ServerArgs: + default=None, + help="The custom dataloader which used to update the model. Should be set with a valid import path, such as my_package.weight_load_func", + ) ++ parser.add_argument( ++ "--custom-pull-weights-pre-read-hook", ++ type=str, ++ default=ServerArgs.custom_pull_weights_pre_read_hook, ++ help="Import path of a hook(source_dir, target_version) that /pull_weights calls before reading the published weights. POSIX shared filesystems need no hook; object-store-backed mounts often lack cross-host read-after-write consistency, so another host's writes only become visible after an explicit refresh.", ++ ) + parser.add_argument( + "--weight-loader-disable-mmap", + action="store_true", +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 0000000..2fc0593 +--- /dev/null ++++ b/python/sglang/srt/weight_sync/local_checkpoint.py +@@ -0,0 +1,349 @@ ++"""Host-local pull of published weights (the /pull_weights endpoint). ++ ++A trainer publishes each weight sync as a version directory ``weight_v{N:06d}/`` ++under a shared ``source_dir``. Each version is a canonical HF checkpoint ++directory of one of two kinds, distinguished by its index metadata: ++ ++- **full**: an ordinary checkpoint. Pulling it copies it into the host-local ++ ``local_checkpoint_dir``, replacing whatever is there — no history needed. ++- **delta** (index metadata carries ``delta_encoding``): safetensors files ++ holding zstd-compressed per-tensor diffs against version N-1, plus per-tensor ++ checksums of the new state. Pulling it patches the local checkpoint in place. ++ ++Version 0 is the engine's own base checkpoint (``model_path``). Every host of a ++(possibly multi-node) deployment runs the same pull; the engine then reloads the ++local checkpoint through the ordinary ``update_weights_from_disk`` path. ++ ++``pull()`` is safe to call concurrently from every scheduler rank on a host: a ++per-host file lock serializes the work and an applied-version marker makes the ++extra calls no-ops. ++""" ++ ++from __future__ import annotations ++ ++import fcntl ++import glob ++import importlib ++import io ++import json ++import logging ++import mmap ++import os ++import shutil ++import struct ++import threading ++import zlib ++from concurrent.futures import ThreadPoolExecutor ++from contextlib import contextmanager ++from typing import Optional ++ ++import numpy as np ++import zstandard ++ ++logger = logging.getLogger(__name__) ++ ++# The delta-apply phases (decompress, XOR/scatter, checksum) are memory-bandwidth ++# bound and release the GIL, so a thread pool over tensors recovers the ++# bandwidth one thread leaves idle. ++NUM_WORKERS = min(32, (os.cpu_count() or 8)) ++ ++# Per-checkpoint dir holding the applied-version marker and the pull lock. ++SYNC_DIR = ".weight_sync" ++ ++ ++def pull( ++ local_checkpoint_dir: str, ++ base_dir: str, ++ source_dir: str, ++ target_version: int, ++ pre_read_hook: Optional[str] = None, ++) -> None: ++ """Bring the host-local checkpoint up to ``target_version``. ++ ++ Seeds from the newest full checkpoint at or below the target — the engine's ++ own base (``base_dir``) for a pure-delta stream, a published full version ++ otherwise — then applies the remaining deltas in order. A local checkpoint ++ already past the seed point just continues its delta chain. Raises on any ++ per-tensor checksum mismatch (fail loud, never serve bad weights). ++ """ ++ # Object-store-backed shared filesystems lack cross-host read-after-write ++ # consistency: the publisher's files only appear here after an explicit ++ # refresh, which the deployment supplies as this hook. POSIX shared ++ # filesystems (NFS, Lustre, ...) need none. ++ if target_version > 0 and pre_read_hook: ++ _load_hook(pre_read_hook)(source_dir, target_version) ++ with _pull_lock(local_checkpoint_dir): ++ applied = _read_applied_version(local_checkpoint_dir) # None on a fresh host ++ # Scan back from the target for the newest full version. Stop at the ++ # local state — below it a reset can never be needed (or, on a fresh ++ # host, at 0 = the engine's base). ++ floor = applied if applied is not None else 0 ++ start = target_version ++ while start > floor and _is_delta(_version_dir(source_dir, start)): ++ start -= 1 ++ if applied is None or start > applied: ++ seed_dir = base_dir if start == 0 else _version_dir(source_dir, start) ++ _reset_checkpoint(seed_dir, local_checkpoint_dir, start) ++ else: ++ start = applied ++ for version in range(start + 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 _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: ++ """A version is a delta iff its index metadata declares an encoding; an ++ ordinary HF checkpoint (with or without an index) is a full version.""" ++ 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: ++ """adler32 behind the incremental .update / .hexdigest interface the hash objects expose.""" ++ ++ 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, buf) -> str: ++ hasher = _new_hasher(algorithm) ++ hasher.update(buf) ++ return hasher.hexdigest() ++ ++ ++@contextmanager ++def _pull_lock(local_checkpoint_dir: str): ++ sync = os.path.join(local_checkpoint_dir, SYNC_DIR) ++ os.makedirs(sync, exist_ok=True) ++ with open(os.path.join(sync, "lock"), "w") as f: ++ fcntl.flock(f, fcntl.LOCK_EX) ++ try: ++ yield ++ finally: ++ fcntl.flock(f, 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") ++ tmp = path + ".tmp" ++ with open(tmp, "w") as f: ++ json.dump({"version": f"{version:06d}"}, f) ++ f.flush() ++ os.fsync(f.fileno()) ++ os.replace(tmp, path) ++ ++ ++def _drop_page_cache(path: str) -> None: ++ """Evict a file from the page cache (POSIX_FADV_DONTNEED).""" ++ 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(src_dir: str, local_checkpoint_dir: str, version: int) -> None: ++ """Make local_checkpoint_dir an exact copy of the full checkpoint in src_dir ++ (files the new checkpoint doesn't have — e.g. differently-sharded old ones — ++ are pruned). Later deltas chain on top of this state.""" ++ logger.info("Pulling full checkpoint v%d %s -> %s", version, src_dir, local_checkpoint_dir) ++ os.makedirs(local_checkpoint_dir, exist_ok=True) ++ src_files = [entry for entry in os.scandir(src_dir) if entry.is_file()] ++ for entry in src_files: ++ shutil.copy2(entry.path, os.path.join(local_checkpoint_dir, entry.name)) ++ # don't let the source evict the local copy we keep resident ++ _drop_page_cache(entry.path) ++ names = {entry.name for entry in src_files} ++ for entry in os.scandir(local_checkpoint_dir): ++ if entry.is_file() and entry.name not in names: ++ os.remove(entry.path) ++ # a truncated copy (e.g. an object-store mount surfacing metadata before ++ # bytes) must fail loud, not serve bad weights ++ for entry in src_files: ++ copied = os.path.getsize(os.path.join(local_checkpoint_dir, entry.name)) ++ if copied != entry.stat().st_size: ++ raise RuntimeError( ++ f"size mismatch copying {entry.name}: src {entry.stat().st_size} != local {copied}" ++ ) ++ _write_applied_version(local_checkpoint_dir, version) ++ ++ ++def _tensor_locations(ckpt_dir: str) -> dict: ++ """Map each tensor name to (file, byte offset, nbytes) by reading every safetensors header.""" ++ locations = {} ++ for path in glob.glob(os.path.join(ckpt_dir, "*.safetensors")): ++ with open(path, "rb") as f: ++ (header_len,) = struct.unpack(" None: ++ """Apply one version's delta in place: decompress + apply + checksum each tensor across a thread ++ pool (each writes a distinct mmap region, so the writes don't conflict). Any mismatch raises.""" ++ with open(os.path.join(version_dir, "model.safetensors.index.json")) as f: ++ meta = json.load(f)["metadata"] ++ applied = _read_applied_version(local_checkpoint_dir) ++ if applied == int(meta["version"]): ++ return ++ if applied != int(meta["base_version"]): ++ raise RuntimeError( ++ f"out-of-order delta: local at {applied}, delta builds on {meta['base_version']}" ++ ) ++ if meta["compression_format"] != "zstd": ++ raise NotImplementedError( ++ f"compression {meta['compression_format']!r} not supported" ++ ) ++ encoding = meta["delta_encoding"] ++ algorithm = meta["checksum_format"] ++ locations = _tensor_locations(local_checkpoint_dir) ++ open_mmaps = {} ++ mismatches = [] ++ lock = threading.Lock() ++ file_bytes = [] # keep alive: items hold zero-copy views into these ++ items = [] # (name, compressed_view, path, offset, nbytes, want_checksum) ++ try: ++ for delta_file in sorted(glob.glob(os.path.join(version_dir, "*.safetensors"))): ++ with open(delta_file, "rb") as f: ++ blob = f.read() ++ file_bytes.append(blob) ++ (header_len,) = struct.unpack(" None: ++ name, compressed, path, offset, nbytes, want = item ++ region = np.ndarray( ++ (nbytes,), dtype=np.uint8, buffer=open_mmaps[path][1], offset=offset ++ ) ++ hasher = _new_hasher(algorithm) ++ reader = zstandard.ZstdDecompressor().stream_reader( ++ io.BytesIO(bytes(compressed)) ++ ) ++ pos = 0 ++ # 2 MB chunks stay L2-resident across decompress -> XOR -> checksum ++ while pos < nbytes: ++ block = reader.read(min(2 << 20, nbytes - pos)) ++ if not block: ++ break ++ chunk = np.frombuffer(block, dtype=np.uint8) ++ region[pos : pos + chunk.size] ^= chunk ++ hasher.update(region[pos : pos + chunk.size]) ++ pos += chunk.size ++ if hasher.hexdigest() != want: ++ with lock: ++ mismatches.append(name) ++ ++ def apply_overwrite(item) -> None: ++ name, compressed, path, offset, nbytes, want = item ++ delta = np.frombuffer( ++ zstandard.ZstdDecompressor().decompress(bytes(compressed)), ++ dtype=np.uint8, ++ ) ++ region = np.ndarray( ++ (nbytes,), dtype=np.uint8, buffer=open_mmaps[path][1], offset=offset ++ ) ++ count = int.from_bytes(delta[:4].tobytes(), "little") ++ positions = np.frombuffer(delta[4 : 4 + 4 * count].tobytes(), dtype=" None +``` + +**Purpose**: Called on each trainer rank after a disk weight sync's files are written +(`--update-weight-transport disk`, full or delta mode), before the engines read them. Use it to +publish the writes on a non-POSIX shared filesystem — e.g. upload pending writes to the +backing object store — where another host cannot see the files without an explicit sync. The hook is called +on every rank and must gate itself (e.g. once per container). + +The read-side counterpart runs inside the inference engine, on every host it spans, and is +therefore an sglang server argument rather than a slime hook: pass +`--sglang-custom-pull-weights-pre-read-hook ` with signature +`hook(source_dir: str, target_version: int)` — called before `/pull_weights` reads the +published weights (e.g. refresh the mount's view). See +[Delta Weight Sync](../advanced/delta-weight-sync.md) for the full mechanism. + ## Testing Custom Function Paths slime also provides CPU-only contract tests for customization interfaces. These tests resolve components through import-path strings, so they can validate both built-in hooks and user-defined implementations passed through the same CLI arguments used by training. diff --git a/docs/zh/advanced/delta-weight-sync.md b/docs/zh/advanced/delta-weight-sync.md index 0aa5434472..8983f28159 100644 --- a/docs/zh/advanced/delta-weight-sync.md +++ b/docs/zh/advanced/delta-weight-sync.md @@ -2,7 +2,7 @@ Delta 权重同步只发送两次同步之间发生变化的字节，而不是每次都写一份完整 checkpoint，以此让非 colocate 的 rollout engine 保持最新。它面向大模型、跨集群或跨数据中心的训推解耦场景——这种场景下每次都写整份 actor 权重是主要开销。 -它**只支持 disk transport**，并且通过**原生**的 `update_weights_from_disk` 端点 reload，因此推理引擎不需要任何 delta 相关的支持。 +它**只支持 disk transport**。训练端把每次同步发布为一份 canonical HF checkpoint 目录；engine 的 `/pull_weights` 端点（随 slime 的 sglang patch 提供）把 apply 扇出到 **engine 覆盖的每一个 host** 并校验，随后 engine 通过**原生**的 `update_weights_from_disk` 端点 reload 打过补丁的本地 checkpoint。slime 对每个 engine 只与一个端点通信，所以多节点 serving 和外部 rollout engine 在 slime 侧都不需要任何额外支持。 ## 配置 @@ -18,7 +18,7 @@ Delta 权重同步只发送两次同步之间发生变化的字节，而不是 | 参数 | 作用 | |---|---| | `--update-weight-disk-dir` | 训练端发布 delta、rollout host 读取 delta 的共享文件系统目录。 | -| `--update-weight-local-checkpoint-dir` | host 本地（如 NVMe）的完整 HF checkpoint，delta 原地 apply 到这里。每个 host 在 engine 启动时由 `--hf-checkpoint` 物化。 | +| `--update-weight-local-checkpoint-dir` | host 本地（如 NVMe）的完整 HF checkpoint，由 `/pull_weights` 保持同步——delta 原地 apply，发布的完整 checkpoint 则整份替换。每个 host 在第一次 `/pull_weights` 时由 engine 的 model path seed。 | | `--update-weight-delta-encoding` | 磁盘上的 delta 编码：`xor`（默认）或 `overwrite`。 | | `--update-weight-delta-checksum` | 逐 tensor 完整性 checksum：`xxh3-128`（默认）、`blake3` 或 `adler32`。 | @@ -26,10 +26,15 @@ delta 始终用 zstd（level 1）压缩；profiling 显示对这类数据它在 ## 工作原理 -1. **Seed。** 第一次同步时，训练端为每个参数捕获一份 CPU snapshot——从 `--hf-checkpoint` seed，而这正是每个 rollout host 物化本地 checkpoint 的来源。此次不发布任何东西；这份 snapshot 就是下一次同步 diff 的基准。 +1. **Seed。** 第一次同步时，训练端为每个参数捕获一份 CPU snapshot——从 `--hf-checkpoint` seed，而这正是每个 rollout host 物化本地 checkpoint 的来源。此次不发布任何东西；这份 snapshot 就是下一次同步 diff 的基准。训练端同时发出 `target_version=0` 的 `/pull_weights`，让每个 host 现在就物化本地 base，与 snapshot 捕获重叠进行。 2. **Publish。** 之后每次同步，训练端把每个 gather 出的 HF tensor 与 snapshot 做 diff，编码、压缩，写到 `--update-weight-disk-dir` 下的新版本目录 `weight_v{N:06d}/`。该目录是一份 canonical HF checkpoint——`model-NNNNN.safetensors` 文件装着压缩后的 diff tensor，外加 `model.safetensors.index.json`（tensor 名 → 文件）承载 apply 元数据——所以这个产物是可移植的，不绑定训练端的并行 layout。随后 snapshot 推进到新值，供下次 diff。 -3. **Apply。** 每个 rollout host 把新版本的 delta 原地 apply 进它的本地 checkpoint。apply 在 tensor 之间并行，并逐 tensor 校验（见“完整性”）。 -4. **Reload。** engine 通过原生 `update_weights_from_disk` 路径 reload 打过补丁的本地 checkpoint——它从不接触 delta 格式。 +3. **Pull。** 训练端对每个 engine 调用 `/pull_weights`。engine 内部把请求广播到每个节点的每个 rank；每个 host 把新版本的 delta 原地 apply 进它的本地 checkpoint（host 级文件锁把同 host 的多个 rank 合并成一次 apply）。apply 在 tensor 之间并行，并逐 tensor 校验（见"完整性"）；只有**每一个 host** 都持有校验通过的 checkpoint，该调用才报告成功。 + + `/pull_weights` 并不绑定 delta：每个发布的版本是自描述的。若某个版本是一份普通的完整 HF + checkpoint（index 中没有 delta 元数据），pull 就直接整份拷贝——同时重置链条，因此晚加入的 + 新 host 从最近的完整版本 seed，而不必回放全部 delta，旧的 delta 也可以被清理。slime 的 + full 模式 disk 同步在设置了 `--update-weight-local-checkpoint-dir` 时正是走这条路径。 +4. **Reload。** engine 通过原生 `update_weights_from_disk` 路径 reload 打过补丁的本地 checkpoint——权重加载代码从不接触 delta 格式。 由于 snapshot 是从 `--hf-checkpoint`（engine 真正的 base）seed，而不是从当前 GPU 权重 seed，即使 Megatron→HF 往返不是逐字节相等（例如 embedding / LM head 中被裁掉的 vocab padding 行），该方案对任意模型也都正确。 @@ -42,13 +47,13 @@ delta 始终用 zstd（level 1）压缩；profiling 显示对这类数据它在 ## 完整性 -训练端把每个 tensor 新状态的逐 tensor checksum 存进版本里。apply 之后每个 host 重新计算 checksum，**任何不匹配都会 raise**，所以损坏的 delta 或错误的 base 会直接报错失败，而不会把坏权重提供出去。apply 还拒绝乱序执行：一个版本只会在它声明的 base 版本之上 apply。 +训练端把每个 tensor 新状态的逐 tensor checksum 存进版本里。apply 之后每个 host 重新计算 checksum，**任何不匹配都会 raise**——失败会通过 `/pull_weights` 的响应传回，所以损坏的 delta 或错误的 base 会直接报错失败，而不会把坏权重提供出去。apply 还拒绝乱序执行：一个版本只会在它声明的 base 版本之上 apply。 `--update-weight-delta-checksum` 选择算法。checksum 不是 apply 的瓶颈（apply 受解压 + XOR 限制），所以这是一个 digest 属性的选择，而非速度选择：`xxh3-128`（默认）是最宽的快速非加密 digest；`blake3` 是加密 digest，用于不可信存储；`adler32` 用于与期望它的系统互操作。 ## 共享文件系统可见性 hook -在 POSIX 共享文件系统（NFS、Lustre……）上不需要额外步骤。对于需要显式 commit/refresh 才能让写入跨 host 可见的对象存储卷，可以提供两个可选 hook（通过 import 路径加载——slime 里不存在任何厂商特定代码）： +在 POSIX 共享文件系统（NFS、Lustre……）上不需要额外步骤。对于需要显式 commit/refresh 才能让写入跨 host 可见的对象存储挂载，可以提供两个可选 hook（通过 import 路径加载——slime 和 sglang 里都不存在任何厂商特定代码）： -- `--custom-delta-pre-push-path`：在一个版本的文件写完之后、通知 engine 读取之前调用（例如 commit volume）。签名：`hook(args, version_dir, rollout_engines)`。 -- `--custom-delta-pre-read-path`：在每个 rollout host 读取 delta 目录之前调用（例如 refresh volume）。签名：`hook(delta_dir, target_version)`。 +- `--custom-update-weight-post-write-path`（slime，训练端）：在一个版本的文件写完之后、通知 engine 读取之前调用（例如把待写入数据上传到底层对象存储）。签名：`hook(args, version_dir, rollout_engines)`。 +- `--sglang-custom-pull-weights-pre-read-hook`（sglang server 参数，engine 端）：在每个 host 上、`/pull_weights` 读取 delta 目录之前于 engine 内部调用（例如刷新挂载视图）。签名：`hook(delta_dir, target_version)`。 diff --git a/docs/zh/advanced/external-rollout-engines.md b/docs/zh/advanced/external-rollout-engines.md index 0007bcff69..c077f82608 100644 --- a/docs/zh/advanced/external-rollout-engines.md +++ b/docs/zh/advanced/external-rollout-engines.md @@ -67,6 +67,8 @@ full checkpoint update from disk 是 external 场景最简单的兜底路径： 每次权重同步时，训练端会在 `--update-weight-disk-dir` 下写一个完整 HF checkpoint 目录，例如 `weight_v000123/`，然后通过 HTTP 调用每个 SGLang engine 的 `update_weights_from_disk`，让 engine 在不重启进程的情况下重新加载 checkpoint。 +额外设置 `--update-weight-local-checkpoint-dir` 后，每个 engine 会先把发布的 checkpoint pull 到它覆盖的每个 host 的本地磁盘（`/pull_weights`，随 slime 的 sglang patch 提供），再从本地（如 NVMe）reload——共享文件系统每个 host 只读一次，而不是每个 rank 读一次；当共享目录是对象存储或 engine 跨多个节点时尤其重要。 + 这个模式的优点是控制面简单：不要求训练器和 engine 建 NCCL group，只要求二者能看到同一个共享文件系统路径。缺点也直接：每次同步都写完整 actor 权重，对大模型和高频同步来说非常重。 调试时可以加： @@ -79,7 +81,7 @@ full checkpoint update from disk 是 external 场景最简单的兜底路径： ## Update With Delta -delta update 面向大模型、跨集群或跨数据中心训推解耦。它不每次都写完整 checkpoint，而是在训练端保留上一次同步的 CPU snapshot，逐参数比对，只发布变化的字节；每个 rollout host 把 delta apply 进自己的本地 checkpoint，再通过原生 `update_weights_from_disk` 端点 reload。 +delta update 面向大模型、跨集群或跨数据中心训推解耦。它不每次都写完整 checkpoint，而是在训练端保留上一次同步的 CPU snapshot，逐参数比对，只发布变化的字节；每个 engine 的 `/pull_weights` 端点（随 slime 的 sglang patch 提供）把 delta apply 进 engine 覆盖的每个 host 的本地 checkpoint，再通过原生 `update_weights_from_disk` 端点 reload。slime 只调用 engine 的 HTTP 端点，所以多节点 external engine 与 slime 拉起的 engine 行为一致。 ```bash --update-weight-mode delta diff --git a/docs/zh/get_started/customization.md b/docs/zh/get_started/customization.md index fd067c04c9..c3cc65e8a3 100644 --- a/docs/zh/get_started/customization.md +++ b/docs/zh/get_started/customization.md @@ -457,6 +457,25 @@ def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler | `--use-routing-replay` | 训练中前向-反向路由一致性。([arXiv:2507.18071](https://arxiv.org/abs/2507.18071)) | | `--use-rollout-routing-replay` | R3：在训练时重放 rollout 阶段的路由。slime 默认的 `sglang_router` 路径支持该功能。([arXiv:2510.11370](https://arxiv.org/abs/2510.11370)) | +--- + +### 19. Disk 权重同步 Post-Write Hook（`--custom-update-weight-post-write-path`） + +**签名**： +```python +def hook(args, version_dir: str, rollout_engines) -> None +``` + +**用途**：在 disk 权重同步（`--update-weight-transport disk`，full 或 delta 模式）的文件写完之后、 +engine 读取之前，在每个训练 rank 上调用。用于在非 POSIX 共享文件系统上发布写入——例如 commit +一个对象存储挂载——否则其他 host 无法看到这些文件。hook 会在每个 rank 上被调用，需要自行去重 +（例如每个容器只执行一次）。 + +读取侧的对应 hook 运行在推理引擎内部、engine 覆盖的每个 host 上，因此它是一个 sglang server +参数而不是 slime hook：传入 `--sglang-custom-pull-weights-pre-read-hook `，签名为 +`hook(source_dir: str, target_version: int)`——在 `/pull_weights` 读取已发布权重之前调用 +（例如刷新挂载视图）。完整机制见 [Delta 权重同步](../advanced/delta-weight-sync.md)。 + ## 自定义函数路径的测试 slime 现在也提供了一组 CPU 契约测试，用于校验这些 customization 接口。测试会通过字符串形式的导入路径来动态加载组件，因此既能回归仓库内置 hook，也能验证用户通过和训练时完全相同的 CLI 参数传入的自定义实现。 diff --git a/examples/delta_weight_sync/README.md b/examples/delta_weight_sync/README.md index 0879ba9fcb..6283900698 100644 --- a/examples/delta_weight_sync/README.md +++ b/examples/delta_weight_sync/README.md @@ -3,9 +3,9 @@ Non-colocated weight sync that ships only the **changed bytes** between two syncs instead of a full checkpoint, for training/inference disaggregation across clusters or datacenters. The trainer publishes per-tensor deltas to a shared filesystem as a canonical HF checkpoint -directory; each rollout host applies them into a host-local checkpoint and the engines reload -through the ordinary `update_weights_from_disk` path — the inference engine needs no -delta-specific support. +directory; each engine's `/pull_weights` applies them into a host-local checkpoint on every +host it spans, and the engines reload through the ordinary `update_weights_from_disk` path — +slime only ever talks to one endpoint per engine. See [Delta Weight Sync](../../docs/en/advanced/delta-weight-sync.md) for the full mechanism, encodings, integrity checks, and shared-filesystem visibility hooks. @@ -31,10 +31,11 @@ at `--update-weight-disk-dir`): - `--update-weight-disk-dir` — shared directory the trainer writes deltas to and the hosts read. - `--update-weight-local-checkpoint-dir` — host-local full HF checkpoint the delta patches in - place; materialized from `--hf-checkpoint` at engine start. + place; materialized from the engine's model path on the first `/pull_weights`. - `--update-weight-delta-encoding` — `xor` (smallest/fastest) or `overwrite` (idempotent). - `--update-weight-delta-checksum` — `xxh3-128` (default), `blake3`, or `adler32`. For object-store-backed volumes that need an explicit commit/refresh to make writes visible -across hosts, supply `--custom-delta-pre-push-path` / `--custom-delta-pre-read-path` (no -vendor-specific code lives in slime; see the doc). +across hosts, supply `--custom-update-weight-post-write-path` (trainer side) / +`--sglang-custom-pull-weights-pre-read-hook` (engine side) — no vendor-specific code lives in slime +or sglang; see the doc. diff --git a/examples/delta_weight_sync/run-glm4.7-30B-A3B-delta.sh b/examples/delta_weight_sync/run-glm4.7-30B-A3B-delta.sh index a399b20bbe..32996f79ad 100644 --- a/examples/delta_weight_sync/run-glm4.7-30B-A3B-delta.sh +++ b/examples/delta_weight_sync/run-glm4.7-30B-A3B-delta.sh @@ -1,8 +1,8 @@ #!/bin/bash # Disk delta weight-sync demo on GLM-4.7-Flash (30B-A3B), non-colocated, 2 nodes x 8 GPU. # The trainer publishes per-tensor deltas to --update-weight-disk-dir as a canonical HF directory; -# each rollout host applies them into --update-weight-local-checkpoint-dir and reloads via the -# vanilla update_weights_from_disk path. +# each engine's /pull_weights applies them into --update-weight-local-checkpoint-dir on every host +# it spans, and the engine reloads via the vanilla update_weights_from_disk path. # # Prerequisites: # - A 2-node (16-GPU) Ray cluster, this script run on the head node. @@ -10,7 +10,7 @@ # - dapo-math-17k.jsonl. # - --update-weight-disk-dir on a filesystem both nodes share. On an object-store-backed volume # that needs an explicit commit/refresh to surface writes across hosts, also pass -# --custom-delta-pre-push-path / --custom-delta-pre-read-path (see the doc). +# --custom-update-weight-post-write-path / --sglang-custom-pull-weights-pre-read-hook (see the doc). set -ex export PYTHONUNBUFFERED=1 diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 155324ffe1..d2ea6bdedb 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -139,8 +139,9 @@ def init( update_weight_transport = self.args.update_weight_transport if update_weight_mode == "delta": - # Delta sync is disk-transport only: each host applies the published deltas into - # its local checkpoint and the engines reload via vanilla update_weights_from_disk. + # Delta sync is disk-transport only: each engine's /pull_weights applies the published + # deltas into a host-local checkpoint on every host it spans, and the engines reload + # via vanilla update_weights_from_disk. assert not self.args.colocate, "--update-weight-mode=delta is not supported with --colocate" assert ( update_weight_transport == "disk" @@ -566,7 +567,6 @@ def update_weights(self) -> None: num_new_engines, engine_gpu_counts, engine_gpu_offsets, - all_engine_actors, ) = ray.get(self.rollout_manager.get_updatable_engines_and_lock.remote()) reconnect_rollout_engines = self.args.offload_train and self.args.use_critic and not self.args.colocate @@ -587,7 +587,6 @@ def update_weights(self) -> None: rollout_engine_lock, engine_gpu_counts=engine_gpu_counts, engine_gpu_offsets=engine_gpu_offsets, - all_engine_actors=all_engine_actors, ) dist.barrier(group=get_gloo_group()) if dist.get_rank() == 0: diff --git a/slime/backends/megatron_utils/hf_checkpoint_saver.py b/slime/backends/megatron_utils/hf_checkpoint_saver.py index 23d618aef8..ec3fe18e13 100644 --- a/slime/backends/megatron_utils/hf_checkpoint_saver.py +++ b/slime/backends/megatron_utils/hf_checkpoint_saver.py @@ -261,14 +261,37 @@ def _finalize_distributed_shards(path: Path, local_state: dict[str, Any]) -> Non else: states = [local_state] - if _is_global_rank_zero(): - _finalize_shard_files(path, states) + _finalize_local_shards(path, local_state, states, write_index=_is_global_rank_zero()) if dist.is_available() and dist.is_initialized(): dist.barrier() -def _finalize_shard_files(path: Path, shard_states: list[dict[str, Any] | None]) -> None: +def _finalize_local_shards( + path: Path, + local_state: dict[str, Any], + shard_states: list[dict[str, Any] | None], + *, + write_index: bool, +) -> None: + """Rename this rank's shard files per the global plan; optionally write the index. + + The plan is deterministic from the gathered states, so each rank renames only + its own files: on a non-POSIX shared filesystem another rank's unpublished + writes are not visible, let alone renamable. + """ + rename_map, index_data = _plan_shard_finalization(shard_states) + for old_name in local_state.get("shard_files", []): + os.replace(path / old_name, path / rename_map[old_name]) + if write_index: + with open(path / "model.safetensors.index.json", "w", encoding="utf-8") as f: + json.dump(index_data, f, indent=2) + + +def _plan_shard_finalization( + shard_states: list[dict[str, Any] | None], +) -> tuple[dict[str, str], dict[str, Any]]: + """Compute the shard rename map and index from every rank's gathered state.""" shard_files = [] total_size = 0 raw_weight_map = {} @@ -295,9 +318,7 @@ def _finalize_shard_files(path: Path, shard_states: list[dict[str, Any] | None]) total_files = len(shard_files) rename_map = {} for idx, old_name in enumerate(shard_files, start=1): - new_name = f"model-{idx:05d}-of-{total_files:05d}.safetensors" - os.replace(path / old_name, path / new_name) - rename_map[old_name] = new_name + rename_map[old_name] = f"model-{idx:05d}-of-{total_files:05d}.safetensors" final_weight_map = {} for name, filename in raw_weight_map.items(): @@ -306,8 +327,7 @@ def _finalize_shard_files(path: Path, shard_states: list[dict[str, Any] | None]) final_weight_map[name] = rename_map[filename] index_data = {"metadata": {"total_size": total_size}, "weight_map": final_weight_map} - with open(path / "model.safetensors.index.json", "w", encoding="utf-8") as f: - json.dump(index_data, f, indent=2) + return rename_map, index_data def _shard_filename_sort_key(filename: str) -> tuple[float, str]: diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_disk.py b/slime/backends/megatron_utils/update_weight/update_weight_from_disk.py index d4c59f0001..06d4068901 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_disk.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_disk.py @@ -35,6 +35,14 @@ def __init__( self.update_weight_metrics: dict[str, float] = {} self.rollout_engines: Sequence[ActorHandle] = [] self.rollout_engine_lock: ActorHandle | None = None + # Post-write hook: object-store-backed shared filesystems lack cross-host + # read-after-write consistency, so written files need an explicit step + # (e.g. uploading them to the backing object store) before the engines can see them. + self._post_write_hook: Callable | None = None + if args.custom_update_weight_post_write_path: + from slime.utils.misc import load_function + + self._post_write_hook = load_function(args.custom_update_weight_post_write_path) def connect_rollout_engines( self, @@ -42,7 +50,6 @@ def connect_rollout_engines( rollout_engine_lock: ActorHandle, engine_gpu_counts: Sequence[int] | None = None, engine_gpu_offsets: Sequence[int] | None = None, - all_engine_actors: Sequence[ActorHandle] | None = None, ) -> None: self.rollout_engines = rollout_engines self.rollout_engine_lock = rollout_engine_lock @@ -63,6 +70,9 @@ def update_weights(self) -> None: shutil.rmtree(version_dir, ignore_errors=True) dist.barrier(group=get_gloo_group()) + # every writing rank creates the dir itself: a non-POSIX shared filesystem may not surface + # one rank's mkdir to another until commit + version_dir.mkdir(parents=True, exist_ok=True) save_hf_model_to_path( self.args, version_dir, @@ -73,6 +83,12 @@ def update_weights(self) -> None: ) dist.barrier(group=get_gloo_group()) + # every rank runs the hook (it gates itself): each container must publish + # its own writes + if self._post_write_hook is not None: + self._post_write_hook(self.args, str(version_dir), list(self.rollout_engines)) + dist.barrier(group=get_gloo_group()) + # SGLang reload is orchestrated by RayTrainGroup after the checkpoint # is fully written, so training-side lifecycle can decide whether # Megatron actors are still alive. diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py b/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py index acb50cf699..4d660a035f 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py @@ -31,8 +31,9 @@ class UpdateWeightFromDiskDelta(UpdateWeightFromDistributed): """ Delta weight sync over a shared filesystem. PP-src ranks diff each gathered HF tensor against a CPU snapshot of the previous sync and publish the changes as a canonical HF checkpoint dir; - every rollout host applies the delta into its local checkpoint and reloads via the ordinary - update_weights_from_disk path, so sglang needs no delta support. + each engine's /pull_weights fans the apply out to every host it spans, then the engine reloads + the patched local checkpoint via the ordinary update_weights_from_disk path. slime only ever + talks to one endpoint per engine, so multi-node serving and external engines need nothing extra. """ def __init__( @@ -51,11 +52,14 @@ def __init__( self.checksum_algorithm = args.update_weight_delta_checksum self._snapshot: dict[str, np.ndarray] = {} self._baseline_captured = False - self._commit_hook: Callable | None = None - if args.custom_delta_pre_push_path: + # Post-write hook: object-store-backed shared filesystems lack cross-host + # read-after-write consistency, so written files need an explicit step + # (e.g. uploading them to the backing object store) before the engines can see them. + self._post_write_hook: Callable | None = None + if args.custom_update_weight_post_write_path: from slime.utils.misc import load_function - self._commit_hook = load_function(args.custom_delta_pre_push_path) + self._post_write_hook = load_function(args.custom_update_weight_post_write_path) def connect_rollout_engines( self, @@ -63,13 +67,10 @@ def connect_rollout_engines( rollout_engine_lock: ActorHandle, engine_gpu_counts: Sequence[int] | None = None, engine_gpu_offsets: Sequence[int] | None = None, - all_engine_actors: Sequence[ActorHandle] | None = None, ) -> None: - # The local checkpoint is host-local, so every host applies its own copy: - # all_engine_actors is one actor per host, vs rollout_engines (node 0 only). The - # rollout_engine_lock the NCCL path uses isn't needed — a per-host flock serializes applies. + # The rollout_engine_lock the NCCL path uses isn't needed — the engine-side apply is + # serialized by a per-host flock. self.rollout_engines = rollout_engines - self.all_engine_actors = list(all_engine_actors or rollout_engines) self._is_pp_src_rank = ( mpu.get_data_parallel_rank(with_context_parallel=True) == 0 and mpu.get_tensor_model_parallel_rank() == 0 ) @@ -95,13 +96,16 @@ def _capture_baseline(self) -> None: stale stream from a prior run. Seeds from hf_checkpoint — what each host materializes its base from — so the invariant ``snapshot == engine base`` holds even where the megatron->HF round-trip trims vocab-padding rows (embed/lm_head). A tensor absent there (rare) falls back - to the gathered value.""" + to the gathered value. pull_weights(0) makes each host materialize its local base now, + overlapped with the snapshot gather, so the first real sync only pays the delta apply.""" # a prior run's versions would apply against the wrong base; start the dir clean + pulls = [] if dist.get_rank() == 0: shutil.rmtree(self.delta_dir, ignore_errors=True) os.makedirs(self.delta_dir, exist_ok=True) - if self._commit_hook is not None: - self._commit_hook(self.args, self.delta_dir, list(self.rollout_engines)) + if self._post_write_hook is not None: + self._post_write_hook(self.args, self.delta_dir, list(self.rollout_engines)) + pulls = [engine.pull_weights.remote(target_version=0) for engine in self.rollout_engines] dist.barrier(group=get_gloo_group()) read_hf = make_tensor_reader(self.args.hf_checkpoint) # index the HF headers once @@ -112,6 +116,7 @@ def _capture_baseline(self) -> None: self._snapshot[name] = tensor.detach().cpu().contiguous().view(torch.uint8).numpy().reshape(-1) logger.warning("seed: %s absent from hf_checkpoint; seeding from current weights", name) if dist.get_rank() == 0: + ray.get(pulls) logger.info( "[disk delta] captured baseline snapshot of %d tensors from %s", len(self._snapshot), @@ -127,7 +132,7 @@ def _publish(self) -> None: def _write_delta_files(self) -> None: """Write this rank's changed tensors as one canonical model-NNNNN.safetensors, and on rank 0 the HF index. The sequential file numbers and the index are coordinated over gloo (small - object gathers), not the filesystem — a shared volume may not surface one rank's writes to + object gathers), not the filesystem — a non-POSIX shared filesystem may not surface one rank's writes to another until commit.""" group = get_gloo_group() world, rank = dist.get_world_size(), dist.get_rank() @@ -162,12 +167,13 @@ def _write_delta_files(self) -> None: dist.barrier(group=group) def _reload_engines(self) -> None: - """Commit the published files, have each host apply the delta, then reload the engines.""" - if self._commit_hook is not None: - self._commit_hook(self.args, self._version_dir, list(self.rollout_engines)) + """Commit the published files, have each engine pull the delta onto every host it spans + (checksum-verified), then reload the engines.""" + if self._post_write_hook is not None: + self._post_write_hook(self.args, self._version_dir, list(self.rollout_engines)) dist.barrier(group=get_gloo_group()) if dist.get_rank() == 0: - ray.get([actor.sync_local_checkpoint.remote(self.weight_version) for actor in self.all_engine_actors]) + ray.get([engine.pull_weights.remote(self.weight_version) for engine in self.rollout_engines]) ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) ray.get( diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 2cede1daac..1da32420c6 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -60,7 +60,6 @@ def connect_rollout_engines( rollout_engine_lock: ActorHandle, engine_gpu_counts: Sequence[int] | None = None, engine_gpu_offsets: Sequence[int] | None = None, - all_engine_actors: Sequence[ActorHandle] | None = None, ) -> None: """ Create NCCL "slime-pp_{pp_rank}" if PP source (DP=TP=0). Lock prevents concurrent broadcasts. diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 606d8aa590..f8efd69c0a 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -65,7 +65,6 @@ def connect_rollout_engines( rollout_engine_lock: ActorHandle, engine_gpu_counts: Sequence[int] | None = None, engine_gpu_offsets: Sequence[int] | None = None, - all_engine_actors: Sequence[ActorHandle] | None = None, ) -> None: """ Split colocated/distributed engines. Global source rank (DP=TP=PP=0) creates NCCL diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index f4636d3b24..5059329995 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -3,7 +3,6 @@ import logging import multiprocessing import os -import threading import time from urllib.parse import quote @@ -168,19 +167,6 @@ def _format_v6_uri(addr): else: self._init_normal(server_args_dict) - # Warm the host-local base off the actor's main thread: sglang serves the first rollout from - # its init-loaded weights, so the materialize (a full base copy) only has to finish before - # the first delta reload. init_local_checkpoint is idempotent and flock-guarded, so the first - # sync_local_checkpoint either finds it done or blocks on the same lock — no join needed. - if self.args.update_weight_mode == "delta" and self.args.update_weight_transport == "disk": - from slime.utils.disk_delta import init_local_checkpoint - - threading.Thread( - target=init_local_checkpoint, - args=(self.args.update_weight_local_checkpoint_dir, self.args.hf_checkpoint), - daemon=True, - ).start() - def _init_external(self, expect_server_args, external_engine_need_check_fields): logger.info(f"Use external SGLang engine (rank={self.rank}, expect_server_args={expect_server_args})") @@ -368,15 +354,6 @@ def get_weight_version(self): response.raise_for_status() return response.json()["weight_version"] - def set_weight_version(self, new_version: str): - """Bump the engine's recorded weight version without changing weights. - - Used by the delta-update path when a sync produced no bytes (e.g. an - all-zero diff): we still need the engine's version to track the - updater's, otherwise the CI version-equality check will trip. - """ - return self._make_request("update_weight_version", {"new_version": str(new_version)}) - def release_memory_occupation(self): self.flush_cache() return self._make_request("release_memory_occupation") @@ -393,23 +370,18 @@ def resume_memory_occupation(self, tags: list[str] = None): def check_weights(self, action: str): return self._make_request("weights_checker", {"action": action}) - def sync_local_checkpoint(self, target_version: int): - """Apply the published deltas into this host's local checkpoint up to target_version; the - engine reloads it afterwards. Assumes this actor shares the checkpoint filesystem with the - sglang it drives (true for slime-launched engines).""" - from slime.utils.disk_delta import apply_deltas, init_local_checkpoint - - init_local_checkpoint(self.args.update_weight_local_checkpoint_dir, self.args.hf_checkpoint) # idempotent - # non-POSIX filesystems lack cross-host read-after-write consistency, so the trainer's - # just-written delta isn't visible on this mount until the hook refreshes it. - if self.args.custom_delta_pre_read_path: - from slime.utils.misc import load_function - - load_function(self.args.custom_delta_pre_read_path)(self.args.update_weight_disk_dir, target_version) - apply_deltas( - self.args.update_weight_local_checkpoint_dir, - self.args.update_weight_disk_dir, - target_version, + def pull_weights(self, target_version: int): + """Have the engine sync every host it spans to target_version: each host pulls the + published weights (a full checkpoint copied as-is, or deltas verified per-tensor and + applied onto the local checkpoint) into its local checkpoint dir. The engine reloads + it afterwards via update_weights_from_disk.""" + return self._make_request( + "pull_weights", + { + "local_checkpoint_dir": self.args.update_weight_local_checkpoint_dir, + "source_dir": self.args.update_weight_disk_dir, + "target_version": target_version, + }, ) def update_weights_from_disk( @@ -417,23 +389,13 @@ def update_weights_from_disk( model_path: str, load_format: str | None = None, weight_version: str | None = None, - files: list[str] | None = None, ): - """Reload weights from *model_path* without restarting the engine. - - Standard HF reload: ``model_path`` is the checkpoint directory. - Delta (``load_format="delta"``): ``model_path`` is the parent of the - per-sync version subdir and ``files`` is the basenames within it to read + - apply. Each delta call is independent — sender owns batching, sync - boundaries, cleanup. - """ + """Reload weights from the checkpoint at *model_path* without restarting the engine.""" payload: dict = {"model_path": model_path} if load_format is not None: payload["load_format"] = load_format if weight_version is not None: payload["weight_version"] = weight_version - if files is not None: - payload["files"] = files return self._make_request("update_weights_from_disk", payload) def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend): diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py index 294381e9ab..4237b80329 100644 --- a/slime/ray/actor_group.py +++ b/slime/ray/actor_group.py @@ -232,12 +232,20 @@ def _reload_rollout_weights_from_disk(self, disk_weight_dir, weight_version): if not self.args.update_weight_disk_keep_files: shutil.rmtree(disk_weight_dir, ignore_errors=True) return + if self.args.update_weight_local_checkpoint_dir: + # each host pulls the published checkpoint onto local disk (e.g. NVMe) and + # the engines reload from there; the pull is disk-only, so it runs before + # pause and overlaps generation + ray.get([engine.pull_weights.remote(int(weight_version)) for engine in engines]) + model_path = self.args.update_weight_local_checkpoint_dir + else: + model_path = str(disk_weight_dir) ray.get([engine.pause_generation.remote() for engine in engines]) ray.get([engine.flush_cache.remote() for engine in engines]) ray.get( [ engine.update_weights_from_disk.remote( - model_path=str(disk_weight_dir), + model_path=model_path, weight_version=weight_version, ) for engine in engines diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index be8f42438c..d0653522f7 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -525,8 +525,7 @@ def get_updatable_engines_and_lock(self): gpu_counts = srv.engine_gpu_counts if srv else [] gpu_offsets = srv.engine_gpu_offsets if srv else [] num_new = srv.num_new_engines if srv else 0 - all_engine_actors = srv.all_engines if srv else [] - return engines, self.rollout_engine_lock, num_new, gpu_counts, gpu_offsets, all_engine_actors + return engines, self.rollout_engine_lock, num_new, gpu_counts, gpu_offsets def get_num_rollout_per_epoch(self): assert self.args.rollout_global_dataset diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index ccb121f35d..0a79b99ded 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -209,37 +209,30 @@ def add_train_arguments(parser): ), ) parser.add_argument( - "--custom-delta-pre-push-path", + "--custom-update-weight-post-write-path", type=str, default=None, help=( - "Path to a custom function called on each trainer rank after its delta files " - "are written, before the engines read them — to publish the writes on a " - "non-POSIX filesystem (no cross-host visibility without an explicit sync). " + "Path to a custom function called on each trainer rank after a disk weight " + "sync's files are written (full or delta), before the engines read them — to " + "publish the writes on a non-POSIX filesystem (no cross-host visibility " + "without an explicit sync). " "Signature: ``def hook(args, version_dir: str, rollout_engines) -> None``; the hook gates itself." ), ) - parser.add_argument( - "--custom-delta-pre-read-path", - type=str, - default=None, - help=( - "Path to a custom function called on each rollout host before it reads the " - "published delta directory — refreshes the mount so the just-published version " - "is visible on a non-POSIX filesystem (no cross-host read-after-write consistency). " - "Signature: ``def hook(delta_dir: str, target_version: int) -> None``." - ), - ) parser.add_argument( "--update-weight-local-checkpoint-dir", type=str, default=None, help=( - "Rollout-host-local directory (NVMe) holding a full HF checkpoint that " - "disk-delta sync patches in place. Each host materializes it from " - "--hf-checkpoint at engine start, applies each version's delta there, and " - "the engines reload from it. Required for --update-weight-mode=delta " - "--update-weight-transport=disk." + "Rollout-host-local directory (NVMe) holding a full HF checkpoint kept in " + "sync by each engine's /pull_weights: every host copies a published full " + "checkpoint as-is or patches published deltas in place, and the engines " + "reload from it. Required for --update-weight-mode=delta " + "--update-weight-transport=disk; optional for full disk sync (engines then " + "pull to local disk instead of reading the shared dir directly). The " + "read-side counterpart of --custom-update-weight-post-write-path is the engine's " + "--sglang-custom-pull-weights-pre-read-hook." ), ) parser.add_argument( diff --git a/slime/utils/disk_delta.py b/slime/utils/disk_delta.py index 3ba8cef12f..50676b6392 100644 --- a/slime/utils/disk_delta.py +++ b/slime/utils/disk_delta.py @@ -1,29 +1,21 @@ from __future__ import annotations -import fcntl import glob -import io import json -import logging -import mmap import os -import shutil import struct -import threading import zlib -from concurrent.futures import ThreadPoolExecutor -from contextlib import contextmanager import numpy as np -import zstandard -logger = logging.getLogger(__name__) - -# The delta phases (XOR/scatter, zstd, checksum) are memory-bandwidth bound and release the GIL, +# The delta phases (diff, zstd, checksum) are memory-bandwidth bound and release the GIL, # so a thread pool over tensors recovers the bandwidth one thread leaves idle. NUM_WORKERS = min(32, (os.cpu_count() or 8)) -SYNC_DIR = ".delta_sync" # per-checkpoint dir holding the applied-version marker and the apply lock +# Trainer-side (publish) helpers for disk-level delta weight sync. The receive side — +# materializing the host-local checkpoint and applying published deltas in place — lives in +# the engine behind its /pull_weights endpoint (sglang.srt.weight_sync.disk_delta), so it +# runs on every host of a multi-node engine while slime only talks to one endpoint. def overwrite_encode(new: np.ndarray, changed_mask: np.ndarray) -> np.ndarray: @@ -66,63 +58,6 @@ def checksum(algorithm: str, buf) -> str: return hasher.hexdigest() -@contextmanager -def _apply_lock(local_ckpt_dir: str): - sync = os.path.join(local_ckpt_dir, SYNC_DIR) - os.makedirs(sync, exist_ok=True) - with open(os.path.join(sync, "lock"), "w") as f: - fcntl.flock(f, fcntl.LOCK_EX) - try: - yield - finally: - fcntl.flock(f, fcntl.LOCK_UN) - - -def _read_applied_version(local_ckpt_dir: str) -> str | None: - try: - with open(os.path.join(local_ckpt_dir, SYNC_DIR, "state.json")) as f: - return json.load(f)["version"] - except FileNotFoundError: - return None - - -def _write_applied_version(local_ckpt_dir: str, version: str) -> None: - path = os.path.join(local_ckpt_dir, SYNC_DIR, "state.json") - tmp = path + ".tmp" - with open(tmp, "w") as f: - json.dump({"version": version}, f) - f.flush() - os.fsync(f.fileno()) - os.replace(tmp, path) - - -def drop_page_cache(path: str) -> None: - """Evict a file from the page cache (POSIX_FADV_DONTNEED).""" - 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 init_local_checkpoint(local_ckpt_dir: str, base_dir: str) -> None: - """Copy the base HF checkpoint into local_ckpt_dir once if absent (run at engine start). Each - later delta is applied on top of this copy in place.""" - with _apply_lock(local_ckpt_dir): - if _read_applied_version(local_ckpt_dir) is not None: - return - logger.info("Materializing base checkpoint %s -> %s", base_dir, local_ckpt_dir) - os.makedirs(local_ckpt_dir, exist_ok=True) - for entry in os.scandir(base_dir): - if entry.is_file(): - shutil.copy2(entry.path, os.path.join(local_ckpt_dir, entry.name)) - drop_page_cache(entry.path) # don't let the source evict the local copy we keep resident - _write_applied_version(local_ckpt_dir, "000000") - - def _tensor_locations(ckpt_dir: str) -> dict[str, tuple[str, int, int]]: """Map each tensor name to (file, byte offset, nbytes) by reading every safetensors header.""" locations: dict[str, tuple[str, int, int]] = {} @@ -150,115 +85,3 @@ def read(name: str) -> np.ndarray: return np.frombuffer(f.read(nbytes), dtype=np.uint8) return read - - -def _apply_version(local_ckpt_dir: str, version_dir: str) -> None: - """Apply one version's delta in place: decompress + apply + checksum each tensor across a thread - pool (each writes a distinct mmap region, so the writes don't conflict). Any mismatch raises.""" - with open(os.path.join(version_dir, "model.safetensors.index.json")) as f: - meta = json.load(f)["metadata"] - applied = _read_applied_version(local_ckpt_dir) - if applied == meta["version"]: - return - if applied != meta["base_version"]: - raise RuntimeError(f"out-of-order delta: local at {applied}, delta builds on {meta['base_version']}") - if meta["compression_format"] != "zstd": - raise NotImplementedError(f"compression {meta['compression_format']!r} not supported") - encoding = meta["delta_encoding"] - algorithm = meta["checksum_format"] - locations = _tensor_locations(local_ckpt_dir) - open_mmaps: dict[str, tuple] = {} - mismatches: list[str] = [] - lock = threading.Lock() - file_bytes: list[bytes] = [] # keep alive: items hold zero-copy views into these - items: list[tuple] = [] # (name, compressed_view, path, offset, nbytes, want_checksum) - try: - for delta_file in sorted(glob.glob(os.path.join(version_dir, "*.safetensors"))): - with open(delta_file, "rb") as f: - blob = f.read() - file_bytes.append(blob) - (header_len,) = struct.unpack(" None: - name, compressed, path, offset, nbytes, want = item - region = np.ndarray((nbytes,), dtype=np.uint8, buffer=open_mmaps[path][1], offset=offset) - hasher = _new_hasher(algorithm) - reader = zstandard.ZstdDecompressor().stream_reader(io.BytesIO(bytes(compressed))) - pos = 0 - while pos < nbytes: # 2 MB chunks stay L2-resident across decompress -> XOR -> checksum - block = reader.read(min(2 << 20, nbytes - pos)) - if not block: - break - chunk = np.frombuffer(block, dtype=np.uint8) - region[pos : pos + chunk.size] ^= chunk - hasher.update(region[pos : pos + chunk.size]) - pos += chunk.size - if hasher.hexdigest() != want: - with lock: - mismatches.append(name) - - def apply_overwrite(item) -> None: - name, compressed, path, offset, nbytes, want = item - delta = np.frombuffer(zstandard.ZstdDecompressor().decompress(bytes(compressed)), dtype=np.uint8) - region = np.ndarray((nbytes,), dtype=np.uint8, buffer=open_mmaps[path][1], offset=offset) - count = int.from_bytes(delta[:4].tobytes(), "little") - positions = np.frombuffer(delta[4 : 4 + 4 * count].tobytes(), dtype=" None: - """Apply the delta chain in order to bring the local checkpoint up to target_version, in place. - A per-tensor checksum guards every write and any mismatch raises (fail loud, never serve bad - weights). Serialized per host by the lock (co-located actors collapse to one apply).""" - with _apply_lock(local_ckpt_dir): - applied = _read_applied_version(local_ckpt_dir) - if applied is None: - raise RuntimeError("local checkpoint not materialized") - for version in range(int(applied) + 1, target_version + 1): - _apply_version(local_ckpt_dir, os.path.join(delta_root, f"weight_v{version:06d}")) diff --git a/tests/utils/test_hf_checkpoint_saver.py b/tests/utils/test_hf_checkpoint_saver.py index e50e17018c..9e0dd85393 100644 --- a/tests/utils/test_hf_checkpoint_saver.py +++ b/tests/utils/test_hf_checkpoint_saver.py @@ -9,7 +9,7 @@ from slime.backends.megatron_utils.hf_checkpoint_saver import ( _clear_existing_hf_weights, _copy_hf_assets, - _finalize_shard_files, + _finalize_local_shards, _SafetensorShardWriter, _write_pending_chunk, save_hf_model_direct_to_path, @@ -88,7 +88,10 @@ def test_finalize_shard_files_merges_node_writer_states(tmp_path: Path): writer0.write([("layers.0.weight", torch.ones(2, 2))], shard_idx=0) writer1.write([("layers.1.weight", torch.zeros(2, 2))], shard_idx=1) - _finalize_shard_files(tmp_path, [writer0.state(), writer1.state()]) + # each rank renames its own files off the shared plan; rank 0 writes the index + states = [writer0.state(), writer1.state()] + for rank, state in enumerate(states): + _finalize_local_shards(tmp_path, state, states, write_index=rank == 0) index = json.loads((tmp_path / "model.safetensors.index.json").read_text(encoding="utf-8")) assert index["metadata"]["total_size"] == 32 @@ -124,7 +127,9 @@ def test_pending_chunk_write_flushes_incomplete_node_group(tmp_path: Path): for i, writer in enumerate(writers): pending_writes[i] = _write_pending_chunk(writer, pending_writes[i]) - _finalize_shard_files(tmp_path, [writer.state() for writer in writers]) + states = [writer.state() for writer in writers] + for rank, state in enumerate(states): + _finalize_local_shards(tmp_path, state, states, write_index=rank == 0) index = json.loads((tmp_path / "model.safetensors.index.json").read_text(encoding="utf-8")) assert index["weight_map"] == {f"layers.{i}.weight": f"model-{i + 1:05d}-of-00005.safetensors" for i in range(5)}