Conversation
Signed-off-by: wangx700 <wangxin700@huawei.com>
Documentation build overview
48 files changed ·
|
![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"}
1 participant
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.