Conversation
![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 disk-based weight synchronization (both full and delta-based) for Ascend NPU environments, allowing weight updates over a shared filesystem using XOR or overwrite delta encodings with zstd compression. The review feedback highlights several critical reliability and robustness improvements: wrapping the delta compression's pinned buffer copy in a try-except block to prevent deadlocks on exception, adding robust error handling for missing or corrupted safetensors files in _tensor_locations, lowering the pinned host buffer memory limit from 32 GB to 4 GB to avoid host instability, and making the patch parser in unit tests more robust against future patch additions.
| 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] |
There was a problem hiding this comment.
If an exception (such as MemoryError during np.empty or np.copyto) occurs inside diff_and_compress before free_q.put(buf) is called, the pinned buffer will be leaked from the pool. If this happens, the main thread's free_q.get() call will eventually block indefinitely, causing a silent deadlock of the entire training process. Wrapping the copy operation in a try...except block ensures that the buffer is always returned to the queue under any failure condition.
def diff_and_compress(name, buf, nbytes, pinned):
try:
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
except Exception:
if pinned:
free_q.put(buf)
raise
old = snapshot[name]| 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("<Q", f.read(8)) | ||
| header = json.loads(f.read(header_len)) | ||
| for name, info in header.items(): | ||
| if name == "__metadata__": | ||
| continue | ||
| begin, end = info["data_offsets"] | ||
| locations[name] = (path, 8 + header_len + begin, end - begin) | ||
| return locations |
There was a problem hiding this comment.
The _tensor_locations function currently assumes that .safetensors files are always present and correctly formatted. If no .safetensors files are found, it silently returns an empty dictionary, which causes downstream components to raise unhelpful KeyError exceptions and flood the logs with warnings. Additionally, if any file is truncated or corrupted, struct.unpack or json.loads will raise unhandled exceptions. Adding explicit validation and robust error handling here will make the initialization process much more reliable.
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 = 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 paths:
try:
with open(path, "rb") as f:
header_bytes = f.read(8)
if len(header_bytes) < 8:
raise ValueError("File header is truncated (less than 8 bytes)")
(header_len,) = struct.unpack("<Q", header_bytes)
header = json.loads(f.read(header_len))
except (struct.error, json.JSONDecodeError, ValueError) as e:
raise RuntimeError(f"Failed to parse safetensors header from {path}: {e}") from e
for name, info in header.items():
if name == "__metadata__":
continue
begin, end = info["data_offsets"]
locations[name] = (path, 8 + header_len + begin, end - begin)
return locations| 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)) |
There was a problem hiding this comment.
Using 32 << 30 (32 GB) as the memory limit for calculating the number of pinned host buffers can lead to extremely large pinned memory allocations on large models (e.g., allocating dozens of buffers of several hundred megabytes each). This can easily exceed the system's memlock limit or cause host Out-Of-Memory (OOM) errors. Lowering this limit to a safer threshold like 4 GB is more than sufficient for pipelining and prevents host instability.
| 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)) | |
| for _ in range(max(4, min(2 * NUM_WORKERS, (4 << 30) // max(max_bytes, 1)))): | |
| free_q.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True)) |
| 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" |
There was a problem hiding this comment.
The test currently parses vllm.patch by splitting at +++ b/vllm/utils/local_checkpoint.py\n and reading all subsequent lines starting with +. If any other file patches are appended to vllm.patch in the future, their diff lines will also be parsed as part of local_checkpoint.py, causing syntax errors or test failures. Splitting the section at the next diff --git line makes the parser robust and future-proof.
| 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" | |
| patch = (ROOT / "docker/npu_patch/vllm.patch").read_text() | |
| section = patch.split("+++ b/vllm/utils/local_checkpoint.py\n", 1)[1] | |
| section = section.split("\ndiff --git", 1)[0] | |
| source = "\n".join(line[1:] for line in section.splitlines() if line.startswith("+")) + "\n" |
Documentation build overview
48 files changed ·
|
| 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": |
There was a problem hiding this comment.
Are these added parameters aligned with the main architecture? If they are extra, what prevents them from being reused?
2 participants
Summary
Ascend currently selects only the colocated tensor updater or distributed collective updater. This adds non-colocated
--update-weight-transport diskwith both--update-weight-mode fulland--update-weight-mode delta, based on ascend at 083aea6.VLLMEngine.pull_weightsuses/collective_rpcto materialize the checkpoint on every NPU worker host before disk reload.docker/patch/latest/vllm-pull_weights.patchinto the existing NPU patch series, and add the NPU worker RPC entry point. Preserve the original model path as the baseline across reloads. Existing Dockerfile/series entries already apply both updated patches.run-qwen3-4B-full-disk.shandrun-qwen3-4B-delta-disk.sh, shared Qwen3-4B settings, configurable paths/Ray ports, and usage/flow documentation. These launchers require explicit device allocation and do not kill existing processes.This adapts the disk protocol to Ascend's current updater lifecycle rather than importing main's newer RayTrainGroup architecture. The default collective path is unchanged; sparse HCCL and profiling changes are outside this PR.
Validation
python3 -m pytest -q tests/unit/backends/megatron_utils/update_weight/test_disk_weight_sync.py tests/unit/utils/test_disk_delta.py— 29 passed. Covers full/delta round trips, unchanged versions and repeated pulls, copy/submission failure injection, pinned allocation limits and cleanup, patch extraction boundaries, and missing/corrupted/truncated checkpoints. Pinned allocation behavior is simulated in CPU tests; this does not measure physical NPU host memory use.