diff --git a/docker/npu_patch/vllm-ascend.patch b/docker/npu_patch/vllm-ascend.patch index 768852a5..e11d6c15 100644 --- a/docker/npu_patch/vllm-ascend.patch +++ b/docker/npu_patch/vllm-ascend.patch @@ -470,3 +470,30 @@ index 6d99ac76a..b7df4df22 100644 # MRV2's scheduler emits new_block_ids_to_zero whenever this flag is # set, so its worker-side consumer must use the same condition. Keep the +diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py +index b7df4df22..b8b9cfbb1 100644 +--- a/vllm_ascend/worker/worker.py ++++ b/vllm_ascend/worker/worker.py +@@ -1287,3 +1287,22 @@ class NPUWorker(WorkerBase): + def reload_weights(self, *args, **kwargs) -> None: + self.model_runner.reload_weights(*args, **kwargs) + ++ def pull_weights( ++ self, ++ local_checkpoint_dir: str, ++ source_dir: str, ++ target_version: int, ++ pre_read_hook: str | None = None, ++ ) -> dict[str, Any]: ++ from vllm.utils.local_checkpoint import pull_checkpoint ++ ++ pull_checkpoint( ++ local_checkpoint_dir=local_checkpoint_dir, ++ base_dir=self.model_config.model, ++ source_dir=source_dir, ++ target_version=target_version, ++ pre_read_hook=pre_read_hook, ++ ) ++ ++ return {"success": True, "weight_version": str(target_version)} ++ diff --git a/docker/npu_patch/vllm.patch b/docker/npu_patch/vllm.patch index 14294152..f4f3110b 100644 --- a/docker/npu_patch/vllm.patch +++ b/docker/npu_patch/vllm.patch @@ -156,3 +156,482 @@ index 1d30a0eaf6..37239f7f62 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=" int: + """Profiles the peak memory usage of the model to determine how much