[Bugfix] [weight-sync] align colocated IPC coordinator/gather with vLLM engine slot

[Meihan-chen]

Summary

Related #56 (comment)

When the Megatron training topology does not match the colocated vLLM engine topology (e.g. TP=4, CP=2, DP=1, EP=8 trainer + one 8-GPU vLLM engine), the colocated CUDA-IPC weight-sync path made two assumptions that only hold when Megatron TP spans the entire vLLM engine slot. Both are now fixed.

• Coordinator selection (connect_rollout_engines): the rank that issues /start_weight_update and /finish_weight_update was chosen by Megatron TP rank 0, which is not unique inside a vLLM engine slot when CP/DP is present (e.g. ranks [0..7] with TP=4, CP=2 have TP ranks [0,1,2,3,0,1,2,3]).
  This caused :

  requests.exceptions.HTTPError: 500 Server Error for url: http://.../start_weight_update
  RuntimeError: start_weight_update called while a weight update is already active.
  Call finish_weight_update first.
  

  The coordinator is now the lowest global rank inside each engine GPU slot (rank == start), which is unique by construction.

• IPC handle gather (_send_hf_chunk_via_ipc): the all-gather of per-rank CUDA IPC payloads ran over Megatron's TP group, so the merged update_info only carried GPU UUIDs for half of the vLLM TP=8 workers, producing

  ValueError: IPC handle not found for GPU UUID ...
  

  A dedicated Gloo process group is now created per engine slot covering the full [start, end) global rank range, and the gather + slot barrier run over that group. The merged payload now contains handles for every physical GPU UUID the colocated vLLM engine owns.

Single-GPU-slot configurations are unaffected (the existing slot_size <= 1 fast path still uses IPCWeightTransferEngine.trainer_send_weights directly).

Test plan

• Unit tests under
  tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py
• End-to-end smoke via the existing tests/test_qwen3_30B_A3B.py with USE_DEEPEP=0

[gemini-code-assist]

Code Review

This pull request updates the weight update logic from tensors to support contexts where CP/DP dimensions are present inside the same vLLM engine slot. It replaces Megatron's TP group gathering with gathering over the complete vLLM engine GPU slot, creating a new Gloo process group for this purpose. The review feedback suggests optimizing this process by conditionally creating the Gloo process group only when there are multiple ranks in the slot, avoiding unnecessary overhead for single-GPU engine slots.

[gemini-code-assist]

For single-GPU engine slots (len(slot_ranks) <= 1), creating a new Gloo process group is unnecessary since the IPC transfer path is bypassed (see the slot_size <= 1 fast path in _send_hf_chunk_via_ipc). Avoiding dist.new_group in this case prevents unnecessary resource allocation and overhead.

We can conditionally create the group only when len(slot_ranks) > 1.

Suggested change

	slot_ranks = list(range(start, end))
	slot_group = dist.new_group(ranks=slot_ranks, backend="gloo")
	slot_ranks = list(range(start, end))
	slot_group = dist.new_group(ranks=slot_ranks, backend="gloo") if len(slot_ranks) > 1 else None

[knlnguyen1802]

@aoshen02 I think this is already handle at #48 right ?

[aoshen02]

@aoshen02 I think this is already handle at #48 right ?

Yah I think so, thus I will close it.