feat: support sparse HCCL weight sync with TP

.buildkite/npu_suites.py

@@ -29,6 +29,7 @@
 SUITES = {
     "smk": [
         ("test_qwen3_4B_npu.py", "npu-8", "", {}),
+        ("test_qwen3_4B_npu.py", "npu-8", "", {"VIME_TEST_UPDATE_MODE": "delta"}),
         ("test_qwen3_30B_A3B_npu.py", "npu-16", "", {}),
         ("test_qwen3_vl_8B_npu.py", "npu-8", "", {}),
     ],

.buildkite/scripts/update-npu-environment.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 # Purpose: Updates an NPU test container to match the requested VIME commit.
 #          - Reads the image's persisted OLD patch series and exact patch bytes
-#          - Updates VIME, then reconciles OLD -> NEW in declared series order
+#          - Updates VIME, then reconciles OLD -> NEW in declared series orde
 #          - Installs the current VIME checkout and normalizes visible devices
 # Usage: Called by Buildkite pipeline during NPU test runs
 set -e -o pipefail
@@ -176,6 +176,10 @@ update_vime_code() {
 
 install_vime_code() {
     pip install -e "$VIME_DIR" --no-deps --break-system-packages || pip install -e "$VIME_DIR" --no-deps
+    # The pre-built smoke image can predate disk-delta dependencies. Keep the
+    # serving and trainer interpreters aligned before loading delta checkpoints.
+    pip install blake3 xxhash zstandard --break-system-packages || \
+        pip install blake3 xxhash zstandard
 }
 
 sort_ascend_visible_devices() {

docker/npu_patch/vllm-ascend.patch

@@ -13,7 +13,7 @@ index 062b3ecd..dd87affb 100644
                  f"Free memory on device "
 @@ -535,15 +535,16 @@ class NPUWorker(WorkerBase):
          self.npugraph_memory_estimate = npugraph_memory_estimate
- 
+
          free_gpu_memory = profile_result.after_profile.free_memory
 -        assert self.init_snapshot.free_memory > free_gpu_memory, (
 -            "Error in memory profiling. "
@@ -50,4 +50,72 @@ index a35d9af8d..66a179cd6 100644
 +        # Source tensors may have been produced asynchronously on the caller's
 +        # stream. Wait before the packing stream reads them in torch.cat.
 +        streams[buffer_idx].wait_stream(source_stream)
-         # Start tasks for the new buffer in a new stream
+        # Start tasks for the new buffer in a new stream
+diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py
+index 062b3ecd..f0d1e4a1 100644
+--- a/vllm_ascend/worker/worker.py
++++ b/vllm_ascend/worker/worker.py
+@@ -333,1 +333,33 @@ class NPUWorker(WorkerBase):
++    def pull_weights(
++        self,
++        local_checkpoint_dir: str,
++        source_dir: str,
++        target_version: int,
++        pre_read_hook: str | None = None,
++    ) -> dict:
++        """Materialize a full or delta checkpoint on this rollout host.
++
++        ``collective_rpc`` invokes this on all NPU workers. The checkpoint
++        helper serializes same-host ranks with a filesystem lock, so each host
++        applies a version exactly once before reload_weights is called.
++        """
++        from vllm.utils.local_checkpoint import pull_checkpoint
++
++        # reload_weights updates model_config.model to the materialized local
++        # checkpoint. Preserve the original model path as the immutable
++        # version-zero seed for retries and restarted training runs.
++        base_dir = getattr(self, "_local_checkpoint_base_dir", None)
++        if base_dir is None:
++            base_dir = self.model_config.model
++            self._local_checkpoint_base_dir = base_dir
++
++        pull_checkpoint(
++            local_checkpoint_dir=local_checkpoint_dir,
++            base_dir=base_dir,
++            source_dir=source_dir,
++            target_version=target_version,
++            pre_read_hook=pre_read_hook,
++        )
++        return {"success": True, "weight_version": str(target_version)}
++
+     def shutdown(self) -> None:
+diff --git a/tests/ut/worker/a2/test_worker_v1.py b/tests/ut/worker/a2/test_worker_v1.py
+index 812b757..06e3875 100644
+--- a/tests/ut/worker/a2/test_worker_v1.py
++++ b/tests/ut/worker/a2/test_worker_v1.py
+@@ -1560,2 +1560,25 @@ class TestNPUWorker(TestBase):
++    def test_pull_weights_preserves_initial_model_path(self):
++        """The version-zero seed must survive reload_weights path updates."""
++        from vllm_ascend.worker.worker import NPUWorker
++
++        with (
++            patch.object(NPUWorker, "__init__", lambda x, **kwargs: None),
++            patch("vllm.utils.local_checkpoint.pull_checkpoint") as pull_checkpoint,
++        ):
++            worker = NPUWorker()
++            worker.model_config = MagicMock()
++            worker.model_config.model = "/models/base"
++
++            worker.pull_weights("/local/checkpoint", "/shared/weights", 0)
++            worker.model_config.model = "/local/checkpoint"
++            worker.pull_weights("/local/checkpoint", "/shared/weights", 1)
++
++        assert pull_checkpoint.call_count == 2
++        assert [call.kwargs["base_dir"] for call in pull_checkpoint.call_args_list] == [
++            "/models/base",
++            "/models/base",
++        ]
++        assert pull_checkpoint.call_args.kwargs["target_version"] == 1
++
+ class TestNPUWorkerWeightUpdate(TestBase):
+     def _make_worker(self, engine=None):

docker/npu_patch/vllm.patch

@@ -22,5 +22,250 @@ index cb61bca..5c076d5 100644
 -        assert request.num_output_placeholders >= 0
 +        request.num_output_placeholders = max(0, request.num_output_placeholders)
 
-         # Cache the new tokens. Preempted requests should be skipped.
-         if status_before_update == RequestStatus.RUNNING:
+        # Cache the new tokens. Preempted requests should be skipped.
+        if status_before_update == RequestStatus.RUNNING:
+diff --git a/vllm/utils/local_checkpoint.py b/vllm/utils/local_checkpoint.py
+new file mode 100644
+--- /dev/null
++++ b/vllm/utils/local_checkpoint.py
+@@ -0,0 +1,240 @@
++# SPDX-License-Identifier: Apache-2.0
++"""Maintain a host-local HF checkpoint from full and delta weight versions."""
++
++from __future__ import annotations
++
++import fcntl
++import glob
++import importlib
++import io
++import json
++import mmap
++import os
++import shutil
++import struct
++import threading
++import zlib
++from concurrent.futures import ThreadPoolExecutor
++from contextlib import ExitStack, contextmanager
++
++import numpy as np
++import zstandard
++
++NUM_WORKERS = min(32, os.cpu_count() or 8)
++SYNC_DIR = ".weight_sync"
++
++
++def pull_checkpoint(local_checkpoint_dir, base_dir, source_dir, target_version, pre_read_hook=None):
++    """Bring a host-local checkpoint to a published full or delta version."""
++    if target_version > 0 and pre_read_hook:
++        module_path, _, function_name = pre_read_hook.rpartition(".")
++        getattr(importlib.import_module(module_path), function_name)(source_dir, target_version)
++    with _pull_lock(local_checkpoint_dir):
++        incomplete = os.path.join(local_checkpoint_dir, SYNC_DIR, "incomplete")
++        # Deltas update mmap'ed safetensor regions in place. A failed apply can
++        # therefore leave bytes changed even though state.json was not advanced;
++        # discard that partial local state on the next retry.
++        applied = None if os.path.exists(incomplete) else _read_applied_version(local_checkpoint_dir)
++        # A pull can legitimately restart a run at version zero or target an
++        # older published version. In that case the current checkpoint cannot
++        # be patched backwards, so search from the base checkpoint instead.
++        floor = applied if applied is not None and applied <= target_version else 0
++        start = target_version
++        while start > floor and _is_delta(_version_dir(source_dir, start)):
++            start -= 1
++        if applied is None or applied > target_version or start > applied:
++            _reset_checkpoint(base_dir if start == 0 else _version_dir(source_dir, start), local_checkpoint_dir, start)
++        else:
++            start = applied
++        try:
++            for version in range(start + 1, target_version + 1):
++                open(incomplete, "a").close()
++                _apply_delta(local_checkpoint_dir, _version_dir(source_dir, version))
++        except BaseException:
++            raise
++        else:
++            if os.path.exists(incomplete):
++                os.remove(incomplete)
++
++
++def _version_dir(source_dir, version):
++    return os.path.join(source_dir, f"weight_v{version:06d}")
++
++
++def _is_delta(version_dir):
++    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 index_file:
++            return "delta_encoding" in json.load(index_file).get("metadata", {})
++    except FileNotFoundError:
++        return False
++
++
++class _Adler32:
++    def __init__(self):
++        self._value = 1
++
++    def update(self, data):
++        self._value = zlib.adler32(data, self._value)
++
++    def hexdigest(self):
++        return f"{self._value:08x}"
++
++
++def _new_hasher(algorithm):
++    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}")
++
++
++@contextmanager
++def _pull_lock(local_checkpoint_dir):
++    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):
++    try:
++        with open(os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json")) as state_file:
++            return int(json.load(state_file)["version"])
++    except FileNotFoundError:
++        return None
++
++
++def _write_applied_version(local_checkpoint_dir, version):
++    path = os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json")
++    temporary = path + ".tmp"
++    with open(temporary, "w") as state_file:
++        json.dump({"version": f"{version:06d}"}, state_file)
++        state_file.flush()
++        os.fsync(state_file.fileno())
++    os.replace(temporary, path)
++
++
++def _reset_checkpoint(source_dir, local_checkpoint_dir, version):
++    os.makedirs(local_checkpoint_dir, exist_ok=True)
++    source_files = [entry for entry in os.scandir(source_dir) if entry.is_file()]
++    for entry in source_files:
++        shutil.copy2(entry.path, os.path.join(local_checkpoint_dir, entry.name))
++    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 = os.path.join(local_checkpoint_dir, entry.name)
++        if os.path.getsize(copied) != entry.stat().st_size:
++            raise RuntimeError(f"Size mismatch copying {entry.name}")
++    _write_applied_version(local_checkpoint_dir, version)
++
++
++def _tensor_locations(checkpoint_dir):
++    locations = {}
++    for path in glob.glob(os.path.join(checkpoint_dir, "*.safetensors")):
++        with open(path, "rb") as tensor_file:
++            (header_length,) = struct.unpack("<Q", tensor_file.read(8))
++            header = json.loads(tensor_file.read(header_length))
++        for name, info in header.items():
++            if name != "__metadata__":
++                begin, end = info["data_offsets"]
++                locations[name] = (path, 8 + header_length + begin, end - begin)
++    return locations
++
++
++@contextmanager
++def _writable_mmap(path):
++    with open(path, "r+b") as file_handle, mmap.mmap(file_handle.fileno(), 0) as mapped:
++        yield mapped
++
++
++def _apply_delta(local_checkpoint_dir, version_dir):
++    with open(os.path.join(version_dir, "model.safetensors.index.json")) as index_file:
++        metadata = json.load(index_file)["metadata"]
++    applied = _read_applied_version(local_checkpoint_dir)
++    version = int(metadata["version"])
++    if applied == version:
++        return
++    if applied != int(metadata["base_version"]):
++        raise RuntimeError(f"Out-of-order delta: local at {applied}, delta builds on {metadata['base_version']}")
++    if metadata["compression_format"] != "zstd":
++        raise NotImplementedError(f"Unsupported compression {metadata['compression_format']!r}")
++
++    locations = _tensor_locations(local_checkpoint_dir)
++    resources, mapped_files, items, mismatches = ExitStack(), {}, [], []
++    mismatch_lock = threading.Lock()
++    try:
++        for delta_path in sorted(glob.glob(os.path.join(version_dir, "*.safetensors"))):
++            with open(delta_path, "rb") as tensor_file:
++                blob = tensor_file.read()
++            (header_length,) = struct.unpack("<Q", blob[:8])
++            header = json.loads(blob[8 : 8 + header_length])
++            metadata_checksums = header.get("__metadata__", {})
++            data_start = 8 + header_length
++            for name, info in header.items():
++                if name == "__metadata__":
++                    continue
++                if name not in locations:
++                    raise KeyError(f"Delta tensor {name!r} is missing from local checkpoint")
++                path, offset, byte_count = locations[name]
++                if path not in mapped_files:
++                    mapped_files[path] = resources.enter_context(_writable_mmap(path))
++                begin, end = info["data_offsets"]
++                items.append((name, memoryview(blob)[data_start + begin : data_start + end], path, offset, byte_count, metadata_checksums.get(name)))
++
++        def mismatch(name):
++            with mismatch_lock:
++                mismatches.append(name)
++
++        def apply_xor(item):
++            name, compressed, path, offset, byte_count, expected = item
++            region = np.ndarray((byte_count,), dtype=np.uint8, buffer=mapped_files[path], offset=offset)
++            reader, position, hasher = zstandard.ZstdDecompressor().stream_reader(io.BytesIO(bytes(compressed))), 0, _new_hasher(metadata["checksum_format"])
++            while position < byte_count:
++                block = reader.read(min(2 << 20, byte_count - 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 position != byte_count or hasher.hexdigest() != expected:
++                mismatch(name)
++
++        def apply_overwrite(item):
++            name, compressed, path, offset, byte_count, expected = item
++            delta = np.frombuffer(zstandard.ZstdDecompressor().decompress(bytes(compressed)), dtype=np.uint8)
++            count = int.from_bytes(delta[:4].tobytes(), "little")
++            positions_end = 4 + 4 * count
++            positions = np.frombuffer(delta[4:positions_end].tobytes(), dtype="<u4")
++            region = np.ndarray((byte_count,), dtype=np.uint8, buffer=mapped_files[path], offset=offset)
++            region[positions] = delta[positions_end:]
++            hasher = _new_hasher(metadata["checksum_format"])
++            hasher.update(region)
++            if hasher.hexdigest() != expected:
++                mismatch(name)
++
++        if metadata["delta_encoding"] == "xor":
++            apply = apply_xor
++        elif metadata["delta_encoding"] == "overwrite":
++            apply = apply_overwrite
++        else:
++            raise NotImplementedError(f"Unsupported delta encoding {metadata['delta_encoding']!r}")
++        with ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor:
++            list(executor.map(apply, items))
++    finally:
++        resources.close()
++    if mismatches:
++        raise RuntimeError(f"Checksum mismatch for {len(mismatches)} tensors after applying {version_dir}: {sorted(mismatches)[:20]}")
++    _write_applied_version(local_checkpoint_dir, version)

examples/README.md

@@ -10,3 +10,6 @@ These examples provide concrete examples to leverage vime in your own RL workflo
 - **[geo3k_vlm_multi_turn](./geo3k_vlm_multi_turn)**: VLM multi-turn training on Geo3k dataset.
 - **[multi_agent](./multi_agent)**: Example of running multi-agent RL with `vime`.
 - **[train_infer_mismatch_helper](./train_infer_mismatch_helper)**: Algorithmic methods for rollout correction (e.g., TIS, MIS).
+# Weight synchronization
+
+- [delta_weight_sync](./delta_weight_sync): Ascend non-colocated disk delta weight sync.

examples/delta_weight_sync/README.md

@@ -0,0 +1,24 @@
+# Ascend Delta Weight Sync
+
+This example enables non-colocated disk delta weight sync for Ascend rollout engines.
+The trainer writes only changed safetensor bytes to a filesystem shared with rollout
+hosts; each host applies them to a local HF safetensors checkpoint and vLLM reloads it.
+
+```bash
+--update-weight-mode delta \
+--update-weight-transport disk \
+--update-weight-disk-dir /shared/vime-delta \
+--update-weight-local-checkpoint-dir /local-nvme/vime-rollout-checkpoint \
+--update-weight-delta-encoding xor \
+--update-weight-delta-checksum xxh3-128
+```
+
+Only non-colocated rollout is supported. `--update-weight-disk-dir` must be a Linux
+POSIX filesystem visible to training and every rollout host. The local checkpoint
+directory is host-local storage and is seeded from `--hf-checkpoint` on the first sync.
+Reserve enough local capacity for one complete HF safetensors checkpoint plus temporary
+copy space during the initial seed. Both the trainer and every rollout image must include
+`blake3`, `xxhash`, and `zstandard`, and use the accompanying vLLM/vLLM-Ascend patches.
+Use `--custom-update-weight-post-write-path` and
+`--custom-update-weight-pre-read-path` for object-storage mounts that need explicit
+publish/refresh operations.