Repository navigation
Conversation
Signed-off-by: wangx700 <wangxin700@huawei.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a mechanism to maintain host-local checkpoints from full and delta weight versions, integrated into the NPU worker. The review feedback highlights critical issues regarding platform compatibility on non-Unix/non-Linux systems, specifically due to the top-level import of fcntl and the use of os.posix_fadvise. Additionally, the feedback identifies a logical bug in rollback scenarios where target_version is less than applied, and suggests adding a unit test to verify rollback behavior.
| + 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 |
There was a problem hiding this comment.
If target_version < applied (a rollback scenario), the current search logic fails because floor is set to applied. Since start is initialized to target_version, the loop while start > floor never runs, leaving start at target_version (which could be a delta version). Then, start > applied is false, so start is set to applied, and no changes are made. This silently leaves the local checkpoint at the newer version while returning success.
To fix this, when target_version < applied, we must search all the way down to 0 to find the correct base full checkpoint, and we must reset the checkpoint if start != applied.
applied = _read_applied_version(local_checkpoint_dir)
if applied is not None and target_version >= applied:
floor = applied
else:
floor = 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
| +import fcntl | ||
| +import glob |
There was a problem hiding this comment.
Remove the top-level import fcntl to prevent ModuleNotFoundError on non-Unix platforms (like Windows) when this module is imported (e.g., during testing or static analysis). We can import it dynamically inside _pull_lock instead.
import glob
| +@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) |
There was a problem hiding this comment.
Import fcntl dynamically here with a fallback to support non-Unix platforms where fcntl is not available.
@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)
try:
import fcntl
except ImportError:
yield
return
with open(os.path.join(sync_dir, "lock"), "w") as lock_file:
fcntl.flock(lock_file, fcntl.LOCK_EX)
try:
yield
finally:
with suppress(OSError):
fcntl.flock(lock_file, fcntl.LOCK_UN)
| +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 |
There was a problem hiding this comment.
On non-Linux platforms (like macOS or Windows), os.posix_fadvise is not available and calling it will raise an AttributeError. We should check if os has the posix_fadvise attribute before attempting to use it.
def _drop_page_cache(path: str) -> None:
if not hasattr(os, "posix_fadvise"):
return
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
| + state = json.loads((local_dir / ".weight_sync" / "state.json").read_text()) | ||
| + assert state == {"version": "000000"} |
There was a problem hiding this comment.
Add a unit test to verify that rolling back to an older version (e.g., from version 1 to version 0) works correctly and restores the baseline weights.
state = json.loads((local_dir / ".weight_sync" / "state.json").read_text())
assert state == {"version": "000000"}
def test_pull_checkpoint_rollback(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 to version 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)
# Rollback to version 0
pull_checkpoint(str(local_dir), str(base_dir), str(source_dir), 0)
actual = safetensors.numpy.load_file(local_dir / "model.safetensors")
np.testing.assert_array_equal(actual["weight"], baseline)
state = json.loads((local_dir / ".weight_sync" / "state.json").read_text())
assert state == {"version": "000000"}
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
88847b3 to
9f71b6a
Compare
| Runs the full disk test configuration for both delta encodings. Two rollout | ||
| rounds exercise inference after a delta reload; recovered tensors are checked | ||
| against the published checksums. | ||
| """ |
There was a problem hiding this comment.
This test isn't added to the NPU suite, so CI never runs it. Could you register it in .buildkite/npu_suites.py?
There was a problem hiding this comment.
Ok, this test was add to .buildkite/npu_suites.py
| dir_dst="/dev/shm", | ||
| hf_checkpoint=HF_CKPT, | ||
| ) | ||
|
|
There was a problem hiding this comment.
The model/dataset paths here don't follow the NPU test convention. /home/vllm/weights looks like a local path, and U.hf_download_dataset hardcodes /root/datasets. The other NPU tests (e.g. test_qwen3_30B_A3B_npu.py) put models and datasets under HF_HOME. Suggest aligning with them.
There was a problem hiding this comment.
ok, the path was changed to hf download path.
| ckpt_args = f"--hf-checkpoint {HF_CKPT} " f"--ref-load {ref_load} " | ||
|
|
||
| rollout_args = ( | ||
| "--prompt-data /root/datasets/gsm8k/train.parquet " |
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
Signed-off-by: wangx700 <wangxin700@huawei.com>
Merge frozen Ascend 5397545 onto the main-aligned branch. Resolve vendor patch conflicts against vLLM 8cc9aa5a and Ascend e00cb64; use main's indexed checkpoint receiver and configured pre-read hook, without vendor test hunks. Adapt the four-NPU XOR/overwrite smoke test to initial full plus subsequent deltas, reuse SafetensorReader and verify all recovered tensors. Extend existing CPU receiver coverage for the NPU patch, checksum failure and full recovery. Static syntax and patch apply/reverse checks passed; targeted pre-commit passed. Unit tests and NPU E2E intentionally not run. Signed-off-by: Meihan-chen <zr010426ztt@outlook.com>
Summary
Add the serving-side patches required by disk-backed full and delta weight synchronization on Ascend:
docker/npu_patch/vllm.patchNPUWorker.pull_weights()todocker/npu_patch/vllm-ascend.patchThe PR is based directly on the current
vllm-project/vime:ascendbranch and changes only these two patch files.Dependency
Depends on #413 for the VIME-side disk full/delta weight synchronization flow and its pinned vLLM/vLLM-Ascend revisions. After #413 merges, this branch can be rebased onto the updated
ascendbranch and merged directly.Validation
Validated on Ascend NPUs 8-15 with Qwen3-4B, disk transport,
num_rollout=10, rollout batch size 4, two samples per prompt, and response length 256. The table excludes the final v10 update because it overlaps Ray/vLLM job teardown; it compares the initial sync and stable v1-v9 updates.E2E uses
perf/step_timeand covers rollout, training wait, training, and the update attributed to the step. Delta transfers about 0.66% of the model on average, but vLLM still reloads a reconstructed 7.49 GiB safetensors checkpoint on this NPU path.Qwen3-0.6B disk delta validation
Added
tests/test_delta_disk_weight_update.pyto run XOR and overwrite sequentially with Qwen3-0.6B. Each encoding runs two rollout/training rounds on four Ascend devices (one training, three rollout), with graph mode enabled. The test validates delta metadata and final recovered tensor checksums.Run from the repository root in a configured environment:
[2026-10-09 10:39:00] update_weight_from_disk_delta.py:286 - [disk delta v=2] density=0.48% wire=0.02 GB [2026-10-09 10:39:00] timer.py:32 - Timer update_weights end (elapsed: 4.4s) ... Disk delta validation passed: {'encoding': 'xor', 'rollout_rounds': 2, 'local_version': '000002', 'checked_tensors': 240}[2026-10-09 10:47:12] update_weight_from_disk_delta.py:286 - [disk delta v=2] density=0.48% wire=0.02 GB [2026-10-09 10:47:12] timer.py:32 - Timer update_weights end (elapsed: 3.9s) ... Disk delta validation passed: {'encoding': 'overwrite', 'rollout_rounds': 2, 'local_version': '000002', 'checked_tensors': 238}