Repository navigation
Conversation
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.
Author
This branch has not been deployed
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
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.
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:
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:
Empty buckets are the common case, not an edge case:
supports_multi_dtypesisTrue,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_chunksleaves mostchunks 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
uint8tensors 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, onlytorch.distributedplussglang's IPC serializer:
Validation
A 744B-A40B MoE model, colocated, 8 nodes × 8 H200, 4 engines of TP16/EP16 each spanning 2
nodes, training TP1/PP16/EP4:
update_weightsalways fails withInvalid device_uuid, for everytraining parallelism we tried.
sleep 35.2 s → update_weights 51.6 s → wake_up 7.4 s, zeroInvalid device_uuid.