feat: support full weight sync over disk by wangx700 · Pull Request #407 · vllm-project/vime

wangx700 · 2026-09-01T02:30:08Z

Summary

add full + disk weight synchronization for non-colocated rollout engines
publish versioned HF Safetensors checkpoints to shared storage
reload vLLM engines through the existing collective reload_weights RPC
add 4B and 30B reproducible disk full/delta benchmark coverage

Validation

argument validation covers --update-weight-mode full --update-weight-transport disk
Qwen3-4B benchmark: full disk steady mean 5.4s
Qwen3-30B-A3B benchmark: full disk steady mean 56.6s

Notes

This PR is based on validated-delta, so it currently includes those prerequisite commits in the diff.

read-the-docs-community · 2026-09-01T02:31:37Z

Documentation build overview

📚 vime | 🛠️ Build #34323704 | 📁 Comparing bfece64 against latest (85ac426)

🔍 Preview build

48 files changed · + 3 added · ± 28 modified · - 17 deleted

+ Added

± Modified

- Deleted

gemini-code-assist

Code Review

This pull request introduces a disk-based delta weight synchronization mechanism for Ascend rollout engines, allowing the trainer to publish byte-level weight deltas or full checkpoints to a shared filesystem and reload them on rollout hosts. Key additions include new weight update classes, utility functions for delta encoding and checksumming, integration with vLLM engines, and associated tests and benchmarks. The review feedback suggests several optimizations: reusing a single pinned host buffer instead of a pool to save memory during delta encoding, avoiding redundant .tobytes() copies in the vLLM patch's overwrite application, and adding early validation to ensure the base checkpoint is in safetensors format when delta mode is enabled.

gemini-code-assist · 2026-09-01T02:32:25Z

+        max_bytes = max((value.nbytes for value in self._snapshot.values()), default=0)
+        free_buffers: queue.Queue[torch.Tensor] = queue.Queue()
+        use_pinned = max_bytes > 0
+        if use_pinned:
+            try:
+                pool_size = max(2, min(2 * NUM_WORKERS, (8 << 30) // max(max_bytes, 1)))
+                for _ in range(pool_size):
+                    free_buffers.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True))
+            except RuntimeError as exc:
+                logger.warning("Pinned host buffers unavailable (%s); using pageable copies", exc)
+                use_pinned = False
+
+        def diff_and_compress(name: str, new: np.ndarray) -> tuple[str, np.ndarray, np.ndarray | None, str | None, int]:
+            old = self._snapshot[name]
+            if new.nbytes != old.nbytes:
+                raise ValueError(f"Delta tensor size changed for {name}: {old.nbytes} != {new.nbytes}")
+            if self.delta_encoding == "xor":
+                diff = new ^ old
+                changed = int(np.count_nonzero(diff))
+            else:
+                mask = new != old
+                changed = int(np.count_nonzero(mask))
+                diff = overwrite_encode(new, mask)
+            if not changed:
+                return name, new, None, None, 0
+            compressed = np.frombuffer(zstandard.ZstdCompressor(level=1).compress(diff), dtype=np.uint8)
+            return name, new, compressed, checksum(self.checksum_algorithm, new), changed
+
+        pool = ThreadPoolExecutor(max_workers=NUM_WORKERS)
+        inflight: deque = deque()
+        try:
+            for name, tensor in self._iter_hf_tensors(progress_desc="Encode disk delta"):
+                new = _tensor_bytes(tensor, free_buffers=free_buffers if use_pinned else None)


Since _tensor_bytes is called sequentially in the main thread before submitting tasks to the ThreadPoolExecutor, there is no concurrent access to the pinned host buffers. Allocating a pool of up to 2 * NUM_WORKERS (up to 64) pinned buffers is redundant and can waste gigabytes of pinned host memory (e.g., up to 8GB), potentially leading to host OOM or swap issues.

We can optimize this by allocating and reusing a single pinned host buffer instead of a queue of buffers.

max_bytes = max((value.nbytes for value in self._snapshot.values()), default=0) pinned_buffer = None if max_bytes > 0: try: pinned_buffer = torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True) except RuntimeError as exc: logger.warning("Pinned host buffer unavailable (%s); using pageable copies", exc) def diff_and_compress(name: str, new: np.ndarray) -> tuple[str, np.ndarray, np.ndarray | None, str | None, int]: old = self._snapshot[name] if new.nbytes != old.nbytes: raise ValueError(f"Delta tensor size changed for {name}: {old.nbytes} != {new.nbytes}") if self.delta_encoding == "xor": diff = new ^ old changed = int(np.count_nonzero(diff)) else: mask = new != old changed = int(np.count_nonzero(mask)) diff = overwrite_encode(new, mask) if not changed: return name, new, None, None, 0 compressed = np.frombuffer(zstandard.ZstdCompressor(level=1).compress(diff), dtype=np.uint8) return name, new, compressed, checksum(self.checksum_algorithm, new), changed pool = ThreadPoolExecutor(max_workers=NUM_WORKERS) inflight: deque = deque() try: for name, tensor in self._iter_hf_tensors(progress_desc="Encode disk delta"): new = _tensor_bytes(tensor, pinned_buffer=pinned_buffer)

gemini-code-assist · 2026-09-01T02:32:25Z

+def _tensor_bytes(tensor: torch.Tensor, *, free_buffers: queue.Queue[torch.Tensor] | None = None) -> np.ndarray:
+    flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1)
+    if flat.device.type == "cpu":
+        return flat.numpy().copy()
+    if free_buffers is None:
+        return flat.cpu().numpy().copy()
+
+    buffer = free_buffers.get()
+    try:
+        buffer[: flat.numel()].copy_(flat, non_blocking=True)
+        if flat.device.type == "npu":
+            torch.npu.current_stream().synchronize()
+        elif flat.device.type == "cuda":
+            torch.cuda.current_stream().synchronize()
+        return buffer[: flat.numel()].numpy().copy()
+    finally:
+        free_buffers.put(buffer)


