diff --git a/docs/en/advanced/delta-weight-sync.md b/docs/en/advanced/delta-weight-sync.md index af0a0e1e8..a309fe475 100644 --- a/docs/en/advanced/delta-weight-sync.md +++ b/docs/en/advanced/delta-weight-sync.md @@ -5,15 +5,24 @@ that changed between two syncs, instead of a full checkpoint each time. It targe training/inference disaggregation across clusters or datacenters, where writing the whole actor every sync is the dominant cost. -It is **disk-transport only**. The trainer publishes each sync as a canonical HF checkpoint +> **Status:** `--update-weight-mode delta` selects the **direct NCCL** delta implementation +> (`--update-weight-transport nccl`): the trainer diffs against a checkpoint-coordinate +> snapshot and ships only changed values over the existing NCCL channel into vLLM's +> checkpoint weight patch API. See +> [`examples/delta_weight_sync/README.md`](https://github.com/vllm-project/vime/blob/main/examples/delta_weight_sync/README.md) +> for requirements and configuration. The **disk-based** design documented +> below is reserved and not implemented: `--update-weight-mode delta +> --update-weight-transport disk` raises `NotImplementedError`. + +The disk-based design is **disk-transport** based. The trainer publishes each sync as a canonical HF checkpoint directory; the engine's `/pull_weights` endpoint (shipped in vime's vllm patch) fans the apply out to **every host the engine spans** and verifies it, then the engine reloads the patched local checkpoint through the **ordinary** `update_weights_from_disk` endpoint. vime only ever talks to one endpoint per engine, so multi-node serving and external rollout engines need nothing extra on the vime side. -Vime currently guards this mechanically synchronized path with a `NotImplementedError` when -`--update-weight-mode=delta` is selected; the implementation below remains upstream reference code. +Vime guards this path with a `NotImplementedError` when `--update-weight-transport=disk` is +combined with delta mode; the implementation below remains upstream reference code. ## Configuration diff --git a/docs/en/get_started/usage.md b/docs/en/get_started/usage.md index d1300ef55..b6b533182 100644 --- a/docs/en/get_started/usage.md +++ b/docs/en/get_started/usage.md @@ -145,7 +145,7 @@ Note: - By default, vLLM reads the maximum context length from the `config.json` in the Hugging Face checkpoint. You can use the `--vllm-max-model-len` parameter to override this value to support longer inference. - During co-located training and inference, although Megatron and vLLM will offload sequentially, they still need to leave some memory for each other. You need to adjust vLLM's total VRAM usage by reducing `--vllm-gpu-memory-utilization`. - vime supports passing through vllm-router parameters by adding a `router` prefix to the original parameter name. For example, vllm-router's `--balance-abs-threshold` parameter should be set as `--router-balance-abs-threshold`. Since vllm-router uses cache-aware routing by default, it may cause uneven request distribution. You can set `--router-balance-abs-threshold 0` to force balanced distribution, but this may affect prefix cache hit rate in multi-turn conversation scenarios. For multi-turn sessions that require session affinity, set `--router-policy consistent_hash` and send a stable `x-session-id` for each session. - - If vLLM engines are pre-launched by an external system, connect to them with `--rollout-external-engine-addrs host1:port host2:port`. When the trainer and engines cannot form an NCCL weight-update group, use `--update-weight-mode full --update-weight-transport disk --update-weight-disk-dir /shared/fs/updates`; vime writes a complete HF checkpoint and asks vLLM to hot-load it through `update_weights_from_disk`. For large models or cross-cluster deployments, use `--update-weight-mode delta --update-weight-transport disk` instead. See [External Rollout Engines Roadmap](../advanced/external-rollout-engines.md) and [Delta Weight Sync](../advanced/delta-weight-sync.md). + - If vLLM engines are pre-launched by an external system, connect to them with `--rollout-external-engine-addrs host1:port host2:port`. When the trainer and engines cannot form an NCCL weight-update group, use `--update-weight-mode full --update-weight-transport disk --update-weight-disk-dir /shared/fs/updates`; vime writes a complete HF checkpoint and asks vLLM to hot-load it through `update_weights_from_disk`. When the trainer and engines can form an NCCL group, `--update-weight-mode delta --update-weight-transport nccl` ships only the weights that changed between syncs (see [Delta Weight Sync](../advanced/delta-weight-sync.md)); disk-based delta sync is a reserved design and not yet implemented. See also [External Rollout Engines Roadmap](../advanced/external-rollout-engines.md). For details on some of vLLM's customizations and the principles behind how vime incorporates vLLM, please see the "How to Use vLLM" section. diff --git a/docs/zh/advanced/delta-weight-sync.md b/docs/zh/advanced/delta-weight-sync.md index 346909fa0..33297e1ba 100644 --- a/docs/zh/advanced/delta-weight-sync.md +++ b/docs/zh/advanced/delta-weight-sync.md @@ -2,9 +2,16 @@ Delta 权重同步只发送两次同步之间发生变化的字节，而不是每次都写一份完整 checkpoint，以此让非 colocate 的 rollout engine 保持最新。它面向大模型、跨集群或跨数据中心的训推解耦场景——这种场景下每次都写整份 actor 权重是主要开销。 -它**只支持 disk transport**。训练端把每次同步发布为一份 canonical HF checkpoint 目录；engine 的 `/pull_weights` 端点（随 vime 的 vllm patch 提供）把 apply 扇出到 **engine 覆盖的每一个 host** 并校验，随后 engine 通过**原生**的 `update_weights_from_disk` 端点 reload 打过补丁的本地 checkpoint。vime 对每个 engine 只与一个端点通信，所以多节点 serving 和外部 rollout engine 在 vime 侧都不需要任何额外支持。 +> **状态：**`--update-weight-mode delta` 现在选择的是**直连 NCCL** 的 delta 实现 +> (`--update-weight-transport nccl`)：训练端基于 checkpoint 坐标快照做 diff，只把变化的 +> 权重通过既有 NCCL 通道送进 vLLM 的 checkpoint weight patch API。环境要求和配置方式见 +> [`examples/delta_weight_sync/README.md`](https://github.com/vllm-project/vime/blob/main/examples/delta_weight_sync/README.md)。 +> 下文描述的**基于磁盘**的设计为保留设计、尚未实现：`--update-weight-mode delta +> --update-weight-transport disk` 会抛出 `NotImplementedError`。 -Vime 当前在选择 `--update-weight-mode=delta` 时会通过 `NotImplementedError` guard 拒绝该路径；下文保留为机械同步的上游参考实现。 +磁盘方案基于 **disk transport**。训练端把每次同步发布为一份 canonical HF checkpoint 目录；engine 的 `/pull_weights` 端点（随 vime 的 vllm patch 提供）把 apply 扇出到 **engine 覆盖的每一个 host** 并校验，随后 engine 通过**原生**的 `update_weights_from_disk` 端点 reload 打过补丁的本地 checkpoint。vime 对每个 engine 只与一个端点通信，所以多节点 serving 和外部 rollout engine 在 vime 侧都不需要任何额外支持。 + +Vime 在 delta 模式与 `--update-weight-transport=disk` 组合时会通过 `NotImplementedError` guard 拒绝该路径；下文保留为上游参考实现。 ## 配置 diff --git a/docs/zh/get_started/usage.md b/docs/zh/get_started/usage.md index 3ecf9fd56..0075ceb75 100644 --- a/docs/zh/get_started/usage.md +++ b/docs/zh/get_started/usage.md @@ -147,7 +147,7 @@ vLLM 的加载非常简单，只需要： - vLLM 默认会从 huggingface ckpt 中 `config.json` 读取模型的最大 context length，可以使用 `--vllm-max-model-len` 参数来对这个值进行覆盖，从而支持进行更长的推理； - 在训推一体的训练过程中，虽然 megatron 和 vLLM 会先后 offload，但是还是需要为对方留有一些空间，需要通过减小 `--vllm-gpu-memory-utilization` 来调整 vLLM 的显存占用总量。 - vime 支持透传 vllm-router 的参数，方式是在原参数名前加上 `router` 前缀。例如，vllm-router 的 `--balance-abs-threshold` 参数需要设置为 `--router-balance-abs-threshold`。由于 vllm-router 默认使用 cache-aware routing，可能会导致请求分配不均衡。可以通过设置 `--router-balance-abs-threshold 0` 来强制均衡分配，但这可能会影响多轮对话场景下 prefix cache 的命中率。对于需要会话亲和的多轮会话，可以设置 `--router-policy consistent_hash`，并为每个会话发送稳定的 `x-session-id`。 -- 如果 vLLM engine 已经由外部系统预启动，可以通过 `--rollout-external-engine-addrs host1:port host2:port` 连接。此时如果训练器和 engine 无法建立 NCCL 权重同步 group，可以使用 `--update-weight-mode full --update-weight-transport disk --update-weight-disk-dir /shared/fs/updates`，vime 会写完整 HF checkpoint 并调用 vLLM 的 `update_weights_from_disk` 热加载；大模型或跨集群场景可进一步使用 `--update-weight-mode delta --update-weight-transport disk`。详见 [External Rollout Engines 配置路线图](../advanced/external-rollout-engines.md) 和 [Delta 权重同步](../advanced/delta-weight-sync.md)。 +- 如果 vLLM engine 已经由外部系统预启动，可以通过 `--rollout-external-engine-addrs host1:port host2:port` 连接。此时如果训练器和 engine 无法建立 NCCL 权重同步 group，可以使用 `--update-weight-mode full --update-weight-transport disk --update-weight-disk-dir /shared/fs/updates`，vime 会写完整 HF checkpoint 并调用 vLLM 的 `update_weights_from_disk` 热加载；当训练器与 engine 可以建立 NCCL group 时，可使用 `--update-weight-mode delta --update-weight-transport nccl` 只传输两次同步之间发生变化的权重（详见 [Delta 权重同步](../advanced/delta-weight-sync.md)）；基于磁盘的 delta 同步为保留设计、尚未实现。另见 [External Rollout Engines 配置路线图](../advanced/external-rollout-engines.md)。 对于一些 vLLM 的自定义以及 vime 引入 vLLM 的原理，请见 vLLM 使用方法一节。 diff --git a/examples/delta_weight_sync/README.md b/examples/delta_weight_sync/README.md index 2f8207a1c..1995a5b71 100644 --- a/examples/delta_weight_sync/README.md +++ b/examples/delta_weight_sync/README.md @@ -1,44 +1,110 @@ # Delta Weight Sync -Non-colocated weight sync that ships only the **changed bytes** between two syncs instead of a -full checkpoint, for training/inference disaggregation across clusters or datacenters. The -trainer publishes per-tensor deltas to a shared filesystem as a canonical HF checkpoint -directory; each engine's `/pull_weights` applies them into a host-local checkpoint on every -host it spans, and the engines reload through the ordinary `update_weights_from_disk` path — -vime only ever talks to one endpoint per engine. +VIME currently provides a direct delta-weight-update (DWU) MVP for a +non-colocated Megatron trainer and VIME-launched vLLM rollout engines. The +trainer exports canonical Hugging Face/checkpoint-coordinate tensors, keeps a +committed CPU snapshot, and sends: -Vime currently rejects `--update-weight-mode delta` with a `NotImplementedError`; this example -is retained as mechanically synchronized upstream reference material. +- a dense BF16 seed for version 1; then +- absolute BF16 values plus flattened `int32` checkpoint indices for elements + whose bit patterns changed after a committed optimizer step. -See [Delta Weight Sync](../../docs/en/advanced/delta-weight-sync.md) for the full mechanism, -encodings, integrity checks, and shared-filesystem visibility hooks. +The vLLM worker applies those patches through the model's native +`load_weights()` mapping. Consequently, VIME does not need to know vLLM's +runtime QKV/gate-up packing or tensor-parallel parameter names. -## Try it +```text +Megatron TP/DP weights + -> canonical HF export + -> dense seed or absolute sparse checkpoint patches + -> NCCL + -> vLLM CheckpointWeightPatch + -> native model.load_weights() +``` + +## vLLM dependency + +Direct DWU requires a vLLM build that contains the checkpoint weight patch +API (`CheckpointWeightPatch` and `load_checkpoint_weight_patches()` in +`vllm.model_executor.model_loader.checkpoint_weight_patch`) from +[vLLM PR #50723](https://github.com/vllm-project/vllm/pull/50723). -`run-glm4.7-30B-A3B-delta.sh` runs the disk delta path on GLM-4.7-Flash, non-colocated across a -2-node (16-GPU) Ray cluster. See its header for prerequisites. +Until #50723 merges, no vLLM release contains that API; build vLLM from the +PR branch. This path was tested at PR commit +`fd07acd5b596c11f949fa71b5f0ee926b9e6bf17`; vime fails fast at engine startup +if the patch API is missing. -## Minimal flags +## Enabling direct DWU -Add to a non-colocated training run (the trainer and engines only need to share the filesystem -at `--update-weight-disk-dir`): +Add to a non-colocated Megatron training run (vime starts and owns the vLLM +engines; do not set `--rollout-external`): ```bash --update-weight-mode delta \ ---update-weight-transport disk \ ---update-weight-disk-dir /shared/fs/delta-updates \ ---update-weight-local-checkpoint-dir /local/nvme/rollout-ckpt \ ---update-weight-delta-encoding xor \ ---update-weight-delta-checksum xxh3-128 +--update-weight-transport nccl ``` -- `--update-weight-disk-dir` — shared directory the trainer writes deltas to and the hosts read. -- `--update-weight-local-checkpoint-dir` — host-local full HF checkpoint the delta patches in - place; materialized from the engine's model path on the first `/pull_weights`. -- `--update-weight-delta-encoding` — `xor` (smallest/fastest) or `overwrite` (idempotent). -- `--update-weight-delta-checksum` — `xxh3-128` (default), `blake3`, or `adler32`. +The first sync ships a mandatory dense seed (start version 0) that aligns the +rollout weights with the trainer checkpoint; every later sync ships only the +weights whose BF16 bits changed. + +## Current MVP boundary + +The direct path currently requires: + +- `--train-backend megatron`; +- non-colocated, VIME-launched rollout engines; +- BF16, unquantized weights; +- Megatron PP=1 and VPP=1; +- vLLM PP=1 and DP=1; +- no rollout offload, speculative decoding, MTP draft update, fault-tolerant + worker replacement, or fully-async rollout; and +- version 0 startup followed by a mandatory dense seed. + +The source currently performs a full canonical-HF export and keeps the +committed weights in CPU memory. Steady-state network traffic is sparse, but +source traversal and snapshot memory are not yet sparse or sharded. + +Delta over disk is only a reserved interface. The current argument validator +rejects: + +```bash +--update-weight-mode delta --update-weight-transport disk +``` + +with `NotImplementedError`. The existing GLM disk script is retained as +historical interface material; it is not a runnable path for the current MVP. + +## Verifying a run + +A successful direct-DWU run must show all of the following in the Ray job +log: + +1. Version 1 logs `dense_seed=True` with `changed == total`. +2. After a real optimizer step, version 2 or later logs + `dense_seed=False`, `0 < changed < total`, and a smaller `wire_bytes` than + the dense seed. +3. Training reports a finite, nonzero `train/grad_norm`. +4. Rollout generation succeeds after the sparse commit and no worker reports a + base-version, sequence, final-manifest, or failed-session error. + +The updater exports these step metrics after a committed update: + +```text +weight_sync/is_dense_seed +weight_sync/total_elements +weight_sync/changed_elements +weight_sync/delta_density +weight_sync/wire_bytes +weight_sync/seconds +``` + +Its summary line has this form: + +```text +Direct DWU committed version= dense_seed= changed=/ \ +density= wire_bytes= seconds= +``` -For object-store-backed volumes that need an explicit commit/refresh to make writes visible -across hosts, supply `--custom-update-weight-post-write-path` (trainer side) / -`--vllm-custom-pull-weights-pre-read-hook` (engine side) — no vendor-specific code lives in vime -or vllm; see the doc. +Process launch, a dense seed alone, or static tests do not by themselves +demonstrate a working delta path; check all four criteria above. diff --git a/tests/test_empty_colocated_weight_bucket.py b/tests/test_empty_colocated_weight_bucket.py index d2051fe02..a1b793e4c 100644 --- a/tests/test_empty_colocated_weight_bucket.py +++ b/tests/test_empty_colocated_weight_bucket.py @@ -112,6 +112,8 @@ def gather_object(obj, object_gather_list, dst, group): update_from_distributed_mod.disconnect_rollout_engines_from_distributed = lambda *args, **kwargs: None update_from_distributed_mod.post_process_weights = lambda *args, **kwargs: None update_from_distributed_mod.update_weights_from_distributed = lambda *args, **kwargs: [] + coordinator_mod = types.ModuleType("vime.backends.megatron_utils.update_weight.coordinator") + coordinator_mod.WeightUpdateCoordinator = object monkeypatch.setitem(sys.modules, "vime", vime_pkg) monkeypatch.setitem(sys.modules, "vime.backends", vime_backends_pkg) @@ -132,6 +134,11 @@ def gather_object(obj, object_gather_list, dst, group): "vime.backends.megatron_utils.update_weight.update_weight_from_distributed", update_from_distributed_mod, ) + monkeypatch.setitem( + sys.modules, + "vime.backends.megatron_utils.update_weight.coordinator", + coordinator_mod, + ) return dist_state diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py index 0e96b1405..b4b9f6f3b 100644 --- a/tests/test_megatron_argument_validation.py +++ b/tests/test_megatron_argument_validation.py @@ -307,18 +307,94 @@ def test_vime_validate_args_preserves_zero_rollout_gpus_without_colocate(monkeyp @pytest.mark.unit -def test_update_weight_delta_disabled(monkeypatch): +def test_update_weight_direct_delta_allowed(monkeypatch): module = load_vime_arguments_module(monkeypatch) - for transport, colocate in (("nccl", False), ("tensor", False), ("nccl", True)): - args = types.SimpleNamespace( - update_weight_mode="delta", - update_weight_transport=transport, - update_weight_disk_dir=None, - update_weight_delta_dir=None, - colocate=colocate, - ) - with pytest.raises(NotImplementedError, match="unverified on vime"): - module._validate_update_weight_args(args) + args = types.SimpleNamespace( + update_weight_mode="delta", + update_weight_transport="nccl", + colocate=False, + rollout_external=False, + train_backend="megatron", + offload_rollout=False, + enable_mtp_training=False, + vllm_speculative_config=None, + fp16=False, + pipeline_model_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + num_layers_per_virtual_pipeline_stage=None, + vllm_pipeline_parallel_size=1, + vllm_data_parallel_size=1, + update_weight_start_version=0, + vllm_worker_extension_cls="", + ) + + module._validate_update_weight_args(args) + + +@pytest.mark.unit +def test_update_weight_delta_disk_is_reserved(monkeypatch): + module = load_vime_arguments_module(monkeypatch) + args = types.SimpleNamespace( + update_weight_mode="delta", + update_weight_transport="disk", + ) + + with pytest.raises(NotImplementedError, match="reserved but not implemented"): + module._validate_update_weight_args(args) + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"colocate": True}, "non-colocated"), + ({"rollout_external": True}, "VIME-launched"), + ({"use_fault_tolerance": True}, "replacement rollout workers"), + ({"offload_rollout": True}, "rollout offload"), + ({"fp16": True}, "BF16"), + ({"pipeline_model_parallel_size": 2}, "Megatron PP=1"), + ({"vllm_pipeline_parallel_size": 2}, "vLLM PP=1"), + ({"vllm_data_parallel_size": 2}, "vLLM DP=1"), + ( + {"vllm_enable_deterministic_inference": True}, + "batch-invariant deterministic inference", + ), + ({"update_weight_start_version": 4}, "version 0"), + ({"vllm_speculative_config": {"method": "mtp"}}, "speculative"), + ( + {"rollout_function_path": "vime.rollout.fully_async_rollout.generate_rollout_fully_async"}, + "fully-async", + ), + ], +) +def test_update_weight_direct_delta_rejects_unsupported_mvp_modes( + monkeypatch, + override, + message, +): + module = load_vime_arguments_module(monkeypatch) + values = dict( + update_weight_mode="delta", + update_weight_transport="nccl", + colocate=False, + rollout_external=False, + train_backend="megatron", + offload_rollout=False, + enable_mtp_training=False, + vllm_speculative_config=None, + fp16=False, + pipeline_model_parallel_size=1, + virtual_pipeline_model_parallel_size=None, + num_layers_per_virtual_pipeline_stage=None, + vllm_pipeline_parallel_size=1, + vllm_data_parallel_size=1, + update_weight_start_version=0, + vllm_worker_extension_cls="", + ) + values.update(override) + + with pytest.raises((ValueError, NotImplementedError), match=message): + module._validate_update_weight_args(types.SimpleNamespace(**values)) if __name__ == "__main__": diff --git a/tests/utils/test_direct_dwu_roundtrip.py b/tests/utils/test_direct_dwu_roundtrip.py new file mode 100644 index 000000000..a87393ef3 --- /dev/null +++ b/tests/utils/test_direct_dwu_roundtrip.py @@ -0,0 +1,413 @@ +"""The comprehensive CPU test suite for VIME direct DWU. + +It drives the REAL sender (CheckpointDeltaSource) into the REAL receiver +(VimeDeltaNCCLWeightTransferEngine): manifests travel as +``chunk.update_info()`` and weights as ``chunk.wire_tensors()`` in sender +order, and the receiver drains them through its own broadcast sequence into a +fake model, so schema, framing, lifecycle, and value semantics are all +checked against each other in one place. Sender-side refusals and +receiver-side protocol rejections live here too, against the same round-trip +harness, instead of in per-side stub suites. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from dataclasses import dataclass +from pathlib import Path + +import pytest +import torch + +from vime.backends.megatron_utils.update_weight.checkpoint_delta import CheckpointDeltaSource + +NUM_GPUS = 0 + + +@dataclass +class _Patch: + name: str + shape: tuple[int, ...] + dtype: torch.dtype + values: torch.Tensor + indices: torch.Tensor | None = None + + +@pytest.fixture +def receiver_module(): + module_names = [ + "vllm", + "vllm.distributed", + "vllm.distributed.weight_transfer", + "vllm.distributed.weight_transfer.base", + "vllm.distributed.weight_transfer.nccl_engine", + "vllm.model_executor", + "vllm.model_executor.model_loader", + "vllm.model_executor.model_loader.reload", + ] + saved = {name: sys.modules.get(name) for name in module_names} + + class Factory: + registry = {} + + @classmethod + def register_engine(cls, name, engine_cls): + cls.registry[name] = engine_cls + + class WeightTransferUpdateInfo: + pass + + class NCCLWeightTransferEngine: + def __init__(self, config, vllm_config, device, model): + self.config = config + self.vllm_config = vllm_config + self.parallel_config = vllm_config.parallel_config + self.model_config = vllm_config.model_config + self.device = device + self.model = model + self.model_update_group = None + + def update_weights(self, update_info): + typed = self.update_info_cls(**update_info) + self.receive_weights(typed) + if torch.accelerator.is_available(): + torch.accelerator.synchronize() + + def shutdown(self): + self.model_update_group = None + + vllm = types.ModuleType("vllm") + vllm.__path__ = [] + distributed = types.ModuleType("vllm.distributed") + distributed.__path__ = [] + weight_transfer = types.ModuleType("vllm.distributed.weight_transfer") + weight_transfer.__path__ = [] + weight_transfer.WeightTransferEngineFactory = Factory + base = types.ModuleType("vllm.distributed.weight_transfer.base") + base.WeightTransferUpdateInfo = WeightTransferUpdateInfo + nccl = types.ModuleType("vllm.distributed.weight_transfer.nccl_engine") + nccl.NCCLWeightTransferEngine = NCCLWeightTransferEngine + + model_executor = types.ModuleType("vllm.model_executor") + model_executor.__path__ = [] + model_loader = types.ModuleType("vllm.model_executor.model_loader") + model_loader.__path__ = [] + reload_module = types.ModuleType("vllm.model_executor.model_loader.reload") + reload_module.events = [] + reload_module.initialize_layerwise_reload = lambda model: reload_module.events.append("initialize") + reload_module.finalize_layerwise_reload = lambda model, config: reload_module.events.append("finalize") + + sys.modules.update( + { + "vllm": vllm, + "vllm.distributed": distributed, + "vllm.distributed.weight_transfer": weight_transfer, + "vllm.distributed.weight_transfer.base": base, + "vllm.distributed.weight_transfer.nccl_engine": nccl, + "vllm.model_executor": model_executor, + "vllm.model_executor.model_loader": model_loader, + "vllm.model_executor.model_loader.reload": reload_module, + } + ) + + module_path = Path(__file__).resolve().parents[2] / "vime" / "backends" / "vllm_utils" / "checkpoint_delta.py" + module_name = "test_vime_direct_dwu_roundtrip_module" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec.loader is not None + spec.loader.exec_module(module) + module._test_factory = Factory + module._test_reload = reload_module + + try: + yield module + finally: + sys.modules.pop(module_name, None) + for name, original in saved.items(): + if original is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = original + + +class _WireGroup: + """Fake NCCL group fed with the sender's wire tensors in sender order. + + Assumes what production relies on: vLLM's ``trainer_send_weights(..., + packed=False)`` decomposes into one broadcast per wire tensor in iteration + order, which the receiver drains with one matching broadcast each.""" + + def __init__(self): + self.payloads: list[torch.Tensor] = [] + + def send_chunk(self, chunk) -> None: + self.payloads.extend(tensor for _, tensor in chunk.wire_tensors()) + + def broadcast(self, destination, *, src, stream): + assert src == 0 + assert stream == "current-stream" + if destination.numel() == 0: + return + payload = self.payloads.pop(0) + assert payload.dtype == destination.dtype + assert payload.numel() == destination.numel() + destination.copy_(payload) + + +def _apply_patches(model, patches, *, max_chunk_bytes, validate_unique_indices): + """Faithful single-copy stand-in for vLLM's checkpoint patch API on an + unsharded model: dense patches replace the runtime tensor, sparse patches + scatter absolute values into flat checkpoint positions.""" + assert max_chunk_bytes > 0 + assert validate_unique_indices is False + applied = set() + for patch in patches: + destination = model.runtime[patch.name] + assert tuple(destination.shape) == tuple(patch.shape) + values = patch.values.to(patch.dtype) + if patch.indices is None: + assert values.numel() == destination.numel() + destination.copy_(values.reshape(patch.shape)) + else: + assert patch.indices.dtype == torch.int32 + assert not torch.isnan(values).any() + destination.reshape(-1)[patch.indices.to(torch.long)] = values + applied.add(patch.name) + model.load_calls.append(patch.name) + return applied + + +def _make_engine(module, monkeypatch, model): + monkeypatch.setattr(module, "_checkpoint_patch_api", lambda: (_Patch, _apply_patches)) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: "current-stream") + engine = object.__new__(module.VimeDeltaNCCLWeightTransferEngine) + engine.config = types.SimpleNamespace() + engine.vllm_config = types.SimpleNamespace( + parallel_config=types.SimpleNamespace(), + model_config=types.SimpleNamespace(dtype=torch.bfloat16), + ) + engine.parallel_config = engine.vllm_config.parallel_config + engine.model_config = engine.vllm_config.model_config + engine.device = torch.device("cpu") + engine.model = model + engine.model_update_group = _WireGroup() + engine._committed_version = 0 + engine._session_base_version = None + engine._session_target_version = None + engine._session_encoding = None + engine._next_sequence_no = 0 + engine._reload_initialized = False + engine._final_received = False + engine._update_failed = False + return engine + + +def _bf16(values) -> torch.Tensor: + return torch.tensor(values, dtype=torch.bfloat16) + + +def _encode_session(source, buckets, *, base_version, target_version): + """Run one trainer update and return its data chunks plus final manifest.""" + source.begin_update(base_version=base_version, target_version=target_version) + chunks = [chunk for bucket in buckets for chunk in source.encode_chunk(bucket)] + return chunks, source.finish_update() + + +def _run_session(source, engine, buckets, *, base_version, target_version): + """Push one full trainer update through the receiver, sender-framed.""" + chunks, final = _encode_session(source, buckets, base_version=base_version, target_version=target_version) + engine.start_weight_update() + for chunk in [*chunks, final]: + engine.model_update_group.send_chunk(chunk) + engine.update_weights(chunk.update_info()) + engine.finish_weight_update() + source.commit() + assert engine.model_update_group.payloads == [] + return chunks + + +@pytest.mark.unit +def test_seed_sparse_and_noop_sessions_apply_exact_values(receiver_module, monkeypatch): + model = types.SimpleNamespace( + runtime={ + "model.a": torch.zeros(2, 2, dtype=torch.bfloat16), + "model.b": torch.zeros(3, dtype=torch.bfloat16), + }, + load_calls=[], + ) + engine = _make_engine(receiver_module, monkeypatch, model) + source = CheckpointDeltaSource() + reload_events = receiver_module._test_reload.events + + # Dense seed across two exporter buckets: the rollout starts from + # different weights and must converge. + seed = [ + [("model.a", _bf16([[1, 2], [3, 4]]))], + [("model.b", _bf16([0.0, 20, 30]))], + ] + chunks = _run_session(source, engine, seed, base_version=0, target_version=1) + assert [chunk.encoding for chunk in chunks] == ["dense", "dense"] + assert engine._committed_version == 1 + assert reload_events == ["initialize", "finalize"] + for bucket in seed: + for name, expected in bucket: + assert torch.equal(model.runtime[name], expected) + + # Sparse update: absolute values at changed positions, including a bitwise + # 0.0 -> -0.0 flip that a value compare would miss. + step_two = [ + [("model.a", _bf16([[1, 9], [3, -4]]))], + [("model.b", _bf16([-0.0, 20, 30]))], + ] + dense_calls = len(model.load_calls) + chunks = _run_session(source, engine, step_two, base_version=1, target_version=2) + assert [chunk.encoding for chunk in chunks] == ["indices", "indices"] + assert engine._committed_version == 2 + assert reload_events == ["initialize", "finalize"] # sparse sessions skip reload + for bucket in step_two: + for name, expected in bucket: + assert torch.equal(model.runtime[name], expected) + assert torch.signbit(model.runtime["model.b"][0]) + assert len(model.load_calls) > dense_calls + assert 0 < source.changed_elements < source.total_elements + assert source.wire_bytes > 0 + + # No-op update: nothing changed, only the final manifest travels, the + # version still advances, and no weights are loaded. + load_calls = len(model.load_calls) + chunks = _run_session(source, engine, step_two, base_version=2, target_version=3) + assert chunks == [] + assert engine._committed_version == 3 + assert len(model.load_calls) == load_calls + for bucket in step_two: + for name, expected in bucket: + assert torch.equal(model.runtime[name], expected) + + +@pytest.mark.unit +def test_abandoned_session_poisons_receiver_until_restart(receiver_module, monkeypatch): + model = types.SimpleNamespace( + runtime={"model.a": torch.zeros(2, dtype=torch.bfloat16)}, + load_calls=[], + ) + engine = _make_engine(receiver_module, monkeypatch, model) + source = CheckpointDeltaSource() + _run_session(source, engine, [[("model.a", _bf16([1, 2]))]], base_version=0, target_version=1) + + # The trainer dies after shipping data but before the final manifest. + source.begin_update(base_version=1, target_version=2) + chunks = source.encode_chunk([("model.a", _bf16([1, 5]))]) + assert chunks + engine.start_weight_update() + for chunk in chunks: + engine.model_update_group.send_chunk(chunk) + engine.update_weights(chunk.update_info()) + source.abort() + + with pytest.raises(RuntimeError, match="without a final manifest"): + engine.finish_weight_update() + with pytest.raises(RuntimeError, match="previous direct DWU session failed"): + engine.start_weight_update() + + # The aborted source keeps its committed baseline: the next update diffs + # against version 1, not the half-shipped version 2 state. + source.begin_update(base_version=1, target_version=2) + retry = source.encode_chunk([("model.a", _bf16([1, 5]))]) + assert len(retry) == 1 + source.abort() + + +@pytest.mark.unit +def test_source_refuses_invalid_updates(): + nan = float("nan") + source = CheckpointDeltaSource() + with pytest.raises(ValueError, match="must start at version 0"): + source.begin_update(base_version=1, target_version=2) + + source.begin_update(base_version=0, target_version=1) + source.encode_chunk([("model.a", _bf16([1, nan]))]) + source.encode_chunk([("model.b", _bf16([3, 4]))]) + source.finish_update() + source.commit() + + # Version guards against the committed snapshot. + with pytest.raises(RuntimeError, match="snapshot=1, update=0"): + source.begin_update(base_version=0, target_version=1) + with pytest.raises(ValueError, match=r"base_version \+ 1"): + source.begin_update(base_version=1, target_version=3) + + # Inventory drift within a bucket fails instead of silently diffing. + source.begin_update(base_version=1, target_version=2) + with pytest.raises(RuntimeError, match="tensor count changed"): + source.encode_chunk([("model.a", _bf16([1, nan])), ("model.x", _bf16([9.0]))]) + source.abort() + + # A NaN that survives training bit-identically is not a change; a weight + # flipping TO NaN is refused at the source, naming the tensor, because NaN + # is the patch API's unchanged-value sentinel. + source.begin_update(base_version=1, target_version=2) + assert source.encode_chunk([("model.a", _bf16([1, nan]))]) == [] + with pytest.raises(ValueError, match=r"model\.b: training produced NaN"): + source.encode_chunk([("model.b", _bf16([3, nan]))]) + source.abort() + + +@pytest.mark.unit +def test_receiver_rejects_protocol_drift(receiver_module, monkeypatch): + """Every malformed session poisons the worker instead of corrupting it.""" + + def fresh(): + model = types.SimpleNamespace( + runtime={"model.a": torch.zeros(2, dtype=torch.bfloat16)}, + load_calls=[], + ) + return _make_engine(receiver_module, monkeypatch, model) + + source = CheckpointDeltaSource() + chunks, final = _encode_session(source, [[("model.a", _bf16([1, 2]))]], base_version=0, target_version=1) + data_info = chunks[0].update_info() + + def deliver(engine, chunk, **overrides): + engine.model_update_group.send_chunk(chunk) + engine.update_weights({**chunk.update_info(), **overrides}) + + # Wrong base version: the wire payload is still drained first so all TP + # ranks complete the same collectives, then the worker fails stop. + engine = fresh() + engine._committed_version = 4 + engine.start_weight_update() + with pytest.raises(RuntimeError, match="base version mismatch"): + deliver(engine, chunks[0]) + assert engine.model_update_group.payloads == [] + assert engine._update_failed is True + + # Out-of-order sequence number. + engine = fresh() + engine.start_weight_update() + with pytest.raises(ValueError, match="sequence mismatch"): + deliver(engine, chunks[0], sequence_no=5) + + # Non-BF16 wire dtype. + engine = fresh() + engine.start_weight_update() + with pytest.raises(ValueError, match="must be BF16"): + deliver(engine, chunks[0], value_dtype_name="float16") + + # A dense session whose only manifest is the final one shipped no weights. + engine = fresh() + engine.start_weight_update() + with pytest.raises(ValueError, match="did not carry any weights"): + engine.update_weights({**final.update_info(), "sequence_no": 0}) + + # Data arriving after the final manifest. + engine = fresh() + engine.start_weight_update() + deliver(engine, chunks[0]) + engine.update_weights(final.update_info()) + engine.model_update_group.send_chunk(chunks[0]) + with pytest.raises(ValueError, match="after the final manifest"): + engine.update_weights({**data_info, "sequence_no": 2}) + diff --git a/tests/utils/test_update_weight_from_distributed.py b/tests/utils/test_update_weight_from_distributed.py index ccf9c323a..cb47487f9 100644 --- a/tests/utils/test_update_weight_from_distributed.py +++ b/tests/utils/test_update_weight_from_distributed.py @@ -20,9 +20,11 @@ from vime.utils.types import ParamInfo MODULE_PATH = "vime.backends.megatron_utils.update_weight.update_weight_from_distributed" +COORDINATOR_MODULE = "vime.backends.megatron_utils.update_weight.coordinator" COMMON_MODULE = "vime.backends.megatron_utils.update_weight.common" DIRECT_MODULE = "vime.backends.megatron_utils.update_weight.hf_weight_iterator_direct" CONVERTER_MODULE = "vime.backends.megatron_utils.megatron_to_hf" +CHECKPOINT_DELTA_UPDATER_MODULE = "vime.backends.megatron_utils.update_weight.update_weight_from_checkpoint_delta" NUM_GPUS = 0 @@ -54,6 +56,8 @@ "vllm.distributed.weight_transfer.nccl_engine", "triton", "triton.language", + COORDINATOR_MODULE, + CHECKPOINT_DELTA_UPDATER_MODULE, ) @@ -137,7 +141,15 @@ class RecordingEngine: default_factory=lambda: RecordingRemoteMethod("destroy_ref") ) start_weight_update: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("start_ref")) + start_draft_weight_update: RecordingRemoteMethod = field( + default_factory=lambda: RecordingRemoteMethod("draft_ref") + ) finish_weight_update: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("finish_ref")) + pause_generation: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("pause_ref")) + flush_cache: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("flush_ref")) + post_process_weights: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("post_ref")) + set_weight_version: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("version_ref")) + continue_generation: RecordingRemoteMethod = field(default_factory=lambda: RecordingRemoteMethod("continue_ref")) @dataclass @@ -199,7 +211,15 @@ def _patch_trainer_send(monkeypatch, upw, seen: list[dict]) -> None: def _make_instance(upw): obj = object.__new__(upw.UpdateWeightFromDistributed) - obj.args = type("Args", (), {"update_weight_buffer_size": 1 << 30})() + obj.args = type( + "Args", + (), + { + "update_weight_buffer_size": 1 << 30, + "enable_mtp_training": False, + "vllm_speculative_config": None, + }, + )() obj.model = [] obj.weights_getter = lambda: {} obj.model_name = "test" @@ -292,6 +312,7 @@ def test_remote_kwargs_are_always_packed(upw, monkeypatch): assert "packed" not in kw assert "group_name" not in kw assert kw["weight_version"] == "42" + assert kw["record_weight_version"] is True assert kw["names"] == ["layer.0.weight"] assert kw["shapes"] == [torch.Size([2, 2])] assert kw["dtypes"] == [torch.float32] @@ -686,22 +707,178 @@ def fake_barrier(*, group=None, **kwargs): upw._begin_vllm_weight_update_session(engines) upw._end_vllm_weight_update_session(engines) + upw._begin_vllm_weight_update_session(engines, draft=True) assert len(engines[0].start_weight_update.calls) == 1 assert engines[0].start_weight_update.calls[0].kwargs["is_checkpoint_format"] is True assert len(engines[1].start_weight_update.calls) == 1 assert len(engines[0].finish_weight_update.calls) == 1 assert len(engines[1].finish_weight_update.calls) == 1 - assert barrier_calls == ["dummy-gloo-group", "dummy-gloo-group"] + assert len(engines[0].start_draft_weight_update.calls) == 1 + assert len(engines[1].start_draft_weight_update.calls) == 1 + assert barrier_calls == ["dummy-gloo-group", "dummy-gloo-group", "dummy-gloo-group"] @pytest.mark.unit -def test_source_wraps_sync_with_weight_update_session(upw): - src = inspect.getsource(upw.UpdateWeightFromDistributed.update_weights) - assert "_begin_vllm_weight_update_session" in src - assert "start_draft_weight_update" in src - assert "_end_vllm_weight_update_session" in src - assert src.count("_send_weights_to_rollout_engines") == 2 +def test_distributed_adapter_follows_coordinator_outcome(upw, monkeypatch): + obj = _make_instance(upw) + sessions = [] + + class RecordingCoordinator: + def __init__(self, rollout_engines, quantization_config): + assert rollout_engines is obj.rollout_engines + assert quantization_config is None + + def run(self, *, current_version, transfer_target, transfer_draft, commit): + assert current_version == 0 + assert transfer_draft is None + assert commit is None + transfer_target(1) + return 1 + + monkeypatch.setattr(upw, "WeightUpdateCoordinator", RecordingCoordinator) + monkeypatch.setattr( + obj, + "_run_weight_update_session", + lambda *, draft=False: sessions.append((obj.weight_version, draft)), + ) + + obj.update_weights() + + assert sessions == [(1, False)] + assert obj.weight_version == 1 + + # A later failing update restores the adapter's committed version marker. + class FailingCoordinator: + def __init__(self, rollout_engines, quantization_config): + pass + + def run(self, *, current_version, transfer_target, transfer_draft, commit): + transfer_target(current_version + 1) + raise RuntimeError("transfer failed") + + monkeypatch.setattr(upw, "WeightUpdateCoordinator", FailingCoordinator) + + with pytest.raises(RuntimeError, match="transfer failed"): + obj.update_weights() + + assert obj.weight_version == 1 + + +@pytest.mark.unit +def test_distributed_adapter_runs_target_then_mtp_draft(upw, monkeypatch): + obj = _make_instance(upw) + obj.args.enable_mtp_training = True + obj.args.vllm_speculative_config = {"method": "mtp"} + sessions = [] + + class RecordingCoordinator: + def __init__(self, rollout_engines, quantization_config): + pass + + def run(self, *, current_version, transfer_target, transfer_draft, commit): + assert transfer_draft is not None + assert commit is None + transfer_target(current_version + 1) + transfer_draft(current_version + 1) + return current_version + 1 + + monkeypatch.setattr(upw, "WeightUpdateCoordinator", RecordingCoordinator) + monkeypatch.setattr( + obj, + "_run_weight_update_session", + lambda *, draft=False: sessions.append((obj.weight_version, draft)), + ) + + obj.update_weights() + + assert sessions == [(1, False), (1, True)] + assert obj.weight_version == 1 + + +@pytest.fixture +def checkpoint_delta_updater(upw): + del upw + sys.modules.pop(CHECKPOINT_DELTA_UPDATER_MODULE, None) + try: + yield importlib.import_module(CHECKPOINT_DELTA_UPDATER_MODULE) + finally: + sys.modules.pop(CHECKPOINT_DELTA_UPDATER_MODULE, None) + + +@pytest.mark.unit +def test_checkpoint_delta_snapshot_follows_coordinator_outcome( + upw, + checkpoint_delta_updater, + monkeypatch, +): + events = [] + + class Source: + is_seed_update = True + total_elements = 4 + changed_elements = 4 + wire_bytes = 8 + + def begin_update(self, *, base_version, target_version): + events.append(("begin", base_version, target_version)) + + def commit(self): + events.append(("commit",)) + + def abort(self): + events.append(("abort",)) + + obj = object.__new__(checkpoint_delta_updater.UpdateWeightFromCheckpointDelta) + obj.weight_version = 0 + obj._is_pp_src_rank = True + obj._delta_source = Source() + obj.update_weight_metrics = {} + + def coordinator_success(_self): + events.append(("coordinator-enter",)) + commit = _self._get_weight_update_commit() + assert commit is not None + commit() + events.append(("coordinator-success",)) + obj.weight_version = 1 + + monkeypatch.setattr( + upw.UpdateWeightFromDistributed, + "update_weights", + coordinator_success, + ) + + obj.update_weights() + + assert events == [ + ("begin", 0, 1), + ("coordinator-enter",), + ("commit",), + ("coordinator-success",), + ] + + # Coordinator failure aborts the pending snapshot instead of committing it. + events.clear() + + def coordinator_failure(_self): + events.append(("coordinator-failure",)) + raise RuntimeError("publish failed") + + monkeypatch.setattr( + upw.UpdateWeightFromDistributed, + "update_weights", + coordinator_failure, + ) + + with pytest.raises(RuntimeError, match="publish failed"): + obj.update_weights() + + assert events == [ + ("begin", 1, 2), + ("coordinator-failure",), + ("abort",), + ] @pytest.mark.unit diff --git a/tests/utils/test_update_weight_from_tensor.py b/tests/utils/test_update_weight_from_tensor.py index 58833c1be..6606766ba 100644 --- a/tests/utils/test_update_weight_from_tensor.py +++ b/tests/utils/test_update_weight_from_tensor.py @@ -20,6 +20,7 @@ import torch MODULE_PATH = "vime.backends.megatron_utils.update_weight.update_weight_from_tensor" +COORDINATOR_MODULE = "vime.backends.megatron_utils.update_weight.coordinator" NUM_GPUS = 0 @@ -37,11 +38,13 @@ def _install_stubs(): dist_stub.get_process_group_ranks.return_value = [0, 1] dist_stub.barrier = MagicMock() dist_stub.all_gather_object = MagicMock() + dist_stub.broadcast_object_list = MagicMock() _dist.get_rank = dist_stub.get_rank _dist.get_world_size = dist_stub.get_world_size _dist.get_process_group_ranks = dist_stub.get_process_group_ranks _dist.barrier = dist_stub.barrier _dist.all_gather_object = dist_stub.all_gather_object + _dist.broadcast_object_list = dist_stub.broadcast_object_list hf_iter_stub = MagicMock() hf_iter_stub.get_hf_weight_chunks.return_value = iter([]) @@ -81,10 +84,18 @@ def _install_stubs(): "ray", "ray.actor", "vime.utils.distributed_utils", + COORDINATOR_MODULE, "vime.backends.megatron_utils.update_weight.hf_weight_iterator_base", "vime.backends.megatron_utils.update_weight.update_weight_from_distributed", ) -_DIST_ATTRS = ("get_rank", "get_world_size", "get_process_group_ranks", "barrier", "all_gather_object") +_DIST_ATTRS = ( + "get_rank", + "get_world_size", + "get_process_group_ranks", + "barrier", + "all_gather_object", + "broadcast_object_list", +) @pytest.fixture(scope="module") @@ -133,6 +144,7 @@ class RecordingVLLMEngine: update_weights_from_tensor: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) pause_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) flush_cache: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) + set_weight_version: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) continue_generation: RecordingRemoteMethod = field(default_factory=RecordingRemoteMethod) @@ -175,6 +187,7 @@ def _make_instance(upw_vllm, args=None): def _bind_single_slot(obj, engine, *, src=0): """Bind ``obj`` to one colocated engine forming a slot whose leader rank is ``src``.""" obj.rollout_engines = [engine] + obj._all_rollout_engines = [engine] obj._ipc_engine = engine obj._ipc_gather_group = "slot_group" obj._ipc_gather_src = src @@ -239,7 +252,10 @@ def test_colocated_lifecycle_uses_pause_flush_and_weight_transfer_apis(upw_vllm) assert len(engine.start_weight_update.calls) == 1 assert engine.start_weight_update.calls[0].kwargs.get("is_checkpoint_format") is True assert len(engine.finish_weight_update.calls) == 1 + assert len(engine.set_weight_version.calls) == 1 + assert engine.set_weight_version.calls[0].args == ("1",) assert len(engine.continue_generation.calls) == 1 + assert obj.weight_version == 1 # Both chunks are kept alive until the bounded in-flight batch drains. assert counters["ipc_collect"] == 2 # lifecycle barriers (no per-chunk barrier). @@ -273,6 +289,63 @@ def test_colocated_mtp_updates_target_then_draft_from_fresh_weight_stream(upw_vl assert len(engine.finish_weight_update.calls) == 2 assert len(engine.update_weights_from_tensor.calls) == 4 assert obj._hf_weight_iterator.get_hf_weight_chunks.call_count == 2 + assert obj.weight_version == 1 + + +@pytest.mark.unit +def test_mixed_colocated_and_remote_engines_share_control_plane(upw_vllm): + obj = _make_instance(upw_vllm) + colocated = RecordingVLLMEngine() + remote = RecordingVLLMEngine() + obj.rollout_engines = [colocated] + obj.distributed_rollout_engines = [remote] + obj._all_rollout_engines = [colocated, remote] + obj.use_distribute = True + obj._is_distributed_src_rank = True + obj._ipc_engine = colocated + obj._ipc_gather_group = "slot_group" + obj._ipc_gather_src = 0 + obj._send_weight_chunks = MagicMock() + + _run_update(obj) + + for engine in (colocated, remote): + assert len(engine.pause_generation.calls) == 1 + assert len(engine.flush_cache.calls) == 1 + assert len(engine.start_weight_update.calls) == 1 + assert len(engine.finish_weight_update.calls) == 1 + assert len(engine.set_weight_version.calls) == 1 + assert engine.set_weight_version.calls[0].args == ("1",) + assert len(engine.continue_generation.calls) == 1 + assert obj.weight_version == 1 + + # The remote leg defers version publication to the coordinator: the NCCL + # send itself must not record a version on the engines. + sender = _make_instance(upw_vllm) + deferred_remote = RecordingVLLMEngine() + sender.use_distribute = True + sender._is_distributed_src_rank = True + sender._model_update_groups = "groups" + sender.distributed_rollout_engines = [deferred_remote] + tensors = _chunks(1)[0] + with patch( + f"{MODULE_PATH}._send_to_colocated_engine", + return_value=(["ipc-ref"], ["keepalive"]), + ), patch( + f"{MODULE_PATH}.update_weights_from_distributed", + return_value=["nccl-ref"], + ) as send_distributed: + refs, keepalive = sender._send_hf_params(tensors) + + assert refs == ["ipc-ref", "nccl-ref"] + assert keepalive == ["keepalive"] + send_distributed.assert_called_once_with( + "groups", + 0, + [deferred_remote], + tensors, + record_weight_version=False, + ) @pytest.mark.unit @@ -306,8 +379,10 @@ def test_send_via_ipc_dispatches_update_weights_from_tensor_with_version(upw_vll assert kwargs["dtype_names"] == dummy_info["dtype_names"] assert kwargs["shapes"] == dummy_info["shapes"] assert kwargs["ipc_handles"] is dummy_info["ipc_handles"] - # weight_version is the trainer's post-increment version (0 + 1 = 1) as a str + # Chunks carry the candidate version but do not publish it before the + # coordinator commits the complete update. assert kwargs["weight_version"] == "1" + assert kwargs["record_weight_version"] is False # finish_weight_update is a stateless bookend now — no kwargs assert len(engine.finish_weight_update.calls) == 1 assert engine.finish_weight_update.calls[0].kwargs == {} @@ -359,6 +434,7 @@ def fake_gather_object(payload, object_gather_list=None, dst=None, group=None): assert kwargs["shapes"] == dummy_info_0["shapes"] assert set(kwargs["ipc_handles"]) == {"uuid-gpu0", "uuid-gpu1"} assert kwargs["weight_version"] == "1" + assert kwargs["record_weight_version"] is False @pytest.mark.unit @@ -469,6 +545,7 @@ def test_connect_binds_engine_and_slot_leader_per_gpu_slot(upw_vllm): assert obj._ipc_gather_src == expected_src is_coordinator = rank == obj._ipc_gather_src assert is_coordinator is (rank in (0, 2)) + assert obj._all_rollout_engines == engines assert obj.use_distribute is False assert obj.distributed_rollout_engines == [] # vLLM #39212: init_weight_transfer_engine fires once during connect (rank 0 only). diff --git a/tests/utils/test_vllm_engine.py b/tests/utils/test_vllm_engine.py index 8e919320a..8c3db550a 100644 --- a/tests/utils/test_vllm_engine.py +++ b/tests/utils/test_vllm_engine.py @@ -52,6 +52,9 @@ def vllm_args() -> SimpleNamespace: vllm_pipeline_parallel_size=1, vllm_data_parallel_size=1, vllm_dp_size=1, + update_weight_mode="full", + update_weight_transport="nccl", + vllm_worker_extension_cls="", ) @@ -247,6 +250,51 @@ def test_compute_server_args_external_check_fields_skip_orchestration_fields(vll assert sa["weight_transfer_config"] == {"backend": "nccl"} +@pytest.mark.unit +def test_direct_delta_server_config_and_env(vllm_args, monkeypatch): + vllm_args.rollout_external = False + vllm_args.update_weight_mode = "delta" + + sa, check_fields = mod._compute_server_args( + vllm_args, + rank=0, + dist_init_addr=None, + host="127.0.0.1", + port=8000, + ) + assert sa["weight_transfer_config"] == {"backend": "vime_delta_nccl"} + assert sa["worker_extension_cls"] == ( + "vime.backends.vllm_utils.checkpoint_delta.VimeDeltaWorkerExtension" + ) + assert "weight_transfer_config" in check_fields + assert "worker_extension_cls" in check_fields + + # Delta mode owns the worker extension slot. + vllm_args.vllm_worker_extension_cls = "custom.WorkerExtension" + with pytest.raises(ValueError, match="owns --vllm-worker-extension-cls"): + mod._compute_server_args( + vllm_args, + rank=0, + dist_init_addr=None, + host="127.0.0.1", + port=8000, + ) + vllm_args.vllm_worker_extension_cls = "" + + # The worker subprocess must see vime on PYTHONPATH to import the extension. + monkeypatch.delenv("PYTHONPATH", raising=False) + monkeypatch.delenv("VLLM_ALLOW_INSECURE_SERIALIZATION", raising=False) + env = mod._build_subprocess_env( + { + "_args": vllm_args, + "_visible_devices": "6,7", + } + ) + assert "PYTHONPATH" in env + assert env["CUDA_VISIBLE_DEVICES"] == "6,7" + assert "VLLM_ALLOW_INSECURE_SERIALIZATION" not in env + + @pytest.mark.unit def test_build_vllm_subprocess_env_colocate(vllm_args, monkeypatch): vllm_args.colocate = True @@ -441,6 +489,54 @@ def test_get_weight_version_worker_rank_returns_none_without_raise(vllm_engine): assert vllm_engine.get_weight_version() is None +@pytest.mark.unit +def test_set_weight_version_commits_engine_core_before_wrapper(vllm_engine, monkeypatch): + calls = [] + monkeypatch.setattr( + vllm_engine, + "_make_request", + lambda endpoint, payload: calls.append((endpoint, payload)) or {"ok": True}, + ) + vllm_engine._weight_version = "old" + + assert vllm_engine.set_weight_version("8") == {"ok": True} + assert calls == [("update_weight_version", {"new_version": "8"})] + assert vllm_engine._weight_version == "8" + + # A server-side failure must not advance the local marker. + monkeypatch.setattr( + vllm_engine, + "_make_request", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("core failed")), + ) + with pytest.raises(RuntimeError, match="core failed"): + vllm_engine.set_weight_version("9") + assert vllm_engine._weight_version == "8" + + +@pytest.mark.unit +def test_set_weight_version_missing_endpoint_degrades_except_for_direct_delta(vllm_engine, monkeypatch): + """vLLM < 0.27.2 has no update_weight_version endpoint: non-delta modes fall + back to local version tracking; direct DWU publishes through the server tag + and must fail-stop.""" + + def raise_404(*_args, **_kwargs): + _MockResponse(status_code=404).raise_for_status() + + monkeypatch.setattr(vllm_engine, "_make_request", raise_404) + vllm_engine._weight_version = "old" + + assert vllm_engine.set_weight_version("8") is None + assert vllm_engine.set_weight_version("9") is None + assert vllm_engine._weight_version == "9" + + vllm_engine.args.update_weight_mode = "delta" + vllm_engine.args.update_weight_transport = "nccl" + with pytest.raises(requests.exceptions.HTTPError): + vllm_engine.set_weight_version("10") + assert vllm_engine._weight_version == "9" + + @pytest.mark.unit def test_update_weights_from_distributed_posts_update_weights_without_checkpoint_flag(vllm_engine, monkeypatch): calls: list[dict] = [] diff --git a/tests/utils/test_weight_update_coordinator.py b/tests/utils/test_weight_update_coordinator.py new file mode 100644 index 000000000..c08ce1c5f --- /dev/null +++ b/tests/utils/test_weight_update_coordinator.py @@ -0,0 +1,273 @@ +"""CPU unit tests for the shared weight update coordinator.""" + +from __future__ import annotations + +import importlib +import sys +from dataclasses import dataclass, field +from pathlib import Path + +_tests_root = Path(__file__).resolve().parents[1] +if str(_tests_root) not in sys.path: + sys.path.insert(0, str(_tests_root)) + +import _unit_stubs +import pytest + +MODULE_PATH = "vime.backends.megatron_utils.update_weight.coordinator" +_STUBBED_MODULES = ("ray", "ray.actor", "vime.utils.distributed_utils", MODULE_PATH) + + +@pytest.fixture(scope="module") +def coordinator_module(): + saved = _unit_stubs.save_sys_modules(_STUBBED_MODULES) + for name in _STUBBED_MODULES: + sys.modules.pop(name, None) + _unit_stubs.install_ray_stub() + _unit_stubs.install_vime_distributed_utils_stub() + try: + yield importlib.import_module(MODULE_PATH) + finally: + _unit_stubs.restore_sys_modules(saved) + + +@dataclass +class _RemoteCall: + args: tuple + kwargs: dict + + +class RecordingRemoteMethod: + def __init__(self, name: str, events: list[tuple]) -> None: + self.name = name + self.events = events + self.calls: list[_RemoteCall] = [] + + def remote(self, *args, **kwargs): + self.calls.append(_RemoteCall(args=args, kwargs=kwargs)) + self.events.append((self.name, args, kwargs)) + return self.name + + +@dataclass +class RecordingEngine: + name: str + events: list[tuple] + pause_generation: RecordingRemoteMethod = field(init=False) + flush_cache: RecordingRemoteMethod = field(init=False) + post_process_weights: RecordingRemoteMethod = field(init=False) + set_weight_version: RecordingRemoteMethod = field(init=False) + continue_generation: RecordingRemoteMethod = field(init=False) + + def __post_init__(self) -> None: + self.pause_generation = RecordingRemoteMethod(f"{self.name}.pause", self.events) + self.flush_cache = RecordingRemoteMethod(f"{self.name}.flush", self.events) + self.post_process_weights = RecordingRemoteMethod(f"{self.name}.post_process", self.events) + self.set_weight_version = RecordingRemoteMethod(f"{self.name}.set_version", self.events) + self.continue_generation = RecordingRemoteMethod(f"{self.name}.continue", self.events) + + +def _patch_runtime( + monkeypatch, + coordinator_module, + *, + rank: int, + events: list[tuple], + ray_get=None, + broadcast=None, +): + monkeypatch.setattr(coordinator_module.dist, "get_rank", lambda: rank) + monkeypatch.setattr( + coordinator_module.dist, + "barrier", + lambda *, group=None: events.append(("barrier", group)), + ) + monkeypatch.setattr( + coordinator_module.dist, + "broadcast_object_list", + broadcast + or ( + lambda status, *, src, group: events.append( + ("broadcast", src, group, status[0]), + ) + ), + ) + monkeypatch.setattr(coordinator_module, "get_gloo_group", lambda: "gloo") + monkeypatch.setattr(coordinator_module.ray, "get", ray_get or (lambda refs: refs)) + + +@pytest.mark.unit +def test_success_orders_control_plane_and_returns_candidate(coordinator_module, monkeypatch): + events: list[tuple] = [] + engines = [RecordingEngine("e0", events), RecordingEngine("e1", events)] + _patch_runtime(monkeypatch, coordinator_module, rank=0, events=events) + coordinator = coordinator_module.WeightUpdateCoordinator( + engines, + {"quant_method": "compressed-tensors"}, + ) + + version = coordinator.run( + current_version=7, + transfer_target=lambda candidate: events.append(("target", candidate)), + transfer_draft=lambda candidate: events.append(("draft", candidate)), + commit=lambda: events.append(("commit",)), + ) + + assert version == 8 + names = [event[0] for event in events] + assert names == [ + "e0.pause", + "e1.pause", + "e0.flush", + "e1.flush", + "e0.post_process", + "e1.post_process", + "broadcast", + "target", + "draft", + "barrier", + "e0.post_process", + "e1.post_process", + "commit", + "e0.set_version", + "e1.set_version", + "e0.continue", + "e1.continue", + "broadcast", + ] + assert engines[0].post_process_weights.calls[0].kwargs == { + "restore_weights_before_load": True, + "post_process_quantization": False, + } + assert engines[0].post_process_weights.calls[1].kwargs == { + "restore_weights_before_load": False, + "post_process_quantization": True, + } + assert engines[0].set_weight_version.calls[0].args == ("8",) + + +@pytest.mark.unit +@pytest.mark.parametrize("failing_phase", ["target", "draft", "quant_post_process", "commit"]) +def test_failure_before_publish_does_not_publish_or_resume( + coordinator_module, + monkeypatch, + failing_phase, +): + events: list[tuple] = [] + engines = [RecordingEngine("e0", events)] + quant = {"quant_method": "compressed-tensors"} if failing_phase == "quant_post_process" else None + + def ray_get(refs): + if ( + failing_phase == "quant_post_process" + and refs == ["e0.post_process"] + and len(engines[0].post_process_weights.calls) == 2 + ): + raise RuntimeError("quant_post_process failed") + return refs + + def phase(name): + def run(*args): + events.append((name, *args)) + if failing_phase == name: + raise RuntimeError(f"{name} failed") + + return run + + _patch_runtime(monkeypatch, coordinator_module, rank=0, events=events, ray_get=ray_get) + coordinator = coordinator_module.WeightUpdateCoordinator(engines, quant) + + with pytest.raises(RuntimeError, match=failing_phase): + coordinator.run( + current_version=3, + transfer_target=phase("target"), + transfer_draft=phase("draft"), + commit=phase("commit"), + ) + + assert engines[0].set_weight_version.calls == [] + assert engines[0].continue_generation.calls == [] + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("failing_call", "resumed"), + [("e0.set_version", False), ("e0.continue", True)], +) +def test_failure_after_publish_restores_version_marker( + coordinator_module, + monkeypatch, + failing_call, + resumed, +): + events: list[tuple] = [] + engines = [RecordingEngine("e0", events)] + + def ray_get(refs): + if refs == [failing_call]: + raise RuntimeError(f"{failing_call} failed") + return refs + + _patch_runtime(monkeypatch, coordinator_module, rank=0, events=events, ray_get=ray_get) + coordinator = coordinator_module.WeightUpdateCoordinator(engines, None) + + with pytest.raises(RuntimeError, match=failing_call): + coordinator.run(current_version=4, transfer_target=lambda _candidate: None) + + # The candidate was published, then the committed marker was restored. + assert [call.args for call in engines[0].set_weight_version.calls] == [ + ("5",), + ("4",), + ] + if resumed: + assert len(engines[0].pause_generation.calls) == 2 + assert len(engines[0].continue_generation.calls) == 1 + else: + assert engines[0].continue_generation.calls == [] + + +@pytest.mark.unit +def test_nonzero_rank_follows_rank_zero(coordinator_module, monkeypatch): + # Success: nonzero ranks only run transfers and collectives, never engine RPCs. + events: list[tuple] = [] + engines = [RecordingEngine("e0", events)] + _patch_runtime(monkeypatch, coordinator_module, rank=1, events=events) + coordinator = coordinator_module.WeightUpdateCoordinator(engines, None) + + version = coordinator.run( + current_version=4, + transfer_target=lambda candidate: events.append(("target", candidate)), + transfer_draft=lambda candidate: events.append(("draft", candidate)), + commit=lambda: events.append(("commit",)), + ) + + assert version == 5 + assert events == [ + ("broadcast", 0, "gloo", None), + ("target", 5), + ("draft", 5), + ("barrier", "gloo"), + ("broadcast", 0, "gloo", None), + ] + + # Failure: a rank-0 control-plane error propagates through the status broadcast. + fail_events: list[tuple] = [] + fail_engines = [RecordingEngine("e0", fail_events)] + + def broadcast(status, *, src, group): + fail_events.append(("broadcast", src, group, status[0])) + status[0] = "RuntimeError: pause failed" + + _patch_runtime( + monkeypatch, + coordinator_module, + rank=1, + events=fail_events, + broadcast=broadcast, + ) + failing = coordinator_module.WeightUpdateCoordinator(fail_engines, None) + + with pytest.raises(RuntimeError, match="quiesce.*pause failed"): + failing.run(current_version=4, transfer_target=lambda _candidate: None) + + assert fail_engines[0].pause_generation.calls == [] diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index 8a757dfe0..811ff39d9 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -151,16 +151,15 @@ def init( update_weight_transport = self.args.update_weight_transport if update_weight_mode == "delta": - # Delta sync is disk-transport only: each engine's /pull_weights applies the published - # deltas into a host-local checkpoint on every host it spans, and the engines reload - # via vanilla update_weights_from_disk. assert not self.args.colocate, "--update-weight-mode=delta is not supported with --colocate" assert ( - update_weight_transport == "disk" - ), "--update-weight-mode=delta requires --update-weight-transport=disk" - from .update_weight.update_weight_from_disk_delta import UpdateWeightFromDiskDelta + update_weight_transport == "nccl" + ), "--update-weight-mode=delta currently requires --update-weight-transport=nccl" + from .update_weight.update_weight_from_checkpoint_delta import ( + UpdateWeightFromCheckpointDelta, + ) - update_weight_cls = UpdateWeightFromDiskDelta + update_weight_cls = UpdateWeightFromCheckpointDelta elif update_weight_transport == "disk": update_weight_cls = UpdateWeightFromDisk elif self.args.colocate: diff --git a/vime/backends/megatron_utils/update_weight/checkpoint_delta.py b/vime/backends/megatron_utils/update_weight/checkpoint_delta.py new file mode 100644 index 000000000..975ffe402 --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/checkpoint_delta.py @@ -0,0 +1,464 @@ +from __future__ import annotations + +from collections.abc import Iterable, Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Literal, Protocol + +import torch + +DeltaEncoding = Literal["dense", "indices"] + + +@dataclass(frozen=True) +class CheckpointPatchSpec: + """Metadata for one checkpoint-coordinate patch inside a wire chunk.""" + + name: str + shape: tuple[int, ...] + dtype_name: str + value_start: int + value_end: int + position_start: int + position_end: int + + def to_dict(self) -> dict: + return { + "name": self.name, + "shape": list(self.shape), + "dtype_name": self.dtype_name, + "value_start": self.value_start, + "value_end": self.value_end, + "position_start": self.position_start, + "position_end": self.position_end, + } + + +@dataclass +class CheckpointDeltaChunk: + """One dense-seed or sparse checkpoint-coordinate transfer chunk.""" + + schema_version: int + base_version: int + target_version: int + sequence_no: int + is_final: bool + encoding: DeltaEncoding + patches: list[CheckpointPatchSpec] + positions: torch.Tensor + values: torch.Tensor + + @property + def wire_bytes(self) -> int: + return ( + self.positions.numel() * self.positions.element_size() + self.values.numel() * self.values.element_size() + ) + + def update_info(self) -> dict: + return { + "schema_version": self.schema_version, + "base_version": self.base_version, + "target_version": self.target_version, + "sequence_no": self.sequence_no, + "is_final": self.is_final, + "encoding": self.encoding, + "patches": [patch.to_dict() for patch in self.patches], + "position_count": self.positions.numel(), + "value_count": self.values.numel(), + "value_dtype_name": str(self.values.dtype).removeprefix("torch."), + } + + def wire_tensors(self) -> Iterator[tuple[str, torch.Tensor]]: + if self.positions.numel(): + yield "__positions__", self.positions + if self.values.numel(): + yield "__values__", self.values + + +class DeltaWeightSource(Protocol): + """Source-side transaction boundary shared by future delta transports.""" + + def begin_update(self, *, base_version: int, target_version: int) -> None: ... + + def encode_chunk(self, named_tensors: Iterable[tuple[str, torch.Tensor]]) -> list[CheckpointDeltaChunk]: ... + + def finish_update(self) -> CheckpointDeltaChunk: ... + + def commit(self) -> None: ... + + def abort(self) -> None: ... + + +@dataclass(frozen=True) +class _CheckpointTensorLayout: + name: str + shape: tuple[int, ...] + dtype: torch.dtype + start: int + end: int + + +@dataclass +class _BucketSnapshot: + layout: tuple[_CheckpointTensorLayout, ...] + values: torch.Tensor + + +@dataclass +class _PendingBucketUpdate: + bucket_index: int + indices: torch.Tensor + values: torch.Tensor + + +class CheckpointDeltaSource: + """Create checkpoint-coordinate patches from deterministic HF buckets. + + The first update sends dense weights and saves one flat CPU snapshot for + each exporter bucket. Later updates compare one complete bucket at a time + on the export device. This keeps the existing HF names and native vLLM + loader while avoiding one GPU synchronization and CPU snapshot update per + checkpoint tensor. Snapshots advance only after every rollout worker has + applied the update. + """ + + def __init__(self, wire_dtype: torch.dtype = torch.bfloat16) -> None: + if wire_dtype != torch.bfloat16: + raise NotImplementedError("The direct DWU MVP supports BF16 only") + self.wire_dtype = wire_dtype + self._snapshot: list[_BucketSnapshot] = [] + self._pending_seed: list[_BucketSnapshot] = [] + self._pending_updates: list[_PendingBucketUpdate] = [] + self._seen_names: set[str] = set() + self._active = False + self._finished = False + self._seed_update = True + self._committed_version = 0 + self._base_version = 0 + self._target_version = 0 + self._next_sequence_no = 0 + self._next_bucket_index = 0 + self.total_elements = 0 + self.changed_elements = 0 + self.wire_bytes = 0 + + @property + def is_seed_update(self) -> bool: + return self._seed_update + + def begin_update(self, *, base_version: int, target_version: int) -> None: + if self._active: + raise RuntimeError("A checkpoint delta update is already active") + if target_version != base_version + 1: + raise ValueError("Checkpoint delta target_version must be base_version + 1") + if not self._snapshot and base_version != 0: + raise ValueError("The first direct DWU transaction must start at version 0") + if self._snapshot and base_version != self._committed_version: + raise RuntimeError( + f"Checkpoint delta source version mismatch: snapshot={self._committed_version}, update={base_version}" + ) + self._active = True + self._finished = False + self._seed_update = not self._snapshot + self._base_version = base_version + self._target_version = target_version + self._next_sequence_no = 0 + self._next_bucket_index = 0 + self._pending_seed = [] + self._pending_updates = [] + self._seen_names = set() + self.total_elements = 0 + self.changed_elements = 0 + self.wire_bytes = 0 + + @torch.no_grad() + def encode_chunk( + self, + named_tensors: Iterable[tuple[str, torch.Tensor]], + ) -> list[CheckpointDeltaChunk]: + self._require_active() + if self._finished: + raise RuntimeError("Cannot encode checkpoint deltas after finish_update()") + layout, current = self._flatten_bucket(named_tensors) + if not layout: + return [] + + bucket_index = self._next_bucket_index + self._next_bucket_index += 1 + if self._seed_update: + self._pending_seed.append( + _BucketSnapshot( + layout=layout, + values=current.to(device="cpu", copy=True), + ) + ) + chunk = self._build_dense_chunk(layout, current) + else: + if bucket_index >= len(self._snapshot): + raise RuntimeError(f"HF exporter added bucket {bucket_index}") + snapshot = self._snapshot[bucket_index] + self._validate_bucket_layout(bucket_index, expected=snapshot.layout, actual=layout) + chunk = self._build_sparse_chunk(bucket_index, snapshot, current) + if chunk is None: + return [] + + self.wire_bytes += chunk.wire_bytes + return [chunk] + + def finish_update(self) -> CheckpointDeltaChunk: + self._require_active() + if self._finished: + raise RuntimeError("finish_update() was already called") + if not self._seed_update and self._next_bucket_index != len(self._snapshot): + raise RuntimeError( + f"HF exporter bucket count changed: expected {len(self._snapshot)}, got {self._next_bucket_index}" + ) + self._finished = True + chunk = CheckpointDeltaChunk( + schema_version=1, + base_version=self._base_version, + target_version=self._target_version, + sequence_no=self._next_sequence_no, + is_final=True, + encoding="dense" if self._seed_update else "indices", + patches=[], + positions=torch.empty(0, dtype=torch.int32), + values=torch.empty(0, dtype=self.wire_dtype), + ) + self._next_sequence_no += 1 + return chunk + + def commit(self) -> None: + self._require_active() + if not self._finished: + raise RuntimeError("finish_update() must be called before commit()") + if self._seed_update: + self._snapshot = self._pending_seed + else: + worker_count = min(16, len(self._pending_updates)) + if worker_count == 1: + self._commit_bucket(self._pending_updates[0]) + elif worker_count > 1: + # Each update writes to an independent bucket snapshot. Parallel + # CPU writes shorten a commit that otherwise extends the rollout pause. + with ThreadPoolExecutor(max_workers=worker_count) as executor: + list(executor.map(self._commit_bucket, self._pending_updates)) + self._committed_version = self._target_version + self._clear_transaction() + + def abort(self) -> None: + if self._active: + self._clear_transaction() + + def _flatten_bucket( + self, + named_tensors: Iterable[tuple[str, torch.Tensor]], + ) -> tuple[tuple[_CheckpointTensorLayout, ...], torch.Tensor]: + tensors: list[torch.Tensor] = [] + layout: list[_CheckpointTensorLayout] = [] + offset = 0 + device: torch.device | None = None + for name, tensor in named_tensors: + if name in self._seen_names: + raise ValueError(f"Duplicate HF tensor in one update: {name!r}") + self._seen_names.add(name) + + current = tensor.detach() + if not current.is_floating_point(): + raise NotImplementedError(f"{name}: direct DWU supports floating-point weights only") + if current.dtype != self.wire_dtype: + current = current.to(self.wire_dtype) + current = current.contiguous() + if current.numel() >= 1 << 31: + raise NotImplementedError(f"{name}: direct DWU int32 indices require fewer than 2^31 elements") + if device is None: + device = current.device + elif current.device != device: + raise ValueError("All tensors in one HF exporter bucket must use the same device") + + end = offset + current.numel() + layout.append( + _CheckpointTensorLayout( + name=name, + shape=tuple(current.shape), + dtype=current.dtype, + start=offset, + end=end, + ) + ) + tensors.append(current.reshape(-1)) + offset = end + self.total_elements += current.numel() + + if not tensors: + return (), torch.empty(0, dtype=self.wire_dtype) + return tuple(layout), torch.cat(tensors) + + def _build_dense_chunk( + self, + layout: tuple[_CheckpointTensorLayout, ...], + values: torch.Tensor, + ) -> CheckpointDeltaChunk: + patches = [ + CheckpointPatchSpec( + name=tensor.name, + shape=tensor.shape, + dtype_name=str(tensor.dtype).removeprefix("torch."), + value_start=tensor.start, + value_end=tensor.end, + position_start=0, + position_end=0, + ) + for tensor in layout + ] + self.changed_elements += values.numel() + return self._new_chunk( + encoding="dense", + patches=patches, + positions=torch.empty(0, dtype=torch.int32, device=values.device), + values=values, + ) + + def _build_sparse_chunk( + self, + bucket_index: int, + snapshot: _BucketSnapshot, + current: torch.Tensor, + ) -> CheckpointDeltaChunk | None: + previous = snapshot.values.to(device=current.device) + changed = torch.nonzero( + current.view(torch.int16) != previous.view(torch.int16), + as_tuple=False, + ).reshape(-1) + if not changed.numel(): + return None + + values = current.index_select(0, changed) + ends = torch.tensor( + [tensor.end for tensor in snapshot.layout], + dtype=torch.int64, + device=current.device, + ) + # The checkpoint patch API reserves NaN as its unchanged-value sentinel + # and rejects NaN wire values on every receiver; fail here on the source + # rank instead, where the offending tensor can still be named. + nan_mask = torch.isnan(values) + if bool(nan_mask.any()): + first_flat = changed[nan_mask.nonzero(as_tuple=False).reshape(-1)[0]] + tensor_index = int(torch.searchsorted(ends, first_flat, right=True).item()) + raise ValueError( + f"{snapshot.layout[tensor_index].name}: training produced NaN weight " + "values; refusing to ship a sparse delta" + ) + cumulative_counts = torch.searchsorted(changed, ends) + counts = cumulative_counts.clone() + counts[1:] -= cumulative_counts[:-1] + starts = torch.tensor( + [tensor.start for tensor in snapshot.layout], + dtype=torch.int64, + device=current.device, + ) + positions = ( + changed + - torch.repeat_interleave( + starts, + counts, + output_size=changed.numel(), + ) + ).to(torch.int32) + + counts_cpu = counts.to(device="cpu").tolist() + patches: list[CheckpointPatchSpec] = [] + offset = 0 + for tensor, count in zip(snapshot.layout, counts_cpu, strict=True): + if not count: + continue + end = offset + count + patches.append( + CheckpointPatchSpec( + name=tensor.name, + shape=tensor.shape, + dtype_name=str(tensor.dtype).removeprefix("torch."), + value_start=offset, + value_end=end, + position_start=offset, + position_end=end, + ) + ) + offset = end + + self._pending_updates.append( + _PendingBucketUpdate( + bucket_index=bucket_index, + indices=changed.to(device="cpu", copy=True), + values=values.to(device="cpu", copy=True), + ) + ) + self.changed_elements += values.numel() + return self._new_chunk( + encoding="indices", + patches=patches, + positions=positions, + values=values, + ) + + def _new_chunk( + self, + *, + encoding: DeltaEncoding, + patches: list[CheckpointPatchSpec], + positions: torch.Tensor, + values: torch.Tensor, + ) -> CheckpointDeltaChunk: + chunk = CheckpointDeltaChunk( + schema_version=1, + base_version=self._base_version, + target_version=self._target_version, + sequence_no=self._next_sequence_no, + is_final=False, + encoding=encoding, + patches=patches, + positions=positions, + values=values, + ) + self._next_sequence_no += 1 + return chunk + + @staticmethod + def _validate_bucket_layout( + bucket_index: int, + *, + expected: tuple[_CheckpointTensorLayout, ...], + actual: tuple[_CheckpointTensorLayout, ...], + ) -> None: + if expected == actual: + return + if len(expected) != len(actual): + raise RuntimeError( + f"HF exporter bucket {bucket_index} tensor count changed: expected {len(expected)}, got {len(actual)}" + ) + for tensor_index, (old, new) in enumerate(zip(expected, actual, strict=True)): + if old != new: + raise RuntimeError( + f"HF exporter bucket {bucket_index} tensor {tensor_index} changed: expected {old}, got {new}" + ) + raise AssertionError("Mismatched HF bucket layouts did not contain a mismatched tensor") + + def _require_active(self) -> None: + if not self._active: + raise RuntimeError("begin_update() must be called first") + + def _commit_bucket(self, update: _PendingBucketUpdate) -> None: + self._snapshot[update.bucket_index].values.index_copy_( + 0, + update.indices, + update.values, + ) + + def _clear_transaction(self) -> None: + self._pending_seed = [] + self._pending_updates = [] + self._seen_names = set() + self._next_bucket_index = 0 + self._active = False + self._finished = False diff --git a/vime/backends/megatron_utils/update_weight/coordinator.py b/vime/backends/megatron_utils/update_weight/coordinator.py new file mode 100644 index 000000000..d8e07119f --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/coordinator.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +import logging +from collections.abc import Callable, Sequence +from typing import TYPE_CHECKING + +import ray +import torch.distributed as dist + +from vime.utils.distributed_utils import get_gloo_group + +if TYPE_CHECKING: + from ray.actor import ActorHandle + +logger = logging.getLogger(__name__) + +WeightTransfer = Callable[[int], None] +WeightCommit = Callable[[], None] + + +def post_process_weights( + restore_weights_before_load: bool, + post_process_quantization: bool, + rollout_engines: Sequence[ActorHandle], +) -> None: + """Run compressed-weight pre/post processing on every rollout engine.""" + ray.get( + [ + engine.post_process_weights.remote( + restore_weights_before_load=restore_weights_before_load, + post_process_quantization=post_process_quantization, + ) + for engine in rollout_engines + ] + ) + + +class WeightUpdateCoordinator: + """Coordinate one committed weight update across rollout engines. + + Transport-specific work stays in ``transfer_target``/``transfer_draft``. + This class owns only the shared control plane: + + pause -> flush -> quant pre-process -> transfer -> quant post-process + -> source commit -> publish committed version -> resume. + + The candidate version is returned only after every phase succeeds. Before + resume, failures leave generation paused. If a batched resume partially + succeeds, the coordinator best-effort pauses every engine again and restores + its version metadata; external recovery is still required because vLLM applies + weight chunks in place and cannot roll them back atomically. Transfer + callbacks are distributed operations and must surface failures collectively; + this coordinator does not turn a rank-local exception into a collective one. + """ + + def __init__( + self, + rollout_engines: Sequence[ActorHandle], + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + self.rollout_engines = tuple(rollout_engines) + self.quantization_config = quantization_config + + def run( + self, + *, + current_version: int, + transfer_target: WeightTransfer, + transfer_draft: WeightTransfer | None = None, + commit: WeightCommit | None = None, + ) -> int: + candidate_version = current_version + 1 + try: + self._quiesce() + transfer_target(candidate_version) + if transfer_draft is not None: + transfer_draft(candidate_version) + + # Keep all trainer ranks aligned before rank 0 publishes the new + # committed version to the rollout engines. + self._barrier() + self._publish_and_resume( + candidate_version=candidate_version, + current_version=current_version, + commit=commit, + ) + except BaseException: + logger.exception( + "Weight update to version %s failed; caller retains version %s " + "and rollout engines require fail-stop recovery.", + candidate_version, + current_version, + ) + raise + return candidate_version + + def _quiesce(self) -> None: + def quiesce_rank_zero() -> None: + ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + if self._uses_compressed_tensors(): + post_process_weights( + restore_weights_before_load=True, + post_process_quantization=False, + rollout_engines=self.rollout_engines, + ) + + self._run_rank_zero_phase("quiesce", quiesce_rank_zero) + + def _publish_and_resume( + self, + *, + candidate_version: int, + current_version: int, + commit: WeightCommit | None, + ) -> None: + def publish_and_resume_rank_zero() -> None: + if self._uses_compressed_tensors(): + post_process_weights( + restore_weights_before_load=False, + post_process_quantization=True, + rollout_engines=self.rollout_engines, + ) + + # Commit transport-side source state while generation is still + # paused. A commit failure is broadcast to every trainer rank before + # the candidate version is published or any engine resumes. + if commit is not None: + commit() + + # Chunk RPCs carry a candidate version for compatibility, but do not + # publish it. Publish once after the source and every rollout worker + # have completed the candidate successfully. + try: + self._set_engine_version(candidate_version) + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + except BaseException: + self._restore_fail_stop_state(current_version) + raise + + self._run_rank_zero_phase("publish/resume", publish_and_resume_rank_zero) + + def _restore_fail_stop_state(self, current_version: int) -> None: + # A list of Ray calls is not an atomic fanout: continue_generation may + # have succeeded on only a subset before ray.get reports an error. + # Re-pause first so partially resumed engines stop serving mixed state. + try: + ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) + except BaseException: + logger.exception("Failed to re-pause every rollout engine after publish/resume error") + + # GPU weights and transport-side source state cannot be rolled back here. + # Restore only the version metadata so no engine advertises a partially + # resumed candidate. Backend-specific fail-stop recovery is still needed. + try: + self._set_engine_version(current_version) + except BaseException: + logger.exception( + "Failed to restore rollout-engine version marker to %s", + current_version, + ) + + def _set_engine_version(self, version: int) -> None: + ray.get([engine.set_weight_version.remote(str(version)) for engine in self.rollout_engines]) + + def _uses_compressed_tensors(self) -> bool: + return bool(self.quantization_config and self.quantization_config.get("quant_method") == "compressed-tensors") + + @staticmethod + def _run_rank_zero_phase(phase: str, operation: Callable[[], None]) -> None: + """Run a control-plane phase once and propagate its result to all ranks.""" + rank_zero_error: BaseException | None = None + status: list[str | None] = [None] + if dist.get_rank() == 0: + try: + operation() + except BaseException as exc: + rank_zero_error = exc + status[0] = f"{type(exc).__name__}: {exc}" + + dist.broadcast_object_list(status, src=0, group=get_gloo_group()) + if rank_zero_error is not None: + raise rank_zero_error + if status[0] is not None: + raise RuntimeError(f"Rank-0 weight-update {phase} failed: {status[0]}") + + @staticmethod + def _barrier() -> None: + dist.barrier(group=get_gloo_group()) diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_checkpoint_delta.py b/vime/backends/megatron_utils/update_weight/update_weight_from_checkpoint_delta.py new file mode 100644 index 000000000..3f6a1e8b6 --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_checkpoint_delta.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import logging +import time + +import ray +import torch +from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerSendWeightsArgs, NCCLWeightTransferEngine + +from .checkpoint_delta import CheckpointDeltaChunk, CheckpointDeltaSource +from .update_weight_from_distributed import UpdateWeightFromDistributed + +logger = logging.getLogger(__name__) + + +class UpdateWeightFromCheckpointDelta(UpdateWeightFromDistributed): + """VIME direct DWU: full HF export, GPU diff, NCCL checkpoint patches.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + if self.quantization_config is not None: + raise NotImplementedError("VIME direct DWU does not support quantized models") + self._delta_source = CheckpointDeltaSource() + + def _get_weight_update_commit(self): + return self._delta_source.commit + + @torch.no_grad() + def update_weights(self) -> None: + is_source = getattr(self, "_is_pp_src_rank", False) + if is_source: + base_version = self.weight_version + self._delta_source.begin_update( + base_version=base_version, + target_version=base_version + 1, + ) + started_at = time.perf_counter() + try: + super().update_weights() + except BaseException: + if is_source: + self._delta_source.abort() + raise + if not is_source: + return + density = ( + self._delta_source.changed_elements / self._delta_source.total_elements + if self._delta_source.total_elements + else 0.0 + ) + elapsed = time.perf_counter() - started_at + self.update_weight_metrics.update( + { + "weight_sync/is_dense_seed": float(self._delta_source.is_seed_update), + "weight_sync/total_elements": float(self._delta_source.total_elements), + "weight_sync/changed_elements": float(self._delta_source.changed_elements), + "weight_sync/delta_density": density, + "weight_sync/wire_bytes": float(self._delta_source.wire_bytes), + "weight_sync/seconds": elapsed, + } + ) + if getattr(self, "_is_pp_src_rank", False): + logger.info( + "Direct DWU committed version=%d dense_seed=%s changed=%d/%d density=%.6f wire_bytes=%d seconds=%.3f", + self.weight_version, + self._delta_source.is_seed_update, + self._delta_source.changed_elements, + self._delta_source.total_elements, + density, + self._delta_source.wire_bytes, + elapsed, + ) + + def _update_bucket_weights_from_distributed( + self, + converted_named_tensors, + pbar=None, + ) -> None: + chunks = self._delta_source.encode_chunk(converted_named_tensors) + converted_named_tensors.clear() + + for chunk in chunks: + while not ray.get(self.rollout_engine_lock.acquire.remote()): + time.sleep(0.1) + try: + refs = self._send_delta_chunk(chunk) + ray.get(refs) + finally: + ray.get(self.rollout_engine_lock.release.remote()) + if pbar is not None: + pbar.update(1) + + def _send_weights_to_rollout_engines(self) -> None: + super()._send_weights_to_rollout_engines() + if not self._is_pp_src_rank: + return + final_chunk = self._delta_source.finish_update() + + while not ray.get(self.rollout_engine_lock.acquire.remote()): + time.sleep(0.1) + try: + ray.get(self._send_delta_chunk(final_chunk)) + finally: + ray.get(self.rollout_engine_lock.release.remote()) + + def _send_delta_chunk(self, chunk: CheckpointDeltaChunk): + refs = [ + engine.update_checkpoint_delta_from_distributed.remote( + update_info=chunk.update_info(), + ) + for engine in self.rollout_engines + ] + NCCLWeightTransferEngine.trainer_send_weights( + chunk.wire_tensors(), + NCCLTrainerSendWeightsArgs( + group=self._model_update_groups, + packed=False, + ), + ) + return refs diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 67da89378..a11eb70a5 100644 --- a/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -21,23 +21,37 @@ from ..megatron_to_hf import convert_to_hf from .common import all_gather_param, named_params_and_buffers +from .coordinator import WeightUpdateCoordinator +from .coordinator import post_process_weights as _post_process_weights from .hf_weight_iterator_base import HfWeightIteratorBase logger = logging.getLogger(__name__) +# Compatibility re-export for callers that imported the old helper location. +post_process_weights = _post_process_weights -def _begin_vllm_weight_update_session(rollout_engines: Sequence[ActorHandle]) -> None: +def _begin_vllm_weight_update_session( + rollout_engines: Sequence[ActorHandle], + *, + draft: bool = False, +) -> None: if dist.get_rank() == 0: - logger.info("vLLM weight update: start_weight_update") - ray.get([engine.start_weight_update.remote(is_checkpoint_format=True) for engine in rollout_engines]) + if draft: + logger.info("vLLM weight update: start_draft_weight_update") + ray.get([engine.start_draft_weight_update.remote() for engine in rollout_engines]) + else: + logger.info("vLLM weight update: start_weight_update") + ray.get([engine.start_weight_update.remote(is_checkpoint_format=True) for engine in rollout_engines]) dist.barrier(group=get_gloo_group()) def _end_vllm_weight_update_session(rollout_engines: Sequence[ActorHandle]) -> None: - if dist.get_rank() == 0: - logger.info("vLLM weight update: finish_weight_update") - ray.get([engine.finish_weight_update.remote() for engine in rollout_engines]) - dist.barrier(group=get_gloo_group()) + try: + if dist.get_rank() == 0: + logger.info("vLLM weight update: finish_weight_update") + ray.get([engine.finish_weight_update.remote() for engine in rollout_engines]) + finally: + dist.barrier(group=get_gloo_group()) class UpdateWeightFromDistributed: @@ -138,52 +152,59 @@ def pop_metrics(self) -> dict[str, float]: out, self.update_weight_metrics = self.update_weight_metrics, {} return out + def _get_weight_update_commit(self) -> Callable[[], None] | None: + """Return an optional source commit run before publish and resume.""" + return None + @torch.no_grad() def update_weights(self) -> None: - """ - Pause → flush → _send_weights → continue. Progress on PP source. - """ - self.weight_version += 1 + """Run one coordinated weight update and commit on success.""" + committed_version = self.weight_version + coordinator = WeightUpdateCoordinator(self.rollout_engines, self.quantization_config) + + def transfer_target(candidate_version: int) -> None: + # Existing bucket senders read ``self.weight_version``. It is + # provisional until the coordinator returns and is restored on any + # failure below. + self.weight_version = candidate_version + self._run_weight_update_session() + + draft_transfer: Callable[[int], None] | None = None + if self.args.enable_mtp_training and (self.args.vllm_speculative_config or {}).get("method") == "mtp": - 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]) - - # int4/fp4 pre_process - if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: - post_process_weights( - restore_weights_before_load=True, - post_process_quantization=False, - rollout_engines=self.rollout_engines, - ) - dist.barrier(group=get_gloo_group()) + def transfer_draft(candidate_version: int) -> None: + assert candidate_version == self.weight_version + self._run_weight_update_session(draft=True) + + draft_transfer = transfer_draft - _begin_vllm_weight_update_session(self.rollout_engines) try: - self._send_weights_to_rollout_engines() - finally: - _end_vllm_weight_update_session(self.rollout_engines) + self.weight_version = coordinator.run( + current_version=committed_version, + transfer_target=transfer_target, + transfer_draft=draft_transfer, + commit=self._get_weight_update_commit(), + ) + except BaseException: + self.weight_version = committed_version + raise - if self.args.enable_mtp_training and (self.args.vllm_speculative_config or {}).get("method") == "mtp": - if dist.get_rank() == 0: - ray.get([engine.start_draft_weight_update.remote() for engine in self.rollout_engines]) - dist.barrier(group=get_gloo_group()) + def _run_weight_update_session(self, *, draft: bool = False) -> None: + _begin_vllm_weight_update_session(self.rollout_engines, draft=draft) + try: + self._send_weights_to_rollout_engines() + except BaseException: + # A failed vLLM update RPC clears its own active session, while a + # local conversion/NCCL failure may leave the session active. Try + # both cases and never replace the primary transfer exception with + # a cleanup error such as "finish without start". try: - self._send_weights_to_rollout_engines() - finally: _end_vllm_weight_update_session(self.rollout_engines) - - dist.barrier(group=get_gloo_group()) - if dist.get_rank() == 0: - # int4/fp4 post_process - if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: - post_process_weights( - restore_weights_before_load=False, - post_process_quantization=True, - rollout_engines=self.rollout_engines, - ) - ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) - dist.barrier(group=get_gloo_group()) + except BaseException: + logger.warning("Failed to finalize vLLM session after weight transfer error", exc_info=True) + raise + else: + _end_vllm_weight_update_session(self.rollout_engines) def _send_weights_to_rollout_engines(self) -> None: if self._uses_persistent_group(): @@ -383,16 +404,18 @@ def _update_bucket_weights_from_distributed( while not ray.get(self.rollout_engine_lock.acquire.remote()): time.sleep(0.1) - refs = update_weights_from_distributed( - self._model_update_groups, - self.weight_version, - self.rollout_engines, - converted_named_tensors, - ) - - ray.get(refs) - converted_named_tensors.clear() - ray.get(self.rollout_engine_lock.release.remote()) + try: + refs = update_weights_from_distributed( + self._model_update_groups, + self.weight_version, + self.rollout_engines, + converted_named_tensors, + record_weight_version=False, + ) + ray.get(refs) + converted_named_tensors.clear() + finally: + ray.get(self.rollout_engine_lock.release.remote()) pbar.update(1) @@ -492,6 +515,8 @@ def update_weights_from_distributed( weight_version: int, rollout_engines: Sequence[ActorHandle], converted_named_tensors: Sequence[tuple[str, torch.Tensor]], + *, + record_weight_version: bool = True, ) -> list[ObjectRef]: """ Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). @@ -505,6 +530,7 @@ def update_weights_from_distributed( dtypes=[param.dtype for _, param in converted_named_tensors], shapes=[param.shape for _, param in converted_named_tensors], weight_version=str(weight_version), + record_weight_version=record_weight_version, ) for engine in rollout_engines ] @@ -519,22 +545,3 @@ def update_weights_from_distributed( ) return refs - - -def post_process_weights( - restore_weights_before_load: bool, - post_process_quantization: bool, - rollout_engines: Sequence[ActorHandle], -): - """ - Trigger post-process for int4/fp4 quantization on all rollout engines. - """ - ray.get( - [ - engine.post_process_weights.remote( - restore_weights_before_load=restore_weights_before_load, - post_process_quantization=post_process_quantization, - ) - for engine in rollout_engines - ] - ) 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 5e3520257..c7c41baa1 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 @@ -13,6 +13,7 @@ from __future__ import annotations +import logging import os from argparse import Namespace from collections.abc import Callable, Iterable, Mapping, Sequence @@ -27,15 +28,16 @@ from vime.utils.distributed_utils import get_gloo_group +from .coordinator import WeightUpdateCoordinator from .hf_weight_iterator_base import HfWeightIteratorBase from .update_weight_from_distributed import ( connect_rollout_engines_from_distributed, disconnect_rollout_engines_from_distributed, - post_process_weights, update_weights_from_distributed, ) _MAX_COLOCATED_UPDATES_INFLIGHT = 4 +logger = logging.getLogger(__name__) def _build_packed_ipc_update_info( @@ -160,6 +162,7 @@ def connect_rollout_engines( Split colocated/distributed engines. Global source rank (DP=TP=PP=0) creates NCCL for distributed. Map ranks to colocated IPC engines. """ + self._all_rollout_engines = list(rollout_engines) self.rollout_engines = rollout_engines if engine_gpu_counts is None: @@ -203,6 +206,8 @@ def connect_rollout_engines( self.distributed_rollout_engines, engine_gpu_counts=distributed_gpu_counts, ) + else: + self.distributed_rollout_engines = [] colocate_gpu_offsets = engine_gpu_offsets[:colocate_engine_nums] colocate_gpu_counts = engine_gpu_counts[:colocate_engine_nums] @@ -243,69 +248,101 @@ def pop_metrics(self) -> dict[str, float]: @torch.no_grad() def update_weights(self) -> None: - """ - version++, flush caches, process buckets. Progress on rank 0. - """ - self.weight_version += 1 - - rank = dist.get_rank() - if 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]) - if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: - post_process_weights( - restore_weights_before_load=True, - post_process_quantization=False, - rollout_engines=self.rollout_engines, - ) - dist.barrier(group=get_gloo_group()) - - # vLLM #39212: enter weight-update mode on each slot leader. - if self._ipc_engine is not None and rank == self._ipc_gather_src: - ray.get(self._ipc_engine.start_weight_update.remote(is_checkpoint_format=True)) - dist.barrier(group=get_gloo_group()) - - megatron_local_weights = self.weights_getter() - self._send_weight_chunks(megatron_local_weights) - - dist.barrier(group=get_gloo_group()) - # After the barrier all engines have returned, so every rank's last-chunk - # IPC handles are now released by the consumers. Clean them up. - torch.cuda.ipc_collect() - - # vLLM #39212: exit weight-update mode. - if self._ipc_engine is not None and rank == self._ipc_gather_src: - ray.get(self._ipc_engine.finish_weight_update.remote()) - dist.barrier(group=get_gloo_group()) - + """Run one coordinated full-weight update and commit on success.""" + committed_version = self.weight_version + rollout_engines = getattr(self, "_all_rollout_engines", self.rollout_engines) + coordinator = WeightUpdateCoordinator(rollout_engines, self.quantization_config) + megatron_local_weights = None + + def transfer_target(candidate_version: int) -> None: + nonlocal megatron_local_weights + self.weight_version = candidate_version + megatron_local_weights = self._run_weight_update_session() + + draft_transfer: Callable[[int], None] | None = None if ( not self.use_distribute and self.args.enable_mtp_training and (self.args.vllm_speculative_config or {}).get("method") == "mtp" ): - if self._ipc_engine is not None and rank == self._ipc_gather_src: - ray.get(self._ipc_engine.start_draft_weight_update.remote()) - dist.barrier(group=get_gloo_group()) - self._send_weight_chunks(megatron_local_weights) + def transfer_draft(candidate_version: int) -> None: + assert candidate_version == self.weight_version + assert megatron_local_weights is not None + self._run_weight_update_session(megatron_local_weights, draft=True) + draft_transfer = transfer_draft + + try: + self.weight_version = coordinator.run( + current_version=committed_version, + transfer_target=transfer_target, + transfer_draft=draft_transfer, + ) + except BaseException: + self.weight_version = committed_version + raise + + def _run_weight_update_session( + self, + megatron_local_weights: Mapping[str, torch.Tensor] | None = None, + *, + draft: bool = False, + ) -> Mapping[str, torch.Tensor]: + self._start_weight_update_session(draft=draft) + try: + # Preserve the colocated memory order: vLLM first enters layerwise + # reload mode, then Megatron materializes the source weights. + if megatron_local_weights is None: + megatron_local_weights = self.weights_getter() + self._send_weight_chunks(megatron_local_weights) dist.barrier(group=get_gloo_group()) + # Every engine has returned, so the consumers have released the + # final chunk's IPC handles. torch.cuda.ipc_collect() - if self._ipc_engine is not None and rank == self._ipc_gather_src: - ray.get(self._ipc_engine.finish_weight_update.remote()) - dist.barrier(group=get_gloo_group()) - - # int4/fp4 post_process - if rank == 0: - if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: - post_process_weights( - restore_weights_before_load=False, - post_process_quantization=True, - rollout_engines=self.rollout_engines, + except BaseException: + try: + self._finish_weight_update_session() + except BaseException: + logger.warning("Failed to finalize vLLM session after weight transfer error", exc_info=True) + raise + else: + self._finish_weight_update_session() + return megatron_local_weights + + def _start_weight_update_session(self, *, draft: bool = False) -> None: + rank = dist.get_rank() + refs = [] + if self._ipc_engine is not None and rank == self._ipc_gather_src: + if draft: + refs.append(self._ipc_engine.start_draft_weight_update.remote()) + else: + refs.append(self._ipc_engine.start_weight_update.remote(is_checkpoint_format=True)) + if self.use_distribute and self._is_distributed_src_rank: + if draft: + refs.extend(engine.start_draft_weight_update.remote() for engine in self.distributed_rollout_engines) + else: + refs.extend( + engine.start_weight_update.remote(is_checkpoint_format=True) + for engine in self.distributed_rollout_engines ) - ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + if refs: + ray.get(refs) dist.barrier(group=get_gloo_group()) + def _finish_weight_update_session(self) -> None: + rank = dist.get_rank() + refs = [] + if self._ipc_engine is not None and rank == self._ipc_gather_src: + refs.append(self._ipc_engine.finish_weight_update.remote()) + if self.use_distribute and self._is_distributed_src_rank: + refs.extend(engine.finish_weight_update.remote() for engine in self.distributed_rollout_engines) + try: + if refs: + ray.get(refs) + finally: + dist.barrier(group=get_gloo_group()) + def _send_weight_chunks(self, megatron_local_weights) -> None: max_inflight = 1 if self.use_distribute else _MAX_COLOCATED_UPDATES_INFLIGHT pending = [] @@ -343,6 +380,7 @@ def _send_hf_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: self.weight_version, self.distributed_rollout_engines, hf_named_tensors, + record_weight_version=False, ) if refs_distributed: all_refs.extend(refs_distributed) @@ -367,7 +405,11 @@ def _send_to_colocated_engine( slot_size = dist.get_world_size(ipc_gather_group) if slot_size <= 1: - ref = ipc_engine.update_weights_from_tensor.remote(**local_info, weight_version=str(weight_version)) + ref = ipc_engine.update_weights_from_tensor.remote( + **local_info, + weight_version=str(weight_version), + record_weight_version=False, + ) return [ref], weight_ref payload = _serialize_ipc_update_info(local_info) @@ -381,6 +423,12 @@ def _send_to_colocated_engine( raise RuntimeError(f"Missing IPC payloads in slot {ipc_gather_src}; got {gathered_payloads!r}") slot_infos = [_deserialize_ipc_update_info(p) for p in gathered_payloads] merged = _merge_ipc_update_infos(slot_infos) - refs.append(ipc_engine.update_weights_from_tensor.remote(**merged, weight_version=str(weight_version))) + refs.append( + ipc_engine.update_weights_from_tensor.remote( + **merged, + weight_version=str(weight_version), + record_weight_version=False, + ) + ) return refs, weight_ref diff --git a/vime/backends/vllm_utils/checkpoint_delta.py b/vime/backends/vllm_utils/checkpoint_delta.py new file mode 100644 index 000000000..75f202d59 --- /dev/null +++ b/vime/backends/vllm_utils/checkpoint_delta.py @@ -0,0 +1,296 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch +from vllm.distributed.weight_transfer import WeightTransferEngineFactory +from vllm.distributed.weight_transfer.base import WeightTransferUpdateInfo +from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +VIME_DELTA_NCCL_BACKEND = "vime_delta_nccl" +VIME_DELTA_WORKER_EXTENSION = "vime.backends.vllm_utils.checkpoint_delta.VimeDeltaWorkerExtension" +_REGISTERED = False + + +def _checkpoint_patch_api(): + try: + from vllm.model_executor.model_loader.checkpoint_weight_patch import ( + CheckpointWeightPatch, + load_checkpoint_weight_patches, + ) + except ModuleNotFoundError as exc: + if exc.name != "vllm.model_executor.model_loader.checkpoint_weight_patch": + raise + raise RuntimeError("VIME direct DWU requires a vLLM build containing PR #50723") from exc + return CheckpointWeightPatch, load_checkpoint_weight_patches + + +def _layerwise_reload_api(): + try: + from vllm.model_executor.model_loader.reload import ( + finalize_layerwise_reload, + initialize_layerwise_reload, + ) + except ModuleNotFoundError as exc: + if exc.name != "vllm.model_executor.model_loader.reload": + raise + raise RuntimeError( + "VIME direct DWU dense seeding requires vLLM's layerwise reload API" + ) from exc + return initialize_layerwise_reload, finalize_layerwise_reload + + +# GPU staging bound for the checkpoint patch API: full checkpoint-shaped +# tensors accumulate up to this target per internal load_weights call. Peak +# transient memory is this target plus the largest single checkpoint tensor, +# because a tensor bigger than the target is still staged whole. +_PATCH_CHUNK_BYTES = 256 << 20 + + +@dataclass +class VimeDeltaNCCLUpdateInfo(WeightTransferUpdateInfo): + schema_version: int + base_version: int + target_version: int + sequence_no: int + is_final: bool + encoding: str + patches: list[dict[str, Any]] + position_count: int + value_count: int + value_dtype_name: str + + def __post_init__(self) -> None: + if self.schema_version != 1: + raise ValueError(f"Unsupported checkpoint delta schema: {self.schema_version}") + if self.target_version != self.base_version + 1: + raise ValueError("target_version must be base_version + 1") + if self.sequence_no < 0: + raise ValueError("sequence_no must be non-negative") + if self.encoding not in {"dense", "indices"}: + raise ValueError(f"Unsupported checkpoint delta encoding: {self.encoding!r}") + if self.position_count < 0 or self.value_count < 0: + raise ValueError("Checkpoint delta tensor sizes must be non-negative") + if self.is_final: + if self.patches or self.position_count or self.value_count: + raise ValueError("A final checkpoint delta manifest cannot carry data") + return + if not self.patches or self.value_count == 0: + raise ValueError("A non-final checkpoint delta chunk must carry patches") + if self.encoding == "dense" and self.position_count: + raise ValueError("Dense checkpoint delta chunks must not contain positions") + if self.encoding == "indices" and self.position_count != self.value_count: + raise ValueError("Sparse checkpoint delta chunks require one position per value") + + +class VimeDeltaNCCLWeightTransferEngine(NCCLWeightTransferEngine): + """Receive VIME checkpoint-coordinate patches over the stock NCCL group.""" + + update_info_cls = VimeDeltaNCCLUpdateInfo + supports_draft_weight_update = False + + def __init__( + self, + config, + vllm_config: VllmConfig, + device: torch.device, + model: torch.nn.Module, + ) -> None: + super().__init__(config, vllm_config, device, model) + # Fail at engine construction, not mid-session, when the vLLM build + # lacks either required model-loader API. + _checkpoint_patch_api() + _layerwise_reload_api() + if self.device.type != "cuda": + raise NotImplementedError("VIME direct DWU requires CUDA") + if self.model_config.dtype != torch.bfloat16: + raise NotImplementedError("VIME direct DWU currently supports BF16 only") + if getattr(self.vllm_config, "quant_config", None) is not None: + raise NotImplementedError("VIME direct DWU does not support quantized models") + if getattr(self.vllm_config, "speculative_config", None) is not None: + raise NotImplementedError("VIME direct DWU does not update speculative draft models") + self._committed_version = 0 + self._session_base_version: int | None = None + self._session_target_version: int | None = None + self._session_encoding: str | None = None + self._next_sequence_no = 0 + self._reload_initialized = False + self._final_received = False + self._update_failed = False + + def init_transfer_engine(self, init_info) -> None: + super().init_transfer_engine(init_info) + + def update_weights(self, update_info: dict[str, Any]) -> None: + try: + super().update_weights(update_info) + except BaseException: + self._update_failed = True + raise + + def start_weight_update(self) -> None: + if self._update_failed: + raise RuntimeError("A previous direct DWU session failed; restart and dense-seed the vLLM workers") + self._session_base_version = None + self._session_target_version = None + self._session_encoding = None + self._next_sequence_no = 0 + self._reload_initialized = False + self._final_received = False + + def receive_weights(self, update_info: VimeDeltaNCCLUpdateInfo) -> None: + if self.model_update_group is None: + raise RuntimeError("VIME direct DWU NCCL group is not initialized") + + try: + if update_info.is_final: + self._accept_chunk_metadata(update_info) + if self._final_received: + raise ValueError("A direct DWU session received two final manifests") + if update_info.encoding == "dense" and not self._reload_initialized: + raise ValueError("A dense direct DWU session did not carry any weights") + self._final_received = True + return + + # Drain this chunk from NCCL before consulting worker-local session + # state. If one TP worker has a stale version or failed earlier, all + # ranks still complete the same collectives and fail-stop cleanly + # instead of stranding the trainer and its healthy peers. + positions = torch.empty( + update_info.position_count, + dtype=torch.int32, + device=self.device, + ) + if positions.numel(): + self.model_update_group.broadcast( + positions, + src=0, + stream=torch.cuda.current_stream(), + ) + + values = torch.empty( + update_info.value_count, + dtype=torch.bfloat16, + device=self.device, + ) + self.model_update_group.broadcast( + values, + src=0, + stream=torch.cuda.current_stream(), + ) + + self._accept_chunk_metadata(update_info) + if self._final_received: + raise ValueError("Checkpoint delta data arrived after the final manifest") + if update_info.value_dtype_name != "bfloat16": + raise ValueError("VIME direct DWU wire values must be BF16") + + if update_info.encoding == "dense" and not self._reload_initialized: + initialize_layerwise_reload, _ = _layerwise_reload_api() + self._reload_initialized = True + try: + initialize_layerwise_reload(self.model) + except BaseException: + self._reload_initialized = False + raise + + CheckpointWeightPatch, load_checkpoint_weight_patches = _checkpoint_patch_api() + patches = [] + for spec in update_info.patches: + patch_indices = None + if update_info.encoding == "indices": + patch_indices = positions[spec["position_start"] : spec["position_end"]] + patches.append( + CheckpointWeightPatch( + name=spec["name"], + shape=tuple(spec["shape"]), + dtype=getattr(torch, spec["dtype_name"]), + values=values[spec["value_start"] : spec["value_end"]], + indices=patch_indices, + ) + ) + # Indices come from one torch.nonzero over a bitwise compare on the + # source, so they are unique by construction; skip the per-patch + # duplicate sort. + load_checkpoint_weight_patches( + self.model, + patches, + max_chunk_bytes=_PATCH_CHUNK_BYTES, + validate_unique_indices=False, + ) + except BaseException: + self._update_failed = True + raise + + def _accept_chunk_metadata(self, update_info: VimeDeltaNCCLUpdateInfo) -> None: + if self._session_base_version is None: + if update_info.base_version != self._committed_version: + raise RuntimeError( + f"Checkpoint delta base version mismatch: worker={self._committed_version}, update={update_info.base_version}" + ) + self._session_base_version = update_info.base_version + self._session_target_version = update_info.target_version + elif ( + update_info.base_version != self._session_base_version + or update_info.target_version != self._session_target_version + ): + raise ValueError("One direct DWU session cannot mix update versions") + + if update_info.sequence_no != self._next_sequence_no: + raise ValueError( + f"Checkpoint delta sequence mismatch: expected {self._next_sequence_no}, got {update_info.sequence_no}" + ) + self._next_sequence_no += 1 + + if self._session_encoding is None: + self._session_encoding = update_info.encoding + elif self._session_encoding != update_info.encoding: + raise ValueError("One direct DWU session cannot mix dense and sparse chunks") + + def finish_weight_update(self) -> None: + try: + if not self._final_received: + raise RuntimeError("Direct DWU session ended without a final manifest") + if self._session_encoding == "dense": + _, finalize_layerwise_reload = _layerwise_reload_api() + finalize_layerwise_reload(self.model, self.model_config) + assert self._session_target_version is not None + self._committed_version = self._session_target_version + except BaseException: + self._update_failed = True + raise + finally: + self._session_base_version = None + self._session_target_version = None + self._session_encoding = None + self._next_sequence_no = 0 + self._reload_initialized = False + self._final_received = False + + def shutdown(self) -> None: + self._session_encoding = None + super().shutdown() + + +def register_vime_delta_weight_transfer_engine() -> None: + global _REGISTERED + if _REGISTERED: + return + WeightTransferEngineFactory.register_engine( + VIME_DELTA_NCCL_BACKEND, + VimeDeltaNCCLWeightTransferEngine, + ) + _REGISTERED = True + + +class VimeDeltaWorkerExtension: + """Register the VIME WTE before the vLLM worker loads its model.""" + + +# Resolving ``worker_extension_cls`` imports this module before GPUWorker loads +# the model and asks the factory to create its configured transfer engine. +register_vime_delta_weight_transfer_engine() diff --git a/vime/backends/vllm_utils/vllm_engine.py b/vime/backends/vllm_utils/vllm_engine.py index ec302c0f2..a473606a6 100644 --- a/vime/backends/vllm_utils/vllm_engine.py +++ b/vime/backends/vllm_utils/vllm_engine.py @@ -19,6 +19,17 @@ logger = logging.getLogger(__name__) _VLLM_WAKE_TAGS = frozenset({"weights", "kv_cache"}) +_DIRECT_DELTA_BACKEND = "vime_delta_nccl" +_DIRECT_DELTA_WORKER_EXTENSION = ( + "vime.backends.vllm_utils.checkpoint_delta.VimeDeltaWorkerExtension" +) + + +def _uses_direct_delta(args) -> bool: + return ( + getattr(args, "update_weight_mode", "full") == "delta" + and getattr(args, "update_weight_transport", "nccl") == "nccl" + ) def get_base_gpu_id(args, rank): @@ -66,13 +77,14 @@ def _build_subprocess_env(server_args_dict: dict[str, Any]) -> dict[str, str]: env.setdefault("VLLM_SERVER_DEV_MODE", "1") if getattr(args, "vllm_enable_deterministic_inference", False): env["VLLM_BATCH_INVARIANT"] = "1" - if getattr(args, "colocate", False): + if getattr(args, "colocate", False) or _uses_direct_delta(args): import vime vime_root = os.path.dirname(os.path.dirname(os.path.abspath(vime.__file__))) existing_pp = env.get("PYTHONPATH", "") if vime_root not in {p for p in existing_pp.split(os.pathsep) if p}: env["PYTHONPATH"] = os.pathsep.join(filter(None, [vime_root, existing_pp])) + if getattr(args, "colocate", False): env.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") worker_type = server_args_dict.get("_worker_type", "regular") @@ -130,6 +142,7 @@ def __init__( self.vllm_overrides = vllm_overrides or {} self.num_gpus_per_engine = num_gpus_per_engine self._weight_version: str | None = None + self._server_version_endpoint_missing = False def init( self, @@ -270,6 +283,7 @@ def update_weights_from_tensor( tensor_sizes: list[int], weight_version: str, flush_cache: bool = False, + record_weight_version: bool = True, ): payload: dict = { "names": names, @@ -282,7 +296,8 @@ def update_weights_from_tensor( if flush_cache: self.flush_cache() result = self._make_request("update_weights", {"update_info": payload}) - self._weight_version = str(weight_version) + if record_weight_version: + self._weight_version = str(weight_version) return result def flush_cache(self): @@ -331,7 +346,28 @@ def get_weight_version(self): return self._weight_version def set_weight_version(self, new_version: str): - self._weight_version = str(new_version) + new_version = str(new_version) + result = None + try: + result = self._make_request( + "update_weight_version", + {"new_version": new_version}, + ) + except requests.exceptions.HTTPError as exc: + status = exc.response.status_code if exc.response is not None else None + # Direct DWU publishes versions through the server tag as part of its + # commit protocol; every other mode only mirrors the local marker, so + # a vLLM without the endpoint (< 0.27.2) must not break weight sync. + if _uses_direct_delta(self.args) or status not in (404, 405): + raise + if not self._server_version_endpoint_missing: + self._server_version_endpoint_missing = True + logger.warning( + "vLLM server has no update_weight_version endpoint; " + "tracking weight versions locally only" + ) + self._weight_version = new_version + return result def release_memory_occupation(self, level: int = 2): self.flush_cache() @@ -424,6 +460,7 @@ def update_weights_from_distributed( *, flush_cache=False, weight_version: str, + record_weight_version: bool = True, ): if flush_cache: self.flush_cache() @@ -435,9 +472,14 @@ def update_weights_from_distributed( "packed": True, } result = self._make_request("update_weights", {"update_info": update_info}) - self._weight_version = str(weight_version) + if record_weight_version: + self._weight_version = str(weight_version) return result + def update_checkpoint_delta_from_distributed(self, *, update_info: dict): + """Send one direct-DWU metadata chunk without publishing a version.""" + return self._make_request("update_weights", {"update_info": update_info}) + def pause_generation(self): if self.node_rank != 0: return @@ -622,7 +664,25 @@ def _compute_server_args( ): kwargs["max_model_len"] = args.rollout_max_context_len - if args.colocate: + direct_delta = _uses_direct_delta(args) + if direct_delta and args.colocate: + raise ValueError("VIME direct DWU requires non-colocated rollout engines") + + configured_extension = getattr(args, "vllm_worker_extension_cls", "") + if direct_delta and configured_extension not in ( + "", + None, + _DIRECT_DELTA_WORKER_EXTENSION, + ): + raise ValueError( + "VIME direct DWU owns --vllm-worker-extension-cls; " + f"got {configured_extension!r}" + ) + + if direct_delta: + kwargs["weight_transfer_config"] = {"backend": _DIRECT_DELTA_BACKEND} + kwargs["worker_extension_cls"] = _DIRECT_DELTA_WORKER_EXTENSION + elif args.colocate: kwargs["weight_transfer_config"] = {"backend": "ipc"} else: kwargs["weight_transfer_config"] = {"backend": "nccl"} @@ -663,6 +723,19 @@ def _compute_server_args( ) if normalized_key in ("model_path",) or normalized_key.startswith("disaggregation"): continue + if direct_delta and normalized_key in { + "weight_transfer_config", + "worker_extension_cls", + }: + required = { + "weight_transfer_config": {"backend": _DIRECT_DELTA_BACKEND}, + "worker_extension_cls": _DIRECT_DELTA_WORKER_EXTENSION, + }[normalized_key] + if value != required: + raise ValueError( + f"VIME direct DWU requires {normalized_key}={required!r}; " + f"got {value!r}" + ) if normalized_key in kwargs: logger.info( f"vllm_overrides: overriding {normalized_key}={kwargs[normalized_key]} -> {value} (rank={rank})" diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 25fd9cc93..72579a48b 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -139,8 +139,9 @@ def add_train_arguments(parser): default="full", help=( "Weight sync strategy. 'full' (default) broadcasts every parameter " - "every sync. 'delta' diffs each sync against a pinned-CPU snapshot of the " - "previous one and ships only the changed bytes (disk transport only)." + "every sync. 'delta' emits checkpoint-coordinate absolute patches: the " + "first sync is a dense seed and later syncs contain only BF16 elements " + "whose bit patterns changed since the last committed update." ), ) parser.add_argument( @@ -150,8 +151,8 @@ def add_train_arguments(parser): help=( "Carrier for weight sync. In full mode, 'nccl' broadcasts chunks and " "'disk' writes a complete HF checkpoint under --update-weight-disk-dir " - "before engines reload it. Delta mode is 'disk' only: each host applies the " - "published deltas into its local checkpoint and reloads via update_weights_from_disk." + "before engines reload it. Delta mode currently supports direct NCCL only; " + "the disk choice is reserved for a future sink using the same delta protocol." ), ) parser.add_argument( @@ -169,8 +170,8 @@ def add_train_arguments(parser): default=None, help=( "Filesystem directory for disk-backed weight sync. In --update-weight-mode=full, " - "one complete HF checkpoint directory is written per sync. In delta mode, " - "one delta directory (changed tensors only) is written per sync." + "one complete HF checkpoint directory is written per sync. Delta disk output " + "is reserved and is not implemented yet." ), ) parser.add_argument( @@ -187,7 +188,7 @@ def add_train_arguments(parser): choices=["xor", "overwrite"], default="xor", help=( - "On-disk delta encoding for --update-weight-mode=delta --update-weight-transport=disk. " + "Reserved encoding for a future delta disk sink; direct NCCL delta sync ignores it. " "'xor' (default): new ^ old — smallest wire and fastest, but an involution that must be " "applied exactly once against the correct base (applying it twice reverts). 'overwrite': " "changed positions + new absolute values — larger, but idempotent (re-applicable any " @@ -200,7 +201,7 @@ def add_train_arguments(parser): choices=["xxh3-128", "blake3", "adler32"], default="xxh3-128", help=( - "Per-tensor integrity checksum for disk delta apply. The checksum is not the " + "Reserved checksum for a future delta disk sink. The checksum is not the " "apply bottleneck (the apply is decompress + XOR bound), so this is a digest-" "property choice, not a speed one. 'xxh3-128' (default): widest fast non-" "cryptographic digest, negligible accidental-corruption collisions. 'blake3': " @@ -229,8 +230,8 @@ def add_train_arguments(parser): "Rollout-host-local directory (NVMe) holding a full HF checkpoint kept in " "sync by each engine's pull_weights: every host copies a published full " "checkpoint as-is or patches published deltas in place, and the engines " - "reload from it. Required for --update-weight-mode=delta " - "--update-weight-transport=disk; optional for full disk sync (engines then " + "reload from it. Reserved for the future delta disk sink; optional for full " + "disk sync (engines then " "pull to local disk instead of reading the shared dir directly). The " "read-side counterpart of --custom-update-weight-post-write-path is " "--vllm-custom-pull-weights-pre-read-hook." @@ -1734,17 +1735,58 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: def _validate_update_weight_args(args) -> None: + if args.update_weight_mode == "delta": + if args.update_weight_transport == "disk": + raise NotImplementedError( + "Delta disk transport is reserved but not implemented; " + "use --update-weight-mode=delta --update-weight-transport=nccl." + ) + if args.update_weight_transport != "nccl": + raise ValueError("Direct delta weight sync requires NCCL transport") + if getattr(args, "train_backend", "megatron") != "megatron": + raise NotImplementedError("Direct delta weight sync currently requires the Megatron trainer") + if getattr(args, "colocate", False): + raise NotImplementedError("Direct delta weight sync currently requires non-colocated rollout engines") + if getattr(args, "rollout_external", False): + raise NotImplementedError("Direct delta weight sync currently requires VIME-launched rollout engines") + if getattr(args, "use_fault_tolerance", False): + raise NotImplementedError("Direct delta weight sync does not yet support replacement rollout workers") + if "fully_async" in getattr(args, "rollout_function_path", ""): + raise NotImplementedError("Direct delta weight sync does not yet support fully-async rollout") + if getattr(args, "offload_rollout", False): + raise NotImplementedError("Direct delta weight sync does not yet support rollout offload") + if getattr(args, "enable_mtp_training", False) or getattr(args, "vllm_speculative_config", None): + raise NotImplementedError("Direct delta weight sync does not yet support speculative or MTP models") + if getattr(args, "fp16", False): + raise NotImplementedError("Direct delta weight sync currently supports BF16 only") + if getattr(args, "pipeline_model_parallel_size", 1) != 1: + raise NotImplementedError("Direct delta weight sync currently requires Megatron PP=1") + if getattr(args, "virtual_pipeline_model_parallel_size", None) not in (None, 1): + raise NotImplementedError("Direct delta weight sync currently requires Megatron VPP=1") + if getattr(args, "num_layers_per_virtual_pipeline_stage", None) is not None: + raise NotImplementedError("Direct delta weight sync currently requires Megatron VPP=1") + if getattr(args, "vllm_pipeline_parallel_size", 1) != 1: + raise NotImplementedError("Direct delta weight sync currently requires vLLM PP=1") + if getattr(args, "vllm_data_parallel_size", 1) != 1: + raise NotImplementedError("Direct delta weight sync currently requires vLLM DP=1") + if getattr(args, "vllm_enable_deterministic_inference", False): + raise NotImplementedError( + "Direct delta weight sync does not support batch-invariant deterministic inference" + ) + if getattr(args, "update_weight_start_version", 0) != 0: + raise ValueError("Direct delta weight sync must start from version 0 and dense-seed the workers") + worker_extension = getattr(args, "vllm_worker_extension_cls", "") + expected_extension = "vime.backends.vllm_utils.checkpoint_delta.VimeDeltaWorkerExtension" + if worker_extension not in ("", None, expected_extension): + raise ValueError("Direct delta weight sync owns --vllm-worker-extension-cls") + return + if args.update_weight_transport == "disk" and not args.update_weight_disk_dir: raise ValueError( "--update-weight-transport=disk requires --update-weight-disk-dir to point at " "a filesystem shared between the trainer and the rollout engines." ) - if args.update_weight_mode == "delta": - raise NotImplementedError( - "--update-weight-mode=delta is unverified on vime+vLLM and is disabled; " "use --update-weight-mode=full." - ) - def vime_validate_args(args): args.eval_datasets = _resolve_eval_datasets(args)