diff --git a/docs/features/preload.md b/docs/features/preload.md index 5787fdb51c1c..f8a51bb08fbe 100644 --- a/docs/features/preload.md +++ b/docs/features/preload.md @@ -128,8 +128,10 @@ state rides the data plane. meta device (no storage) and requests the tensors from the daemon on its GPU. 3. Before serving anything, the engine and daemon compare a fingerprint of the - cached weights: checkpoint content (hashed from safetensors metadata, so - identical weights in different directories still match), model + cached weights: checkpoint content (hashed from safetensors metadata plus + the first 1 KiB of each tensor, so identical weights in different + directories still match and a change in any sampled region also changes + the key), model architecture, TP/DP size/rank, dtype, quantization method and config, text RoPE parameters, model revision, KV cache dtype and quantization skip layers, and vLLM version. On any mismatch the engine falls back to disk loading diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index e9ff7dc17c4e..6b6a2070f956 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -775,6 +775,81 @@ def test_ipc_rope_matching( server_task.result(timeout=5) +def test_hash_checkpoint_distinguishes_weight_values(tmp_path): + # Same layout, different values: a base checkpoint and its fine-tune must + # not share a key, or an ipc_cache engine silently maps the daemon's stale + # weights (#59647). + import torch + from safetensors.torch import save_file + + from vllm.model_executor.model_loader.weight_cache.protocol import ( + hash_checkpoint, + ) + + for name, value in (("base", 0.0), ("finetuned", 1.0)): + (tmp_path / name).mkdir() + save_file( + {"w": torch.full((4, 4), value)}, + str(tmp_path / name / "model.safetensors"), + ) + assert hash_checkpoint(str(tmp_path / "base")) != hash_checkpoint( + str(tmp_path / "finetuned") + ) + + +def test_hash_checkpoint_stable_across_directory_copies(tmp_path): + # The #54921 property: byte-identical checkpoints under different paths + # share a key, so a daemon preloaded from one directory can serve an + # engine started on the copy. + import torch + from safetensors.torch import save_file + + from vllm.model_executor.model_loader.weight_cache.protocol import ( + hash_checkpoint, + ) + + (tmp_path / "a").mkdir() + save_file( + {"w": torch.arange(4096, dtype=torch.float32).reshape(64, 64)}, + str(tmp_path / "a" / "model.safetensors"), + ) + shutil.copytree(tmp_path / "a", tmp_path / "b") + assert hash_checkpoint(str(tmp_path / "a")) == hash_checkpoint(str(tmp_path / "b")) + + +def test_hash_checkpoint_catches_single_tensor_change_among_many(tmp_path): + # Every tensor's head contributes to the key, so changing one tensor out + # of 64 flips it. The sampled reads run through a thread pool, and 64 + # tensors exceed the worker count, so this also pins down that the digest + # does not depend on read scheduling. + import torch + from safetensors.torch import save_file + + from vllm.model_executor.model_loader.weight_cache.protocol import ( + hash_checkpoint, + ) + + tensors = {f"w{i:03d}": torch.full((8, 8), float(i)) for i in range(64)} + (tmp_path / "base").mkdir() + save_file(tensors, str(tmp_path / "base" / "model.safetensors")) + changed = dict(tensors) + changed["w037"] = torch.full((8, 8), -1.0) + (tmp_path / "changed").mkdir() + save_file(changed, str(tmp_path / "changed" / "model.safetensors")) + assert hash_checkpoint(str(tmp_path / "base")) != hash_checkpoint( + str(tmp_path / "changed") + ) + + +def test_hash_checkpoint_without_safetensors_returns_none(tmp_path): + from vllm.model_executor.model_loader.weight_cache.protocol import ( + hash_checkpoint, + ) + + assert hash_checkpoint(str(tmp_path)) is None + assert hash_checkpoint(str(tmp_path / "not-a-dir")) is None + + def test_ipc_loader_copy_mode_reports_no_external_weight_memory(): """Copy mode clones the weights into the engine, so nothing is external (vllm_config is never touched).""" diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 334c13e7b0d6..5161686f5fff 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -20,6 +20,7 @@ import struct import tempfile import time +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, fields from typing import Any, NamedTuple @@ -280,13 +281,67 @@ def _safetensors_header(path: str) -> bytes: return size_bytes + f.read(header_len) -def hash_checkpoint(model: str) -> str | None: - """Fingerprint checkpoint content from local safetensors metadata. +_TENSOR_HEAD_BYTES = 1024 +# Bounded concurrency for the sampled reads. On NFS each pread costs one RPC +# RTT, so sequential sampling of a K3-scale checkpoint (~70k tensors) would +# take over a minute at 1 ms RTT; 32 in-flight reads keep that at a few +# seconds without hammering the server. +_FINGERPRINT_READ_WORKERS = 32 + + +def _mix_safetensors_fingerprint( + path: str, hasher: Any, pool: ThreadPoolExecutor +) -> None: + """Mix one shard's layout and sampled weight bytes into ``hasher``. - Hashes each shard's safetensors header so a daemon and an engine pointing - at identical weights in different directories produce the same key. Returns - None when local safetensors files can't be located (e.g. an undownloaded - Hugging Face repo id), leaving the caller to fall back to the model path. + The header carries tensor names, dtypes, shapes and offsets but no weight + values, so hashing it alone gives a base model and its fine-tune the same + key (#59647). Sampling the head of every tensor's data region catches + value-only changes within that region at the cost of one small read per + tensor, and stays stable for byte-identical copies in different + directories. The reads go + through ``pool`` but are mixed in sorted-name order, so the digest is + identical to sequential sampling. + """ + hasher.update(str(os.path.getsize(path)).encode()) + with open(path, "rb") as f: + size_bytes = f.read(8) + (header_len,) = struct.unpack(" str | None: + """Fingerprint checkpoint content from local safetensors files. + + Hashes each shard's metadata plus sampled tensor bytes so a daemon and an + engine pointing at identical weights in different directories produce the + same key, while a difference within any tensor's sampled head produces a + different key. + Returns None when local safetensors files can't be located (e.g. an + undownloaded Hugging Face repo id), leaving the caller to fall back to the + model path. """ if not os.path.isdir(model): return None @@ -298,9 +353,10 @@ def hash_checkpoint(model: str) -> str | None: if not files: return None hasher = safe_hash(b"", usedforsecurity=False) - for path in sorted(files, key=os.path.basename): - hasher.update(os.path.basename(path).encode()) - hasher.update(_safetensors_header(path)) + with ThreadPoolExecutor(max_workers=_FINGERPRINT_READ_WORKERS) as pool: + for path in sorted(files, key=os.path.basename): + hasher.update(os.path.basename(path).encode()) + _mix_safetensors_fingerprint(path, hasher, pool) return hasher.hexdigest() @@ -361,7 +417,7 @@ def from_model_config( may mutate hf_config.quantization_config, which would change the hash between the daemon and the engine. - The checkpoint is identified by a hash of its safetensors metadata when + The checkpoint is identified by a hash of its safetensors files when the weights are available locally, so a daemon and engine referencing identical weights in different directories still match; otherwise it falls back to the model path.