fix(weight-sync): record IPC weight version on engines so ci_test passes - #45
Merged
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request threads the weight_version through the finish_weight_update method to ensure that colocated engines correctly record the updated weight version during IPC weight updates, preventing a version mismatch in tests. The reviewer suggests updating self._weight_version only after the HTTP request to /finish_weight_update successfully completes to ensure transactional correctness in case of network or timeout failures.
Comment on lines
825
to
829
| if weight_version is not None: | ||
| self._weight_version = str(weight_version) | ||
| update_timeout_s = float(os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900")) | ||
| response = self._post_json("finish_weight_update", {}, timeout=update_timeout_s) | ||
| response.raise_for_status() |
There was a problem hiding this comment.
To ensure transactional correctness, consider updating self._weight_version only after the HTTP request to /finish_weight_update has successfully completed. If the request fails (e.g., due to a timeout or network error), the local version tracker should not be updated to the new version, as the engine failed to successfully exit the weight-update mode.
Suggested change
| if weight_version is not None: | |
| self._weight_version = str(weight_version) | |
| update_timeout_s = float(os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900")) | |
| response = self._post_json("finish_weight_update", {}, timeout=update_timeout_s) | |
| response.raise_for_status() | |
| update_timeout_s = float(os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900")) | |
| response = self._post_json("finish_weight_update", {}, timeout=update_timeout_s) | |
| response.raise_for_status() | |
| if weight_version is not None: | |
| self._weight_version = str(weight_version) |
After PR #22 added the colocated CUDA IPC weight sync path, the IPC branch bypasses ``VLLMEngine.update_weights_from_tensor`` entirely (data moves through ``IPCWeightTransferEngine.trainer_send_weights`` directly, not the ``/update_weights`` RPC). That RPC is the normal place where the engine's ``self._weight_version`` gets recorded; the IPC path leaves it at the initial ``None``. The fail mode shows up the first time IPC weight sync runs with ``--ci-test``: ``actor.update_weights`` at slime/backends/megatron_utils/actor.py picks a random rollout engine, calls ``get_weight_version()``, and compares against the trainer-side ``weight_updater.weight_version``. Because the engine's ``_weight_version`` is still ``None``, ``get_weight_version`` falls back to ``GET /v1/models`` and returns the model path string (e.g. ``/root/models/Qwen2.5-0.5B-Instruct``); the updater returns the integer version string (``"1"``, ``"2"``, ...). They never match: RuntimeError: Weight version mismatch! Engine: /root/models/Qwen2.5-0.5B-Instruct, Updater: 1 Reproducer: ``tests/test_qwen2.5_0.5B_short.py`` and ``tests/test_qwen2.5_0.5B_ppo_critic_only_short.py`` once they reach ``update_weights`` — both colocate + IPC, both have ``--ci-test`` set. Thread the version through: - ``VLLMEngine.finish_weight_update(weight_version: str | None = None)`` now sets ``self._weight_version = str(weight_version)`` before issuing the ``POST /finish_weight_update``. - ``UpdateWeightFromTensor.update_weights`` (step 5) passes ``str(self.weight_version)`` when the coordinator rank fires the per-engine ``finish_weight_update`` RPC. Each colocated engine's coordinator rank fires its own ``finish_weight_update``, so every engine ends up with the new version recorded; the distributed path is unchanged. Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02
force-pushed
the
aoshen/fix-ipc-weight-version
branch
from
May 27, 2026 02:40
d884078 to
86b808a
Compare
Collaborator
Author
|
Thanks, applied — moved the |
Collaborator
|
LGTM |
3 tasks
aoshen02
added a commit
that referenced
this pull request
May 27, 2026
…version-with-data) Background ---------- After PR #22 introduced the colocated CUDA IPC path, vime ended up with three ``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was inconsistent: - ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version`` inside the RPC, version travels with data — slime-style. - ``update_weights`` (IPC path, called from ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never recorded ``_weight_version`` — vLLM's payload schema does not carry it. - ``update_weights_from_tensor`` (PR #18 legacy entry): kept a SGLang-ish ``serialized_named_tensors`` payload and wrote ``_weight_version``, but had no callers in main. The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine: /root/models/<...>, Updater: N" failure on every colocated test with ``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto ``finish_weight_update``). slime's design avoids this entirely: both IPC and distributed call ``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same RPC name across both repos, with version travelling alongside the data in the same RPC. This PR ------- Rewire vime's IPC path to match slime's interface: 1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point. Signature ``(update_info: dict, weight_version: str | None, flush_cache)``; payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names / shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after the POST succeeds — mirrors ``update_weights_from_distributed``'s post-POST ordering so a failed transfer never advances the engine's tracked version. 2. Delete ``vllm_engine.update_weights`` — was the vLLM ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used after step 4 below. 3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only ``update_weights_from_tensor``'s old SGLang-ish path called. 4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version`` now lives in the data-carrying RPC, so the bookend no longer needs to piggyback a kwarg. (Undoes the kwarg added in #45.) 5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in ``_send_hf_chunk_via_ipc`` with direct ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)`` for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing vLLM's ``reduce_tensor`` for IPC handle creation (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is ours — so we don't fork the vLLM IPC protocol, just route through our own RPC surface. Why not just keep #45's piggyback? - #45 worked but coupled version bookkeeping to the lifecycle hook ``finish_weight_update`` instead of the data RPC. The wire shape doesn't match slime's, and a new IPC-style entry point added later would have to remember to also write ``_weight_version`` — exactly the trap PR #22 fell into. Centralising the write inside the data RPC removes the trap. Unit tests ---------- - ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC call recording stays complete. - Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` -> ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``, asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and that ``finish_weight_update`` is now stateless (no kwargs). - Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``: asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata fields pass through, ``_weight_version`` advances on POST success. - Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``: asserts POST failure does not advance ``_weight_version`` (matches the same post-POST ordering review note from #45). Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py (``_weight_transfer_http_timeout``, ``_response_json_or_fallback``, ``server_host``) are unchanged from main — main has 8 failed / 24 passed, this PR has 8 failed / 26 passed (the two added tests). Not in scope here. Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02
added a commit
that referenced
this pull request
May 27, 2026
…version-with-data) Background ---------- After PR #22 introduced the colocated CUDA IPC path, vime ended up with three ``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was inconsistent: - ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version`` inside the RPC, version travels with data — slime-style. - ``update_weights`` (IPC path, called from ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never recorded ``_weight_version`` — vLLM's payload schema does not carry it. - ``update_weights_from_tensor`` (PR #18 legacy entry): kept a SGLang-ish ``serialized_named_tensors`` payload and wrote ``_weight_version``, but had no callers in main. The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine: /root/models/<...>, Updater: N" failure on every colocated test with ``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto ``finish_weight_update``). slime's design avoids this entirely: both IPC and distributed call ``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same RPC name across both repos, with version travelling alongside the data in the same RPC. This PR ------- Rewire vime's IPC path to match slime's interface: 1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point. Signature ``(update_info: dict, weight_version: str | None, flush_cache)``; payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names / shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after the POST succeeds — mirrors ``update_weights_from_distributed``'s post-POST ordering so a failed transfer never advances the engine's tracked version. 2. Delete ``vllm_engine.update_weights`` — was the vLLM ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used after step 4 below. 3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only ``update_weights_from_tensor``'s old SGLang-ish path called. 4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version`` now lives in the data-carrying RPC, so the bookend no longer needs to piggyback a kwarg. (Undoes the kwarg added in #45.) 5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in ``_send_hf_chunk_via_ipc`` with direct ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)`` for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing vLLM's ``reduce_tensor`` for IPC handle creation (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is ours — so we don't fork the vLLM IPC protocol, just route through our own RPC surface. Why not just keep #45's piggyback? - #45 worked but coupled version bookkeeping to the lifecycle hook ``finish_weight_update`` instead of the data RPC. The wire shape doesn't match slime's, and a new IPC-style entry point added later would have to remember to also write ``_weight_version`` — exactly the trap PR #22 fell into. Centralising the write inside the data RPC removes the trap. Unit tests ---------- - ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC call recording stays complete. - Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` -> ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``, asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and that ``finish_weight_update`` is now stateless (no kwargs). - Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``: asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata fields pass through, ``_weight_version`` advances on POST success. - Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``: asserts POST failure does not advance ``_weight_version`` (matches the same post-POST ordering review note from #45). Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py (``_weight_transfer_http_timeout``, ``_response_json_or_fallback``, ``server_host``) are unchanged from main — main has 8 failed / 24 passed, this PR has 8 failed / 26 passed (the two added tests). Not in scope here. Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02
added a commit
that referenced
this pull request
May 27, 2026
…version-with-data) Background ---------- After PR #22 introduced the colocated CUDA IPC path, vime ended up with three ``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was inconsistent: - ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version`` inside the RPC, version travels with data — slime-style. - ``update_weights`` (IPC path, called from ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never recorded ``_weight_version`` — vLLM's payload schema does not carry it. - ``update_weights_from_tensor`` (PR #18 legacy entry): kept a SGLang-ish ``serialized_named_tensors`` payload and wrote ``_weight_version``, but had no callers in main. The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine: /root/models/<...>, Updater: N" failure on every colocated test with ``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto ``finish_weight_update``). slime's design avoids this entirely: both IPC and distributed call ``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same RPC name across both repos, with version travelling alongside the data in the same RPC. This PR ------- Rewire vime's IPC path to match slime's interface: 1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point. Signature ``(update_info: dict, weight_version: str | None, flush_cache)``; payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names / shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after the POST succeeds — mirrors ``update_weights_from_distributed``'s post-POST ordering so a failed transfer never advances the engine's tracked version. 2. Delete ``vllm_engine.update_weights`` — was the vLLM ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used after step 4 below. 3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only ``update_weights_from_tensor``'s old SGLang-ish path called. 4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version`` now lives in the data-carrying RPC, so the bookend no longer needs to piggyback a kwarg. (Undoes the kwarg added in #45.) 5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in ``_send_hf_chunk_via_ipc`` with direct ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)`` for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing vLLM's ``reduce_tensor`` for IPC handle creation (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is ours — so we don't fork the vLLM IPC protocol, just route through our own RPC surface. Why not just keep #45's piggyback? - #45 worked but coupled version bookkeeping to the lifecycle hook ``finish_weight_update`` instead of the data RPC. The wire shape doesn't match slime's, and a new IPC-style entry point added later would have to remember to also write ``_weight_version`` — exactly the trap PR #22 fell into. Centralising the write inside the data RPC removes the trap. Unit tests ---------- - ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC call recording stays complete. - Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` -> ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``, asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and that ``finish_weight_update`` is now stateless (no kwargs). - Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``: asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata fields pass through, ``_weight_version`` advances on POST success. - Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``: asserts POST failure does not advance ``_weight_version`` (matches the same post-POST ordering review note from #45). Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py (``_weight_transfer_http_timeout``, ``_response_json_or_fallback``, ``server_host``) are unchanged from main — main has 8 failed / 24 passed, this PR has 8 failed / 26 passed (the two added tests). Not in scope here. Signed-off-by: aoshen02 <aoshen@inferact.ai>
CalvinXKY
pushed a commit
that referenced
this pull request
May 27, 2026
…version-with-data) Background ---------- After PR #22 introduced the colocated CUDA IPC path, vime ended up with three ``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was inconsistent: - ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version`` inside the RPC, version travels with data — slime-style. - ``update_weights`` (IPC path, called from ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never recorded ``_weight_version`` — vLLM's payload schema does not carry it. - ``update_weights_from_tensor`` (PR #18 legacy entry): kept a SGLang-ish ``serialized_named_tensors`` payload and wrote ``_weight_version``, but had no callers in main. The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine: /root/models/<...>, Updater: N" failure on every colocated test with ``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto ``finish_weight_update``). slime's design avoids this entirely: both IPC and distributed call ``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same RPC name across both repos, with version travelling alongside the data in the same RPC. This PR ------- Rewire vime's IPC path to match slime's interface: 1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point. Signature ``(update_info: dict, weight_version: str | None, flush_cache)``; payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names / shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after the POST succeeds — mirrors ``update_weights_from_distributed``'s post-POST ordering so a failed transfer never advances the engine's tracked version. 2. Delete ``vllm_engine.update_weights`` — was the vLLM ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used after step 4 below. 3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only ``update_weights_from_tensor``'s old SGLang-ish path called. 4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version`` now lives in the data-carrying RPC, so the bookend no longer needs to piggyback a kwarg. (Undoes the kwarg added in #45.) 5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in ``_send_hf_chunk_via_ipc`` with direct ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)`` for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing vLLM's ``reduce_tensor`` for IPC handle creation (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is ours — so we don't fork the vLLM IPC protocol, just route through our own RPC surface. Why not just keep #45's piggyback? - #45 worked but coupled version bookkeeping to the lifecycle hook ``finish_weight_update`` instead of the data RPC. The wire shape doesn't match slime's, and a new IPC-style entry point added later would have to remember to also write ``_weight_version`` — exactly the trap PR #22 fell into. Centralising the write inside the data RPC removes the trap. Unit tests ---------- - ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC call recording stays complete. - Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` -> ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``, asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and that ``finish_weight_update`` is now stateless (no kwargs). - Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``: asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata fields pass through, ``_weight_version`` advances on POST success. - Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``: asserts POST failure does not advance ``_weight_version`` (matches the same post-POST ordering review note from #45). Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py (``_weight_transfer_http_timeout``, ``_response_json_or_fallback``, ``server_host``) are unchanged from main — main has 8 failed / 24 passed, this PR has 8 failed / 26 passed (the two added tests). Not in scope here. Signed-off-by: aoshen02 <aoshen@inferact.ai>
CalvinXKY
pushed a commit
that referenced
this pull request
May 28, 2026
…version-with-data) (#48) * refactor(weight-sync): align IPC RPC contract with slime (single-RPC version-with-data) Background ---------- After PR #22 introduced the colocated CUDA IPC path, vime ended up with three ``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was inconsistent: - ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version`` inside the RPC, version travels with data — slime-style. - ``update_weights`` (IPC path, called from ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never recorded ``_weight_version`` — vLLM's payload schema does not carry it. - ``update_weights_from_tensor`` (PR #18 legacy entry): kept a SGLang-ish ``serialized_named_tensors`` payload and wrote ``_weight_version``, but had no callers in main. The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine: /root/models/<...>, Updater: N" failure on every colocated test with ``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto ``finish_weight_update``). slime's design avoids this entirely: both IPC and distributed call ``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same RPC name across both repos, with version travelling alongside the data in the same RPC. This PR ------- Rewire vime's IPC path to match slime's interface: 1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point. Signature ``(update_info: dict, weight_version: str | None, flush_cache)``; payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names / shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after the POST succeeds — mirrors ``update_weights_from_distributed``'s post-POST ordering so a failed transfer never advances the engine's tracked version. 2. Delete ``vllm_engine.update_weights`` — was the vLLM ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used after step 4 below. 3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only ``update_weights_from_tensor``'s old SGLang-ish path called. 4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version`` now lives in the data-carrying RPC, so the bookend no longer needs to piggyback a kwarg. (Undoes the kwarg added in #45.) 5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in ``_send_hf_chunk_via_ipc`` with direct ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)`` for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing vLLM's ``reduce_tensor`` for IPC handle creation (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is ours — so we don't fork the vLLM IPC protocol, just route through our own RPC surface. Why not just keep #45's piggyback? - #45 worked but coupled version bookkeeping to the lifecycle hook ``finish_weight_update`` instead of the data RPC. The wire shape doesn't match slime's, and a new IPC-style entry point added later would have to remember to also write ``_weight_version`` — exactly the trap PR #22 fell into. Centralising the write inside the data RPC removes the trap. Unit tests ---------- - ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC call recording stays complete. - Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` -> ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``, asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and that ``finish_weight_update`` is now stateless (no kwargs). - Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``: asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata fields pass through, ``_weight_version`` advances on POST success. - Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``: asserts POST failure does not advance ``_weight_version`` (matches the same post-POST ordering review note from #45). Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py (``_weight_transfer_http_timeout``, ``_response_json_or_fallback``, ``server_host``) are unchanged from main — main has 8 failed / 24 passed, this PR has 8 failed / 26 passed (the two added tests). Not in scope here. Signed-off-by: aoshen02 <aoshen@inferact.ai> * fix(weight-sync): correct IPC slot leader gating and gather group Two regressions in the previous commit only fire when Megatron TP != rollout-num-gpus-per-engine (e.g. parallel-check sweeps Megatron TP=1 with rollout TP=2). Both surface via ``tests/test_qwen3_0.6B_parallel_check.py``. Bug 1: leader gating wrong reference group ------------------------------------------ ``connect_rollout_engines`` used:: if mpu.get_tensor_model_parallel_rank() == 0: self._ipc_engine_coordinator = True The intent was "TP rank 0 within the engine GPU slot", but ``mpu.get_tensor_model_parallel_rank()`` is the Megatron TP rank, not the engine-slot rank. When Megatron TP=1, every trainer rank sees ``tp_rank=0`` and becomes a coordinator. For slot_size > 1, both ranks in the slot then call ``start_weight_update`` → the second call explodes:: Worker failed with error 'start_weight_update called while a weight update is already active. Call finish_weight_update first.' Fix: gate on ``rank == start`` (lowest trainer rank in the engine GPU range). Unique per slot regardless of Megatron parallelism. Bug 2: gather group wrong scope ------------------------------- ``_send_hf_chunk_via_ipc`` used ``mpu.get_tensor_model_parallel_group()`` to all_gather IPC payloads from peers in the engine slot. Again Megatron TP group ≠ engine slot when their world sizes differ. The merge then only had the coordinator's own UUID; downstream workers reading a different physical GPU got:: ValueError: IPC handle not found for GPU UUID <peer>. Available UUIDs: ['<coordinator>'] Fix: build per-slot process groups in ``connect_rollout_engines`` collectively (every trainer rank calls ``dist.new_group(slot_ranks)`` for every engine slot, keeps the one it belongs to). Use that group instead of Megatron's TP group for the gather and the trailing barrier. Validation ---------- Ran ``tests/test_qwen3_0.6B_parallel_check.py`` on 8×H200 with ``--num-rollout 2``. The test sweeps tp_size ∈ {1, 2, 4, 8} × pp_size ∈ {1, 2, 4} × cp_size ∈ {1, 2, 4, 8} for num_gpus ∈ {8, 4, 2} — every leg also uses ``--rollout-num-gpus-per-engine 2``, so the Megatron-TP=1 cases now exercise the new slot-group path. Pre-fix: fails on the very first Megatron-TP=1 config with bug 1 above; after bug 1 is patched, the next iteration fails with bug 2. Post-fix: the entire ~2-hour sweep completes with "Job succeeded" on every leg. Other tests already exercising IPC at Megatron-TP=2 == rollout-TP=2 are unaffected by this fix (``rank == start`` is equivalent to ``tp_rank == 0`` when the slot fits a single Megatron TP group). The following also pass post-fix as a sanity check: ``test_qwen3_4B_ppo``, ``test_qwen3_4B_ppo_train_critic_only``, ``test_qwen3_4B_ppo_disaggregate``, ``test_mimo_7B_mtp_only_grad``, ``test_moonlight_16B_A3B``, ``test_quick_start_glm4_9B``, ``test_qwen2.5_0.5B_{short,async_short,debug_rollout_then_train, ppo_critic_only_short}``, ``test_qwen3.5_0.8B_gsm8k_{short,async_short}``. Signed-off-by: aoshen02 <aoshen@inferact.ai> 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix(weight-sync): use gloo backend for ipc payload gather, add multi-gpu test Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: SamitHuang <285365963@qq.com> Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(weight-sync): drop dead _apply_monkey_patch_torch_reductions calls The two _apply_monkey_patch_torch_reductions() call sites in this file (trainer send + vLLM worker hijack) are no-ops on the IPC path here: 1. We route IPC handles by physical GPU UUID dict key (set on the trainer side at _build_ipc_update_info_from_named_tensors via _current_gpu_uuid() from torch.cuda.get_device_properties().uuid). The receiver looks up by its own UUID, independent of args[6]. 2. vLLM's IPCWeightTransferEngine.receive_weights unconditionally overwrites args[6] with the receiver's local device_index before calling rebuild_cuda_tensor. Whatever a torch reductions patch encodes into args[6] is therefore discarded. The patch was the historical mechanism (sglang upstream) for translating device indices across CUDA_VISIBLE_DEVICES boundaries by stuffing UUID strings into args[6]. Our UUID-keyed dict + vLLM's explicit device_index override accomplish the same thing without the global torch reductions mutation. Also expand the _build_ipc_update_info_from_named_tensors docstring to spell out the UUID-keyed routing contract so future readers don't have to chase this through git history. Side effect: hf_weight_iterator_direct.py also calls monkey_patch_torch_reductions() at module-collective time. That call site is similarly decorative (only NCCL broadcast / all_gather collectives run there, no cross-process pickling) but lives outside this file's scope and is not touched here. Tracked alongside #29. * refactor(weight-sync): finish removing monkey_patch_torch_reductions dead code Follow-up to 39bf899 (deleted _apply_monkey_patch_torch_reductions from update_weight_from_tensor.py). With that helper gone, two more references are now dead in the vime IPC weight-transfer path: 1. hf_weight_iterator_direct.py:48 called monkey_patch_torch_reductions() at the top of _get_megatron_full_params(). On vime this never has effect on the IPC handle path: _get_megatron_full_params only runs NCCL broadcast/all_gather collectives (no cross-process pickling), and the chunks it returns are subsequently sent via PR #48's UUID-keyed {gpu_uuid: reduce_tensor(weight)} dict that vLLM's receiver routes by physical UUID + explicit args[6] overwrite. The call survives in slime/ miles upstream because their downstream path pickles tensors through sglang's MultiprocessingSerializer.serialize (ForkingPickler → reduce_tensor), where the patched encoding/decoding does real work; PR #48 does not use that pipeline, so the call here was incidentally inherited rather than functionally required. 2. slime/backends/megatron_utils/sglang.py's monkey_patch_torch_reductions re-export + __all__ entry now have no remaining importers in vime. Remove them. Also (this commit, B): 3. vllm_engine.py's update_weights_from_tensor docstring referred to "closures injected by _apply_monkey_patch_torch_reductions" as the reason for cloudpickle. That helper is gone; the cloudpickle is still correct because reduce_tensor returns a (rebuild_fn, args) tuple where the rebuild_fn is a module-level callable that JSON can't serialise. Update the docstring to reflect the actual reason. Net behaviour: identical — the deletions remove dead code paths. The patch_torch shim itself is still importable for callers outside vime (none currently in this tree). Cross-references: - #29 — issue documenting the no-op stub - 39bf899 — prior commit that deleted the helper from update_weight_from_tensor.py * docs(weight-sync): drop sglang refs + add IPC call-stack notes - update_weight_from_tensor.py module docstring: rewrite from a pure vLLM perspective. Step (2) now spells out the merged {uuid_G0: handle, uuid_G1: handle, ...} dict + collective_rpc fan-out inside the vLLM server, and step (3) makes version-with-data atomicity explicit. Drop the "match slime's sglang_engine signature" framing. - vllm_engine.update_weights_from_tensor: collapse the long sglang-vs-vLLM compare block into a focused docstring describing what the POST does and why ipc_handles needs cloudpickle. Add a Chinese call-stack walkthrough (trainer → ★this method★ → server collective_rpc → per-TP worker receive_weights) and note that `node_rank != 0` is a dead branch since VLLMEngine pins node_rank=0 (see PR #48 review comment for the follow-up cleanup). - Hoist base64/cloudpickle imports to module scope so the hot path no longer pays the per-call import overhead. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> --------- Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: SamitHuang <285365963@qq.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: SamitHuang <285365963@qq.com> Co-authored-by: Cursor <cursoragent@cursor.com>
momo609
pushed a commit
that referenced
this pull request
Jun 8, 2026
…ses (#45) After PR #22 added the colocated CUDA IPC weight sync path, the IPC branch bypasses ``VLLMEngine.update_weights_from_tensor`` entirely (data moves through ``IPCWeightTransferEngine.trainer_send_weights`` directly, not the ``/update_weights`` RPC). That RPC is the normal place where the engine's ``self._weight_version`` gets recorded; the IPC path leaves it at the initial ``None``. The fail mode shows up the first time IPC weight sync runs with ``--ci-test``: ``actor.update_weights`` at slime/backends/megatron_utils/actor.py picks a random rollout engine, calls ``get_weight_version()``, and compares against the trainer-side ``weight_updater.weight_version``. Because the engine's ``_weight_version`` is still ``None``, ``get_weight_version`` falls back to ``GET /v1/models`` and returns the model path string (e.g. ``/root/models/Qwen2.5-0.5B-Instruct``); the updater returns the integer version string (``"1"``, ``"2"``, ...). They never match: RuntimeError: Weight version mismatch! Engine: /root/models/Qwen2.5-0.5B-Instruct, Updater: 1 Reproducer: ``tests/test_qwen2.5_0.5B_short.py`` and ``tests/test_qwen2.5_0.5B_ppo_critic_only_short.py`` once they reach ``update_weights`` — both colocate + IPC, both have ``--ci-test`` set. Thread the version through: - ``VLLMEngine.finish_weight_update(weight_version: str | None = None)`` now sets ``self._weight_version = str(weight_version)`` before issuing the ``POST /finish_weight_update``. - ``UpdateWeightFromTensor.update_weights`` (step 5) passes ``str(self.weight_version)`` when the coordinator rank fires the per-engine ``finish_weight_update`` RPC. Each colocated engine's coordinator rank fires its own ``finish_weight_update``, so every engine ends up with the new version recorded; the distributed path is unchanged. Signed-off-by: aoshen02 <aoshen@inferact.ai>
momo609
pushed a commit
that referenced
this pull request
Jun 8, 2026
…version-with-data) (#48) * refactor(weight-sync): align IPC RPC contract with slime (single-RPC version-with-data) Background ---------- After PR #22 introduced the colocated CUDA IPC path, vime ended up with three ``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was inconsistent: - ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version`` inside the RPC, version travels with data — slime-style. - ``update_weights`` (IPC path, called from ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never recorded ``_weight_version`` — vLLM's payload schema does not carry it. - ``update_weights_from_tensor`` (PR #18 legacy entry): kept a SGLang-ish ``serialized_named_tensors`` payload and wrote ``_weight_version``, but had no callers in main. The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine: /root/models/<...>, Updater: N" failure on every colocated test with ``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto ``finish_weight_update``). slime's design avoids this entirely: both IPC and distributed call ``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same RPC name across both repos, with version travelling alongside the data in the same RPC. This PR ------- Rewire vime's IPC path to match slime's interface: 1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point. Signature ``(update_info: dict, weight_version: str | None, flush_cache)``; payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names / shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after the POST succeeds — mirrors ``update_weights_from_distributed``'s post-POST ordering so a failed transfer never advances the engine's tracked version. 2. Delete ``vllm_engine.update_weights`` — was the vLLM ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used after step 4 below. 3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only ``update_weights_from_tensor``'s old SGLang-ish path called. 4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version`` now lives in the data-carrying RPC, so the bookend no longer needs to piggyback a kwarg. (Undoes the kwarg added in #45.) 5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in ``_send_hf_chunk_via_ipc`` with direct ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)`` for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing vLLM's ``reduce_tensor`` for IPC handle creation (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is ours — so we don't fork the vLLM IPC protocol, just route through our own RPC surface. Why not just keep #45's piggyback? - #45 worked but coupled version bookkeeping to the lifecycle hook ``finish_weight_update`` instead of the data RPC. The wire shape doesn't match slime's, and a new IPC-style entry point added later would have to remember to also write ``_weight_version`` — exactly the trap PR #22 fell into. Centralising the write inside the data RPC removes the trap. Unit tests ---------- - ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC call recording stays complete. - Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` -> ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``, asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and that ``finish_weight_update`` is now stateless (no kwargs). - Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``: asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata fields pass through, ``_weight_version`` advances on POST success. - Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``: asserts POST failure does not advance ``_weight_version`` (matches the same post-POST ordering review note from #45). Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py (``_weight_transfer_http_timeout``, ``_response_json_or_fallback``, ``server_host``) are unchanged from main — main has 8 failed / 24 passed, this PR has 8 failed / 26 passed (the two added tests). Not in scope here. Signed-off-by: aoshen02 <aoshen@inferact.ai> * fix(weight-sync): correct IPC slot leader gating and gather group Two regressions in the previous commit only fire when Megatron TP != rollout-num-gpus-per-engine (e.g. parallel-check sweeps Megatron TP=1 with rollout TP=2). Both surface via ``tests/test_qwen3_0.6B_parallel_check.py``. Bug 1: leader gating wrong reference group ------------------------------------------ ``connect_rollout_engines`` used:: if mpu.get_tensor_model_parallel_rank() == 0: self._ipc_engine_coordinator = True The intent was "TP rank 0 within the engine GPU slot", but ``mpu.get_tensor_model_parallel_rank()`` is the Megatron TP rank, not the engine-slot rank. When Megatron TP=1, every trainer rank sees ``tp_rank=0`` and becomes a coordinator. For slot_size > 1, both ranks in the slot then call ``start_weight_update`` → the second call explodes:: Worker failed with error 'start_weight_update called while a weight update is already active. Call finish_weight_update first.' Fix: gate on ``rank == start`` (lowest trainer rank in the engine GPU range). Unique per slot regardless of Megatron parallelism. Bug 2: gather group wrong scope ------------------------------- ``_send_hf_chunk_via_ipc`` used ``mpu.get_tensor_model_parallel_group()`` to all_gather IPC payloads from peers in the engine slot. Again Megatron TP group ≠ engine slot when their world sizes differ. The merge then only had the coordinator's own UUID; downstream workers reading a different physical GPU got:: ValueError: IPC handle not found for GPU UUID <peer>. Available UUIDs: ['<coordinator>'] Fix: build per-slot process groups in ``connect_rollout_engines`` collectively (every trainer rank calls ``dist.new_group(slot_ranks)`` for every engine slot, keeps the one it belongs to). Use that group instead of Megatron's TP group for the gather and the trailing barrier. Validation ---------- Ran ``tests/test_qwen3_0.6B_parallel_check.py`` on 8×H200 with ``--num-rollout 2``. The test sweeps tp_size ∈ {1, 2, 4, 8} × pp_size ∈ {1, 2, 4} × cp_size ∈ {1, 2, 4, 8} for num_gpus ∈ {8, 4, 2} — every leg also uses ``--rollout-num-gpus-per-engine 2``, so the Megatron-TP=1 cases now exercise the new slot-group path. Pre-fix: fails on the very first Megatron-TP=1 config with bug 1 above; after bug 1 is patched, the next iteration fails with bug 2. Post-fix: the entire ~2-hour sweep completes with "Job succeeded" on every leg. Other tests already exercising IPC at Megatron-TP=2 == rollout-TP=2 are unaffected by this fix (``rank == start`` is equivalent to ``tp_rank == 0`` when the slot fits a single Megatron TP group). The following also pass post-fix as a sanity check: ``test_qwen3_4B_ppo``, ``test_qwen3_4B_ppo_train_critic_only``, ``test_qwen3_4B_ppo_disaggregate``, ``test_mimo_7B_mtp_only_grad``, ``test_moonlight_16B_A3B``, ``test_quick_start_glm4_9B``, ``test_qwen2.5_0.5B_{short,async_short,debug_rollout_then_train, ppo_critic_only_short}``, ``test_qwen3.5_0.8B_gsm8k_{short,async_short}``. Signed-off-by: aoshen02 <aoshen@inferact.ai> 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix(weight-sync): use gloo backend for ipc payload gather, add multi-gpu test Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: SamitHuang <285365963@qq.com> Co-authored-by: Cursor <cursoragent@cursor.com> * refactor(weight-sync): drop dead _apply_monkey_patch_torch_reductions calls The two _apply_monkey_patch_torch_reductions() call sites in this file (trainer send + vLLM worker hijack) are no-ops on the IPC path here: 1. We route IPC handles by physical GPU UUID dict key (set on the trainer side at _build_ipc_update_info_from_named_tensors via _current_gpu_uuid() from torch.cuda.get_device_properties().uuid). The receiver looks up by its own UUID, independent of args[6]. 2. vLLM's IPCWeightTransferEngine.receive_weights unconditionally overwrites args[6] with the receiver's local device_index before calling rebuild_cuda_tensor. Whatever a torch reductions patch encodes into args[6] is therefore discarded. The patch was the historical mechanism (sglang upstream) for translating device indices across CUDA_VISIBLE_DEVICES boundaries by stuffing UUID strings into args[6]. Our UUID-keyed dict + vLLM's explicit device_index override accomplish the same thing without the global torch reductions mutation. Also expand the _build_ipc_update_info_from_named_tensors docstring to spell out the UUID-keyed routing contract so future readers don't have to chase this through git history. Side effect: hf_weight_iterator_direct.py also calls monkey_patch_torch_reductions() at module-collective time. That call site is similarly decorative (only NCCL broadcast / all_gather collectives run there, no cross-process pickling) but lives outside this file's scope and is not touched here. Tracked alongside #29. * refactor(weight-sync): finish removing monkey_patch_torch_reductions dead code Follow-up to 39bf899 (deleted _apply_monkey_patch_torch_reductions from update_weight_from_tensor.py). With that helper gone, two more references are now dead in the vime IPC weight-transfer path: 1. hf_weight_iterator_direct.py:48 called monkey_patch_torch_reductions() at the top of _get_megatron_full_params(). On vime this never has effect on the IPC handle path: _get_megatron_full_params only runs NCCL broadcast/all_gather collectives (no cross-process pickling), and the chunks it returns are subsequently sent via PR #48's UUID-keyed {gpu_uuid: reduce_tensor(weight)} dict that vLLM's receiver routes by physical UUID + explicit args[6] overwrite. The call survives in slime/ miles upstream because their downstream path pickles tensors through sglang's MultiprocessingSerializer.serialize (ForkingPickler → reduce_tensor), where the patched encoding/decoding does real work; PR #48 does not use that pipeline, so the call here was incidentally inherited rather than functionally required. 2. slime/backends/megatron_utils/sglang.py's monkey_patch_torch_reductions re-export + __all__ entry now have no remaining importers in vime. Remove them. Also (this commit, B): 3. vllm_engine.py's update_weights_from_tensor docstring referred to "closures injected by _apply_monkey_patch_torch_reductions" as the reason for cloudpickle. That helper is gone; the cloudpickle is still correct because reduce_tensor returns a (rebuild_fn, args) tuple where the rebuild_fn is a module-level callable that JSON can't serialise. Update the docstring to reflect the actual reason. Net behaviour: identical — the deletions remove dead code paths. The patch_torch shim itself is still importable for callers outside vime (none currently in this tree). Cross-references: - #29 — issue documenting the no-op stub - 39bf899 — prior commit that deleted the helper from update_weight_from_tensor.py * docs(weight-sync): drop sglang refs + add IPC call-stack notes - update_weight_from_tensor.py module docstring: rewrite from a pure vLLM perspective. Step (2) now spells out the merged {uuid_G0: handle, uuid_G1: handle, ...} dict + collective_rpc fan-out inside the vLLM server, and step (3) makes version-with-data atomicity explicit. Drop the "match slime's sglang_engine signature" framing. - vllm_engine.update_weights_from_tensor: collapse the long sglang-vs-vLLM compare block into a focused docstring describing what the POST does and why ipc_handles needs cloudpickle. Add a Chinese call-stack walkthrough (trainer → ★this method★ → server collective_rpc → per-TP worker receive_weights) and note that `node_rank != 0` is a dead branch since VLLMEngine pins node_rank=0 (see PR #48 review comment for the follow-up cleanup). - Hoist base64/cloudpickle imports to module scope so the hot path no longer pays the per-call import overhead. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> --------- Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: SamitHuang <285365963@qq.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: SamitHuang <285365963@qq.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
After #22 added the colocated CUDA IPC weight sync path, the IPC branch in
UpdateWeightFromTensor.update_weightsbypassesVLLMEngine.update_weights_from_tensorentirely — data moves throughIPCWeightTransferEngine.trainer_send_weightsdirectly, not the/update_weightsRPC. That RPC is the normal place where the engine'sself._weight_versionis recorded; the IPC path leaves it at the initialNone.Failure mode
The first time IPC weight sync runs with
--ci-test,actor.update_weights(slime/backends/megatron_utils/actor.py) picks a random rollout engine, callsget_weight_version(), and compares against the trainer-sideweight_updater.weight_version. Because the engine's_weight_versionis stillNone,get_weight_versionfalls back toGET /v1/modelsand returns the model path string (e.g./root/models/Qwen2.5-0.5B-Instruct); the updater returns the integer version string ("1","2", …). They never match:Reproducer: any colocated test with
--ci-testthat reaches the firstupdate_weights, e.g.tests/test_qwen2.5_0.5B_short.pyortests/test_qwen2.5_0.5B_ppo_critic_only_short.py.Fix
Thread the version through the IPC
finish_weight_updatehook:VLLMEngine.finish_weight_update(weight_version: str | None = None)now setsself._weight_version = str(weight_version)before issuing thePOST /finish_weight_update.UpdateWeightFromTensor.update_weights(step 5) passesstr(self.weight_version)when the coordinator rank fires the per-enginefinish_weight_updateRPC.Each colocated engine's coordinator rank fires its own
finish_weight_update, so every engine ends up with the new version recorded. The distributed/NCCL path is unchanged — it already writes_weight_versionviaupdate_weights_from_tensor.Test plan
RuntimeError: Weight version mismatch!ontests/test_qwen2.5_0.5B_short.py(colocate + IPC + ci_test) before the patchtests/test_qwen2.5_0.5B_ppo_critic_only_short.py(which adds PPO+critic on top of colocate+IPC+ci_test) also runs end-to-end after this patch (paired with the fix(args): drop incorrect critic GPU add to rollout_num_gpus in colocate #41 colocate-critic fix)