Update _tensor_bytes to accept and use the single pinned_buffer instead of the free_buffers queue.

def _tensor_bytes(tensor: torch.Tensor, *, pinned_buffer: torch.Tensor | None = None) -> np.ndarray: flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1) if flat.device.type == "cpu": return flat.numpy().copy() if pinned_buffer is None: return flat.cpu().numpy().copy() pinned_buffer[: flat.numel()].copy_(flat, non_blocking=True) if flat.device.type == "npu": torch.npu.current_stream().synchronize() elif flat.device.type == "cuda": torch.cuda.current_stream().synchronize() return pinned_buffer[: flat.numel()].numpy().copy()

gemini-code-assist · 2026-09-01T02:32:25Z

+            count = int.from_bytes(delta[:4].tobytes(), "little")
+            positions_end = 4 + 4 * count
+            positions = np.frombuffer(delta[4:positions_end].tobytes(), dtype="<u4")


We can optimize apply_overwrite by avoiding the .tobytes() copies. Since delta is a 1D uint8 array, and the offsets are 4-byte aligned, we can use np.frombuffer directly on the slices of delta without copying the underlying memory.

count = int(np.frombuffer(delta[:4], dtype="<u4")[0]) positions_end = 4 + 4 * count positions = np.frombuffer(delta[4:positions_end], dtype="<u4")

gemini-code-assist · 2026-09-01T02:32:26Z

+    if args.update_weight_mode == "delta":
+        if args.update_weight_transport != "disk":
+            raise ValueError("--update-weight-mode=delta requires --update-weight-transport=disk.")
+        if args.colocate:
+            raise ValueError(
+                "--update-weight-mode=delta is not supported with --colocate; "
+                "colocated NPU IPC already avoids full-model copies."
+            )
+        if not args.update_weight_disk_dir:
+            raise ValueError("--update-weight-mode=delta requires --update-weight-disk-dir.")
+        if not args.update_weight_local_checkpoint_dir:
+            raise ValueError("--update-weight-mode=delta requires --update-weight-local-checkpoint-dir.")


Since delta weight synchronization relies on patching a local safetensors checkpoint, the base checkpoint specified by --hf-checkpoint must be in safetensors format (containing .safetensors files). If it is in another format (like .bin), the delta sync will fail with a cryptic KeyError during rollout.

Adding an early validation check here will provide a clear, user-friendly error message at startup.

Suggested change

if args.update_weight_mode == "delta":

if args.update_weight_transport != "disk":

raise ValueError("--update-weight-mode=delta requires --update-weight-transport=disk.")

if args.colocate:

raise ValueError(

"--update-weight-mode=delta is not supported with --colocate; "

"colocated NPU IPC already avoids full-model copies."

)

if not args.update_weight_disk_dir:

raise ValueError("--update-weight-mode=delta requires --update-weight-disk-dir.")

if not args.update_weight_local_checkpoint_dir:

raise ValueError("--update-weight-mode=delta requires --update-weight-local-checkpoint-dir.")

if args.update_weight_mode == "delta":

if args.update_weight_transport != "disk":

raise ValueError("--update-weight-mode=delta requires --update-weight-transport=disk.")

if args.colocate:

raise ValueError(

"--update-weight-mode=delta is not supported with --colocate; "

"colocated NPU IPC already avoids full-model copies."

)

if not args.update_weight_disk_dir:

raise ValueError("--update-weight-mode=delta requires --update-weight-disk-dir.")

if not args.update_weight_local_checkpoint_dir:

raise ValueError("--update-weight-mode=delta requires --update-weight-local-checkpoint-dir.")

if args.hf_checkpoint and os.path.exists(args.hf_checkpoint):

import glob

if not glob.glob(os.path.join(args.hf_checkpoint, "*.safetensors")):

raise ValueError(

f"--update-weight-mode=delta requires the HF checkpoint at {args.hf_checkpoint} "

"to be in safetensors format (containing .safetensors files)."

)

wangx700 and others added 8 commits August 26, 2026 00:38

feat: add Ascend disk delta weight sync

e8d18a0

fix: complete Ascend disk delta pull flow

ae68ae4

fix: isolate Ray ports for NPU delta run

3acfbf3

fix: use NPU workspace Python paths

1fe9a20

fix-ascend-full-weight-update-signatures

1196563

add-disk-delta-stage-timing-metrics

e1dca3e

feat: support full weight sync over disk

d7ca57a

perf: add disk weight sync benchmark

bfece64

gemini-code-assist Bot reviewed Sep 1, 2026

View reviewed changes

wangx700 marked this pull request as draft September 1, 2026 06:38

Navigation Menu

Uh oh!

feat: support full weight sync over disk - #407

feat: support full weight sync over disk#407
wangx700 wants to merge 8 commits into
vllm-project:ascendfrom
wangx700:feature/disk-full-weight-sync

wangx700 commented Sep 1, 2026

Uh oh!

read-the-docs-community Bot commented Sep 1, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Sep 1, 2026

Uh oh!

gemini-code-assist Bot Sep 1, 2026

Uh oh!

gemini-code-assist Bot Sep 1, 2026

Uh oh!

gemini-code-assist Bot Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

wangx700 commented Sep 1, 2026

Summary

Validation

Notes

Uh oh!

read-the-docs-community Bot commented Sep 1, 2026

Documentation build overview

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 1, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant