From ea7ec1423f236c69185ecf7ef71a647ab22c4729 Mon Sep 17 00:00:00 2001 From: wangx700 Date: Tue, 15 Sep 2026 10:03:11 +0800 Subject: [PATCH 1/3] feat(ascend): add disk checkpoint NPU patches --- docker/npu_patch/vllm-ascend.patch | 30 ++ docker/npu_patch/vllm.patch | 479 +++++++++++++++++++++++++++++ 2 files changed, 509 insertions(+) diff --git a/docker/npu_patch/vllm-ascend.patch b/docker/npu_patch/vllm-ascend.patch index 768852a5..3d318950 100644 --- a/docker/npu_patch/vllm-ascend.patch +++ b/docker/npu_patch/vllm-ascend.patch @@ -470,3 +470,33 @@ 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,6 +1287,25 @@ 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)} ++ + def check_health(self) -> None: + import subprocess + 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 From 1cc3a0ec82cb4dced49726919be49f568ed97956 Mon Sep 17 00:00:00 2001 From: wangx700 Date: Tue, 15 Sep 2026 10:39:59 +0800 Subject: [PATCH 2/3] feat(scripts): add Qwen3-4B GPU disk sync examples --- scripts/run-qwen3-4B-gpu-delta-disk.sh | 13 +++ scripts/run-qwen3-4B-gpu-disk-common.sh | 147 ++++++++++++++++++++++++ scripts/run-qwen3-4B-gpu-full-disk.sh | 11 ++ 3 files changed, 171 insertions(+) create mode 100755 scripts/run-qwen3-4B-gpu-delta-disk.sh create mode 100755 scripts/run-qwen3-4B-gpu-disk-common.sh create mode 100755 scripts/run-qwen3-4B-gpu-full-disk.sh diff --git a/scripts/run-qwen3-4B-gpu-delta-disk.sh b/scripts/run-qwen3-4B-gpu-delta-disk.sh new file mode 100755 index 00000000..327e0305 --- /dev/null +++ b/scripts/run-qwen3-4B-gpu-delta-disk.sh @@ -0,0 +1,13 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +export UPDATE_WEIGHT_MODE=delta +# Shared trainer/rollout path containing the published delta stream. +export UPDATE_WEIGHT_DISK_DIR="${UPDATE_WEIGHT_DISK_DIR:-/tmp/vime-gpu-delta-weights}" +# Rollout-host-local NVMe path containing the materialized checkpoint. +export UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR="${UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR:-/tmp/vime-gpu-rollout-checkpoint}" +export RAY_GCS_PORT="${RAY_GCS_PORT:-6399}" +export RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8267}" +export RAY_TEMP_DIR="${RAY_TEMP_DIR:-/tmp/ray-vime-gpu-delta-disk}" +exec bash "${SCRIPT_DIR}/run-qwen3-4B-gpu-disk-common.sh" "$@" diff --git a/scripts/run-qwen3-4B-gpu-disk-common.sh b/scripts/run-qwen3-4B-gpu-disk-common.sh new file mode 100755 index 00000000..43c8c27b --- /dev/null +++ b/scripts/run-qwen3-4B-gpu-disk-common.sh @@ -0,0 +1,147 @@ +#!/bin/bash +set -euo pipefail + +export PYTHONUNBUFFERED=1 +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/models/qwen3-4B.sh" + +DATA_ROOT="${DATA_ROOT:-/root}" +HF_CHECKPOINT="${HF_CHECKPOINT:-${DATA_ROOT}/Qwen3-4B}" +MEGATRON_CHECKPOINT="${MEGATRON_CHECKPOINT:-${DATA_ROOT}/Qwen3-4B_vime}" +REF_CHECKPOINT="${REF_CHECKPOINT:-${DATA_ROOT}/Qwen3-4B_torch_dist}" +PROMPT_DATA="${PROMPT_DATA:-${DATA_ROOT}/dapo-math-17k/dapo-math-17k.jsonl}" +MEGATRON_LM_PATH="${MEGATRON_LM_PATH:-${DATA_ROOT}/Megatron-LM}" + +ACTOR_GPUS="${ACTOR_GPUS:-4}" +ROLLOUT_GPUS="${ROLLOUT_GPUS:-4}" +ROLLOUT_GPUS_PER_ENGINE="${ROLLOUT_GPUS_PER_ENGINE:-2}" + +if [[ "${UPDATE_WEIGHT_MODE}" == delta ]]; then + python3 - <<'PY' +from vllm.v1.worker.gpu_worker import Worker + +if not hasattr(Worker, "pull_weights"): + raise RuntimeError( + "disk + delta requires vllm.v1.worker.gpu_worker.Worker.pull_weights; " + "apply the GPU worker disk checkpoint patch to the installed vLLM source" + ) +PY +fi + +CKPT_ARGS=( + --hf-checkpoint "${HF_CHECKPOINT}" + --ref-load "${REF_CHECKPOINT}" + --load "${MEGATRON_CHECKPOINT}" + --save "${MEGATRON_CHECKPOINT}" + --save-interval "${SAVE_INTERVAL:-20}" +) + +ROLLOUT_ARGS=( + --prompt-data "${PROMPT_DATA}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type "${RM_TYPE:-deepscaler}" + --num-rollout "${NUM_ROLLOUT:-3000}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE:-32}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT:-8}" + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN:-8192}" + --rollout-temperature "${ROLLOUT_TEMPERATURE:-1}" + --global-batch-size "${GLOBAL_BATCH_SIZE:-256}" + --balance-data +) + +PERF_ARGS=( + --tensor-model-parallel-size "${TRAIN_TP_SIZE:-2}" + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --use-dynamic-batch-size + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU:-9216}" +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr "${LR:-1e-6}" + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +VLLM_ARGS=( + --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}" + --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION:-0.7}" +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +SYNC_ARGS=( + --update-weight-mode "${UPDATE_WEIGHT_MODE}" + --update-weight-transport disk + --update-weight-disk-dir "${UPDATE_WEIGHT_DISK_DIR}" +) +if [[ "${UPDATE_WEIGHT_MODE}" == delta ]]; then + SYNC_ARGS+=( + --update-weight-local-checkpoint-dir "${UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR}" + --update-weight-delta-encoding "${UPDATE_WEIGHT_DELTA_ENCODING:-xor}" + ) +fi + +cd "${SCRIPT_DIR}/.." +ray start --head \ + --port="${RAY_GCS_PORT}" \ + --temp-dir="${RAY_TEMP_DIR}" \ + --node-ip-address="${MASTER_ADDR:-127.0.0.1}" \ + --num-gpus="$((ACTOR_GPUS + ROLLOUT_GPUS))" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port="${RAY_DASHBOARD_PORT}" + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"${MEGATRON_LM_PATH}:${PYTHONPATH:-}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\" + } +}" + +ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT}" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node "${ACTOR_GPUS}" \ + --rollout-num-gpus "${ROLLOUT_GPUS}" \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${VLLM_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "${SYNC_ARGS[@]}" \ + "$@" diff --git a/scripts/run-qwen3-4B-gpu-full-disk.sh b/scripts/run-qwen3-4B-gpu-full-disk.sh new file mode 100755 index 00000000..5a6953ac --- /dev/null +++ b/scripts/run-qwen3-4B-gpu-full-disk.sh @@ -0,0 +1,11 @@ +#!/bin/bash +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +export UPDATE_WEIGHT_MODE=full +# Shared trainer/rollout path containing each published full checkpoint. +export UPDATE_WEIGHT_DISK_DIR="${UPDATE_WEIGHT_DISK_DIR:-/tmp/vime-gpu-full-weights}" +export RAY_GCS_PORT="${RAY_GCS_PORT:-6400}" +export RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8268}" +export RAY_TEMP_DIR="${RAY_TEMP_DIR:-/tmp/ray-vime-gpu-full-disk}" +exec bash "${SCRIPT_DIR}/run-qwen3-4B-gpu-disk-common.sh" "$@" From c18186d3d8a2471bbe5ccffa56552d2641ad1239 Mon Sep 17 00:00:00 2001 From: wangx700 Date: Tue, 15 Sep 2026 16:37:53 +0800 Subject: [PATCH 3/3] chore: defer GPU disk sync examples --- docker/npu_patch/vllm-ascend.patch | 5 +- scripts/run-qwen3-4B-gpu-delta-disk.sh | 13 --- scripts/run-qwen3-4B-gpu-disk-common.sh | 147 ------------------------ scripts/run-qwen3-4B-gpu-full-disk.sh | 11 -- 4 files changed, 1 insertion(+), 175 deletions(-) delete mode 100755 scripts/run-qwen3-4B-gpu-delta-disk.sh delete mode 100755 scripts/run-qwen3-4B-gpu-disk-common.sh delete mode 100755 scripts/run-qwen3-4B-gpu-full-disk.sh diff --git a/docker/npu_patch/vllm-ascend.patch b/docker/npu_patch/vllm-ascend.patch index 3d318950..e11d6c15 100644 --- a/docker/npu_patch/vllm-ascend.patch +++ b/docker/npu_patch/vllm-ascend.patch @@ -474,7 +474,7 @@ 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,6 +1287,25 @@ class NPUWorker(WorkerBase): +@@ -1287,3 +1287,22 @@ class NPUWorker(WorkerBase): def reload_weights(self, *args, **kwargs) -> None: self.model_runner.reload_weights(*args, **kwargs) @@ -497,6 +497,3 @@ index b7df4df22..b8b9cfbb1 100644 + + return {"success": True, "weight_version": str(target_version)} + - def check_health(self) -> None: - import subprocess - diff --git a/scripts/run-qwen3-4B-gpu-delta-disk.sh b/scripts/run-qwen3-4B-gpu-delta-disk.sh deleted file mode 100755 index 327e0305..00000000 --- a/scripts/run-qwen3-4B-gpu-delta-disk.sh +++ /dev/null @@ -1,13 +0,0 @@ -#!/bin/bash -set -euo pipefail - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -export UPDATE_WEIGHT_MODE=delta -# Shared trainer/rollout path containing the published delta stream. -export UPDATE_WEIGHT_DISK_DIR="${UPDATE_WEIGHT_DISK_DIR:-/tmp/vime-gpu-delta-weights}" -# Rollout-host-local NVMe path containing the materialized checkpoint. -export UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR="${UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR:-/tmp/vime-gpu-rollout-checkpoint}" -export RAY_GCS_PORT="${RAY_GCS_PORT:-6399}" -export RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8267}" -export RAY_TEMP_DIR="${RAY_TEMP_DIR:-/tmp/ray-vime-gpu-delta-disk}" -exec bash "${SCRIPT_DIR}/run-qwen3-4B-gpu-disk-common.sh" "$@" diff --git a/scripts/run-qwen3-4B-gpu-disk-common.sh b/scripts/run-qwen3-4B-gpu-disk-common.sh deleted file mode 100755 index 43c8c27b..00000000 --- a/scripts/run-qwen3-4B-gpu-disk-common.sh +++ /dev/null @@ -1,147 +0,0 @@ -#!/bin/bash -set -euo pipefail - -export PYTHONUNBUFFERED=1 -export CUDA_DEVICE_MAX_CONNECTIONS=1 - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -source "${SCRIPT_DIR}/models/qwen3-4B.sh" - -DATA_ROOT="${DATA_ROOT:-/root}" -HF_CHECKPOINT="${HF_CHECKPOINT:-${DATA_ROOT}/Qwen3-4B}" -MEGATRON_CHECKPOINT="${MEGATRON_CHECKPOINT:-${DATA_ROOT}/Qwen3-4B_vime}" -REF_CHECKPOINT="${REF_CHECKPOINT:-${DATA_ROOT}/Qwen3-4B_torch_dist}" -PROMPT_DATA="${PROMPT_DATA:-${DATA_ROOT}/dapo-math-17k/dapo-math-17k.jsonl}" -MEGATRON_LM_PATH="${MEGATRON_LM_PATH:-${DATA_ROOT}/Megatron-LM}" - -ACTOR_GPUS="${ACTOR_GPUS:-4}" -ROLLOUT_GPUS="${ROLLOUT_GPUS:-4}" -ROLLOUT_GPUS_PER_ENGINE="${ROLLOUT_GPUS_PER_ENGINE:-2}" - -if [[ "${UPDATE_WEIGHT_MODE}" == delta ]]; then - python3 - <<'PY' -from vllm.v1.worker.gpu_worker import Worker - -if not hasattr(Worker, "pull_weights"): - raise RuntimeError( - "disk + delta requires vllm.v1.worker.gpu_worker.Worker.pull_weights; " - "apply the GPU worker disk checkpoint patch to the installed vLLM source" - ) -PY -fi - -CKPT_ARGS=( - --hf-checkpoint "${HF_CHECKPOINT}" - --ref-load "${REF_CHECKPOINT}" - --load "${MEGATRON_CHECKPOINT}" - --save "${MEGATRON_CHECKPOINT}" - --save-interval "${SAVE_INTERVAL:-20}" -) - -ROLLOUT_ARGS=( - --prompt-data "${PROMPT_DATA}" - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - --rm-type "${RM_TYPE:-deepscaler}" - --num-rollout "${NUM_ROLLOUT:-3000}" - --rollout-batch-size "${ROLLOUT_BATCH_SIZE:-32}" - --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT:-8}" - --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN:-8192}" - --rollout-temperature "${ROLLOUT_TEMPERATURE:-1}" - --global-batch-size "${GLOBAL_BATCH_SIZE:-256}" - --balance-data -) - -PERF_ARGS=( - --tensor-model-parallel-size "${TRAIN_TP_SIZE:-2}" - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - --use-dynamic-batch-size - --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU:-9216}" -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr "${LR:-1e-6}" - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -VLLM_ARGS=( - --rollout-num-gpus-per-engine "${ROLLOUT_GPUS_PER_ENGINE}" - --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION:-0.7}" -) - -MISC_ARGS=( - --attention-dropout 0.0 - --hidden-dropout 0.0 - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - --attention-backend flash -) - -SYNC_ARGS=( - --update-weight-mode "${UPDATE_WEIGHT_MODE}" - --update-weight-transport disk - --update-weight-disk-dir "${UPDATE_WEIGHT_DISK_DIR}" -) -if [[ "${UPDATE_WEIGHT_MODE}" == delta ]]; then - SYNC_ARGS+=( - --update-weight-local-checkpoint-dir "${UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR}" - --update-weight-delta-encoding "${UPDATE_WEIGHT_DELTA_ENCODING:-xor}" - ) -fi - -cd "${SCRIPT_DIR}/.." -ray start --head \ - --port="${RAY_GCS_PORT}" \ - --temp-dir="${RAY_TEMP_DIR}" \ - --node-ip-address="${MASTER_ADDR:-127.0.0.1}" \ - --num-gpus="$((ACTOR_GPUS + ROLLOUT_GPUS))" \ - --disable-usage-stats \ - --dashboard-host=0.0.0.0 \ - --dashboard-port="${RAY_DASHBOARD_PORT}" - -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"${MEGATRON_LM_PATH}:${PYTHONPATH:-}\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\" - } -}" - -ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT}" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node "${ACTOR_GPUS}" \ - --rollout-num-gpus "${ROLLOUT_GPUS}" \ - "${MODEL_ARGS[@]}" \ - "${CKPT_ARGS[@]}" \ - "${ROLLOUT_ARGS[@]}" \ - "${OPTIMIZER_ARGS[@]}" \ - "${GRPO_ARGS[@]}" \ - "${PERF_ARGS[@]}" \ - "${VLLM_ARGS[@]}" \ - "${MISC_ARGS[@]}" \ - "${SYNC_ARGS[@]}" \ - "$@" diff --git a/scripts/run-qwen3-4B-gpu-full-disk.sh b/scripts/run-qwen3-4B-gpu-full-disk.sh deleted file mode 100755 index 5a6953ac..00000000 --- a/scripts/run-qwen3-4B-gpu-full-disk.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash -set -euo pipefail - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -export UPDATE_WEIGHT_MODE=full -# Shared trainer/rollout path containing each published full checkpoint. -export UPDATE_WEIGHT_DISK_DIR="${UPDATE_WEIGHT_DISK_DIR:-/tmp/vime-gpu-full-weights}" -export RAY_GCS_PORT="${RAY_GCS_PORT:-6400}" -export RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8268}" -export RAY_TEMP_DIR="${RAY_TEMP_DIR:-/tmp/ray-vime-gpu-full-disk}" -exec bash "${SCRIPT_DIR}/run-qwen3-4B-gpu-disk-common.sh" "$@"