diff --git a/.buildkite/npu_suites.py b/.buildkite/npu_suites.py index 275afd125..d80059192 100644 --- a/.buildkite/npu_suites.py +++ b/.buildkite/npu_suites.py @@ -29,6 +29,7 @@ SUITES = { "smk": [ ("test_qwen3_4B_npu.py", "npu-8", "", {}), + ("test_qwen3_4B_npu.py", "npu-8", "", {"VIME_TEST_UPDATE_MODE": "delta"}), ("test_qwen3_30B_A3B_npu.py", "npu-16", "", {}), ("test_qwen3_vl_8B_npu.py", "npu-8", "", {}), ], diff --git a/.buildkite/scripts/update-npu-environment.sh b/.buildkite/scripts/update-npu-environment.sh index 761f2406d..1779e2d8b 100644 --- a/.buildkite/scripts/update-npu-environment.sh +++ b/.buildkite/scripts/update-npu-environment.sh @@ -1,7 +1,7 @@ #!/bin/bash # Purpose: Updates an NPU test container to match the requested VIME commit. # - Reads the image's persisted OLD patch series and exact patch bytes -# - Updates VIME, then reconciles OLD -> NEW in declared series order +# - Updates VIME, then reconciles OLD -> NEW in declared series orde # - Installs the current VIME checkout and normalizes visible devices # Usage: Called by Buildkite pipeline during NPU test runs set -e -o pipefail @@ -176,6 +176,10 @@ update_vime_code() { install_vime_code() { pip install -e "$VIME_DIR" --no-deps --break-system-packages || pip install -e "$VIME_DIR" --no-deps + # The pre-built smoke image can predate disk-delta dependencies. Keep the + # serving and trainer interpreters aligned before loading delta checkpoints. + pip install blake3 xxhash zstandard --break-system-packages || \ + pip install blake3 xxhash zstandard } sort_ascend_visible_devices() { diff --git a/docker/npu_patch/vllm-ascend.patch b/docker/npu_patch/vllm-ascend.patch index fa89a9129..1efcc5263 100644 --- a/docker/npu_patch/vllm-ascend.patch +++ b/docker/npu_patch/vllm-ascend.patch @@ -13,7 +13,7 @@ index 062b3ecd..dd87affb 100644 f"Free memory on device " @@ -535,15 +535,16 @@ class NPUWorker(WorkerBase): self.npugraph_memory_estimate = npugraph_memory_estimate - + free_gpu_memory = profile_result.after_profile.free_memory - assert self.init_snapshot.free_memory > free_gpu_memory, ( - "Error in memory profiling. " @@ -50,4 +50,72 @@ index a35d9af8d..66a179cd6 100644 + # Source tensors may have been produced asynchronously on the caller's + # stream. Wait before the packing stream reads them in torch.cat. + streams[buffer_idx].wait_stream(source_stream) - # Start tasks for the new buffer in a new stream + # Start tasks for the new buffer in a new stream +diff --git a/vllm_ascend/worker/worker.py b/vllm_ascend/worker/worker.py +index 062b3ecd..f0d1e4a1 100644 +--- a/vllm_ascend/worker/worker.py ++++ b/vllm_ascend/worker/worker.py +@@ -333,1 +333,33 @@ class NPUWorker(WorkerBase): ++ def pull_weights( ++ self, ++ local_checkpoint_dir: str, ++ source_dir: str, ++ target_version: int, ++ pre_read_hook: str | None = None, ++ ) -> dict: ++ """Materialize a full or delta checkpoint on this rollout host. ++ ++ ``collective_rpc`` invokes this on all NPU workers. The checkpoint ++ helper serializes same-host ranks with a filesystem lock, so each host ++ applies a version exactly once before reload_weights is called. ++ """ ++ from vllm.utils.local_checkpoint import pull_checkpoint ++ ++ # reload_weights updates model_config.model to the materialized local ++ # checkpoint. Preserve the original model path as the immutable ++ # version-zero seed for retries and restarted training runs. ++ base_dir = getattr(self, "_local_checkpoint_base_dir", None) ++ if base_dir is None: ++ base_dir = self.model_config.model ++ self._local_checkpoint_base_dir = base_dir ++ ++ pull_checkpoint( ++ local_checkpoint_dir=local_checkpoint_dir, ++ base_dir=base_dir, ++ source_dir=source_dir, ++ target_version=target_version, ++ pre_read_hook=pre_read_hook, ++ ) ++ return {"success": True, "weight_version": str(target_version)} ++ + def shutdown(self) -> None: +diff --git a/tests/ut/worker/a2/test_worker_v1.py b/tests/ut/worker/a2/test_worker_v1.py +index 812b757..06e3875 100644 +--- a/tests/ut/worker/a2/test_worker_v1.py ++++ b/tests/ut/worker/a2/test_worker_v1.py +@@ -1560,2 +1560,25 @@ class TestNPUWorker(TestBase): ++ def test_pull_weights_preserves_initial_model_path(self): ++ """The version-zero seed must survive reload_weights path updates.""" ++ from vllm_ascend.worker.worker import NPUWorker ++ ++ with ( ++ patch.object(NPUWorker, "__init__", lambda x, **kwargs: None), ++ patch("vllm.utils.local_checkpoint.pull_checkpoint") as pull_checkpoint, ++ ): ++ worker = NPUWorker() ++ worker.model_config = MagicMock() ++ worker.model_config.model = "/models/base" ++ ++ worker.pull_weights("/local/checkpoint", "/shared/weights", 0) ++ worker.model_config.model = "/local/checkpoint" ++ worker.pull_weights("/local/checkpoint", "/shared/weights", 1) ++ ++ assert pull_checkpoint.call_count == 2 ++ assert [call.kwargs["base_dir"] for call in pull_checkpoint.call_args_list] == [ ++ "/models/base", ++ "/models/base", ++ ] ++ assert pull_checkpoint.call_args.kwargs["target_version"] == 1 ++ + class TestNPUWorkerWeightUpdate(TestBase): + def _make_worker(self, engine=None): diff --git a/docker/npu_patch/vllm.patch b/docker/npu_patch/vllm.patch index f1576e389..35bc7556d 100644 --- a/docker/npu_patch/vllm.patch +++ b/docker/npu_patch/vllm.patch @@ -22,5 +22,250 @@ index cb61bca..5c076d5 100644 - assert request.num_output_placeholders >= 0 + request.num_output_placeholders = max(0, request.num_output_placeholders) - # Cache the new tokens. Preempted requests should be skipped. - if status_before_update == RequestStatus.RUNNING: + # Cache the new tokens. Preempted requests should be skipped. + if status_before_update == RequestStatus.RUNNING: +diff --git a/vllm/utils/local_checkpoint.py b/vllm/utils/local_checkpoint.py +new file mode 100644 +--- /dev/null ++++ b/vllm/utils/local_checkpoint.py +@@ -0,0 +1,240 @@ ++# SPDX-License-Identifier: Apache-2.0 ++"""Maintain a host-local HF 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 ++ ++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, base_dir, source_dir, target_version, pre_read_hook=None): ++ """Bring a host-local checkpoint to a published full or delta version.""" ++ if target_version > 0 and pre_read_hook: ++ module_path, _, function_name = pre_read_hook.rpartition(".") ++ getattr(importlib.import_module(module_path), function_name)(source_dir, target_version) ++ with _pull_lock(local_checkpoint_dir): ++ incomplete = os.path.join(local_checkpoint_dir, SYNC_DIR, "incomplete") ++ # Deltas update mmap'ed safetensor regions in place. A failed apply can ++ # therefore leave bytes changed even though state.json was not advanced; ++ # discard that partial local state on the next retry. ++ applied = None if os.path.exists(incomplete) else _read_applied_version(local_checkpoint_dir) ++ # A pull can legitimately restart a run at version zero or target an ++ # older published version. In that case the current checkpoint cannot ++ # be patched backwards, so search from the base checkpoint instead. ++ floor = applied if applied is not None and applied <= target_version else 0 ++ start = target_version ++ while start > floor and _is_delta(_version_dir(source_dir, start)): ++ start -= 1 ++ if applied is None or applied > target_version or start > applied: ++ _reset_checkpoint(base_dir if start == 0 else _version_dir(source_dir, start), local_checkpoint_dir, start) ++ else: ++ start = applied ++ try: ++ for version in range(start + 1, target_version + 1): ++ open(incomplete, "a").close() ++ _apply_delta(local_checkpoint_dir, _version_dir(source_dir, version)) ++ except BaseException: ++ raise ++ else: ++ if os.path.exists(incomplete): ++ os.remove(incomplete) ++ ++ ++def _version_dir(source_dir, version): ++ return os.path.join(source_dir, f"weight_v{version:06d}") ++ ++ ++def _is_delta(version_dir): ++ 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): ++ self._value = 1 ++ ++ def update(self, data): ++ self._value = zlib.adler32(data, self._value) ++ ++ def hexdigest(self): ++ return f"{self._value:08x}" ++ ++ ++def _new_hasher(algorithm): ++ 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}") ++ ++ ++@contextmanager ++def _pull_lock(local_checkpoint_dir): ++ 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): ++ 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, version): ++ path = os.path.join(local_checkpoint_dir, SYNC_DIR, "state.json") ++ temporary = 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 _reset_checkpoint(source_dir, local_checkpoint_dir, version): ++ 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)) ++ 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 = os.path.join(local_checkpoint_dir, entry.name) ++ if os.path.getsize(copied) != entry.stat().st_size: ++ raise RuntimeError(f"Size mismatch copying {entry.name}") ++ _write_applied_version(local_checkpoint_dir, version) ++ ++ ++def _tensor_locations(checkpoint_dir): ++ locations = {} ++ for path in glob.glob(os.path.join(checkpoint_dir, "*.safetensors")): ++ with open(path, "rb") as tensor_file: ++ (header_length,) = struct.unpack("=0.1.14 wandb +xxhash +zstandard diff --git a/scripts/run-qwen3-4B-npu-sparse-hccl.sh b/scripts/run-qwen3-4B-npu-sparse-hccl.sh new file mode 100755 index 000000000..883f02003 --- /dev/null +++ b/scripts/run-qwen3-4B-npu-sparse-hccl.sh @@ -0,0 +1,10 @@ +#!/bin/bash +set -e + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +export VLLM_VERSION="${VLLM_VERSION:-0.26.0}" +export UPDATE_WEIGHT_MODE=sparse +export UPDATE_WEIGHT_TRANSPORT=nccl +export VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}" + +exec bash "${SCRIPT_DIR}/run-qwen3-4B-npu.sh" "$@" diff --git a/scripts/run-qwen3-4B-npu.sh b/scripts/run-qwen3-4B-npu.sh index 73b672866..b46f2550c 100644 --- a/scripts/run-qwen3-4B-npu.sh +++ b/scripts/run-qwen3-4B-npu.sh @@ -24,7 +24,9 @@ export HYDRA_FULL_ERROR=1 export DISABLE_L2_CACHE=1 export VLLM_ASCEND_ENABLE_NZ=0 export VLLM_USE_AOT_COMPILE=0 -export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:${PYTHONPATH:-}" +export VLLM_VERSION="${VLLM_VERSION:-0.26.0}" +VIME_WORKSPACE_ROOT="${VIME_WORKSPACE_ROOT:-/home/vllm/c00944022/0623}" +export PYTHONPATH="${VIME_WORKSPACE_ROOT}/Megatron-Bridge/src:${VIME_WORKSPACE_ROOT}/Megatron-LM:${PYTHONPATH:-}" unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY @@ -32,27 +34,44 @@ SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" source "${SCRIPT_DIR}/models/qwen3-4B.sh" DATA_ROOT="${DATA_ROOT:-/root}" +MODEL_PATH="${MODEL_PATH:-/home/vllm/weights/Qwen3-4B}" +PROMPT_DATA_PATH="${PROMPT_DATA_PATH:-/home/vllm/c00944022/datasets/dapo-math-17k/dapo-math-17k.jsonl}" +UPDATE_WEIGHT_DISK_DIR="${UPDATE_WEIGHT_DISK_DIR:-/home/vllm/c00944022/0623/vime-delta-weights}" +UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR="${UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR:-/tmp/vime-rollout-checkpoint}" +UPDATE_WEIGHT_MODE="${UPDATE_WEIGHT_MODE:-delta}" +UPDATE_WEIGHT_TRANSPORT="${UPDATE_WEIGHT_TRANSPORT:-disk}" +VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.6}" +NUM_ROLLOUT="${NUM_ROLLOUT:-200}" +ROLLOUT_BATCH_SIZE="${ROLLOUT_BATCH_SIZE:-32}" +N_SAMPLES_PER_PROMPT="${N_SAMPLES_PER_PROMPT:-8}" +ROLLOUT_MAX_RESPONSE_LEN="${ROLLOUT_MAX_RESPONSE_LEN:-2048}" +GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-256}" +TRAIN_LR="${TRAIN_LR:-1e-6}" +ENTROPY_COEF="${ENTROPY_COEF:-0.0}" +RAY_GCS_PORT="${RAY_GCS_PORT:-6399}" +RAY_DASHBOARD_PORT="${RAY_DASHBOARD_PORT:-8267}" +RAY_TEMP_DIR="${RAY_TEMP_DIR:-/tmp/ray-vime-delta}" CKPT_ARGS=( - --hf-checkpoint ${DATA_ROOT}/models/Qwen3-4B/ - --load ${DATA_ROOT}/models/Qwen3-4B/ - --ref-load ${DATA_ROOT}/models/Qwen3-4B/ + --hf-checkpoint "${MODEL_PATH}" + --load "${MODEL_PATH}" + --ref-load "${MODEL_PATH}" --megatron-to-hf-mode bridge ) ROLLOUT_ARGS=( - --prompt-data ${DATA_ROOT}/datasets/dapo-math-17k/dapo-math-17k.jsonl + --prompt-data "${PROMPT_DATA_PATH}" --input-key prompt --label-key label --apply-chat-template --rollout-shuffle --rm-type math - --num-rollout 200 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 2048 + --num-rollout "${NUM_ROLLOUT}" + --rollout-batch-size "${ROLLOUT_BATCH_SIZE}" + --n-samples-per-prompt "${N_SAMPLES_PER_PROMPT}" + --rollout-max-response-len "${ROLLOUT_MAX_RESPONSE_LEN}" --rollout-temperature 1 - --global-batch-size 256 + --global-batch-size "${GLOBAL_BATCH_SIZE}" --balance-data ) @@ -75,14 +94,14 @@ GRPO_ARGS=( --kl-loss-coef 0.0 --kl-loss-type low_var_kl --kl-coef 0.00 - --entropy-coef 0.0 + --entropy-coef "${ENTROPY_COEF}" --eps-clip 0.2 --eps-clip-high 0.28 ) OPTIMIZER_ARGS=( --optimizer adam - --lr 1e-6 + --lr "${TRAIN_LR}" --lr-decay-style constant --weight-decay 0.1 --adam-beta1 0.9 @@ -94,9 +113,22 @@ OPTIMIZER_ARGS=( VLLM_ARGS=( --rollout-num-gpus-per-engine 4 - --vllm-gpu-memory-utilization 0.6 + --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION}" ) +UPDATE_WEIGHT_ARGS=( + --update-weight-mode "${UPDATE_WEIGHT_MODE}" + --update-weight-transport "${UPDATE_WEIGHT_TRANSPORT}" +) +if [[ "${UPDATE_WEIGHT_MODE}" == "delta" ]]; then + UPDATE_WEIGHT_ARGS+=( + --update-weight-disk-dir "${UPDATE_WEIGHT_DISK_DIR}" + --update-weight-local-checkpoint-dir "${UPDATE_WEIGHT_LOCAL_CHECKPOINT_DIR}" + --update-weight-delta-encoding xor + --update-weight-delta-checksum xxh3-128 + ) +fi + MISC_ARGS=( --attention-dropout 0.0 --hidden-dropout 0.0 @@ -107,9 +139,11 @@ MISC_ARGS=( --use-flash-attn ) -ray start --head --node-ip-address 127.0.0.1 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 +ray start --head --port="${RAY_GCS_PORT}" --temp-dir="${RAY_TEMP_DIR}" \ +--node-ip-address 127.0.0.1 --disable-usage-stats \ +--dashboard-host=0.0.0.0 --dashboard-port="${RAY_DASHBOARD_PORT}" -ray job submit --address="http://127.0.0.1:8265" \ +ray job submit --address="http://127.0.0.1:${RAY_DASHBOARD_PORT}" \ -- python3 train.py \ --actor-num-nodes 1 \ --actor-num-gpus-per-node 4 \ @@ -121,4 +155,5 @@ ${OPTIMIZER_ARGS[@]} \ ${GRPO_ARGS[@]} \ ${PERF_ARGS[@]} \ ${VLLM_ARGS[@]} \ +${UPDATE_WEIGHT_ARGS[@]} \ ${MISC_ARGS[@]} diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index af20fa71c..a57ef0b35 100644 --- a/tests/test_megatron_argument_validation.py +++ b/tests/test_megatron_argument_validation.py @@ -39,6 +39,64 @@ def load_arguments_module(monkeypatch): return module +def load_vime_arguments_module(monkeypatch): + """Load VIME argument validation without importing the full runtime stack.""" + router_pkg_mod = types.ModuleType("vllm_router") + router_launch_mod = types.ModuleType("vllm_router.launch_router") + vllm_arguments_mod = types.ModuleType("vime.backends.vllm_utils.arguments") + common_mod = types.ModuleType("vime.utils.common") + logging_utils_mod = types.ModuleType("vime.utils.logging_utils") + router_launch_mod.RouterArgs = object + vllm_arguments_mod.vllm_parse_args = lambda *args, **kwargs: None + vllm_arguments_mod.validate_args = lambda args: args + common_mod.is_npu = lambda: True + logging_utils_mod.configure_logger = lambda *args, **kwargs: None + monkeypatch.setitem(sys.modules, "vllm_router", router_pkg_mod) + monkeypatch.setitem(sys.modules, "vllm_router.launch_router", router_launch_mod) + monkeypatch.setitem(sys.modules, "vime.backends.vllm_utils.arguments", vllm_arguments_mod) + monkeypatch.setitem(sys.modules, "vime.utils.common", common_mod) + monkeypatch.setitem(sys.modules, "vime.utils.logging_utils", logging_utils_mod) + module_path = Path(__file__).resolve().parents[1] / "vime" / "utils" / "arguments.py" + module_name = "test_vime_argument_validation_module" + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def make_vime_validate_args(**overrides): + values = dict( + eval_config=None, eval_prompt_data=None, kl_coef=0, use_kl_loss=False, + ref_load=None, use_opd=False, opd_type=None, opd_teacher_load=None, + load=None, hf_checkpoint="/tmp/hf", megatron_to_hf_mode="bridge", ref_ckpt_step=None, ckpt_step=None, + no_load_optim=False, no_load_rng=False, finetune=False, start_rollout_id=None, + eval_interval=None, save_interval=None, save=None, kl_loss_coef=0, + advantage_estimator="grpo", normalize_advantages=False, use_rollout_logprobs=False, + use_tis=False, get_mismatch_metrics=False, custom_tis_function_path=None, + use_dynamic_batch_size=False, max_tokens_per_gpu=None, log_probs_max_tokens_per_gpu=None, + balance_by_flops=False, balance_data=False, eps_clip_high=None, eps_clip=0.2, + eval_reward_key=None, reward_key="reward", dump_details=None, + save_debug_rollout_data=None, save_debug_train_data=None, load_debug_rollout_data=None, + rollout_external_engine_addrs=None, debug_train_only=False, actor_num_gpus_per_node=8, + actor_num_nodes=1, num_gpus_per_node=8, offload=False, offload_train=None, + offload_rollout=None, debug_rollout_only=False, colocate=False, rollout_num_gpus=8, + eval_function_path=None, rollout_function_path="custom.rollout", num_steps_per_rollout=None, + rollout_batch_size=1, n_samples_per_prompt=1, global_batch_size=None, + grpo_std_normalization=True, over_sampling_batch_size=None, num_epoch=None, + num_rollout=1, rollout_global_dataset=False, enable_mtp_training=False, + mtp_num_layers=None, use_rollout_routing_replay=False, use_routing_replay=False, + custom_config_path=None, eval_max_context_len=None, rollout_max_context_len=None, + rollout_max_prompt_len=None, train_backend="megatron", release_train=False, + keep_old_actor=False, only_train_params_name_list=None, freeze_params_name_list=None, + update_weight_transport="nccl", update_weight_disk_dir=None, + update_weight_local_checkpoint_dir=None, update_weight_mode="full", qkv_format="sbhd", + ) + values.update(overrides) + return types.SimpleNamespace(**values) + + def make_qwen3_6_args(**overrides): values = dict( hidden_size=2048, @@ -139,5 +197,72 @@ def test_allgather_cp_ignores_cp_size_one(monkeypatch): module._validate_allgather_cp_supported(args) +@pytest.mark.unit +def test_update_weight_delta_disk_is_valid(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + module.vime_validate_args( + make_vime_validate_args( + update_weight_mode="delta", + update_weight_transport="disk", + update_weight_disk_dir="/shared/delta", + update_weight_local_checkpoint_dir="/local/delta", + ) + ) + + +@pytest.mark.unit +def test_update_weight_sparse_hccl_is_valid(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + module.vime_validate_args( + make_vime_validate_args(update_weight_mode="sparse", update_weight_transport="nccl") + ) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("overrides", "error"), + [ + ({"update_weight_mode": "sparse", "update_weight_transport": "disk"}, "requires.*nccl"), + ({"update_weight_mode": "sparse", "colocate": True}, "non-colocated"), + ({"update_weight_mode": "sparse", "megatron_to_hf_mode": "raw"}, "requires.*bridge"), + ], +) +def test_update_weight_sparse_rejects_invalid_combinations(monkeypatch, overrides, error): + module = load_vime_arguments_module(monkeypatch) + with pytest.raises(ValueError, match=error): + module.vime_validate_args(make_vime_validate_args(**overrides)) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("overrides", "error"), + [ + ({"update_weight_mode": "delta"}, "requires --update-weight-transport=disk"), + ( + { + "update_weight_mode": "delta", + "update_weight_transport": "disk", + "update_weight_disk_dir": "/shared/delta", + "colocate": True, + }, + "not supported with --colocate", + ), + ( + { + "update_weight_mode": "delta", + "update_weight_transport": "disk", + "update_weight_disk_dir": "/shared/delta", + }, + "requires --update-weight-local-checkpoint-dir", + ), + ({"update_weight_transport": "disk"}, "supported only with --update-weight-mode=delta"), + ], +) +def test_update_weight_disk_rejects_invalid_combinations(monkeypatch, overrides, error): + module = load_vime_arguments_module(monkeypatch) + with pytest.raises(ValueError, match=error): + module.vime_validate_args(make_vime_validate_args(**overrides)) + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_qwen3_4B_npu.py b/tests/test_qwen3_4B_npu.py index 1b7d7cf3b..9a1542d30 100644 --- a/tests/test_qwen3_4B_npu.py +++ b/tests/test_qwen3_4B_npu.py @@ -1,5 +1,7 @@ import os import shlex +import tempfile +from pathlib import Path import vime.utils.external_utils.command_utils as U @@ -23,6 +25,7 @@ def prepare(): def execute(): model_dir = shlex.quote(MODEL_DIR) prompt_data = shlex.quote(f"{DATASET_DIR}/dapo-math-17k.jsonl") + update_mode = os.environ.get("VIME_TEST_UPDATE_MODE", "full") checkpoint_args = ( f"--hf-checkpoint {model_dir} " @@ -66,7 +69,7 @@ def execute(): "--kl-loss-coef 0.0 " "--kl-loss-type low_var_kl " "--kl-coef 0.00 " - "--entropy-coef 0.0 " + f"--entropy-coef {'0.01' if update_mode == 'delta' else '0.0'} " "--eps-clip 0.2 " "--eps-clip-high 0.28 " ) @@ -109,22 +112,44 @@ def execute(): "--ci-test " ) - train_args = ( - checkpoint_args - + rollout_args - + parallel_args - + grpo_args - + optimizer_args - + vllm_args - + model_args - + runtime_args - ) - U.execute_train( - train_args=train_args, - num_gpus_per_node=8, - megatron_model_type="qwen3-4B", - extra_env_vars={}, - ) + with tempfile.TemporaryDirectory(prefix="vime_npu_delta_shared_") as shared_dir, tempfile.TemporaryDirectory( + prefix="vime_npu_delta_local_" + ) as local_dir: + disk_args = "" + if update_mode == "delta": + disk_args = ( + "--update-weight-mode delta " + "--update-weight-transport disk " + f"--update-weight-disk-dir {shlex.quote(shared_dir)} " + f"--update-weight-local-checkpoint-dir {shlex.quote(local_dir)} " + "--update-weight-delta-encoding xor " + "--update-weight-delta-checksum xxh3-128 " + ) + + train_args = ( + checkpoint_args + + rollout_args + + parallel_args + + grpo_args + + optimizer_args + + vllm_args + + model_args + + runtime_args + + disk_args + ) + U.execute_train( + train_args=train_args, + num_gpus_per_node=8, + megatron_model_type="qwen3-4B", + extra_env_vars={}, + ) + + if update_mode == "delta": + versions = sorted(Path(shared_dir).glob("weight_v*")) + assert len(versions) >= 2, f"Expected two delta updates under {shared_dir}, got {versions}" + assert all((version / "model.safetensors.index.json").exists() for version in versions) + assert any("delta_encoding" in (version / "model.safetensors.index.json").read_text() for version in versions) + assert (Path(local_dir) / ".weight_sync" / "state.json").exists() def main(): diff --git a/tests/unit/backends/megatron_utils/update_weight/test_megatron_sparse_export.py b/tests/unit/backends/megatron_utils/update_weight/test_megatron_sparse_export.py new file mode 100644 index 000000000..87f083b66 --- /dev/null +++ b/tests/unit/backends/megatron_utils/update_weight/test_megatron_sparse_export.py @@ -0,0 +1,223 @@ +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace + +import torch + +MODULE_PATH = ( + Path(__file__).parents[5] + / "vime/backends/megatron_utils/update_weight/megatron_sparse_export.py" +) +SPEC = importlib.util.spec_from_file_location("test_megatron_sparse_export_module", MODULE_PATH) +assert SPEC is not None and SPEC.loader is not None +export = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = export +SPEC.loader.exec_module(export) + + +def test_local_bit_exact_diff_uses_storage_bits() -> None: + old_bits = torch.tensor([0x0000, 0x8000, 0x7FC1, 0x3F80], dtype=torch.uint16) + new_bits = torch.tensor([0x8000, 0x8000, 0x7FC2, 0x3F80], dtype=torch.uint16) + old = old_bits.view(torch.bfloat16) + new = new_bits.view(torch.bfloat16) + + indices, values = export.local_bit_exact_diff(new, old) + + assert indices.tolist() == [0, 2] + assert values.view(torch.uint16).tolist() == [0x8000, 0x7FC2] + + +def test_clone_cpu_snapshot_does_not_alias_mutable_backup() -> None: + backup = torch.tensor([1.0, 2.0], dtype=torch.bfloat16) + + snapshot = export.clone_cpu_snapshot(backup) + backup.copy_(torch.tensor([3.0, 4.0], dtype=torch.bfloat16)) + + indices, values = export.local_bit_exact_diff(backup, snapshot) + assert indices.tolist() == [0, 1] + assert values.tolist() == [3.0, 4.0] + + +class _SplitProbe: + def megatron_to_hf(self, tensor, _module): + # Mimic a fused Megatron parameter split into two final HF tensors. + return {"q.weight": tensor[:2], "k.weight": tensor[2:]} + + +def test_sparse_hf_entry_splits_fused_parameter() -> None: + record = export.SparseExportRecord( + megatron_name="decoder.qkv.weight", + weight_key="vp_stages.0.decoder.qkv.weight", + param=torch.empty(4, dtype=torch.bfloat16), + gather_group=None, + contributes=True, + probe=_SplitProbe(), + module=object(), + ) + cache = {} + + slots, counts, indices, values = export.sparse_hf_entry( + record, + torch.tensor([1, 3]), + torch.tensor([2.0, 4.0], dtype=torch.bfloat16), + cache, + ) + + assert slots == [("q.weight", (2,)), ("k.weight", (2,))] + assert counts.tolist() == [1, 1] + assert indices.tolist() == [1, 1] + assert values.tolist() == [2.0, 4.0] + + +class _GlobalOffsetProbe: + def megatron_to_hf(self, tensor, _module): + missing = torch.full_like(tensor, float("nan")) + return {"proj.weight": torch.cat((missing, tensor))} + + +def test_sparse_hf_entry_preserves_final_hf_global_indices() -> None: + record = export.SparseExportRecord( + megatron_name="decoder.proj.weight", + weight_key="vp_stages.0.decoder.proj.weight", + param=torch.empty(3, dtype=torch.bfloat16), + gather_group=object(), + contributes=True, + probe=_GlobalOffsetProbe(), + module=object(), + ) + + slots, counts, indices, values = export.sparse_hf_entry( + record, + torch.tensor([0, 2]), + torch.tensor([5.0, 7.0], dtype=torch.bfloat16), + {}, + ) + + assert slots == [("proj.weight", (6,))] + assert counts.tolist() == [2] + assert indices.tolist() == [3, 5] + assert values.tolist() == [5.0, 7.0] + + +def _mapping(class_name, **attributes): + mapping = type(class_name, (), {})() + for name, value in attributes.items(): + setattr(mapping, name, value) + return mapping + + +def _fast_record(mapping, shape): + return SimpleNamespace( + mapping=mapping, + param=torch.empty(shape, dtype=torch.bfloat16), + module=SimpleNamespace(), + megatron_name="weight", + slots=None, + ) + + +def test_fast_column_and_row_coordinate_mapping() -> None: + column = _fast_record( + _mapping("ColumnParallelMapping", hf_param="column", tp_rank=2), + (2, 3), + ) + column.slots = [("column", (8, 3))] + _slots, counts, indices, values = export.sparse_hf_entry( + column, + torch.tensor([0, 5]), + torch.tensor([1.0, 2.0], dtype=torch.bfloat16), + {}, + ) + assert counts.tolist() == [2] + assert indices.tolist() == [12, 17] + assert values.tolist() == [1.0, 2.0] + + row = _fast_record( + _mapping("RowParallelMapping", hf_param="row", tp_rank=1), + (2, 3), + ) + row.slots = [("row", (2, 12))] + _slots, counts, indices, _values = export.sparse_hf_entry( + row, + torch.tensor([0, 5]), + torch.tensor([1.0, 2.0], dtype=torch.bfloat16), + {}, + ) + assert counts.tolist() == [2] + assert indices.tolist() == [3, 17] + + +def test_fast_gated_mlp_coordinate_mapping() -> None: + mapping = _mapping( + "GatedMLPMapping", + hf_param={"gate": "gate", "up": "up"}, + tp_rank=1, + ) + record = _fast_record(mapping, (4, 3)) + record.slots = [("gate", (4, 3)), ("up", (4, 3))] + + _slots, counts, indices, values = export.sparse_hf_entry( + record, + torch.tensor([0, 7, 11]), + torch.tensor([1.0, 2.0, 3.0], dtype=torch.bfloat16), + {}, + ) + + assert counts.tolist() == [1, 2] + assert indices.tolist() == [6, 7, 11] + assert values.tolist() == [1.0, 2.0, 3.0] + + +def test_fast_qkv_coordinate_mapping() -> None: + config = SimpleNamespace( + num_attention_heads=4, + num_query_groups=2, + kv_channels=2, + hidden_size=4, + attention_output_gate=False, + ) + mapping = _mapping( + "QKVMapping", + hf_param={"q": "q", "k": "k", "v": "v"}, + tp_rank=1, + _get_config=lambda _module: config, + ) + # Global packed rows are [q0, q1, k0, v0, q2, q3, k1, v1], with + # two scalar rows per head. TP rank 1 owns the second four heads. + record = _fast_record(mapping, (8, 4)) + record.slots = [("q", (8, 4)), ("k", (4, 4)), ("v", (4, 4))] + local_rows = torch.tensor([0, 2, 4, 6]) + local_indices = local_rows * 4 + 1 + + _slots, counts, indices, values = export.sparse_hf_entry( + record, + local_indices, + torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.bfloat16), + {}, + ) + + assert counts.tolist() == [2, 1, 1] + assert indices.tolist() == [17, 25, 9, 9] + assert values.tolist() == [1.0, 2.0, 3.0, 4.0] + + +def test_exchange_slot_tables_uses_gloo_control_group(monkeypatch) -> None: + gloo_group = object() + seen = {} + record = SimpleNamespace(megatron_name="weight", slots=None) + cache = {"weight": [("model.weight", (2, 2))]} + + monkeypatch.setattr(export, "get_gloo_group", lambda: gloo_group) + monkeypatch.setattr(export.dist, "get_world_size", lambda: 1) + + def fake_all_gather_object(output, value, group=None): + seen["group"] = group + output[0] = value + + monkeypatch.setattr(export.dist, "all_gather_object", fake_all_gather_object) + + export._exchange_slot_tables([record], cache) + + assert seen["group"] is gloo_group + assert record.slots == [("model.weight", (2, 2))] diff --git a/tests/unit/backends/megatron_utils/update_weight/test_sparse_gather.py b/tests/unit/backends/megatron_utils/update_weight/test_sparse_gather.py new file mode 100644 index 000000000..8c1d6b4b6 --- /dev/null +++ b/tests/unit/backends/megatron_utils/update_weight/test_sparse_gather.py @@ -0,0 +1,61 @@ +import importlib.util +from pathlib import Path + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +MODULE_PATH = ( + Path(__file__).parents[5] + / "vime/backends/megatron_utils/update_weight/sparse_gather.py" +) +SPEC = importlib.util.spec_from_file_location("test_sparse_gather_module", MODULE_PATH) +assert SPEC is not None and SPEC.loader is not None +sparse_gather = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(sparse_gather) + + +def _gather_worker(rank: int, world: int, init_file: str) -> None: + dist.init_process_group( + "gloo", init_method=f"file://{init_file}", rank=rank, world_size=world + ) + inputs = [ + ([1, 0, 1], [0, 2], [10.0, 12.0]), + ([0, 1, 0], [1], [21.0]), + ([0, 0, 0], [], []), + ([1, 1, 1], [3, 4, 5], [33.0, 34.0, 35.0]), + ] + count_values, index_values, value_values = inputs[rank] + + result = sparse_gather.gather_slot_entries_to_rank0( + torch.tensor(index_values, dtype=torch.int64), + torch.tensor(value_values, dtype=torch.float32), + torch.tensor(count_values, dtype=torch.int64), + dist.group.WORLD, + ) + + if rank == 0: + assert result is not None + assert [part[0].tolist() for part in result] == [ + [0, 3], + [1, 4], + [2, 5], + ] + assert [part[1].tolist() for part in result] == [ + [10.0, 33.0], + [21.0, 34.0], + [12.0, 35.0], + ] + else: + assert result is None + dist.destroy_process_group() + + +def test_variable_length_p2p_gather_only_materializes_on_rank0(tmp_path) -> None: + init_file = tmp_path / "gloo-init" + mp.spawn( + _gather_worker, + args=(4, str(init_file)), + nprocs=4, + join=True, + ) diff --git a/tests/unit/backends/vllm_utils/conftest.py b/tests/unit/backends/vllm_utils/conftest.py index ec0d8dcb8..d1ab01589 100644 --- a/tests/unit/backends/vllm_utils/conftest.py +++ b/tests/unit/backends/vllm_utils/conftest.py @@ -23,6 +23,9 @@ def vllm_args() -> SimpleNamespace: use_critic=False, critic_num_gpus_per_node=0, critic_num_nodes=0, + update_weight_disk_dir="/shared/weights", + update_weight_local_checkpoint_dir="/local/weights", + custom_update_weight_pre_read_path=None, ) diff --git a/tests/unit/backends/vllm_utils/test_vllm_engine.py b/tests/unit/backends/vllm_utils/test_vllm_engine.py index 579da57ee..b356cdda2 100644 --- a/tests/unit/backends/vllm_utils/test_vllm_engine.py +++ b/tests/unit/backends/vllm_utils/test_vllm_engine.py @@ -353,6 +353,85 @@ def fake_post_vllm(update_info: dict) -> dict: assert vllm_engine._weight_version == "7" +@pytest.mark.unit +def test_update_sparse_weights_from_distributed_posts_counts(vllm_engine, monkeypatch): + calls = [] + monkeypatch.setattr(vllm_engine, "_post_vllm_update_weights_http", lambda info: calls.append(info) or {"ok": True}) + + vllm_engine.update_sparse_weights_from_distributed( + ["model.embed_tokens.weight"], + [torch.bfloat16], + [torch.Size([8, 4])], + [3], + group_name="vime-sparse-hccl", + weight_version="9", + ) + + assert calls == [ + { + "names": ["model.embed_tokens.weight"], + "dtype_names": ["bfloat16"], + "shapes": [[8, 4]], + "num_updates_list": [3], + } + ] + assert vllm_engine._weight_version == "9" + + +@pytest.mark.unit +def test_pull_weights_posts_collective_rpc_and_records_version(vllm_engine, monkeypatch): + calls = [] + + def fake_post(url, *, json=None, timeout=None): + calls.append((url, json, timeout)) + return _MockResponse(json_data={"success": True}) + + monkeypatch.setattr(mod.requests, "post", fake_post) + + assert vllm_engine.pull_weights(8) == {"success": True} + assert calls == [ + ( + "http://127.0.0.1:8765/collective_rpc", + { + "method": "pull_weights", + "kwargs": { + "local_checkpoint_dir": "/local/weights", + "source_dir": "/shared/weights", + "target_version": 8, + "pre_read_hook": None, + }, + }, + 600, + ) + ] + assert vllm_engine._weight_version == "8" + + +@pytest.mark.unit +def test_pull_weights_failure_does_not_advance_version(vllm_engine, monkeypatch): + vllm_engine._weight_version = "old" + + def fake_post(url, *, json=None, timeout=None): + del url, json, timeout + return _MockResponse(status_code=500, text="failed") + + monkeypatch.setattr(mod.requests, "post", fake_post) + with pytest.raises(requests.exceptions.HTTPError): + vllm_engine.pull_weights(8) + assert vllm_engine._weight_version == "old" + + +@pytest.mark.unit +def test_disk_reload_records_version_only_after_success(vllm_engine, monkeypatch): + def fake_post(url, *, json=None, timeout=None): + del url, json, timeout + return _MockResponse(json_data={"reloaded": True}) + + monkeypatch.setattr(mod.requests, "post", fake_post) + assert vllm_engine.update_weights_from_disk("/local/weights", weight_version="9") == {"reloaded": True} + assert vllm_engine._weight_version == "9" + + @pytest.mark.unit def test_post_vllm_update_weights_http_wraps_update_info(vllm_engine, monkeypatch): seen: list[tuple] = [] @@ -680,5 +759,7 @@ def _boom(*a, **k): assert vllm_engine.finish_weight_update() is None assert vllm_engine.init_weights_update_group("addr", 1, 0, 4, "g", "nccl") is None assert vllm_engine.update_weights_from_distributed(["w"], [torch.float32], [[1]], "g") is None + assert vllm_engine.pull_weights(1) is None + assert vllm_engine.update_weights_from_disk("/local/weights", weight_version="1") is None assert vllm_engine.release_memory_occupation() is None assert vllm_engine.resume_memory_occupation() is None diff --git a/tests/unit/utils/test_disk_delta.py b/tests/unit/utils/test_disk_delta.py new file mode 100644 index 000000000..f9f288f9b --- /dev/null +++ b/tests/unit/utils/test_disk_delta.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import numpy as np +import safetensors.numpy + +from vime.utils.disk_delta import checksum, make_tensor_reader, overwrite_encode + + +def test_overwrite_encode_contains_changed_positions_and_values(): + old = np.array([1, 2, 3, 4], dtype=np.uint8) + new = np.array([1, 9, 3, 8], dtype=np.uint8) + + encoded = overwrite_encode(new, new != old) + + count = int.from_bytes(encoded[:4].tobytes(), "little") + positions = np.frombuffer(encoded[4 : 4 + count * 4].tobytes(), dtype=" None: + calls = [] + + if supports_keyword: + + def target(self, is_checkpoint_format: bool = True) -> None: + calls.append((self, is_checkpoint_format)) + + else: + + def target(self) -> None: + calls.append((self, "without_keyword")) + + receiver = object() + call_with_optional_keyword( + target, + receiver, + keyword="is_checkpoint_format", + value=False, + ) + + assert calls == [(receiver, False if supports_keyword else "without_keyword")] diff --git a/tests/unit/utils/test_tensor_backper.py b/tests/unit/utils/test_tensor_backper.py new file mode 100644 index 000000000..890bfaab3 --- /dev/null +++ b/tests/unit/utils/test_tensor_backper.py @@ -0,0 +1,33 @@ +import torch + +from vime.utils.tensor_backper import TensorBackuper + + +def test_double_buffer_keeps_previous_backup_immutable(monkeypatch): + source = torch.tensor([1.0, 2.0]) + original_empty_like = torch.empty_like + monkeypatch.setattr( + torch, + "empty_like", + lambda tensor, **kwargs: original_empty_like(tensor, device="cpu"), + ) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + backuper = TensorBackuper.create( + source_getter=lambda: [("weight", source)], single_tag=None + ) + backuper.enable_double_buffer("actor") + + backuper.backup("actor") + first = backuper.get("actor")["weight"] + source.add_(10) + backuper.backup("actor") + second = backuper.get("actor")["weight"] + + assert first.tolist() == [1.0, 2.0] + assert second.tolist() == [11.0, 12.0] + assert first.data_ptr() != second.data_ptr() + + source.add_(10) + backuper.backup("actor") + assert second.tolist() == [11.0, 12.0] + assert backuper.get("actor")["weight"].tolist() == [21.0, 22.0] diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index e8781cb09..3845ffc96 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -52,7 +52,9 @@ def _safe_empty_cache(): from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values from .model import forward_only, initialize_model_and_optimizer, save, train from .update_weight.common import named_params_and_buffers +from .update_weight.update_weight_from_disk_delta import UpdateWeightFromDiskDelta from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed +from .update_weight.update_weight_from_sparse_distributed import UpdateWeightFromSparseDistributed from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor logging.getLogger("megatron").setLevel(logging.WARNING) @@ -140,6 +142,11 @@ def init( ), single_tag=None, ) + if getattr(self.args, "update_weight_mode", "full") == "sparse": + # Sparse diff keeps the previous actor backup as its transactional + # baseline. Alternate two pinned CPU buffers so the next backup + # cannot mutate that baseline and no full-model clone is needed. + self.weights_backuper.enable_double_buffer("actor") self._active_model_tag: str | None = "actor" self.weights_backuper.backup("actor") @@ -163,7 +170,11 @@ def init( hf_vocab = getattr(self.hf_config, "vocab_size", None) self.args.vocab_size = hf_vocab if hf_vocab is not None else self.tokenizer.vocab_size - if self.args.colocate: + if getattr(self.args, "update_weight_mode", "full") == "delta": + update_weight_cls = UpdateWeightFromDiskDelta + elif self.args.update_weight_mode == "sparse": + update_weight_cls = UpdateWeightFromSparseDistributed + elif self.args.colocate: update_weight_cls = UpdateWeightFromTensor else: update_weight_cls = UpdateWeightFromDistributed diff --git a/vime/backends/megatron_utils/update_weight/megatron_sparse_export.py b/vime/backends/megatron_utils/update_weight/megatron_sparse_export.py new file mode 100644 index 000000000..dc246b05e --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/megatron_sparse_export.py @@ -0,0 +1,463 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Shard-local sparse export through Megatron-Bridge parameter mappings. + +This follows verl's Megatron delta exporter: each rank runs a communication- +free copy of the real Bridge mapping with its local shard inserted among NaN +placeholders. Non-NaN survivors are that rank's contribution in final HF +coordinates, including QKV and gate/up rearrangements. +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +from megatron.core import mpu + +from vime.utils.distributed_utils import get_gloo_group + + +class _ProbeGroup: + """Process-group stand-in that preserves only size and rank.""" + + def __init__(self, size: int, rank: int): + self._size = int(size) + self._rank = int(rank) + + def size(self) -> int: + return self._size + + def rank(self) -> int: + return self._rank + + def __getattr__(self, name): + raise RuntimeError( + f"Sparse Bridge probe attempted unstubbed communication {name!r}" + ) + + +_NAN_POOL: dict[tuple, torch.Tensor] = {} + + +def _nan_block(shape, dtype, device) -> torch.Tensor: + key = (tuple(shape), dtype, str(device)) + tensor = _NAN_POOL.get(key) + if tensor is None: + tensor = torch.full( + tuple(shape), float("nan"), dtype=dtype, device=device + ) + _NAN_POOL[key] = tensor + return tensor + + +def _warm_lazy_mapping(mapping, module) -> None: + if ( + hasattr(mapping, "_detect_parallelism_type") + and getattr(mapping, "_mapping", None) is None + ): + parallelism = mapping._detect_parallelism_type(module) + mapping._mapping = mapping._get_or_create_mapping(parallelism) + mapping._detected_type = parallelism + + +def make_probe(mapping, module): + """Copy a Bridge mapping and replace its collectives with local synthesis.""" + from megatron.bridge.models.conversion.param_mapping import ( + MegatronParamMapping, + ) + from megatron.core.utils import get_pg_rank, get_pg_size + + _warm_lazy_mapping(mapping, module) + + def _stub(mapping_copy): + mapping_copy.pp_group = _ProbeGroup(1, 0) + for attr in ("ep_group", "_tp_group", "_etp_group"): + group = getattr(mapping_copy, attr, None) + if group is None: + setattr(mapping_copy, attr, _ProbeGroup(1, 0)) + else: + setattr( + mapping_copy, + attr, + _ProbeGroup(get_pg_size(group), get_pg_rank(group)), + ) + + def _gather_tp(tensor, owner=mapping_copy): + missing = _nan_block(tensor.shape, tensor.dtype, tensor.device) + gathered = [missing] * owner.tp_size + gathered[owner.tp_rank] = tensor + return gathered + + mapping_copy.gather_from_tp_ranks = _gather_tp + mapping_copy.gather_from_ep_ranks = ( + lambda weight, _module, name: {str(name): weight} + ) + mapping_copy.gather_from_ep_ranks_scale = ( + lambda weight, _module, name: { + str(name): weight.unsqueeze(0).squeeze().unsqueeze(-1) + } + ) + return mapping_copy + + def _inject(node): + result = _stub(copy.copy(node)) + for attr, value in list(vars(result).items()): + if isinstance(value, MegatronParamMapping): + _warm_lazy_mapping(value, module) + setattr(result, attr, _inject(value)) + return result + + return _inject(mapping) + + +@dataclass +class SparseExportRecord: + megatron_name: str + weight_key: str + param: torch.Tensor + gather_group: dist.ProcessGroup | None + contributes: bool + probe: Any + mapping: Any = None + module: Any = None + slots: list[tuple[str, tuple[int, ...]]] | None = None + + +def build_sparse_export_index( + bridge, + model, + local_weights: dict[str, torch.Tensor], + slot_cache: dict[str, list[tuple[str, tuple[int, ...]]]], +) -> list[SparseExportRecord]: + """Build the static Bridge directory and local probe for every parameter.""" + if mpu.get_pipeline_model_parallel_world_size() != 1: + raise NotImplementedError( + "Sparse shard export currently supports Megatron PP=1" + ) + if mpu.get_expert_model_parallel_world_size() != 1: + raise NotImplementedError( + "Sparse shard export currently supports Megatron EP=1" + ) + + tp_group = mpu.get_tensor_model_parallel_group() + tp_world = dist.get_world_size(group=tp_group) + tp_rank = mpu.get_tensor_model_parallel_rank() + dp_rank = mpu.get_data_parallel_rank(with_context_parallel=True) + records: list[SparseExportRecord] = [] + for task in bridge.get_conversion_tasks(model): + if task.param_weight is None: + continue + weight_key = f"vp_stages.{task.vp_stage}.{task.param_name}" + if weight_key not in local_weights: + raise KeyError( + f"Megatron sparse export weight {weight_key!r} is unavailable" + ) + param = task.param_weight + tp_sharded = ( + getattr(param, "tensor_model_parallel", False) and tp_world > 1 + ) + records.append( + SparseExportRecord( + megatron_name=task.global_param_name, + weight_key=weight_key, + param=param, + gather_group=tp_group if tp_sharded else None, + contributes=dp_rank == 0 and (tp_sharded or tp_rank == 0), + probe=make_probe(task.mapping, task.megatron_module), + mapping=task.mapping, + module=task.megatron_module, + ) + ) + + _exchange_slot_tables(records, slot_cache) + return records + + +def _exchange_slot_tables( + records: list[SparseExportRecord], + slot_cache: dict[str, list[tuple[str, tuple[int, ...]]]], +) -> None: + """Make every directory row use an identical final-HF slot table.""" + local_rows = [] + for record in records: + if record.megatron_name not in slot_cache: + empty_idx = torch.empty( + 0, dtype=torch.int64, device=record.param.device + ) + empty_val = torch.empty( + 0, dtype=record.param.dtype, device=record.param.device + ) + sparse_hf_entry(record, empty_idx, empty_val, slot_cache) + local_rows.append(slot_cache[record.megatron_name]) + + gathered: list = [None] * dist.get_world_size() + # The directory contains Python strings/shapes and belongs on the control + # plane. Do not let ``all_gather_object`` fall back to the HCCL world + # group: besides staging a potentially large object through NPU memory, + # some torch/HCCL combinations do not support object collectives at all. + dist.all_gather_object(gathered, local_rows, group=get_gloo_group()) + if not all(len(rows) == len(local_rows) for rows in gathered): + raise RuntimeError("Megatron sparse Bridge directories differ by rank") + for row_index, record in enumerate(records): + union: dict[tuple[str, tuple[int, ...]], None] = {} + for rows in gathered: + for name, shape in rows[row_index]: + union[(name, tuple(shape))] = None + record.slots = list(union) + + +def sparse_hf_entry( + record: SparseExportRecord, + local_indices: torch.Tensor, + local_values: torch.Tensor, + slot_cache: dict[str, list[tuple[str, tuple[int, ...]]]], +) -> tuple[ + list[tuple[str, tuple[int, ...]]], + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Map one local sparse shard patch to final HF names and coordinates.""" + slots = record.slots or slot_cache.get(record.megatron_name) + if local_indices.numel() == 0 and slots is not None: + return ( + slots, + torch.zeros(len(slots), dtype=torch.int64), + torch.empty( + 0, dtype=torch.int32, device=local_values.device + ), + torch.empty( + 0, dtype=local_values.dtype, device=local_values.device + ), + ) + + if slots is not None: + fast_entry = _sparse_hf_entry_fast( + record, slots, local_indices, local_values + ) + if fast_entry is not None: + return fast_entry + + buffer = torch.full( + tuple(record.param.shape), + float("nan"), + dtype=local_values.dtype, + device=local_values.device, + ) + if local_indices.numel(): + buffer.view(-1)[local_indices] = local_values + outputs = record.probe.megatron_to_hf(buffer, record.module) + + if slots is None: + slots = [ + (name, tuple(int(dim) for dim in tensor.shape)) + for name, tensor in outputs.items() + ] + slot_cache[record.megatron_name] = slots + unknown = set(outputs) - {name for name, _ in slots} + if unknown: + raise RuntimeError( + f"Bridge probe emitted unknown HF slots: {sorted(unknown)}" + ) + + counts = torch.zeros(len(slots), dtype=torch.int64) + index_parts: list[torch.Tensor] = [] + value_parts: list[torch.Tensor] = [] + for slot_index, (name, _shape) in enumerate(slots): + output = outputs.get(name) + if output is None: + continue + flat = output.reshape(-1) + indices = (~torch.isnan(flat)).nonzero(as_tuple=False).view(-1) + if indices.numel(): + counts[slot_index] = indices.numel() + index_parts.append(indices.to(torch.int32)) + value_parts.append(flat[indices]) + + if index_parts: + return slots, counts, torch.cat(index_parts), torch.cat(value_parts) + return ( + slots, + counts, + torch.empty(0, dtype=torch.int32, device=local_values.device), + torch.empty( + 0, dtype=local_values.dtype, device=local_values.device + ), + ) + + +def _sparse_hf_entry_fast( + record: SparseExportRecord, + slots: list[tuple[str, tuple[int, ...]]], + local_indices: torch.Tensor, + local_values: torch.Tensor, +) -> tuple[ + list[tuple[str, tuple[int, ...]]], + torch.Tensor, + torch.Tensor, + torch.Tensor, +] | None: + """Map common dense-Megatron layouts without materializing NaN tensors. + + Qwen3 uses only Auto(Column/Row/Replicated), QKV and GatedMLP mappings. + These transforms rearrange coordinates, so changed entries can be routed + directly in O(number of changes), instead of building and scanning dense + full-HF probe outputs on every TP rank. Unknown Bridge mappings retain the + generic probe fallback above. + """ + mapping = getattr(record, "mapping", None) + if mapping is None: + return None + mapping_name = type(mapping).__name__ + if mapping_name == "AutoMapping" or hasattr(mapping, "_mapping"): + concrete = getattr(mapping, "_mapping", None) + if concrete is not None: + mapping = concrete + mapping_name = type(mapping).__name__ + + slot_by_name = {str(name): index for index, (name, _shape) in enumerate(slots)} + tp_rank = int(getattr(mapping, "tp_rank", 0)) + contributions: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + + if mapping_name in {"DirectMapping", "ReplicatedMapping"}: + slot = slot_by_name.get(str(mapping.hf_param)) + if slot is None: + return None + contributions[slot] = (local_indices.to(torch.int32), local_values) + + elif mapping_name == "ColumnParallelMapping": + slot = slot_by_name.get(str(mapping.hf_param)) + if slot is None: + return None + global_indices = local_indices + tp_rank * record.param.numel() + contributions[slot] = (global_indices.to(torch.int32), local_values) + + elif mapping_name == "RowParallelMapping": + slot = slot_by_name.get(str(mapping.hf_param)) + if slot is None: + return None + if record.param.ndim <= 1: + global_indices = local_indices + else: + local_width = record.param.shape[1] + full_width = slots[slot][1][1] + row = torch.div(local_indices, local_width, rounding_mode="floor") + column = local_indices.remainder(local_width) + global_indices = ( + row * full_width + tp_rank * local_width + column + ) + contributions[slot] = (global_indices.to(torch.int32), local_values) + + elif mapping_name == "GatedMLPMapping": + if record.param.shape[0] % 2: + return None + gate_slot = slot_by_name.get(str(mapping.hf_param["gate"])) + up_slot = slot_by_name.get(str(mapping.hf_param["up"])) + if gate_slot is None or up_slot is None: + return None + width = record.param.numel() // record.param.shape[0] + half_rows = record.param.shape[0] // 2 + row = torch.div(local_indices, width, rounding_mode="floor") + inner = local_indices.remainder(width) + gate_mask = row < half_rows + for slot, mask, local_row in ( + (gate_slot, gate_mask, row), + (up_slot, ~gate_mask, row - half_rows), + ): + global_indices = ( + (tp_rank * half_rows + local_row[mask]) * width + inner[mask] + ) + contributions[slot] = ( + global_indices.to(torch.int32), + local_values[mask], + ) + + elif mapping_name == "QKVMapping": + config = mapping._get_config(record.module) + if getattr(config, "attention_output_gate", False): + return None + q_slot = slot_by_name.get(str(mapping.hf_param["q"])) + k_slot = slot_by_name.get(str(mapping.hf_param["k"])) + v_slot = slot_by_name.get(str(mapping.hf_param["v"])) + if q_slot is None or k_slot is None or v_slot is None: + return None + head_count = int(config.num_attention_heads) + group_count = int(config.num_query_groups) + heads_per_group = head_count // group_count + head_size = int(config.kv_channels or (config.hidden_size // head_count)) + width = record.param.numel() // record.param.shape[0] + local_row = torch.div(local_indices, width, rounding_mode="floor") + inner = local_indices.remainder(width) + packed_row = tp_rank * record.param.shape[0] + local_row + packed_head = torch.div(packed_row, head_size, rounding_mode="floor") + head_inner = packed_row.remainder(head_size) + group_width = heads_per_group + 2 + group = torch.div(packed_head, group_width, rounding_mode="floor") + position = packed_head.remainder(group_width) + q_mask = position < heads_per_group + k_mask = position == heads_per_group + v_mask = position == heads_per_group + 1 + q_row = (group * heads_per_group + position) * head_size + head_inner + kv_row = group * head_size + head_inner + for slot, mask, output_row in ( + (q_slot, q_mask, q_row), + (k_slot, k_mask, kv_row), + (v_slot, v_mask, kv_row), + ): + contributions[slot] = ( + (output_row[mask] * width + inner[mask]).to(torch.int32), + local_values[mask], + ) + else: + return None + + counts = torch.zeros(len(slots), dtype=torch.int64) + index_parts: list[torch.Tensor] = [] + value_parts: list[torch.Tensor] = [] + for slot in range(len(slots)): + indices, values = contributions.get( + slot, + ( + torch.empty(0, dtype=torch.int32, device=local_indices.device), + torch.empty(0, dtype=local_values.dtype, device=local_values.device), + ), + ) + counts[slot] = indices.numel() + index_parts.append(indices) + value_parts.append(values) + return slots, counts, torch.cat(index_parts), torch.cat(value_parts) + + +_INTEGER_DTYPE = { + 1: torch.uint8, + 2: torch.int16, + 4: torch.int32, + 8: torch.int64, +} + + +def clone_cpu_snapshot(tensor: torch.Tensor) -> torch.Tensor: + """Create an owning CPU snapshot that cannot alias a mutable weight backup.""" + return tensor.detach().cpu().contiguous().clone() + + +def local_bit_exact_diff( + current: torch.Tensor, snapshot: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """Return changed local flat indices and current values, losslessly.""" + current = current.detach().contiguous().view(-1) + snapshot = snapshot.detach().contiguous().view(-1) + if current.shape != snapshot.shape or current.dtype != snapshot.dtype: + raise RuntimeError("Megatron sparse local snapshot shape/dtype mismatch") + integer_dtype = _INTEGER_DTYPE.get(current.element_size()) + if integer_dtype is None: + raise NotImplementedError( + f"Unsupported sparse local dtype {current.dtype}" + ) + changed = current.view(integer_dtype) != snapshot.view(integer_dtype) + indices = changed.nonzero(as_tuple=False).view(-1) + return indices, current[indices] diff --git a/vime/backends/megatron_utils/update_weight/sparse_gather.py b/vime/backends/megatron_utils/update_weight/sparse_gather.py new file mode 100644 index 000000000..96bc0d65e --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/sparse_gather.py @@ -0,0 +1,161 @@ +# Copyright 2025 Bytedance Ltd. and/or its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Variable-length sparse gather adapted from verl delta weight sync.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + + +def gather_slot_entries_to_rank0( + indices: torch.Tensor, + values: torch.Tensor, + counts: torch.Tensor, + group: dist.ProcessGroup, + max_round_bytes: int | None = None, +) -> list[tuple[torch.Tensor, torch.Tensor]] | None: + """Gather variable-length sparse slots only to group rank 0. + + HCCL implements ``dist.gather`` with an all-gather fallback. That copies every + rank's sparse payload to every training rank and also requires padding all + payloads to the largest rank. Exchange the small per-slot counts collectively, + then use exact-sized point-to-point transfers for the payload instead. + """ + rank = dist.get_rank(group) + world = dist.get_world_size(group) + destination = dist.get_global_rank(group, 0) + device = indices.device + slot_count = int(counts.numel()) + + all_counts = [torch.zeros_like(counts) for _ in range(world)] + dist.all_gather(all_counts, counts.to(device), group=group) + counts_cpu = torch.stack(all_counts).cpu().tolist() + + if max_round_bytes is not None and slot_count > 1: + bytes_per_entry = indices.element_size() + values.element_size() + budget = max(int(max_round_bytes) // bytes_per_entry, 1) + cuts = [0] + running = [0] * world + for slot_index in range(slot_count): + running = [ + running[r] + counts_cpu[r][slot_index] + for r in range(world) + ] + if max(running) > budget and cuts[-1] != slot_index: + cuts.append(slot_index) + running = [ + counts_cpu[r][slot_index] for r in range(world) + ] + cuts.append(slot_count) + if len(cuts) > 2: + offsets = [0] + for count in counts_cpu[rank]: + offsets.append(offsets[-1] + count) + output = [] + for start, end in zip(cuts[:-1], cuts[1:], strict=True): + sub_counts = torch.tensor( + counts_cpu[rank][start:end], + dtype=torch.int64, + device=device, + ) + gathered = gather_slot_entries_to_rank0( + indices[offsets[start] : offsets[end]], + values[offsets[start] : offsets[end]], + sub_counts, + group, + ) + if rank == 0: + output.extend(gathered) + return output if rank == 0 else None + + totals = [sum(row) for row in counts_cpu] + max_entries = max(totals) if totals else 0 + if max_entries == 0: + if rank != 0: + return None + return [ + ( + torch.empty(0, dtype=indices.dtype, device=device), + torch.empty(0, dtype=values.dtype, device=device), + ) + for _ in range(slot_count) + ] + + if rank == 0: + index_list = [indices] + value_list = [values] + p2p_ops = [] + for group_rank in range(1, world): + entry_count = totals[group_rank] + if entry_count == 0: + index_list.append( + torch.empty(0, dtype=indices.dtype, device=device) + ) + value_list.append( + torch.empty(0, dtype=values.dtype, device=device) + ) + continue + index_buffer = torch.empty( + entry_count, dtype=indices.dtype, device=device + ) + value_buffer = torch.empty( + entry_count, dtype=values.dtype, device=device + ) + index_list.append(index_buffer) + value_list.append(value_buffer) + peer = dist.get_global_rank(group, group_rank) + p2p_ops.extend( + [ + dist.P2POp(dist.irecv, index_buffer, peer, group), + dist.P2POp(dist.irecv, value_buffer, peer, group), + ] + ) + elif totals[rank] > 0: + p2p_ops = [ + dist.P2POp(dist.isend, indices, destination, group), + dist.P2POp(dist.isend, values, destination, group), + ] + else: + p2p_ops = [] + + if p2p_ops: + for request in dist.batch_isend_irecv(p2p_ops): + request.wait() + if rank != 0: + return None + + offsets = [[0] * (slot_count + 1) for _ in range(world)] + for rank_index in range(world): + for slot_index in range(slot_count): + offsets[rank_index][slot_index + 1] = ( + offsets[rank_index][slot_index] + + counts_cpu[rank_index][slot_index] + ) + + output = [] + for slot_index in range(slot_count): + index_parts = [ + index_list[r][ + offsets[r][slot_index] : offsets[r][slot_index + 1] + ] + for r in range(world) + if counts_cpu[r][slot_index] + ] + value_parts = [ + value_list[r][ + offsets[r][slot_index] : offsets[r][slot_index + 1] + ] + for r in range(world) + if counts_cpu[r][slot_index] + ] + if index_parts: + output.append((torch.cat(index_parts), torch.cat(value_parts))) + else: + output.append( + ( + torch.empty(0, dtype=indices.dtype, device=device), + torch.empty(0, dtype=values.dtype, device=device), + ) + ) + return output diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py b/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py new file mode 100644 index 000000000..47e587227 --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_disk_delta.py @@ -0,0 +1,318 @@ +from __future__ import annotations + +import json +import logging +import os +import queue +import shutil +import time +from argparse import Namespace +from collections import deque +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import ray +import safetensors.numpy +import torch +import torch.distributed as dist +import zstandard +from ray.actor import ActorHandle + +from vime.utils.disk_delta import NUM_WORKERS, checksum, make_tensor_reader, overwrite_encode +from vime.utils.distributed_utils import get_gloo_group + +from .hf_weight_iterator_base import HfWeightIteratorBase + +logger = logging.getLogger(__name__) + + +class UpdateWeightFromDiskDelta: + """Publish byte-level HF weight deltas and reload them through local checkpoints. + + This is intentionally independent of ``UpdateWeightFromDistributed``: all + training ranks still participate in the Ascend PP/TP/EP conversion iterator, + while only global rank zero publishes the canonical result. No HCCL group is + created for this transport. + """ + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + self.args = args + self.weights_getter = weights_getter + self.weight_version = 0 + self.update_weight_metrics: dict[str, float] = {} + self.rollout_engines: list[ActorHandle] = [] + self.delta_dir = args.update_weight_disk_dir + self.delta_encoding = args.update_weight_delta_encoding + self.checksum_algorithm = args.update_weight_delta_checksum + self._iterator = HfWeightIteratorBase.create( + args=args, + model=model, + model_name=model_name, + quantization_config=quantization_config, + ) + self._snapshot: dict[str, np.ndarray] = {} + self._baseline_captured = False + self._post_write_hook: Callable | None = None + if args.custom_update_weight_post_write_path: + from vime.utils.misc import load_function + + self._post_write_hook = load_function(args.custom_update_weight_post_write_path) + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + ) -> None: + del rollout_engine_lock, engine_gpu_counts, engine_gpu_offsets + self.rollout_engines = list(rollout_engines) + + def disconnect_rollout_engines(self) -> None: + return + + def pop_metrics(self) -> dict[str, float]: + metrics, self.update_weight_metrics = self.update_weight_metrics, {} + return metrics + + @torch.no_grad() + def update_weights(self) -> None: + if not self._baseline_captured: + self._capture_baseline() + self._baseline_captured = True + return + + self.weight_version += 1 + self._stage_metrics: dict[str, float] = {} + self._publish() + self._reload_engines() + self._record_metrics() + + def _capture_baseline(self) -> None: + """Use the serving checkpoint as the byte-exact version-zero baseline.""" + pulls = [] + if dist.get_rank() == 0: + shutil.rmtree(self.delta_dir, ignore_errors=True) + os.makedirs(self.delta_dir, exist_ok=True) + if self._post_write_hook is not None: + self._post_write_hook(self.args, self.delta_dir, self.rollout_engines) + pulls = [engine.pull_weights.remote(target_version=0) for engine in self.rollout_engines] + dist.barrier(group=get_gloo_group()) + + read_hf = make_tensor_reader(self.args.hf_checkpoint) + for name, tensor in self._iter_hf_tensors(progress_desc="Capture disk delta baseline"): + try: + self._snapshot[name] = read_hf(name) + except KeyError: + self._snapshot[name] = _tensor_bytes(tensor) + logger.warning("delta baseline: %s absent from hf_checkpoint; using current converted weight", name) + + if dist.get_rank() == 0: + ray.get(pulls) + logger.info("[disk delta] captured baseline for %d tensors", len(self._snapshot)) + + def _publish(self) -> None: + started = time.perf_counter() + self._encode_delta() + encoded = time.perf_counter() + dist.barrier(group=get_gloo_group()) + if dist.get_rank() == 0: + write_started = time.perf_counter() + self._write_delta_files() + self._stage_metrics["perf/update_weights_delta_write_time"] = time.perf_counter() - write_started + dist.barrier(group=get_gloo_group()) + if dist.get_rank() == 0: + self._stage_metrics["perf/update_weights_delta_encode_time"] = encoded - started + self._stage_metrics["perf/update_weights_delta_publish_time"] = time.perf_counter() - started + + def _iter_hf_tensors(self, *, progress_desc: str): + """All ranks execute conversion collectives; only rank zero publishes tensors.""" + for chunk in self._iterator.get_hf_weight_chunks(self.weights_getter(), progress_desc=progress_desc): + if dist.get_rank() == 0: + yield from chunk + + def _encode_delta(self) -> None: + self._version_dir = os.path.join(self.delta_dir, f"weight_v{self.weight_version:06d}") + self._delta: dict[str, np.ndarray] = {} + self._checksums: dict[str, str] = {} + self.changed_bytes = 0 + self.total_bytes = 0 + self.wire_bytes = 0 + + if dist.get_rank() != 0: + # Do not return: every rank must execute the conversion iterator's collectives. + for _name, _tensor in self._iter_hf_tensors(progress_desc="Encode disk delta"): + pass + return + + os.makedirs(self._version_dir, exist_ok=True) + max_bytes = max((value.nbytes for value in self._snapshot.values()), default=0) + free_buffers: queue.Queue[torch.Tensor] = queue.Queue() + use_pinned = max_bytes > 0 + if use_pinned: + try: + pool_size = max(2, min(2 * NUM_WORKERS, (8 << 30) // max(max_bytes, 1))) + for _ in range(pool_size): + free_buffers.put(torch.empty(max_bytes, dtype=torch.uint8, pin_memory=True)) + except RuntimeError as exc: + logger.warning("Pinned host buffers unavailable (%s); using pageable copies", exc) + use_pinned = False + + def diff_and_compress(name: str, new: np.ndarray) -> tuple[str, np.ndarray, np.ndarray | None, str | None, int]: + old = self._snapshot[name] + if new.nbytes != old.nbytes: + raise ValueError(f"Delta tensor size changed for {name}: {old.nbytes} != {new.nbytes}") + if self.delta_encoding == "xor": + diff = new ^ old + changed = int(np.count_nonzero(diff)) + else: + mask = new != old + changed = int(np.count_nonzero(mask)) + diff = overwrite_encode(new, mask) + if not changed: + return name, new, None, None, 0 + compressed = np.frombuffer(zstandard.ZstdCompressor(level=1).compress(diff), dtype=np.uint8) + return name, new, compressed, checksum(self.checksum_algorithm, new), changed + + pool = ThreadPoolExecutor(max_workers=NUM_WORKERS) + inflight: deque = deque() + try: + for name, tensor in self._iter_hf_tensors(progress_desc="Encode disk delta"): + new = _tensor_bytes(tensor, free_buffers=free_buffers if use_pinned else None) + self.total_bytes += new.nbytes + inflight.append(pool.submit(diff_and_compress, name, new)) + if len(inflight) >= 2 * NUM_WORKERS: + self._collect_encoded(inflight.popleft()) + while inflight: + self._collect_encoded(inflight.popleft()) + finally: + pool.shutdown() + + def _collect_encoded(self, future) -> None: + name, new, compressed, digest, changed = future.result() + self._snapshot[name] = new + if changed: + self.changed_bytes += changed + assert compressed is not None and digest is not None + self._delta[name] = compressed + self._checksums[name] = digest + + def _write_delta_files(self) -> None: + if self._delta: + filename = "model-00000-of-00001.safetensors" + blob = safetensors.numpy.save(self._delta, metadata=self._checksums) + self.wire_bytes = len(blob) + _atomic_write(os.path.join(self._version_dir, filename), blob) + else: + filename = None + index = { + "metadata": { + "version": f"{self.weight_version:06d}", + "base_version": f"{self.weight_version - 1:06d}", + "delta_encoding": self.delta_encoding, + "compression_format": "zstd", + "checksum_format": self.checksum_algorithm, + }, + "weight_map": {name: filename for name in self._delta}, + } + _atomic_write( + os.path.join(self._version_dir, "model.safetensors.index.json"), + json.dumps(index).encode(), + ) + + def _reload_engines(self) -> None: + if self._post_write_hook is not None: + self._post_write_hook(self.args, self._version_dir, self.rollout_engines) + dist.barrier(group=get_gloo_group()) + if dist.get_rank() == 0: + started = time.perf_counter() + ray.get([engine.pull_weights.remote(self.weight_version) for engine in self.rollout_engines]) + pulled = time.perf_counter() + ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) + paused = time.perf_counter() + try: + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + flushed = time.perf_counter() + ray.get( + [ + engine.update_weights_from_disk.remote( + self.args.update_weight_local_checkpoint_dir, + weight_version=str(self.weight_version), + ) + for engine in self.rollout_engines + ] + ) + reloaded = time.perf_counter() + finally: + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + resumed = time.perf_counter() + self._stage_metrics.update( + { + "perf/update_weights_delta_materialize_time": pulled - started, + "perf/update_weights_delta_pause_time": paused - pulled, + "perf/update_weights_delta_flush_time": flushed - paused, + "perf/update_weights_delta_reload_time": reloaded - flushed, + "perf/update_weights_delta_resume_time": resumed - reloaded, + "perf/update_weights_delta_rollout_time": resumed - started, + } + ) + dist.barrier(group=get_gloo_group()) + + def _record_metrics(self) -> None: + device = _metric_device() + counts = torch.tensor( + [self.changed_bytes, self.total_bytes, self.wire_bytes], dtype=torch.int64, device=device + ) + dist.all_reduce(counts) + changed, total, wire = counts.tolist() + self.update_weight_metrics["perf/update_weights_density"] = changed / max(total, 1) + self.update_weight_metrics["perf/update_weights_wire_bytes"] = wire + if dist.get_rank() == 0: + self.update_weight_metrics.update(self._stage_metrics) + logger.info("[disk delta timings v=%s] %s", self.weight_version, self._stage_metrics) + logger.info("[disk delta v=%s] density=%.2f%% wire=%.2f GB", self.weight_version, 100 * changed / max(total, 1), wire / 1e9) + + +def _tensor_bytes(tensor: torch.Tensor, *, free_buffers: queue.Queue[torch.Tensor] | None = None) -> np.ndarray: + flat = tensor.detach().contiguous().view(torch.uint8).reshape(-1) + if flat.device.type == "cpu": + return flat.numpy().copy() + if free_buffers is None: + return flat.cpu().numpy().copy() + + buffer = free_buffers.get() + try: + buffer[: flat.numel()].copy_(flat, non_blocking=True) + if flat.device.type == "npu": + torch.npu.current_stream().synchronize() + elif flat.device.type == "cuda": + torch.cuda.current_stream().synchronize() + return buffer[: flat.numel()].numpy().copy() + finally: + free_buffers.put(buffer) + + +def _metric_device() -> torch.device: + if hasattr(torch, "npu") and torch.npu.is_available(): + return torch.device("npu", torch.npu.current_device()) + if torch.cuda.is_available(): + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + + +def _atomic_write(path: str, data: bytes) -> None: + temporary = path + ".tmp" + with open(temporary, "wb") as output: + output.write(data) + output.flush() + os.fsync(output.fileno()) + os.replace(temporary, path) diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_sparse_distributed.py b/vime/backends/megatron_utils/update_weight/update_weight_from_sparse_distributed.py new file mode 100644 index 000000000..f4c8fbb7a --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_sparse_distributed.py @@ -0,0 +1,436 @@ +from __future__ import annotations + +import hashlib +import logging +import os +import time +from argparse import Namespace +from collections.abc import Callable, Mapping, Sequence + +import ray +import torch +import torch.distributed as dist +from megatron.core import mpu +from ray.actor import ActorHandle +from vllm_ascend.distributed.weight_transfer.hccl_engine import HCCLTrainerSendWeightsArgs +from vllm_ascend.distributed.weight_transfer.sparse_hccl_engine import SparseHCCLWeightTransferEngine +from vllm_ascend.distributed.weight_transfer.sparse_weight_patch import SparseWeightPatch + +from vime.utils import megatron_bridge_utils +from vime.utils.distributed_utils import get_gloo_group + +from ..misc_utils import strip_param_name_prefix +from .hf_weight_iterator_base import HfWeightIteratorBase +from .megatron_sparse_export import ( + build_sparse_export_index, + local_bit_exact_diff, + sparse_hf_entry, +) +from .sparse_gather import gather_slot_entries_to_rank0 +from .update_weight_from_distributed import ( + connect_rollout_engines_from_distributed, + disconnect_rollout_engines_from_distributed, +) + +logger = logging.getLogger(__name__) + + +class _GatherQueue: + """Count-triggered queues keep collective ordering identical on all ranks.""" + + def __init__(self, batch_size: int, max_round_bytes: int, is_source: bool, consume): + self.batch_size = max(int(batch_size), 1) + self.max_round_bytes = int(max_round_bytes) + self.is_source = is_source + self.consume = consume + self.queues: dict[int, tuple] = {} + self.gather_seconds = 0.0 + + def put(self, group, slots, counts, indices, values) -> None: + _group, entries = self.queues.setdefault(id(group), (group, [])) + entries.append((slots, counts, indices, values)) + if len(entries) >= self.batch_size: + self._flush(group, entries) + + def flush_all(self) -> None: + for group, entries in self.queues.values(): + self._flush(group, entries) + + def _flush(self, group, entries) -> None: + if not entries: + return + batch = list(entries) + entries.clear() + if group is None: + if self.is_source: + for slots, counts, indices, values in batch: + offset = 0 + for (name, shape), count in zip(slots, counts.tolist(), strict=True): + self.consume(name, shape, indices[offset : offset + count], values[offset : offset + count]) + offset += count + return + + device = batch[0][2].device + counts = torch.cat([entry[1] for entry in batch]).to(device) + indices = torch.cat([entry[2] for entry in batch]) + values = torch.cat([entry[3] for entry in batch]) + gather_started = time.perf_counter() + try: + gathered = gather_slot_entries_to_rank0( + indices, + values, + counts, + group=group, + max_round_bytes=self.max_round_bytes, + ) + finally: + self.gather_seconds += time.perf_counter() - gather_started + if self.is_source and gathered is not None: + slot_index = 0 + for slots, _counts, _indices, _values in batch: + for name, shape in slots: + merged_indices, merged_values = gathered[slot_index] + slot_index += 1 + self.consume(name, shape, merged_indices, merged_values) + + +class UpdateWeightFromSparseDistributed: + """Diff Megatron shards locally and gather only final-HF sparse entries.""" + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + if quantization_config: + raise NotImplementedError("Sparse HCCL weight sync currently supports unquantized rollout weights only") + self.args = args + self.model = model + self.weights_getter = weights_getter + sparse_cpu_threads = max( + int(os.getenv("VIME_SPARSE_CPU_THREADS", "4")), 1 + ) + if torch.get_num_threads() != sparse_cpu_threads: + torch.set_num_threads(sparse_cpu_threads) + self.weight_version = 0 + self.update_weight_metrics: dict[str, float] = {} + self._snapshot: dict[str, torch.Tensor] = {} + self._slot_cache: dict[str, list[tuple[str, tuple[int, ...]]]] = {} + self._export_index = None + self._baseline_captured = False + self._model_update_groups = None + self._iterator = HfWeightIteratorBase.create( + args=args, model=model, model_name=model_name, quantization_config=quantization_config + ) + self._is_src_rank = ( + mpu.get_data_parallel_rank(with_context_parallel=True) == 0 + and mpu.get_tensor_model_parallel_rank() == 0 + and mpu.get_pipeline_model_parallel_rank() == 0 + ) + self._verify_full_diff = os.getenv("VIME_SPARSE_HCCL_VERIFY_FULL_DIFF", "0").lower() in { + "1", "true", "yes" + } + if dist.get_rank() == 0: + logger.info( + "[sparse HCCL] CPU diff threads per training rank: %d", + sparse_cpu_threads, + ) + self._legacy_snapshot: dict[str, torch.Tensor] = {} + self._distributed_signatures: dict[str, tuple] = {} + if self._is_src_rank: + self._group_name = "vime-sparse-hccl" + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + ) -> None: + del engine_gpu_offsets + self.rollout_engines = list(rollout_engines) + self.rollout_engine_lock = rollout_engine_lock + if self._is_src_rank: + self._model_update_groups = connect_rollout_engines_from_distributed( + self.args, self._group_name, self.rollout_engines, engine_gpu_counts=engine_gpu_counts + ) + + def disconnect_rollout_engines(self) -> None: + if self._is_src_rank and self._model_update_groups is not None: + disconnect_rollout_engines_from_distributed( + self.args, self._group_name, self._model_update_groups, self.rollout_engines + ) + self._model_update_groups = None + + def pop_metrics(self) -> dict[str, float]: + metrics, self.update_weight_metrics = self.update_weight_metrics, {} + return metrics + + def _local_weights(self) -> dict[str, torch.Tensor]: + return {strip_param_name_prefix(name): tensor for name, tensor in self.weights_getter().items()} + + def _iter_hf_tensors(self): + for chunk in self._iterator.get_hf_weight_chunks( + self.weights_getter(), progress_desc="Sparse HCCL full-diff verification" + ): + if self._is_src_rank: + yield from chunk + + def _capture_baseline(self) -> None: + local_weights = self._local_weights() + with megatron_bridge_utils.patch_megatron_model(self.model): + self._export_index = build_sparse_export_index( + self._iterator._bridge, self.model, local_weights, self._slot_cache + ) + for record in self._export_index: + self._snapshot[record.weight_key] = local_weights[record.weight_key] + + if self._verify_full_diff: + for name, tensor in self._iter_hf_tensors(): + self._legacy_snapshot[name] = tensor.detach().cpu().contiguous().clone() + self._baseline_captured = True + dist.barrier(group=get_gloo_group()) + local_bytes = sum(t.numel() * t.element_size() for t in self._snapshot.values()) + if self._is_src_rank: + logger.info( + "[sparse HCCL] captured rank-local baseline: %d shards, %.2f MB%s", + len(self._snapshot), local_bytes / 1e6, + " (full-diff verification enabled)" if self._verify_full_diff else "", + ) + + @torch.no_grad() + def update_weights(self) -> None: + if not self._baseline_captured: + self._capture_baseline() + return + + next_weight_version = self.weight_version + 1 + started = time.perf_counter() + if dist.get_rank() == 0: + ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + ray.get([ + engine.start_weight_update.remote(is_checkpoint_format=False) + for engine in self.rollout_engines + ]) + dist.barrier(group=get_gloo_group()) + + statistics = {"changed": 0, "total": 0, "wire": 0} + transfer_statistics = {"patches": 0, "batches": 0, "seconds": 0.0} + next_snapshot: dict[str, torch.Tensor] = {} + seen_names: set[str] = set() + self._distributed_signatures.clear() + pending_patches: list[tuple[SparseWeightPatch, list[int]]] = [] + pending_wire_bytes = 0 + + def flush_pending_patches() -> None: + nonlocal pending_wire_bytes + if not pending_patches: + return + transfer_started = time.perf_counter() + self._send_patches(pending_patches, next_weight_version) + transfer_statistics["seconds"] += time.perf_counter() - transfer_started + transfer_statistics["batches"] += 1 + pending_patches.clear() + pending_wire_bytes = 0 + + def consume(name, shape, indices, values) -> None: + nonlocal pending_wire_bytes + if name in seen_names: + raise RuntimeError(f"Sparse Bridge emitted duplicate HF tensor {name!r}") + seen_names.add(name) + total = 1 + for dimension in shape: + total *= dimension + statistics["total"] += total + statistics["changed"] += indices.numel() + if self._verify_full_diff: + self._distributed_signatures[name] = self._patch_signature(shape, indices, values) + if indices.numel() == 0: + return + patch_wire_bytes = ( + indices.numel() * indices.element_size() + values.numel() * values.element_size() + ) + statistics["wire"] += patch_wire_bytes + patch = SparseWeightPatch( + name=name, indices=indices.to(torch.int32).contiguous(), values=values.contiguous() + ) + # Match the bucketed/flush design used by verl weight sync: keep + # payloads bounded by the configured communication buffer, while + # amortizing Ray RPC and HCCL launch latency across many tensors. + # The sparse HCCL receiver already accepts multiple tensor entries + # in one update request and consumes their broadcasts in order. + if pending_patches and ( + pending_wire_bytes + patch_wire_bytes > self.args.update_weight_buffer_size + ): + flush_pending_patches() + pending_patches.append((patch, list(shape))) + pending_wire_bytes += patch_wire_bytes + transfer_statistics["patches"] += 1 + + queue = _GatherQueue( + batch_size=32, + max_round_bytes=self.args.update_weight_buffer_size, + is_source=self._is_src_rank, + consume=consume, + ) + local_pipeline_seconds = 0.0 + try: + local_pipeline_started = time.perf_counter() + local_weights = self._local_weights() + with megatron_bridge_utils.patch_megatron_model(self.model): + for record in self._export_index: + current = local_weights[record.weight_key].detach().cpu().contiguous() + snapshot = self._snapshot[record.weight_key] + if current.data_ptr() == snapshot.data_ptr(): + raise RuntimeError( + "Sparse actor backups must be double-buffered; " + f"current weight aliases its snapshot: {record.weight_key}" + ) + local_indices, local_values = local_bit_exact_diff(current, snapshot) + # Commit snapshots only after every collective and rollout + # update has succeeded. A failed update can then be + # retried without silently dropping its local changes. + # TensorBackuper alternates two pinned CPU actor buffers. + # The current buffer therefore stays immutable until the + # next diff completes and can become the baseline without + # another full-model clone. + next_snapshot[record.weight_key] = current + if not record.contributes: + local_indices = local_indices[:0] + local_values = local_values[:0] + device = record.param.device + slots, counts, hf_indices, hf_values = sparse_hf_entry( + record, + local_indices.to(device=device, non_blocking=False), + local_values.to(device=device, non_blocking=False), + self._slot_cache, + ) + queue.put(record.gather_group, slots, counts, hf_indices, hf_values) + queue.flush_all() + flush_pending_patches() + local_pipeline_seconds = ( + time.perf_counter() + - local_pipeline_started + - transfer_statistics["seconds"] + ) + if self._verify_full_diff: + next_legacy_snapshot = self._verify_against_full_diff() + else: + next_legacy_snapshot = None + if self._is_src_rank: + torch.npu.synchronize() + finally: + dist.barrier(group=get_gloo_group()) + if dist.get_rank() == 0: + try: + ray.get([engine.finish_weight_update.remote() for engine in self.rollout_engines]) + finally: + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + dist.barrier(group=get_gloo_group()) + + self._snapshot.update(next_snapshot) + if next_legacy_snapshot is not None: + self._legacy_snapshot = next_legacy_snapshot + self.weight_version = next_weight_version + + if self._is_src_rank: + elapsed = time.perf_counter() - started + changed, total, wire = statistics["changed"], statistics["total"], statistics["wire"] + local_compute_seconds = max( + local_pipeline_seconds - queue.gather_seconds, 0.0 + ) + self.update_weight_metrics.update({ + "perf/update_weights_density": changed / max(total, 1), + "perf/update_weights_wire_bytes": wire, + "perf/update_weights_sparse_hccl_time": elapsed, + "perf/update_weights_sparse_hccl_transfer_time": transfer_statistics["seconds"], + "perf/update_weights_sparse_local_process_time": max( + local_pipeline_seconds, 0.0 + ), + "perf/update_weights_sparse_tp_gather_time": queue.gather_seconds, + "perf/update_weights_sparse_local_compute_time": local_compute_seconds, + "perf/update_weights_sparse_hccl_patches": transfer_statistics["patches"], + "perf/update_weights_sparse_hccl_batches": transfer_statistics["batches"], + }) + logger.info( + "[sparse HCCL v=%d] density=%.4f%% wire=%.2f MB " + "patches=%d batches=%d local=%.3fs (compute=%.3fs gather=%.3fs) " + "transfer=%.3fs elapsed=%.3fs", + self.weight_version, 100 * changed / max(total, 1), wire / 1e6, + transfer_statistics["patches"], transfer_statistics["batches"], + max(local_pipeline_seconds, 0.0), local_compute_seconds, + queue.gather_seconds, transfer_statistics["seconds"], elapsed, + ) + + @staticmethod + def _patch_signature(shape, indices, values) -> tuple: + indices = indices.detach().cpu().to(torch.int64) + values = values.detach().cpu().contiguous() + if indices.numel(): + order = torch.argsort(indices) + indices = indices[order] + values = values[order] + digest = hashlib.sha256() + digest.update(indices.numpy().tobytes()) + digest.update(values.view(torch.uint8).numpy().tobytes()) + return tuple(shape), indices.numel(), digest.hexdigest() + + def _verify_against_full_diff(self) -> dict[str, torch.Tensor]: + expected: dict[str, tuple] = {} + next_snapshot: dict[str, torch.Tensor] = {} + for name, tensor in self._iter_hf_tensors(): + current = tensor.detach().cpu().contiguous() + snapshot = self._legacy_snapshot.get(name) + if snapshot is None: + raise RuntimeError(f"Full-diff baseline is missing {name!r}") + indices, values = local_bit_exact_diff(current, snapshot) + expected[name] = self._patch_signature(current.shape, indices, values) + next_snapshot[name] = current + if self._is_src_rank and expected != self._distributed_signatures: + missing = sorted(expected.keys() - self._distributed_signatures.keys()) + extra = sorted(self._distributed_signatures.keys() - expected.keys()) + mismatched = sorted( + name for name in expected.keys() & self._distributed_signatures.keys() + if expected[name] != self._distributed_signatures[name] + ) + raise AssertionError( + "Sparse shard export differs from full HF bit-exact diff: " + f"missing={missing[:8]}, extra={extra[:8]}, mismatched={mismatched[:8]}" + ) + if self._is_src_rank: + logger.info( + "[sparse HCCL] TP shard export matched full HF bit-exact diff for %d tensors", len(expected) + ) + return next_snapshot + + def _send_patches( + self, + patches_with_shapes: Sequence[tuple[SparseWeightPatch, list[int]]], + weight_version: int, + ) -> None: + patches = [patch for patch, _shape in patches_with_shapes] + while not ray.get(self.rollout_engine_lock.acquire.remote()): + time.sleep(0.1) + try: + refs = [ + engine.update_sparse_weights_from_distributed.remote( + names=[patch.name for patch in patches], + dtypes=[patch.values.dtype for patch in patches], + shapes=[shape for _patch, shape in patches_with_shapes], + num_updates_list=[patch.indices.numel() for patch in patches], + group_name=self._group_name, + weight_version=str(weight_version), + ) + for engine in self.rollout_engines + ] + SparseHCCLWeightTransferEngine.trainer_send_weights( + iter(patches), HCCLTrainerSendWeightsArgs(group=self._model_update_groups, packed=False) + ) + ray.get(refs) + finally: + ray.get(self.rollout_engine_lock.release.remote()) diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index f3ec65292..10b9b1798 100644 --- a/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -16,6 +16,7 @@ from ray.actor import ActorHandle from vime.utils.common import is_npu +from vime.utils.function_compat import call_with_optional_keyword from vime.utils.distributed_utils import get_gloo_group from .hf_weight_iterator_base import HfWeightIteratorBase @@ -349,7 +350,12 @@ def _patched_start_weight_update( self, is_checkpoint_format: bool = True, _orig=_orig_start_weight_update ) -> None: _VLLMHijack.patch_moe_weight_loader(self.model_runner.model) - _orig(self, is_checkpoint_format=is_checkpoint_format) + call_with_optional_keyword( + _orig, + self, + keyword="is_checkpoint_format", + value=is_checkpoint_format, + ) _VLLMHijack._invalidate_moe_alltoall_expert_ids() def _patched_wake_up(self, tags=None, _orig=_orig_wake_up) -> None: diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index 69de3a1b7..5296e886f 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/vime/backends/vllm_utils/vllm_engine.py @@ -458,6 +458,8 @@ def build_vllm_cmd_and_env(server_args: dict[str, Any]) -> tuple[list[str], dict "--weight-transfer-config", _serialize_weight_transfer_config(args.vllm_weight_transfer_config), ] + elif getattr(args, "update_weight_mode", "full") == "sparse": + cmd += ["--weight-transfer-config", '{"backend":"sparse_nccl"}'] elif getattr(args, "colocate", False): cmd += ["--weight-transfer-config", '{"backend":"ipc"}'] else: @@ -951,7 +953,57 @@ def update_weights_from_distributed( } return self._post_vllm_update_weights_http(update_info) - def update_weights_from_disk(self, model_path: str, load_format: str | None = None): + def update_sparse_weights_from_distributed( + self, + names, + dtypes, + shapes, + num_updates_list, + group_name, + flush_cache=False, + weight_version: str | None = None, + ): + """Sparse HCCL path: send metadata while indices/values use HCCL.""" + del group_name + if weight_version is not None: + self._weight_version = str(weight_version) + if flush_cache: + self.flush_cache() + update_info = { + "names": names, + "dtype_names": [str(dtype).replace("torch.", "") for dtype in dtypes], + "shapes": [list(shape) for shape in shapes], + "num_updates_list": list(num_updates_list), + } + return self._post_vllm_update_weights_http(update_info) + + def pull_weights(self, target_version: int): + """Materialize a published disk version on every host of this engine.""" + if self.node_rank != 0: + return None + response = requests.post( + f"{self._http_base()}/collective_rpc", + json={ + "method": "pull_weights", + "kwargs": { + "local_checkpoint_dir": self.args.update_weight_local_checkpoint_dir, + "source_dir": self.args.update_weight_disk_dir, + "target_version": target_version, + "pre_read_hook": self.args.custom_update_weight_pre_read_path, + }, + }, + timeout=600, + ) + result = _response_json(response) + self._weight_version = str(target_version) + return result + + def update_weights_from_disk( + self, + model_path: str, + load_format: str | None = None, + weight_version: str | None = None, + ): """``POST /collective_rpc`` with ``reload_weights`` and ``weights_path``.""" if self.node_rank != 0: return @@ -964,7 +1016,10 @@ def update_weights_from_disk(self, model_path: str, load_format: str | None = No }, timeout=600, ) - return _response_json(response) + result = _response_json(response) + if weight_version is not None: + self._weight_version = str(weight_version) + return result def pause_generation(self): """``POST /pause`` with mode="keep"; returns the ``requests.Response``.""" diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 7b3ae8968..63161434f 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -132,6 +132,69 @@ def add_train_arguments(parser): default=1024**3, help="Add margin for train memory allocation. By default we will reserve 1GB as margin.", ) + parser.add_argument( + "--update-weight-mode", + choices=["full", "delta", "sparse"], + default="full", + help=( + "Weight sync strategy. 'full' keeps the existing accelerator-native " + "HCCL/NPU-IPC path. 'delta' publishes only changed weight bytes " + "through a shared filesystem. 'sparse' sends changed BF16 elements " + "directly through the accelerator collective transport." + ), + ) + parser.add_argument( + "--update-weight-transport", + choices=["nccl", "disk"], + default="nccl", + help=( + "Weight sync transport. 'nccl' is retained for CLI compatibility and " + "selects the existing accelerator-native HCCL/NPU-IPC path on Ascend; " + "'disk' is supported only with --update-weight-mode=delta." + ), + ) + parser.add_argument( + "--update-weight-disk-dir", + type=str, + default=None, + help="Shared filesystem directory where delta weight versions are published.", + ) + parser.add_argument( + "--update-weight-local-checkpoint-dir", + type=str, + default=None, + help="Host-local HF checkpoint directory patched in place by rollout workers.", + ) + parser.add_argument( + "--update-weight-delta-encoding", + choices=["xor", "overwrite"], + default="xor", + help="Delta encoding: xor (compact) or overwrite (idempotent).", + ) + parser.add_argument( + "--update-weight-delta-checksum", + choices=["xxh3-128", "blake3", "adler32"], + default="xxh3-128", + help="Per-tensor checksum algorithm for delta application.", + ) + parser.add_argument( + "--custom-update-weight-post-write-path", + type=str, + default=None, + help=( + "Optional trainer-side hook after a delta version is written. Signature: " + "hook(args, version_dir: str, rollout_engines) -> None." + ), + ) + parser.add_argument( + "--custom-update-weight-pre-read-path", + type=str, + default=None, + help=( + "Optional rollout-host hook before a published version is read. Signature: " + "hook(source_dir: str, target_version: int) -> None." + ), + ) try: default_megatron_to_hf_mode = "bridge" if is_npu() else "raw" except RuntimeError: @@ -1688,6 +1751,31 @@ def vime_validate_args(args): 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": + if args.update_weight_transport != "disk": + raise ValueError("--update-weight-mode=delta requires --update-weight-transport=disk.") + if args.colocate: + raise ValueError( + "--update-weight-mode=delta is not supported with --colocate; " + "colocated NPU IPC already avoids full-model copies." + ) + if not args.update_weight_disk_dir: + raise ValueError("--update-weight-mode=delta requires --update-weight-disk-dir.") + if not args.update_weight_local_checkpoint_dir: + raise ValueError("--update-weight-mode=delta requires --update-weight-local-checkpoint-dir.") + elif args.update_weight_mode == "sparse": + if args.update_weight_transport != "nccl": + raise ValueError("--update-weight-mode=sparse requires --update-weight-transport=nccl.") + if args.colocate: + raise ValueError("--update-weight-mode=sparse is supported only for non-colocated rollout.") + if args.megatron_to_hf_mode != "bridge": + raise ValueError("--update-weight-mode=sparse requires --megatron-to-hf-mode=bridge.") + elif args.update_weight_transport == "disk": + raise ValueError( + "--update-weight-transport=disk is currently supported only with " + "--update-weight-mode=delta on Ascend." + ) + assert not (args.kl_coef != 0 and args.kl_loss_coef != 0), "Only one of kl_coef and kl_loss_coef can be set" if args.advantage_estimator in ["reinforce_plus_plus", "reinforce_plus_plus_baseline"]: diff --git a/vime/utils/disk_delta.py b/vime/utils/disk_delta.py new file mode 100644 index 000000000..e4a4c84a4 --- /dev/null +++ b/vime/utils/disk_delta.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import glob +import json +import os +import struct +import zlib + +import numpy as np + + +NUM_WORKERS = min(32, os.cpu_count() or 8) + + +def overwrite_encode(new: np.ndarray, changed_mask: np.ndarray) -> np.ndarray: + """Encode changed byte positions and their replacement values.""" + positions = np.flatnonzero(changed_mask).astype(" 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, buffer) -> str: + hasher = _new_hasher(algorithm) + hasher.update(buffer) + return hasher.hexdigest() + + +def _tensor_locations(checkpoint_dir: str) -> dict[str, tuple[str, int, int]]: + locations: dict[str, tuple[str, int, int]] = {} + for path in glob.glob(os.path.join(checkpoint_dir, "*.safetensors")): + with open(path, "rb") as tensor_file: + (header_len,) = struct.unpack(" np.ndarray: + path, offset, nbytes = locations[name] + with open(path, "rb") as tensor_file: + tensor_file.seek(offset) + return np.frombuffer(tensor_file.read(nbytes), dtype=np.uint8) + + return read diff --git a/vime/utils/function_compat.py b/vime/utils/function_compat.py new file mode 100644 index 000000000..576e4b3db --- /dev/null +++ b/vime/utils/function_compat.py @@ -0,0 +1,12 @@ +import inspect +from collections.abc import Callable +from typing import Any + + +def call_with_optional_keyword( + func: Callable[..., Any], *args: Any, keyword: str, value: Any +) -> Any: + """Call ``func`` with a keyword only when its installed version supports it.""" + if keyword in inspect.signature(func).parameters: + return func(*args, **{keyword: value}) + return func(*args) diff --git a/vime/utils/tensor_backper.py b/vime/utils/tensor_backper.py index 2fc2a6359..8f378cfb8 100644 --- a/vime/utils/tensor_backper.py +++ b/vime/utils/tensor_backper.py @@ -34,6 +34,9 @@ def backup(self, tag: str): def copy(self, *, src_tag: str, dst_tag: str): raise NotImplementedError + def enable_double_buffer(self, tag: str) -> None: + raise NotImplementedError + @abstractmethod def restore(self, tag: str): raise NotImplementedError @@ -43,6 +46,8 @@ class _TensorBackuperNormal(TensorBackuper): def __init__(self, source_getter): super().__init__(source_getter=source_getter) self._backups: dict[str, dict[str, torch.Tensor]] = defaultdict(dict) + self._spare_backups: dict[str, dict[str, torch.Tensor]] = defaultdict(dict) + self._double_buffer_tags: set[str] = set() @property def backup_tags(self): @@ -53,12 +58,28 @@ def get(self, tag: str): @torch.no_grad() def backup(self, tag: str) -> None: - backup_dict = self._backups[tag] + backup_dict = ( + self._spare_backups[tag] + if tag in self._double_buffer_tags + else self._backups[tag] + ) for name, param in self._source_getter(): if name not in backup_dict: backup_dict[name] = torch.empty_like(param, device=torch.device("cpu"), pin_memory=True) backup_dict[name].copy_(param.detach(), non_blocking=True) torch.cuda.synchronize() + if tag in self._double_buffer_tags: + self._backups[tag], self._spare_backups[tag] = ( + backup_dict, + self._backups[tag], + ) + + def enable_double_buffer(self, tag: str) -> None: + if self._backups[tag]: + raise RuntimeError( + f"Double buffering must be enabled before the first backup of {tag!r}" + ) + self._double_buffer_tags.add(tag) @torch.no_grad() def copy(self, *, src_tag: str, dst_tag: str): @@ -96,6 +117,9 @@ def backup(self, tag: str) -> None: self._backup_hash_dict = _compute_hash_dict(dict(self._source_getter())) torch.cuda.synchronize() + def enable_double_buffer(self, tag: str) -> None: + del tag + def restore(self, tag: str) -> None: assert tag == self._single_tag assert _compute_hash_dict(dict(self._source_getter())) == self._backup_hash_dict