Skip to content

fix: Fix placeholder IPC handles when a rollout engine spans multiple nodes - #2520

Open
leondada wants to merge 1 commit into
THUDM:mainfrom
leondada:fix/placeholder-ipc-device
Open

leondada wants to merge 1 commit into
THUDM:mainfrom
leondada:fix/placeholder-ipc-device

Conversation

@leondada

@leondada leondada commented Nov 5, 2026

Copy link
Copy Markdown

Fix placeholder IPC handles when a rollout engine spans multiple nodes

Bug

In colocated mode, when a rank has no weights for the current bucket, the IPC gather
source fabricates that rank's empty placeholder:

# update_weight_from_tensor.py, _send_to_colocated_engine
if dist.get_rank() == ipc_gather_src:
    ...
    if empty_serialized_tensor is None:
        empty_tensor_data = _empty_flattened_tensor_data()   # allocates on the CALLER's device
        empty_serialized_tensor = MultiprocessingSerializer.serialize(empty_tensor_data, ...)
    serialized_tensors_for_dtype.append(empty_serialized_tensor)   # -> any rank's slot

A serialized placeholder is a CUDA IPC handle tagged with the producing device, and the
caller is always the gather source. So a slot owned by a rank on another node carries a
handle for a GPU that node cannot see, and the engine worker reading it fails:

tp_worker.py:177  deserialize(recv_req.serialized_named_tensors[self.tp_rank])
patch_torch.py:97 _device_from_maybe_uuid
Exception: Invalid device_uuid=54dcf85c-8127-2620-201a-cad88e8fcccf

Empty buckets are the common case, not an edge case: supports_multi_dtypes is True,
so a rank contributes 1 bucket or 0, and both sync paths produce zeros — expert routing
skips ranks that are neither source nor target, and get_hf_weight_chunks leaves most
chunks empty for most ranks under pipeline parallelism.

Only triggers with colocated training + an engine spanning more than one node
(rollout_num_gpus_per_engine > num_gpus_per_node). When an engine fits on one node,
the gather source shares that node with every rank in its group, so the fabricated
placeholder is importable and nothing breaks. We hit it on a 744B-A40B MoE model: ~1.48 TB in
BF16 does not fit on one 8×H200 node, so each engine is TP16 across two nodes.

Fix

Every rank contributes its own placeholder, gathered alongside its real payload; the
gather source then fills a missing slot with the placeholder from that rank. A slot's
handle therefore always refers to a device visible to the worker reading it.

No numerical change: placeholders are zero-length uint8 tensors carrying no weights.
They only need to be deserializable so the receiver can walk the whole list.

Reproducer

repro_placeholder_device.py — single node, ≥2 GPUs, only torch.distributed plus
sglang's IPC serializer:

$ python3 repro_placeholder_device.py          # rank 1 has an empty bucket
  current: slot1 -> cuda:0  WRONG: handle on cuda:0, consumer is cuda:1
$ python3 repro_placeholder_device.py --fixed
  fixed  : slot1 -> cuda:1  ok

Validation

A 744B-A40B MoE model, colocated, 8 nodes × 8 H200, 4 engines of TP16/EP16 each spanning 2
nodes, training TP1/PP16/EP4:

  • Before: first update_weights always fails with Invalid device_uuid, for every
    training parallelism we tried.
  • After: full colocate cycle completes — sleep 35.2 s → update_weights 51.6 s → wake_up 7.4 s, zero Invalid device_uuid.

In colocated mode, when a rank has no weights for the current bucket, the
IPC gather source fabricated that rank's empty placeholder. A serialized
placeholder is a CUDA IPC handle tagged with the producing device, and the
producer was always the gather source, so a slot owned by a rank on another
node carried a handle for a GPU that node cannot see. The engine worker
reading such a slot then failed with sglang's
"Invalid device_uuid=..." in _device_from_maybe_uuid.

Empty buckets are the common case rather than an edge case: with
supports_multi_dtypes a rank contributes one bucket or none, and both sync
paths produce zeros - expert routing skips ranks that are neither source nor
target, and get_hf_weight_chunks leaves most chunks empty for most ranks
under pipeline parallelism.

This only triggers with colocated training plus an engine that spans more
than one node (rollout_num_gpus_per_engine > num_gpus_per_node); when an
engine fits on one node the fabricated placeholder is importable anyway.

Now every rank contributes its own placeholder alongside its real payload,
and the gather source fills a missing slot with the placeholder from that
rank, so a slot's handle always refers to a device visible to the worker
reading it. Placeholders are zero-length uint8 tensors, so the numerical
path is unchanged.
@leondada

Copy link
Copy Markdown
Author

@vpj

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant