From 03a8e1bf984157d137c3bb16d73d7af2a5d9c292 Mon Sep 17 00:00:00 2001 From: wangx700 Date: Fri, 11 Sep 2026 16:44:12 +0800 Subject: [PATCH 1/2] feat(ascend): support full and delta weight sync over disk --- docker/npu_patch/vllm-ascend.patch | 38 ++ docker/npu_patch/vllm.patch | 475 ++++++++++++++++++ requirements.txt | 3 + scripts/README-disk-weight-sync.md | 50 ++ scripts/run-qwen3-4B-delta-disk.sh | 12 + scripts/run-qwen3-4B-disk-common.sh | 120 +++++ scripts/run-qwen3-4B-full-disk.sh | 11 + .../update_weight/test_disk_weight_sync.py | 104 ++++ vime/backends/megatron_utils/actor.py | 8 +- .../update_weight/update_weight_from_disk.py | 176 +++++++ .../update_weight_from_disk_delta.py | 318 ++++++++++++ vime/backends/vllm_utils/vllm_engine.py | 33 +- vime/utils/arguments.py | 85 ++++ vime/utils/disk_delta.py | 86 ++++ 14 files changed, 1516 insertions(+), 3 deletions(-) create mode 100644 scripts/README-disk-weight-sync.md create mode 100644 scripts/run-qwen3-4B-delta-disk.sh create mode 100644 scripts/run-qwen3-4B-disk-common.sh create mode 100644 scripts/run-qwen3-4B-full-disk.sh create mode 100644 tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py create mode 100644 vime/backends/megatron_utils/update_weight/update_weight_from_disk.py create mode 100644 vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py create mode 100644 vime/utils/disk_delta.py diff --git a/docker/npu_patch/vllm-ascend.patch b/docker/npu_patch/vllm-ascend.patch index fa89a9129..44ee63575 100644 --- a/docker/npu_patch/vllm-ascend.patch +++ b/docker/npu_patch/vllm-ascend.patch @@ -51,3 +51,41 @@ index a35d9af8d..66a179cd6 100644 + # 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 +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/docker/npu_patch/vllm.patch b/docker/npu_patch/vllm.patch index 976fea2b9..dac75639f 100644 --- a/docker/npu_patch/vllm.patch +++ b/docker/npu_patch/vllm.patch @@ -40,3 +40,478 @@ index 596cb48..8c30495 100644 inputs_embeds = self.enorm(inputs_embeds) previous_hidden_states = self.hnorm(previous_hidden_states) +diff --git a/tests/model_executor/model_loader/test_local_checkpoint.py b/tests/model_executor/model_loader/test_local_checkpoint.py +new file mode 100644 +index 0000000000..192b7e979a +--- /dev/null ++++ b/tests/model_executor/model_loader/test_local_checkpoint.py +@@ -0,0 +1,144 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++ ++import json ++import zlib ++ ++import numpy as np ++import pytest ++import safetensors.numpy ++import zstandard ++ ++from vllm.utils.local_checkpoint import pull_checkpoint ++ ++ ++def _checksum(data: np.ndarray) -> str: ++ return f"{zlib.adler32(data):08x}" ++ ++ ++def _write_delta( ++ source_dir, ++ version: int, ++ old: np.ndarray, ++ new: np.ndarray, ++ encoding: str, ++) -> None: ++ version_dir = source_dir / f"weight_v{version:06d}" ++ version_dir.mkdir() ++ old_bytes = old.view(np.uint8).reshape(-1) ++ new_bytes = new.view(np.uint8).reshape(-1) ++ if encoding == "xor": ++ payload = old_bytes ^ new_bytes ++ else: ++ positions = np.flatnonzero(old_bytes != new_bytes).astype(" None: ++ version_dir = source_dir / f"weight_v{version:06d}" ++ version_dir.mkdir() ++ safetensors.numpy.save_file({"weight": weight}, version_dir / "model.safetensors") ++ (version_dir / "config.json").write_text("{}") ++ ++ ++@pytest.mark.parametrize("encoding", ["xor", "overwrite"]) ++def test_pull_checkpoint_applies_vime_delta(tmp_path, encoding): ++ base_dir = tmp_path / "base" ++ source_dir = tmp_path / "published" ++ local_dir = tmp_path / "local" ++ base_dir.mkdir() ++ source_dir.mkdir() ++ ++ baseline = np.arange(12, dtype=np.float32).reshape(3, 4) ++ updated = baseline.copy() ++ updated[0, 1] = 100.0 ++ updated[2, 3] = -5.0 ++ safetensors.numpy.save_file({"weight": baseline}, base_dir / "model.safetensors") ++ (base_dir / "config.json").write_text("{}") ++ _write_delta(source_dir, 1, baseline, updated, encoding) ++ ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 0) ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 1) ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 1) ++ ++ actual = safetensors.numpy.load_file(local_dir / "model.safetensors") ++ np.testing.assert_array_equal(actual["weight"], updated) ++ state = json.loads((local_dir / ".weight_sync" / "state.json").read_text()) ++ assert state == {"version": "000001"} ++ ++ ++def test_pull_checkpoint_resets_to_latest_full_version(tmp_path): ++ base_dir = tmp_path / "base" ++ source_dir = tmp_path / "published" ++ local_dir = tmp_path / "local" ++ base_dir.mkdir() ++ source_dir.mkdir() ++ ++ baseline = np.arange(12, dtype=np.float32).reshape(3, 4) ++ first = baseline + 1 ++ reset = baseline + 10 ++ latest = reset.copy() ++ latest[1, 2] = -7.0 ++ safetensors.numpy.save_file({"weight": baseline}, base_dir / "model.safetensors") ++ (base_dir / "config.json").write_text("{}") ++ _write_delta(source_dir, 1, baseline, first, "xor") ++ _write_full(source_dir, 2, reset) ++ _write_delta(source_dir, 3, reset, latest, "xor") ++ ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 1) ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 3) ++ ++ actual = safetensors.numpy.load_file(local_dir / "model.safetensors") ++ np.testing.assert_array_equal(actual["weight"], latest) ++ state = json.loads((local_dir / ".weight_sync" / "state.json").read_text()) ++ assert state == {"version": "000003"} ++ ++ ++def test_pull_checkpoint_does_not_advance_on_checksum_failure(tmp_path): ++ base_dir = tmp_path / "base" ++ source_dir = tmp_path / "published" ++ local_dir = tmp_path / "local" ++ base_dir.mkdir() ++ source_dir.mkdir() ++ ++ baseline = np.arange(12, dtype=np.float32).reshape(3, 4) ++ updated = baseline + 1 ++ safetensors.numpy.save_file({"weight": baseline}, base_dir / "model.safetensors") ++ (base_dir / "config.json").write_text("{}") ++ _write_delta(source_dir, 1, baseline, updated, "xor") ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 0) ++ safetensors.numpy.save_file( ++ {"weight": baseline + 2}, local_dir / "model.safetensors" ++ ) ++ ++ with pytest.raises(RuntimeError, match="Checksum mismatch"): ++ pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 1) ++ ++ state = json.loads((local_dir / ".weight_sync" / "state.json").read_text()) ++ assert state == {"version": "000000"} +diff --git a/vllm/utils/local_checkpoint.py b/vllm/utils/local_checkpoint.py +new file mode 100644 +index 0000000000..42be58d249 +--- /dev/null ++++ b/vllm/utils/local_checkpoint.py +@@ -0,0 +1,319 @@ ++# SPDX-License-Identifier: Apache-2.0 ++# SPDX-FileCopyrightText: Copyright contributors to the vLLM project ++"""Maintain a host-local 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, suppress ++ ++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: str, ++ base_dir: str, ++ source_dir: str, ++ target_version: int, ++ pre_read_hook: str | None = None, ++) -> None: ++ """Bring a host-local checkpoint to a published weight version.""" ++ if target_version > 0 and pre_read_hook: ++ module_path, _, function_name = pre_read_hook.rpartition(".") ++ hook = getattr(importlib.import_module(module_path), function_name) ++ hook(source_dir, target_version) ++ with _pull_lock(local_checkpoint_dir): ++ applied = _read_applied_version(local_checkpoint_dir) ++ 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 _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: ++ 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) -> None: ++ 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, data) -> str: ++ hasher = _new_hasher(algorithm) ++ hasher.update(data) ++ return hasher.hexdigest() ++ ++ ++@contextmanager ++def _pull_lock(local_checkpoint_dir: str): ++ sync_dir = os.path.join(local_checkpoint_dir, SYNC_DIR) ++ os.makedirs(sync_dir, exist_ok=True) ++ with open(os.path.join(sync_dir, "lock"), "w") as lock_file: ++ fcntl.flock(lock_file, fcntl.LOCK_EX) ++ try: ++ yield ++ finally: ++ fcntl.flock(lock_file, fcntl.LOCK_UN) ++ ++ ++def _read_applied_version(local_checkpoint_dir: str) -> int | None: ++ 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: str, version: int) -> None: ++ path = os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json") ++ temporary = f"{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 _drop_page_cache(path: str) -> None: ++ try: ++ file_descriptor = os.open(path, os.O_RDONLY) ++ try: ++ os.posix_fadvise(file_descriptor, 0, 0, os.POSIX_FADV_DONTNEED) ++ finally: ++ os.close(file_descriptor) ++ except OSError: ++ pass ++ ++ ++def _reset_checkpoint(source_dir: str, local_checkpoint_dir: str, version: int) -> None: ++ 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)) ++ _drop_page_cache(entry.path) ++ ++ source_names = {entry.name for entry in source_files} ++ for entry in os.scandir(local_checkpoint_dir): ++ if entry.is_file() and entry.name not in source_names: ++ os.remove(entry.path) ++ ++ for entry in source_files: ++ copied_size = os.path.getsize(os.path.join(local_checkpoint_dir, entry.name)) ++ if copied_size != entry.stat().st_size: ++ raise RuntimeError( ++ f"Size mismatch copying {entry.name}: " ++ f"source {entry.stat().st_size} != local {copied_size}" ++ ) ++ _write_applied_version(local_checkpoint_dir, version) ++ ++ ++def _tensor_locations(checkpoint_dir: str) -> dict[str, tuple[str, int, int]]: ++ locations = {} ++ for path in glob.glob(os.path.join(checkpoint_dir, "*.safetensors")): ++ with open(path, "rb") as tensor_file: ++ (header_length,) = struct.unpack(" None: ++ 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}, " ++ f"delta builds on {metadata['base_version']}" ++ ) ++ if metadata["compression_format"] != "zstd": ++ raise NotImplementedError( ++ f"Compression {metadata['compression_format']!r} is not supported" ++ ) ++ ++ encoding = metadata["delta_encoding"] ++ checksum_algorithm = metadata["checksum_format"] ++ locations = _tensor_locations(local_checkpoint_dir) ++ open_mmaps = {} ++ resources = ExitStack() ++ mismatches = [] ++ mismatch_lock = threading.Lock() ++ delta_blobs = [] ++ items = [] ++ try: ++ for delta_file in sorted(glob.glob(os.path.join(version_dir, "*.safetensors"))): ++ with open(delta_file, "rb") as tensor_file: ++ blob = tensor_file.read() ++ delta_blobs.append(blob) ++ (header_length,) = struct.unpack(" None: ++ with mismatch_lock: ++ mismatches.append(name) ++ ++ def apply_xor(item) -> None: ++ name, compressed, path, offset, byte_count, expected = item ++ region = np.ndarray( ++ (byte_count,), ++ dtype=np.uint8, ++ buffer=open_mmaps[path], ++ offset=offset, ++ ) ++ hasher = _new_hasher(checksum_algorithm) ++ reader = zstandard.ZstdDecompressor().stream_reader( ++ io.BytesIO(bytes(compressed)) ++ ) ++ position = 0 ++ 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: ++ report_mismatch(name) ++ ++ def apply_overwrite(item) -> None: ++ name, compressed, path, offset, byte_count, expected = item ++ delta = np.frombuffer( ++ zstandard.ZstdDecompressor().decompress(bytes(compressed)), ++ dtype=np.uint8, ++ ) ++ region = np.ndarray( ++ (byte_count,), ++ dtype=np.uint8, ++ buffer=open_mmaps[path], ++ offset=offset, ++ ) ++ count = int.from_bytes(delta[:4].tobytes(), "little") ++ positions_end = 4 + 4 * count ++ positions = np.frombuffer(delta[4:positions_end].tobytes(), dtype="=0.1.14 wandb +xxhash +zstandard diff --git a/scripts/README-disk-weight-sync.md b/scripts/README-disk-weight-sync.md new file mode 100644 index 000000000..2bb986998 --- /dev/null +++ b/scripts/README-disk-weight-sync.md @@ -0,0 +1,50 @@ +# Ascend disk weight synchronization + +Run these smoke cases in an Ascend environment built with this branch's +`docker/npu_patch/vllm.patch` and `vllm-ascend.patch`. The existing Dockerfile +and patch series already apply both files. Python dependencies are declared in +`requirements.txt` (`zstandard`, `xxhash`, and `blake3`). + +Both cases use Qwen3-4B, four training NPUs and four rollout NPUs, with three +rollout iterations by default. Set `DATA_ROOT` to a directory containing +`models/Qwen3-4B` and `datasets/dapo-math-17k/dapo-math-17k.jsonl`, as in the +existing Qwen3-4B NPU example. Run one case at a time on allocated devices. + +```bash +ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 DATA_ROOT=/root \ + bash scripts/run-qwen3-4B-full-disk.sh + +ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 DATA_ROOT=/root \ + bash scripts/run-qwen3-4B-delta-disk.sh +``` + +Set `UPDATE_WEIGHT_DISK_DIR` to a dedicated directory shared at the same path +between trainer and rollout hosts. The `/tmp` default is for a single host. +For delta mode, set `UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR` to a dedicated writable +directory on each rollout host. Existing delta-stream files are cleared when +capturing the initial baseline; do not share these directories between jobs. +`NUM_ROLLOUT`, `RAY_GCS_PORT`, `RAY_DASHBOARD_PORT`, and `RAY_TEMP_DIR` can be +overridden. The launchers do not terminate existing processes. They require a +clean, dedicated Ray runtime; stop the runtime you started after the run. + +Full flow: the Megatron actor selects `UpdateWeightFromDisk`; all ranks take +part in HF conversion, rank zero writes safetensors shards and their index, +then rank zero pauses rollout, flushes its cache and calls +`VLLMEngine.update_weights_from_disk`. `/collective_rpc` dispatches +`reload_weights(weights_path=...)` to the NPU workers before generation resumes. + +Delta flow: the first update captures the original HF checkpoint as version +zero. Later updates publish Zstandard-compressed XOR or overwrite deltas and +checksums. `VLLMEngine.pull_weights` dispatches `/collective_rpc` to every NPU +worker. The mainline local-checkpoint helper locks each host's checkpoint, +applies each version once and verifies the resulting bytes; the engine then +reloads that local checkpoint through the same full reload interface. + +Disk synchronization currently requires non-colocated training and rollout. +The existing default collective transport remains unchanged. These launchers +exercise end-to-end training; protocol unit tests additionally check exported +full shards, both delta encodings, unchanged versions and repeated pulls: + +```bash +python3 -m pytest -q tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py +``` diff --git a/scripts/run-qwen3-4B-delta-disk.sh b/scripts/run-qwen3-4B-delta-disk.sh new file mode 100644 index 000000000..a779da85f --- /dev/null +++ b/scripts/run-qwen3-4B-delta-disk.sh @@ -0,0 +1,12 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +export UPDATE_WEIGHT_MODE=delta +# This directory must be visible at the same path to trainer and rollout hosts. +export UPDATE_WEIGHT_DISK_DIR="${UPDATE_WEIGHT_DISK_DIR:-/tmp/vime-delta-weights}" +export RAY_GCS_PORT="${RAY_GCS_PORT:-6399}" +export RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8267}" +export RAY_TEMP_DIR="${RAY_TEMP_DIR:-/tmp/ray-vime-delta-disk}" +export UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR="${UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR:-/tmp/vime-rollout-checkpoint}" +exec bash "${SCRIPT_DIR}/run-qwen3-4B-disk-common.sh" "$@" diff --git a/scripts/run-qwen3-4B-disk-common.sh b/scripts/run-qwen3-4B-disk-common.sh new file mode 100644 index 000000000..331eb78aa --- /dev/null +++ b/scripts/run-qwen3-4B-disk-common.sh @@ -0,0 +1,120 @@ +#!/bin/bash +set -euo pipefail + +export PYTHONUNBUFFERED=1 +: "${ASCEND_RT_VISIBLE_DEVICES:?Set ASCEND_RT_VISIBLE_DEVICES to the allocated NPUs}" +export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1 +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050 +export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050 +export HYDRA_FULL_ERROR=1 +export DISABLE_L2_CACHE=1 +export VLLM_ASCEND_ENABLE_NZ=0 +export VLLM_USE_AOT_COMPILE=0 +export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:${PYTHONPATH:-}" + +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-4B.sh" + +DATA_ROOT="${DATA_ROOT:-/root}" + +CKPT_ARGS=( + --hf-checkpoint ${DATA_ROOT}/models/Qwen3-4B/ + --load ${DATA_ROOT}/models/Qwen3-4B/ + --ref-load ${DATA_ROOT}/models/Qwen3-4B/ + --megatron-to-hf-mode bridge +) + +ROLLOUT_ARGS=( + --prompt-data ${DATA_ROOT}/datasets/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout "${NUM_ROLLOUT:-3}" + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 2048 + --rollout-temperature 1 + --global-batch-size 256 + --balance-data +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 8192 + --megatron-to-hf-mode bridge +) + +GRPO_ARGS=( + --advantage-estimator grpo + --kl-loss-coef 0.0 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.0 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine 4 + --vllm-gpu-memory-utilization 0.6 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --micro-batch-size 1 + --use-flash-attn +) + +SYNC_ARGS=(--update-weight-mode "${UPDATE_WEIGHT_MODE}" --update-weight-transport disk + --update-weight-disk-dir "${UPDATE_WEIGHT_DISK_DIR}") +if [[ "${UPDATE_WEIGHT_MODE}" == delta ]]; then + SYNC_ARGS+=(--update-weight-local-checkpoint-dir "${UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR}" + --update-weight-delta-encoding "${UPDATE_WEIGHT_DELTA_ENCODING:-xor}") +fi +cd "${SCRIPT_DIR}/.." +ray start --head --port="${RAY_GCS_PORT}" --temp-dir="${RAY_TEMP_DIR}" --node-ip-address 127.0.0.1 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port="${RAY_DASHBOARD_PORT}" + +ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT}" \ +-- python3 train.py \ +--actor-num-nodes 1 \ +--actor-num-gpus-per-node 4 \ +--rollout-num-gpus 4 \ +${MODEL_ARGS[@]} \ +${CKPT_ARGS[@]} \ +${ROLLOUT_ARGS[@]} \ +${OPTIMIZER_ARGS[@]} \ +${GRPO_ARGS[@]} \ +${PERF_ARGS[@]} \ +${VLLM_ARGS[@]} \ +"${MISC_ARGS[@]}" \ +"${SYNC_ARGS[@]}" \ +"$@" diff --git a/scripts/run-qwen3-4B-full-disk.sh b/scripts/run-qwen3-4B-full-disk.sh new file mode 100644 index 000000000..e07abbe89 --- /dev/null +++ b/scripts/run-qwen3-4B-full-disk.sh @@ -0,0 +1,11 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +export UPDATE_WEIGHT_MODE=full +# This directory must be visible at the same path to trainer and rollout hosts. +export UPDATE_WEIGHT_DISK_DIR="${UPDATE_WEIGHT_DISK_DIR:-/tmp/vime-full-weights}" +export RAY_GCS_PORT="${RAY_GCS_PORT:-6400}" +export RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8268}" +export RAY_TEMP_DIR="${RAY_TEMP_DIR:-/tmp/ray-vime-full-disk}" +exec bash "${SCRIPT_DIR}/run-qwen3-4B-disk-common.sh" "$@" diff --git a/tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py b/tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py new file mode 100644 index 000000000..eaafd381f --- /dev/null +++ b/tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py @@ -0,0 +1,104 @@ +"""CPU protocol checks; these do not replace distributed Ascend smoke runs.""" + +import importlib +import importlib.util +import json +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import safetensors.torch +import torch + +ROOT = Path(__file__).resolve().parents[5] + + +@pytest.fixture +def modules(monkeypatch, tmp_path): + # Load the updater package without importing the Megatron training runtime. + package = types.ModuleType("_disk_sync_test") + package.__path__ = [str(ROOT / "vime/backends/megatron_utils/update_weight")] + monkeypatch.setitem(sys.modules, package.__name__, package) + loaded = [] + for name in ("update_weight_from_disk", "update_weight_from_disk_delta"): + module = importlib.import_module(f"{package.__name__}.{name}") + monkeypatch.setattr(module, "get_gloo_group", lambda: None) + monkeypatch.setattr(module.dist, "get_rank", lambda: 0) + monkeypatch.setattr(module.dist, "get_world_size", lambda: 1) + monkeypatch.setattr(module.dist, "barrier", lambda **kwargs: None) + monkeypatch.setattr(module.dist, "all_gather_object", lambda out, value, **kwargs: out.__setitem__(0, value)) + monkeypatch.setattr(module.ray, "get", lambda values: values) + loaded.append(module) + monkeypatch.setattr(loaded[1], "NUM_WORKERS", 1) + empty = torch.empty + monkeypatch.setattr(torch, "empty", lambda *args, **kwargs: empty(*args, **dict(kwargs, pin_memory=False))) + iterator = SimpleNamespace(get_hf_weight_chunks=lambda weights, **kwargs: iter([list(weights.items())])) + monkeypatch.setattr(loaded[0].HfWeightIteratorBase, "create", lambda **kwargs: iterator) + + # Exercise exactly the receiver shipped in the NPU patch. + patch = (ROOT / "docker/npu_patch/vllm.patch").read_text() + section = patch.split("+++ b/vllm/utils/local_checkpoint.py\n", 1)[1] + source = "\n".join(line[1:] for line in section.splitlines() if line.startswith("+")) + "\n" + receiver_path = tmp_path / "receiver.py" + receiver_path.write_text(source) + spec = importlib.util.spec_from_file_location("disk_receiver", receiver_path) + receiver = importlib.util.module_from_spec(spec) + spec.loader.exec_module(receiver) + yield *loaded, receiver + for name in list(sys.modules): + if name.startswith("_disk_sync_test."): + monkeypatch.delitem(sys.modules, name) + + +def make_args(tmp_path, encoding="xor"): + base = tmp_path / "base" + base.mkdir() + (base / "config.json").write_text("{}") + return SimpleNamespace( + hf_checkpoint=str(base), + update_weight_disk_dir=str(tmp_path / "published"), + update_weight_local_checkpoint_dir=str(tmp_path / "local"), + update_weight_delta_encoding=encoding, + update_weight_delta_checksum="adler32", + custom_update_weight_post_write_path=None, + ) + + +def test_full_checkpoint_has_loadable_shards_and_index(modules, tmp_path): + full, _, _ = modules + args = make_args(tmp_path) + weights = {"weight": torch.arange(12, dtype=torch.float32).reshape(3, 4)} + safetensors.torch.save_file(weights, Path(args.hf_checkpoint) / "model.safetensors") + updater = full.UpdateWeightFromDisk(args, [], lambda: weights, model_name="qwen3", quantization_config=None) + updater.update_weights() + version = Path(args.update_weight_disk_dir) / "weight_v000001" + index = json.loads((version / "model.safetensors.index.json").read_text()) + shard = safetensors.torch.load_file(version / index["weight_map"]["weight"]) + torch.testing.assert_close(shard["weight"], weights["weight"]) + assert index["metadata"]["total_size"] == 48 + assert (version / "config.json").exists() + + +@pytest.mark.parametrize("encoding", ["xor", "overwrite"]) +def test_delta_roundtrip_versions_and_repeated_pull(modules, tmp_path, encoding): + _, delta, receiver = modules + args = make_args(tmp_path, encoding) + weights = {"weight": torch.arange(12, dtype=torch.float32).reshape(3, 4)} + safetensors.torch.save_file(weights, Path(args.hf_checkpoint) / "model.safetensors") + updater = delta.UpdateWeightFromDiskDelta(args, [], lambda: weights, model_name="qwen3", quantization_config=None) + updater.update_weights() # Capture the same byte-exact HF base used by rollout. + assert updater.weight_version == 0 + for version in (1, 2, 3): + if version != 2: # Include an unchanged-weight version. + weights["weight"][0, 0] += version + updater.weight_version = version + updater._publish() + for _ in range(2): # Repeating XOR application must not corrupt bytes. + receiver.pull_checkpoint( + args.update_weight_local_checkpoint_dir, args.hf_checkpoint, args.update_weight_disk_dir, version + ) + actual = safetensors.torch.load_file(Path(args.update_weight_local_checkpoint_dir) / "model.safetensors") + np.testing.assert_array_equal(actual["weight"].numpy(), weights["weight"].numpy()) diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index 68aac0f71..0e4b22830 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -64,6 +64,8 @@ def _safe_empty_cache(): from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values from .model import forward_only, initialize_model_and_optimizer, save, train from .update_weight.common import named_params_and_buffers +from .update_weight.update_weight_from_disk import UpdateWeightFromDisk +from .update_weight.update_weight_from_disk_delta import UpdateWeightFromDiskDelta from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor @@ -175,7 +177,11 @@ def init( hf_vocab = getattr(self.hf_config, "vocab_size", None) self.args.vocab_size = hf_vocab if hf_vocab is not None else self.tokenizer.vocab_size - if self.args.colocate: + if self.args.update_weight_mode == "delta": + update_weight_cls = UpdateWeightFromDiskDelta + elif self.args.update_weight_transport == "disk": + update_weight_cls = UpdateWeightFromDisk + elif self.args.colocate: update_weight_cls = UpdateWeightFromTensor else: update_weight_cls = UpdateWeightFromDistributed diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py b/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py new file mode 100644 index 000000000..58deb2ee9 --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_disk.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +import json +import logging +import os +import shutil +from argparse import Namespace +from collections.abc import Callable, Mapping, Sequence +from pathlib import Path + +import ray +import torch +import torch.distributed as dist +from ray.actor import ActorHandle +from safetensors.torch import save_file + +from vime.utils.distributed_utils import get_gloo_group + +from .hf_weight_iterator_base import HfWeightIteratorBase + +logger = logging.getLogger(__name__) + +_HF_WEIGHT_FILE_NAMES = { + "model.safetensors.index.json", + "pytorch_model.bin.index.json", + "tf_model.h5", + "flax_model.msgpack", +} +_HF_WEIGHT_FILE_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".ckpt", ".msgpack") + + +class UpdateWeightFromDisk: + """Publish a full HF checkpoint and reload non-colocated rollout engines from disk.""" + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + self.args = args + self.weights_getter = weights_getter + self.weight_version = 0 + self.update_weight_metrics: dict[str, float] = {} + self.rollout_engines: list[ActorHandle] = [] + self._iterator = HfWeightIteratorBase.create( + args=args, + model=model, + model_name=model_name, + quantization_config=quantization_config, + ) + self._post_write_hook: Callable | None = None + if args.custom_update_weight_post_write_path: + from vime.utils.misc import load_function + + self._post_write_hook = load_function(args.custom_update_weight_post_write_path) + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + ) -> None: + del rollout_engine_lock, engine_gpu_counts, engine_gpu_offsets + self.rollout_engines = list(rollout_engines) + + def disconnect_rollout_engines(self) -> None: + return + + def pop_metrics(self) -> dict[str, float]: + metrics, self.update_weight_metrics = self.update_weight_metrics, {} + return metrics + + @torch.no_grad() + def update_weights(self) -> None: + self.weight_version += 1 + version_dir = Path(self.args.update_weight_disk_dir) / f"weight_v{self.weight_version:06d}" + + self._prepare_version_dir(version_dir) + self._write_full_checkpoint(version_dir) + dist.barrier(group=get_gloo_group()) + + if self._post_write_hook is not None: + self._post_write_hook(self.args, str(version_dir), self.rollout_engines) + dist.barrier(group=get_gloo_group()) + + if dist.get_rank() == 0: + ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) + try: + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + ray.get( + [ + engine.update_weights_from_disk.remote( + model_path=str(version_dir), weight_version=str(self.weight_version) + ) + for engine in self.rollout_engines + ] + ) + finally: + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + dist.barrier(group=get_gloo_group()) + + def _prepare_version_dir(self, version_dir: Path) -> None: + if dist.get_rank() == 0: + shutil.rmtree(version_dir, ignore_errors=True) + version_dir.mkdir(parents=True, exist_ok=True) + source_dir = Path(self.args.hf_checkpoint) + for source in source_dir.iterdir(): + if source.is_file() and not _is_hf_weight_file(source): + shutil.copy2(source, version_dir / source.name) + dist.barrier(group=get_gloo_group()) + + def _write_full_checkpoint(self, version_dir: Path) -> None: + is_writer = dist.get_rank() == 0 + weight_map: dict[str, str] = {} + total_size = 0 + shard_files: list[str] = [] + + chunk_iterator = iter( + self._iterator.get_hf_weight_chunks(self.weights_getter(), progress_desc="Save full disk checkpoint") + ) + chunk_index = 0 + while True: + try: + chunk = next(chunk_iterator) + except StopIteration: + break + chunk_index += 1 + if not is_writer: + continue + state_dict: dict[str, torch.Tensor] = {} + for name, tensor in chunk: + if name in weight_map or name in state_dict: + raise ValueError(f"Duplicate HF tensor while saving full disk checkpoint: {name}") + tensor = tensor.detach() + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + if tensor.device.type != "cpu": + tensor = tensor.cpu() + state_dict[name] = tensor + total_size += tensor.numel() * tensor.element_size() + if not state_dict: + continue + filename = f"model-{chunk_index:05d}.safetensors" + save_file(state_dict, version_dir / filename, metadata={"format": "pt"}) + shard_files.append(filename) + weight_map.update({name: filename for name in state_dict}) + + if is_writer: + if not shard_files: + raise ValueError("No HF tensors were produced for full disk checkpoint") + rename_map: dict[str, str] = {} + total_files = len(shard_files) + for index, old_name in enumerate(shard_files, start=1): + new_name = f"model-{index:05d}-of-{total_files:05d}.safetensors" + os.replace(version_dir / old_name, version_dir / new_name) + rename_map[old_name] = new_name + index_data = { + "metadata": {"total_size": total_size}, + "weight_map": {name: rename_map[filename] for name, filename in weight_map.items()}, + } + index_path = version_dir / "model.safetensors.index.json" + temporary = index_path.with_suffix(index_path.suffix + ".tmp") + with temporary.open("w", encoding="utf-8") as output: + json.dump(index_data, output) + output.flush() + os.fsync(output.fileno()) + os.replace(temporary, index_path) + + +def _is_hf_weight_file(path: Path) -> bool: + return path.name in _HF_WEIGHT_FILE_NAMES or path.name.endswith(_HF_WEIGHT_FILE_SUFFIXES) diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py b/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py new file mode 100644 index 000000000..970761685 --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +import json +import logging +import os +import queue +import shutil +from argparse import Namespace +from collections import deque +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import ray +import safetensors.numpy +import torch +import torch.distributed as dist +import zstandard +from ray.actor import ActorHandle + +from vime.utils.disk_delta import NUM_WORKERS, checksum, make_tensor_reader, overwrite_encode +from vime.utils.distributed_utils import get_gloo_group + +from .hf_weight_iterator_base import HfWeightIteratorBase + +logger = logging.getLogger(__name__) + + +class UpdateWeightFromDiskDelta: + """ + Delta weight sync over a shared filesystem. All ranks participate in HF conversion; + rank zero diffs the gathered tensors against a CPU snapshot and publishes a version. + Each engine fans out pull_weights through collective_rpc to materialize the checkpoint + on every rollout host, then reloads it through update_weights_from_disk. + """ + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + self.args = args + self.weights_getter = weights_getter + self.weight_version = 0 + self.update_weight_metrics: dict[str, float] = {} + self.rollout_engines: list[ActorHandle] = [] + self.delta_dir = args.update_weight_disk_dir + self.delta_encoding = args.update_weight_delta_encoding + self.checksum_algorithm = args.update_weight_delta_checksum + self._iterator = HfWeightIteratorBase.create( + args=args, + model=model, + model_name=model_name, + quantization_config=quantization_config, + ) + self._snapshot: dict[str, np.ndarray] = {} + self._baseline_captured = False + self._post_write_hook: Callable | None = None + if args.custom_update_weight_post_write_path: + from vime.utils.misc import load_function + + self._post_write_hook = load_function(args.custom_update_weight_post_write_path) + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + ) -> None: + del rollout_engine_lock, engine_gpu_counts, engine_gpu_offsets + self.rollout_engines = list(rollout_engines) + + def disconnect_rollout_engines(self) -> None: + return + + def pop_metrics(self) -> dict[str, float]: + metrics, self.update_weight_metrics = self.update_weight_metrics, {} + return metrics + + @torch.no_grad() + def update_weights(self) -> None: + # The first call only captures the baseline snapshot the next sync diffs against. + if not self._baseline_captured: + self._capture_baseline() + self._baseline_captured = True + return + + self.weight_version += 1 + self._publish() + self._reload_engines() + self._record_metrics() + + def _capture_baseline(self) -> None: + """Capture the baseline snapshot the first delta diffs against (no publish), and clear any + 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. 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._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 + for name, tensor in self._iter_hf_tensors(): + try: + self._snapshot[name] = read_hf(name) + except KeyError: + 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), + self.args.hf_checkpoint, + ) + + def _publish(self) -> None: + """Encode this version's changed tensors (PP-src ranks), then write it as a canonical HF dir.""" + self._encode_delta() + dist.barrier(group=get_gloo_group()) + self._write_delta_files() + + 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 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() + + # number the files sequentially across only the ranks that have one (no gaps) + counts: list = [None] * world + dist.all_gather_object(counts, int(bool(self._delta)), group=group) + offset, total = sum(counts[:rank]), sum(counts) + + fname = None + self.wire_bytes = 0 + if self._delta: + fname = f"model-{offset:05d}-of-{total:05d}.safetensors" + blob = safetensors.numpy.save(self._delta, metadata=self._checksums) + self.wire_bytes = len(blob) + _atomic_write(os.path.join(self._version_dir, fname), blob) + + maps: list = [None] * world + dist.all_gather_object(maps, {name: fname for name in self._delta}, group=group) + if rank == 0: + index = { + "metadata": { + "version": f"{self.weight_version:06d}", + "base_version": f"{self.weight_version - 1:06d}", + "delta_encoding": self.delta_encoding, + "compression_format": "zstd", + "checksum_format": self.checksum_algorithm, + }, + "weight_map": {name: f for m in maps for name, f in m.items()}, + } + _atomic_write(os.path.join(self._version_dir, "model.safetensors.index.json"), json.dumps(index).encode()) + dist.barrier(group=group) + + def _reload_engines(self) -> None: + """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([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]) + try: + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + ray.get( + [ + engine.update_weights_from_disk.remote( + model_path=self.args.update_weight_local_checkpoint_dir, + weight_version=str(self.weight_version), + ) + for engine in self.rollout_engines + ] + ) + finally: + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + dist.barrier(group=get_gloo_group()) + + def _iter_hf_tensors(self): + """Every rank participates in HF conversion; rank zero publishes.""" + for chunk in self._iterator.get_hf_weight_chunks(self.weights_getter()): + if dist.get_rank() == 0: + yield from chunk + + def _encode_delta(self) -> None: + """Diff each gathered HF tensor against the snapshot, keeping the changed ones (compressed) + in self._delta with their checksums. The GPU->CPU gather is pipelined into a compute pool: + the main loop copies one tensor to a pinned buffer and submits it; pool workers diff and + compress in parallel (each is a few big GIL-releasing numpy/zstd calls).""" + self._version_dir = os.path.join(self.delta_dir, f"weight_v{self.weight_version:06d}") + if dist.get_rank() == 0: + os.makedirs(self._version_dir, exist_ok=True) + snapshot = self._snapshot + self._delta: dict[str, np.ndarray] = {} # changed tensor name -> compressed diff + self._checksums: dict[str, str] = {} # changed tensor name -> new-state checksum + self.changed_bytes = self.total_bytes = 0 + + # Pinned host-buffer pool: a pinned non_blocking GPU->CPU copy is far faster than .cpu(). + max_bytes = max((int(v.nbytes) for v in snapshot.values()), default=0) + free_q: queue.Queue = queue.Queue() + use_pinned = True + try: + for _ in range(max(4, min(2 * NUM_WORKERS, (32 << 30) // max(max_bytes, 1)))): + free_q.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True)) + except RuntimeError as e: # low memlock limit + logger.warning("pinned host buffers unavailable (%s); using pageable .cpu()", e) + use_pinned = False + + def diff_and_compress(name, buf, nbytes, pinned): + if pinned: # copy out and free the pinned buffer before the heavy diff/compress + new = np.empty(nbytes, dtype=np.uint8) + np.copyto(new, buf.numpy()[:nbytes]) + free_q.put(buf) + else: + new = buf + old = snapshot[name] + if self.delta_encoding == "xor": + diff = new ^ old + changed = int(np.count_nonzero(diff)) + elif self.delta_encoding == "overwrite": + mask = new != old + changed = int(np.count_nonzero(mask)) + diff = overwrite_encode(new, mask) + else: + raise ValueError(f"unknown delta encoding {self.delta_encoding!r}") + 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 + + def collect(fut): + name, new, compressed, digest, changed = fut.result() + snapshot[name] = new # becomes the next sync's base + if changed: + self.changed_bytes += changed + self._delta[name] = compressed + self._checksums[name] = digest + + pool = ThreadPoolExecutor(max_workers=NUM_WORKERS) + inflight: deque = deque() + try: + for name, tensor in self._iter_hf_tensors(): + flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1) + nbytes = int(flat.numel()) + if use_pinned and nbytes <= max_bytes: + buf = free_q.get() # blocks when all buffers are in flight -> backpressures the gather + buf[:nbytes].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() + payload, pinned = buf, True + else: + payload, pinned = flat.cpu().numpy().copy(), False + self.total_bytes += nbytes + inflight.append(pool.submit(diff_and_compress, name, payload, nbytes, pinned)) + if len(inflight) >= 2 * NUM_WORKERS: + collect(inflight.popleft()) + while inflight: + collect(inflight.popleft()) + finally: + pool.shutdown() + + def _record_metrics(self) -> None: + """All-reduce the byte counts and record changed-fraction / wire size; the actor drains + update_weight_metrics onto the step log.""" + counts = torch.tensor( + [self.changed_bytes, self.total_bytes, self.wire_bytes], + dtype=torch.int64, + device=_metric_device(), + ) + dist.all_reduce(counts) + changed, total, wire = counts.tolist() + m = self.update_weight_metrics + m["perf/update_weights_density"] = changed / max(total, 1) + m["perf/update_weights_wire_bytes"] = wire + if dist.get_rank() == 0: + logger.info( + "[disk delta v=%s] density=%.2f%% wire=%.2f GB", + self.weight_version, + 100.0 * changed / max(total, 1), + wire / 1e9, + ) + + +def _atomic_write(path: str, data: bytes) -> None: + tmp = path + ".tmp" + with open(tmp, "wb") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + + +def _metric_device() -> torch.device: + if hasattr(torch, "npu") and torch.npu.is_available(): + return torch.device("npu", torch.npu.current_device()) + if torch.cuda.is_available(): + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index 69de3a1b7..78ee6a320 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/vime/backends/vllm_utils/vllm_engine.py @@ -951,7 +951,33 @@ def update_weights_from_distributed( } return self._post_vllm_update_weights_http(update_info) - def update_weights_from_disk(self, model_path: str, load_format: str | None = None): + def pull_weights(self, target_version: int): + """Materialize a published disk version on every host of this engine.""" + if self.node_rank != 0: + return None + response = requests.post( + f"{self._http_base()}/collective_rpc", + json={ + "method": "pull_weights", + "kwargs": { + "local_checkpoint_dir": self.args.update_weight_local_checkpoint_dir, + "source_dir": self.args.update_weight_disk_dir, + "target_version": target_version, + "pre_read_hook": self.args.custom_update_weight_pre_read_path, + }, + }, + timeout=600, + ) + result = _response_json(response) + self._weight_version = str(target_version) + return result + + def update_weights_from_disk( + self, + model_path: str, + load_format: str | None = None, + weight_version: str | None = None, + ): """``POST /collective_rpc`` with ``reload_weights`` and ``weights_path``.""" if self.node_rank != 0: return @@ -964,7 +990,10 @@ def update_weights_from_disk(self, model_path: str, load_format: str | None = No }, timeout=600, ) - return _response_json(response) + result = _response_json(response) + if weight_version is not None: + self._weight_version = str(weight_version) + return result def pause_generation(self): """``POST /pause`` with mode="keep"; returns the ``requests.Response``.""" diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 0dc1145ff..a1b718e7a 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -132,6 +132,71 @@ def add_train_arguments(parser): default=1024**3, help="Add margin for train memory allocation. By default we will reserve 1GB as margin.", ) + parser.add_argument( + "--update-weight-mode", + choices=["full", "delta"], + default="full", + help=( + "Weight sync strategy. 'full' sends all weights over the selected " + "transport. 'delta' publishes only changed weight bytes " + "through a shared filesystem." + ), + ) + parser.add_argument( + "--update-weight-transport", + choices=["nccl", "disk"], + default="nccl", + help=( + "Weight sync transport. 'nccl' is retained for CLI compatibility and " + "selects the existing accelerator-native HCCL/NPU-IPC path on Ascend; " + "'disk' publishes either full checkpoints or byte deltas through a shared filesystem." + ), + ) + parser.add_argument( + "--update-weight-disk-dir", + type=str, + default=None, + help="Shared filesystem directory where full or delta weight versions are published.", + ) + parser.add_argument( + "--update-weight-local-checkpoint-dir", + type=str, + default=None, + help=( + "Host-local HF checkpoint directory patched in place by rollout workers. " + "Required for delta disk sync and unused by full disk sync." + ), + ) + parser.add_argument( + "--update-weight-delta-encoding", + choices=["xor", "overwrite"], + default="xor", + help="Delta encoding: xor (compact) or overwrite (idempotent).", + ) + parser.add_argument( + "--update-weight-delta-checksum", + choices=["xxh3-128", "blake3", "adler32"], + default="xxh3-128", + help="Per-tensor checksum algorithm for delta application.", + ) + parser.add_argument( + "--custom-update-weight-post-write-path", + type=str, + default=None, + help=( + "Optional trainer-side hook after a disk version is written. Signature: " + "hook(args, version_dir: str, rollout_engines) -> None." + ), + ) + parser.add_argument( + "--custom-update-weight-pre-read-path", + type=str, + default=None, + help=( + "Optional rollout-host hook before a published version is read. Signature: " + "hook(source_dir: str, target_version: int) -> None." + ), + ) parser.add_argument( "--megatron-to-hf-mode", choices=["raw", "bridge"], @@ -1684,6 +1749,26 @@ def vime_validate_args(args): if args.save_interval is not None: assert args.save is not None, "'--save' is required when save_interval is set." + 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.") + elif args.update_weight_transport == "disk": + if args.colocate: + raise ValueError( + "--update-weight-mode=full --update-weight-transport=disk is not supported with --colocate." + ) + if not args.update_weight_disk_dir: + raise ValueError("full disk weight sync requires --update-weight-disk-dir.") + assert not (args.kl_coef != 0 and args.kl_loss_coef != 0), "Only one of kl_coef and kl_loss_coef can be set" if args.advantage_estimator in ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]: diff --git a/vime/utils/disk_delta.py b/vime/utils/disk_delta.py new file mode 100644 index 000000000..c3abc65b1 --- /dev/null +++ b/vime/utils/disk_delta.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import glob +import json +import os +import struct +import zlib + +import numpy as np + +# 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)) + +# Trainer-side helpers for disk-level delta weight sync. The receive side — materializing the +# host-local checkpoint and applying published deltas in place — lives in vLLM behind its +# /pull_weights endpoint, so it runs on every host while Vime only talks to one endpoint per engine. + + +def overwrite_encode(new: np.ndarray, changed_mask: np.ndarray) -> np.ndarray: + """The 'overwrite' delta: changed-position count (u4), positions (u4 each), then new values. + Idempotent to apply, unlike xor (an involution); the trainer picks the encoding per the docs.""" + pos = np.flatnonzero(changed_mask).astype(" 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() + + +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]] = {} + for path in glob.glob(os.path.join(ckpt_dir, "*.safetensors")): + with open(path, "rb") as f: + (header_len,) = struct.unpack(" uint8 bytes`` that seeks straight to the + tensor — for reading many tensors without rescanning every header. KeyError if absent.""" + locations = _tensor_locations(ckpt_dir) + + def read(name: str) -> np.ndarray: + path, offset, nbytes = locations[name] + with open(path, "rb") as f: + f.seek(offset) + return np.frombuffer(f.read(nbytes), dtype=np.uint8) + + return read From febdeb2eb2d3fa338d33eef420cbad6b818db95d Mon Sep 17 00:00:00 2001 From: wangx700 Date: Fri, 11 Sep 2026 17:12:51 +0800 Subject: [PATCH 2/2] fix(ascend): bound disk delta buffers and validate checkpoints --- .../update_weight/test_disk_weight_sync.py | 89 ++++++++++++++++++- tests/unit/utils/test_disk_delta.py | 43 +++++++++ .../update_weight_from_disk_delta.py | 65 +++++++++----- vime/utils/disk_delta.py | 41 +++++++-- 4 files changed, 207 insertions(+), 31 deletions(-) create mode 100644 tests/unit/utils/test_disk_delta.py diff --git a/tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py b/tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py index eaafd381f..5f1078cb7 100644 --- a/tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py +++ b/tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py @@ -3,8 +3,10 @@ import importlib import importlib.util import json +import queue import sys import types +import weakref from pathlib import Path from types import SimpleNamespace @@ -16,6 +18,19 @@ ROOT = Path(__file__).resolve().parents[5] +def _receiver_source(patch): + section = patch.split("+++ b/vllm/utils/local_checkpoint.py\n", 1)[1] + section = section.split("\ndiff --git", 1)[0] + return "\n".join(line[1:] for line in section.splitlines() if line.startswith("+")) + "\n" + + +def test_receiver_source_ignores_following_file(): + patch = (ROOT / "docker/npu_patch/vllm.patch").read_text() + appended = patch + "\ndiff --git a/other.py b/other.py\n+++ b/other.py\n@@ -0,0 +1 @@\n+invalid python!\n" + assert _receiver_source(appended) == _receiver_source(patch) + compile(_receiver_source(appended), "receiver.py", "exec") + + @pytest.fixture def modules(monkeypatch, tmp_path): # Load the updater package without importing the Megatron training runtime. @@ -40,8 +55,7 @@ def modules(monkeypatch, tmp_path): # Exercise exactly the receiver shipped in the NPU patch. patch = (ROOT / "docker/npu_patch/vllm.patch").read_text() - section = patch.split("+++ b/vllm/utils/local_checkpoint.py\n", 1)[1] - source = "\n".join(line[1:] for line in section.splitlines() if line.startswith("+")) + "\n" + source = _receiver_source(patch) receiver_path = tmp_path / "receiver.py" receiver_path.write_text(source) spec = importlib.util.spec_from_file_location("disk_receiver", receiver_path) @@ -102,3 +116,74 @@ def test_delta_roundtrip_versions_and_repeated_pull(modules, tmp_path, encoding) ) actual = safetensors.torch.load_file(Path(args.update_weight_local_checkpoint_dir) / "model.safetensors") np.testing.assert_array_equal(actual["weight"].numpy(), weights["weight"].numpy()) + + +@pytest.mark.parametrize("failure", ["empty", "copyto", "submit", "device_copy"]) +def test_pinned_buffer_returned_on_failure(modules, tmp_path, monkeypatch, failure): + _, delta, _ = modules + args = make_args(tmp_path) + weights = {f"weight{i}": torch.arange(12, dtype=torch.float32) for i in range(3)} + safetensors.torch.save_file(weights, Path(args.hf_checkpoint) / "model.safetensors") + updater = delta.UpdateWeightFromDiskDelta(args, [], lambda: weights, model_name="qwen3", quantization_config=None) + updater.update_weights() + + class BoundedWaitQueue(queue.Queue): + def get(self, block=True, timeout=None): + # Turn a leaked-buffer deadlock into a bounded test failure. + return super().get(block=block, timeout=2 if block and timeout is None else timeout) + + buffers = BoundedWaitQueue() + buffer = torch.empty(48, dtype=torch.uint8) + buffers.put(buffer) + monkeypatch.setattr(delta, "_make_pinned_pool", lambda size: buffers) + + def fail(*args, **kwargs): + raise MemoryError("injected buffer failure") + + if failure in ("empty", "copyto"): + monkeypatch.setattr(delta.np, failure, fail) + elif failure == "submit": + monkeypatch.setattr(delta.ThreadPoolExecutor, "submit", fail) + else: + monkeypatch.setattr(torch.Tensor, "copy_", fail) + with pytest.raises(MemoryError, match="injected buffer failure"): + updater._encode_delta() + assert buffers.qsize() == 1 + assert buffers.get_nowait() is buffer + + +@pytest.mark.parametrize("size,count", [(0, 0), (1 << 30, 8), (3 << 30, 2), (8 << 30, 1), ((8 << 30) + 1, 0)]) +def test_pinned_pool_respects_budget(modules, monkeypatch, size, count): + _, delta, _ = modules + monkeypatch.setattr(delta, "NUM_WORKERS", 16) + allocations = [] + + def allocate(nbytes, **kwargs): + allocations.append(nbytes) + return object() + + monkeypatch.setattr(delta.torch, "empty", allocate) + buffers = delta._make_pinned_pool(size) + assert buffers.qsize() == count + assert allocations == [size] * count + assert sum(allocations) <= 8 << 30 + + +@pytest.mark.parametrize("error", [RuntimeError, MemoryError]) +def test_pinned_pool_releases_partial_allocations(modules, monkeypatch, error): + _, delta, _ = modules + references = [] + + class Buffer: + pass + + def allocate(*args, **kwargs): + if references: + raise error("allocation failed") + buffer = Buffer() + references.append(weakref.ref(buffer)) + return buffer + + monkeypatch.setattr(delta.torch, "empty", allocate) + assert delta._make_pinned_pool(1024).empty() + assert references[0]() is None diff --git a/tests/unit/utils/test_disk_delta.py b/tests/unit/utils/test_disk_delta.py new file mode 100644 index 000000000..1167d2812 --- /dev/null +++ b/tests/unit/utils/test_disk_delta.py @@ -0,0 +1,43 @@ +import json +import struct + +import pytest + +from vime.utils.disk_delta import make_tensor_reader + + +def test_missing_safetensors(tmp_path): + with pytest.raises(FileNotFoundError, match=str(tmp_path)): + make_tensor_reader(str(tmp_path)) + + +@pytest.mark.parametrize( + "content", + [b"short", struct.pack(" None: # Pinned host-buffer pool: a pinned non_blocking GPU->CPU copy is far faster than .cpu(). max_bytes = max((int(v.nbytes) for v in snapshot.values()), default=0) - free_q: queue.Queue = queue.Queue() - use_pinned = True - try: - for _ in range(max(4, min(2 * NUM_WORKERS, (32 << 30) // max(max_bytes, 1)))): - free_q.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True)) - except RuntimeError as e: # low memlock limit - logger.warning("pinned host buffers unavailable (%s); using pageable .cpu()", e) - use_pinned = False + free_q = _make_pinned_pool(max_bytes) + use_pinned = not free_q.empty() def diff_and_compress(name, buf, nbytes, pinned): if pinned: # copy out and free the pinned buffer before the heavy diff/compress - new = np.empty(nbytes, dtype=np.uint8) - np.copyto(new, buf.numpy()[:nbytes]) - free_q.put(buf) + try: + new = np.empty(nbytes, dtype=np.uint8) + np.copyto(new, buf.numpy()[:nbytes]) + finally: + free_q.put(buf) else: new = buf old = snapshot[name] @@ -260,18 +256,26 @@ def collect(fut): for name, tensor in self._iter_hf_tensors(): flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1) nbytes = int(flat.numel()) - if use_pinned and nbytes <= max_bytes: - buf = free_q.get() # blocks when all buffers are in flight -> backpressures the gather - buf[:nbytes].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() - payload, pinned = buf, True - else: - payload, pinned = flat.cpu().numpy().copy(), False - self.total_bytes += nbytes - inflight.append(pool.submit(diff_and_compress, name, payload, nbytes, pinned)) + buf = None + submitted = False + try: + if use_pinned and nbytes <= max_bytes: + buf = free_q.get() # backpressure until a worker returns a buffer + buf[:nbytes].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() + payload, pinned = buf, True + else: + payload, pinned = flat.cpu().numpy().copy(), False + self.total_bytes += nbytes + future = pool.submit(diff_and_compress, name, payload, nbytes, pinned) + submitted = True # the worker now owns returning the buffer + inflight.append(future) + finally: + if buf is not None and not submitted: + free_q.put(buf) if len(inflight) >= 2 * NUM_WORKERS: collect(inflight.popleft()) while inflight: @@ -301,6 +305,21 @@ def _record_metrics(self) -> None: ) +def _make_pinned_pool(max_bytes: int) -> queue.Queue: + """Limit this pool's requested pinned storage to 8 GiB, excluding other CPU state.""" + free_q: queue.Queue = queue.Queue() + num_buffers = min(2 * NUM_WORKERS, (8 << 30) // max_bytes) if max_bytes > 0 else 0 + try: + for _ in range(num_buffers): + free_q.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True)) + except (RuntimeError, MemoryError) as e: + # Release partial allocations before falling back to pageable copies. + while not free_q.empty(): + free_q.get_nowait() + logger.warning("pinned host buffers unavailable (%s); using pageable .cpu()", e) + return free_q + + def _atomic_write(path: str, data: bytes) -> None: tmp = path + ".tmp" with open(tmp, "wb") as f: diff --git a/vime/utils/disk_delta.py b/vime/utils/disk_delta.py index c3abc65b1..e6631b794 100644 --- a/vime/utils/disk_delta.py +++ b/vime/utils/disk_delta.py @@ -59,15 +59,41 @@ def checksum(algorithm: str, buf) -> str: 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.""" + paths = sorted(glob.glob(os.path.join(ckpt_dir, "*.safetensors"))) + if not paths: + raise FileNotFoundError(f"No .safetensors files found in checkpoint directory: {ckpt_dir}") locations: dict[str, tuple[str, int, int]] = {} - for path in glob.glob(os.path.join(ckpt_dir, "*.safetensors")): - with open(path, "rb") as f: - (header_len,) = struct.unpack(" file_size - 8: + raise ValueError("declared header length exceeds file size") + header_bytes = f.read(header_len) + if len(header_bytes) != header_len: + raise ValueError("truncated header") + header = json.loads(header_bytes) + if not isinstance(header, dict): + raise ValueError("header must be a JSON object") + except (ValueError, UnicodeError, struct.error) as e: + raise RuntimeError(f"Failed to parse safetensors header from {path}: {e}") from e + data_size = file_size - 8 - header_len for name, info in header.items(): if name == "__metadata__": continue - begin, end = info["data_offsets"] + offsets = info.get("data_offsets") if isinstance(info, dict) else None + if ( + not isinstance(offsets, list) + or len(offsets) != 2 + or any(type(value) is not int for value in offsets) + or not 0 <= offsets[0] <= offsets[1] <= data_size + ): + raise RuntimeError(f"Invalid data_offsets for tensor {name!r} in {path}: {offsets!r}") + begin, end = offsets locations[name] = (path, 8 + header_len + begin, end - begin) return locations @@ -81,6 +107,9 @@ def read(name: str) -> np.ndarray: path, offset, nbytes = locations[name] with open(path, "rb") as f: f.seek(offset) - return np.frombuffer(f.read(nbytes), dtype=np.uint8) + data = f.read(nbytes) + if len(data) != nbytes: + raise RuntimeError(f"Truncated tensor {name!r} in {path}: expected {nbytes} bytes, read {len(data)}") + return np.frombuffer(data, dtype=np.uint8) return read