[BugFix]:Fix vLLM IPC weight transfer for MiMo MTP training - #55
Conversation
Co-authored-by: Cursor <cursoragent@cursor.com>
![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 modifies the CUDA IPC weight update mechanism to ensure that producer tensors remain alive until all tensor parallel workers have opened the handles and the coordinator's update has completed. Specifically, _build_ipc_update_info_from_named_tensors now returns a list of contiguous tensor references along with the payload, which are explicitly deleted only after the synchronization barrier in _send_hf_chunk_via_ipc. Unit tests have been updated accordingly to mock the new return signature. There are no review comments to address, and I have no additional feedback to provide.
princepride
changed the title
princepride
changed the title
|
Independently arrived at the same One concern about the merge order, documented in #58: The official Dockerfile pins
Since vllm 0.22 hasn't released yet, there's no clean release tag to bump to. The options (full discussion in #58):
Option 1 in this PR is probably the cleanest if reviewers are OK with the broader scope; otherwise option 2 needs to land first. Either way the constraint is the same: this PR shouldn't merge ahead of the image change, or CI on main breaks the moment it does. |
Resolves conflicts with PR #55 (8ffcabf, "Fix vLLM IPC weight transfer for MiMo MTP training") which merged into main while PR #48 was open. PR #55 overlaps PR #48's IPC sender in update_weight_from_tensor.py: 1. _build_ipc_update_info_from_named_tensors return type - PR #48 (39bf899 docstring): returned dict[str, list], explained UUID-keyed routing as the contract that replaces the need for a torch reductions monkey-patch. - PR #55: changed return to tuple[dict[str, list], list[torch.Tensor]], adding weight_refs liveness guard (producer storage must stay alive until receiver opens the IPC handle). - Resolution: keep PR #55's tuple return + weight_refs; merge both docstring paragraphs (the UUID-routing rationale and the liveness rationale are orthogonal and both correct). 2. _send_hf_chunk_via_ipc barrier + cleanup - PR #48 (0d435a0): barrier group changed from tp_group to slot_group as part of the slot-leader gating fix (Bug 1/2 series). - PR #55: added ``del weight_refs`` after the barrier to release sender storage. - Resolution: keep PR #48's slot_group + PR #55's del weight_refs. Both improvements are orthogonal. Also propagate PR #55's tuple return into the slot_size <= 1 fast path which PR #48 introduced after PR #55 was written; that branch was still spreading the bare dict (``**local_info``) which now needs to unpack ``local_info, weight_refs = _build_...`` and ``del weight_refs`` after the ray.get() RPC. And drop the ``with patch(...) ._apply_monkey_patch_torch_reductions`` line in the unit test ``_run_update`` helper — the helper was deleted by my earlier commit 39bf899 on this branch, so the patch context would raise ``AttributeError`` after the merge. Pre-existing failure on this branch: ``test_connect_marks_one_coordinator_per_engine_gpu_slot`` fails on PR #48 HEAD (0d435a0, fd12344, and now the merge result) but passes on main; the failure ("Default process group has not been initialized") is independent of this merge and was already broken before main was pulled in. Not addressed here; should be fixed in a separate commit on PR #48. Verified: ``pytest tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py --deselect ::test_connect_marks_one_coordinator_per_engine_gpu_slot`` → 6 passed.
2 participants
Propose
Related: #11
The failure happened because Vime passed the full reduce_tensor() return pair into vLLM IPC weight transfer, while vLLM expected only the CUDA IPC rebuild args, causing list_args[6] to index into the wrong object.
The fix extracts and sends only ipc_args, while keeping the producer tensors alive until vLLM workers finish opening the CUDA IPC handles.
Test Plan