Repository navigation
[Transport][Bug fix] Key received slices by offset; refuse sharded streams across TP sizes - #623
Conversation
Automated reviewThe PR adds Findings: none. I traced these paths and found no problems:
The description is accurate and proportionate. It lists the commands run and their results, including which tests fail on Advisory, not a merge gate. Rules: |
Live test: [Transport][Bug fix] Key received slices by offset; refuse sharded streams across TP sizesNo regressions, but the served check didn't reach the fix: Qwen3-Omni with the thinker at TP2 serves cleanly, and a probe on the receive cache saw only misses, so no second slice of a tensor arrived on one worker. The PR's tests fail on main and pass here. 3x H200, Claims
Other
|
Thanks for the run and the receive-cache probe. 0 hits matches the description: no shipped config sets |
merceod
left a comment
merceod
left a comment
There was a problem hiding this comment.
Checked on klone (5c780b9, 2 x H200). The startup refusal works: Orpheus with a TP2 LLM feeding a TP1 snac_decoder through a sharded new_token stops at startup with a clear message naming the signal, the walk and both TP sizes. main does not start that config either (worker_1 SIGABRT). TP2 Qwen3.5-0.8B with the vision encoder on rank 0, greedy at c=1: identical to main, 50/50 text and 12/12 image. CPU: test_shm_tensor_comm.py 47 and test_startup_graph_validation.py 93 pass on the head where main with these tests fails 2 and 5, and the arena two_slices cases pass here and fail on main. Not shown served: two different slices of one tensor on one worker (no shipped config sets shard_dim).
One thing seen on the way, against main: the shipped orpheus_tp2.yaml does not start on this stack on either tree (main with a worker_1 SIGABRT, this branch with an NCCL all_gather "CUDA error: invalid argument"), so Orpheus TP2 is broken on main regardless of this PR. Transport receive path, so leaving the merge to a second pair of eyes.
|
Thanks for the runs.
Filed as #706. The same config served 8 of 8 on a branch based on main as of #544 in the #658 review and in that PR's live test, so something that landed after #544 (or the host stack) is the likely cause; the issue lists the candidate merges and a bisect plan. No code change here since the approval. |
|
@Gaurav-Shah05 could you take a look when you get a chance, and merge if it looks good? @merceod approved it and CI's green. it keys received slices by offset and refuses sharded streams across different TP sizes. |
…ss TP sizes Every slice the fanout cuts from one producer tensor carries the producer's uuid, and the receive cache in the file transport, the SHM arena and Mooncake was keyed by the uuid alone, so a second slice on the same worker came back as the first. A slice other than the one stored under the producer uuid now reads into a uuid of the consumer's own, and its ack still names the producer's uuid. A sharded streamed signal across TP sizes cannot be served: a gathered chunk arrives under a uuid the StreamBuffer never registered, the shard buffer is per signal rather than per chunk, and local streaming ignores the slice. The conductor now refuses that config at startup. Fixes mstar-project#549
The startup check looked only at node output edges; a Loop streams its outputs and accumulated outputs on completion, so a TP2 producer inside a Loop feeding a TP1 consumer started and then never released a chunk.
A fanout slice of a tensor whose rows hold zero bytes (a [4, 0] tensor) has offset 0 and nbytes 0 whatever its rows, so with a TP2 node and a TP4 group on one worker the TP4 rank's [1, 0] slice matched the stored [2, 0] one. The receive cache now compares dims as well, and keys the re-pointed slices by them.
A slice re-pointed to a local uuid can sit in one edge beside a tensor reused from an earlier edge (an empty tensor's slices are the same for every rank). The whole edge then passed through as a re-delivery, so the new slice kept its transit ref and never got its shard layout back. Now only an edge whose every tensor was already emitted passes through, and the transit ref is dropped for each tensor read for the first time.
Two nodes on one worker can take the same sharded signal, one gathering every source rank's shard and one taking a slice of a single rank's. The gather buffer was keyed by request and signal alone, so the slice arrived as one of the gather's shards and consolidation failed on the unequal counts. It is now keyed by the consumer node too.
5c780b9 to
2f2db25
Compare
|
@merceod FYI since your approval: rebased onto main, slices now also match on dims (a |
merceod
left a comment
merceod
left a comment
There was a problem hiding this comment.
Approve again on 2f2db25. Three receive-path fixes since my last look, all yours. A cached slice is matched and keyed on dims too (a [4,0] tensor's slices all sit at offset 0 with 0 bytes), an edge passes through as a re-delivery only when all of its tensors were seen and the transit ref is dropped per tensor on its first read, and the gather buffer is keyed by signal and consumer node. The eight new cases in test_shm_tensor_comm.py fail on the head I approved and pass here, 158 over the file. Still held for a second approval (transport).
One note, pre-existing on main and unreachable with the shipped configs (none sets shard_dim). The comment at tensors.py lines 556 to 558 says two nodes on one worker can each gather a signal, but a node whose edge reuses an already-read uuid takes the all-seen pass-through at line 1470 and never reaches its gather buffer. Two TP1 gatherers leave the second with un-gathered shards, and a TP2 whole-tensor consumer arriving before a TP1 gatherer leaves the gatherer waiting. Say if you want the repro scripts. One for a follow-up.
|
@Gaurav-Shah05 merceod re-approved on the current head and it's clean on main. would you mind merging when you get a chance? |
The Mooncake reader now reads with batch_transfer_async_read and polls transfer_check_status, so the fake engine mstar-project#623's slice tests use gets both: each batch is copied on submit and reports completed.
…ed, at TP>1 too (#419) * Fail only the request whose async (Mooncake) tensor read failed, and free its buffers once its transfers stop Fixes #418. get_ready_tensors re-raised a failed async read on every poll and never reached the line that removes it, so every request behind it on that worker or data worker stalled until the failed request was torn down. A failed or cancelled read is now dropped and reported once through take_failed_reads. The worker fails its request through FAIL_REQUESTS; the data worker answers it with one error chunk in the output slot of each tensor the read would have delivered. Tensor reads go through TrackedMooncakeReader: one batch per tensor, read with batch_transfer_async_read and polled with transfer_check_status, with the sync read's retries and deadline. A failed read hands over the batch ids still running; the manager keeps the read's buffers (out of the store, so teardown cannot free them) and frees them once every id has reported COMPLETED or FAILED. A failed submit, a status other than those four, an error that names no transfers, or a failed unregister keeps the buffers for good. KV transfers keep the batched sync read. * Hold READS_DONE while a failed read's transfers run, and read a torn-down read's outcome from its future * Poll a read's future through done() and result() only, as the SHM arena's event future offers no more * test: the output-failure teardown fake tensor manager reports no failed reads * Transfer metrics: a failed read is dropped and reported, not raised, so the read delivered before it counts once * test: the serving-metrics stub tensor manager reports no failed reads * test: the fake Mooncake engine's report waits for the reader to submit the batch it targets On a loaded host the reader thread can submit after the test reports a status, which dropped the report and left the read waiting out its deadline. * Count a read's retries under the device filter the Mooncake engine was initialized with get_local_topology builds its own engine, so without the TCP device it counted HCAs the engine filtered out, and retried a failed tensor more often than the sync read does. * Log a failed output read in one line, as a failed input read is The tensor manager already logs the read's error; the data worker logged a full traceback per tensor slot on top. * Count the HCAs for read retries on the first failed tensor, not at startup get_local_topology starts a second Mooncake engine, about 1 s per manager, and the count is only used to retry a failed tensor. * test: over real Mooncake TCP, a failed read keeps its buffers until a transfer still running reports it stopped * test: acceptance tests for a TP follower whose input read fails after rank 0 read its copy Two real ranks in their own processes over gloo, running the real Worker.run, scheduler, SHM tensor manager, Engine/StepRunner/KVManager and Qwen3.5's prepare_inputs. Checks behavior only: no hang and only the failed request fails, nothing of the failed row routed or sent, its pages never leasable with prefix_cache_tp, and the request beside it bit-identical to a run without the failure; for an image-grid row and a text-only row. * TP: take a row whose follower read failed out of the step on every rank, before preparing it Rank 0 schedules a lockstep step from its own copy of each input, so a follower whose read of one fails after rank 0 sent the step could never build the row, and the instance hung in its next collective. Rank 0 now marks a step whose rows take an input the followers read from elsewhere (ScheduleTPNode.read_check). A follower builds such a step as soon as each row's read landed or failed, with the rows it could not read in it and no inputs, so every rank has rank 0's rows. On the GPU thread, before prepare_inputs, one gloo all-reduce of a flag per row tells every rank which rows to take out; they go to failed_requests and fail through the conductor, as a per-row failure does. They are never prepared, admitted, committed or sent, so their outputs and pages never exist, and a prepare that runs a collective per row (Qwen3.5's vocab-parallel embedding) stays in step. A read-checked step is not pre-planned, so every rank admits it after the check. * TP read check: record every failed read for the check, and share the row vote - A follower's failed read of a request it already failed for another reason now goes to read_failed_rids too (the failed_rids filter stays on the report only). Before, a read-checked head naming it waited on the follower with no warning while rank 0 waited in the check's all-reduce. - The per-row vote moves to distributed/all_reduce.py as flagged_on_any_rank, next to _all_ranks_agree. - The async follower puts its unread rows back through with_unread_rows, as the serial path does; a comment and a test pin why pushing back an unread fresh row on a voided speculation stays harmless. - test_a_step_whose_every_row_failed_runs_empty_on_every_rank bounds its run time against a run with no failure instead of a fixed 8 s, which slow hosts went past on startup alone. * test: the in-process Mooncake engine of the slice tests reads async too The Mooncake reader now reads with batch_transfer_async_read and polls transfer_check_status, so the fake engine #623's slice tests use gets both: each batch is copied on submit and reports completed.
Fixes #549.
What does this PR do?
Two slices of one producer tensor received on one worker aliased (file transport, SHM arena, Mooncake): the receive cache was keyed by uuid alone, so the second slice returned the first. And a streamed signal with a
shard_dimacross TP sizes never released a chunk. Now the slices no longer alias, and such a stream is refused at startup, also when a Loop's outputs carry it.tensors.pyandarena.py, a slice other than the one stored under the producer uuid reads into a local uuid; acks still name the producer uuid. A slice is matched by offset, nbytes and dims, since a[4, 0]tensor's slices all have offset 0 and nbytes 0.distributed/base.pyandconductor.pyrefuse a sharded streamed signal across TP sizes (per emitting walk): gathered chunks arrive under a uuid the StreamBuffer never registered, and local streaming ignores slices.How was it tested?
On the head (2f2db25, on main 294cbab), CPU,
mstar_rustbuilt from the tree.ruff check .: clean.--ignore=test/modular/test_tp_async_lockstep_sim.py): 6345 passed, 1138 skipped, 4 failed; the 4 needftfy/prometheus_client, missing from this host's env, and fail on main too.test/modular/test_tp_async_lockstep_sim.py): 7 passed.test/rust/ test/modular/test_row_output_store.py: 380 passed, 89 skipped.-k two_slicesintest_shm_tensor_comm.pyandtest_arena_transport.py(file, Mooncake on a memcpy engine, arena w/ Python and Rust bookkeeping;[4, 3]and[4, 0]): 8 passed, the 4[4, 0]ones also fail w/o the dims match.-k mixing(reused empty tensor next to a new slice in one edge): 6 passed.-k another_nodes_gather(TP4 slice arriving mid-gather of a TP1 node on the same worker): 2 passed.test_startup_graph_validation.py -k "sharded_stream or sharded_signal": 7 passed.python -m benchmark.eval compare --base 294cbab66 --head 2f2db2575 --models qwen3tts_1p7b_split_shm,qwen3tts_1p7b_split_shm_arena(worker-to-worker SHM over files and over the arena, the receive paths this changes; no sharding in these configs): no/dev/shmleftovers or crash lines after any case on either arm, WER 2.90% and UTMOS 4.30 on both, c=1 latency within 0.2%, c=8 throughput -1.4% and -0.6% (no A/A floor for these rows yet; within the passes' spread).CPU suite on main, same host
Main (294cbab): 6326 passed, 1138 skipped, 6 failed: the same 4, plus 2 signal-timing tests (
test_eval_compare.py::test_ctrl_c_on_compare_stops_the_remote_job_and_prints_one_line,test_eval_framework.py::test_sigterm_to_the_runner_stops_the_managed_server) that passed on the head run; the host ran other jobs. Rust and row store: 374 passed, 89 skipped (the 6 fewer are the new arena cases).Not tested
Multi-GPU served mixed-TP run (needs an 8-GPU VM, and no shipped config sets
shard_dimacross TP sizes; the CPU tests drive the real fanout and the real receive paths for that topology). Mooncake over a real TCP or RDMA link (the wheel needs libcuda).Not here: serving sharded streams across TP sizes (needs per-chunk fan-in and slice offsets in both runtimes' local-streaming output).
Depends on and merge order
Depends on: none.
Merge order: conflicts only with #419 (tail of
cleanup_request, keep both sides); trial merges with #637 and #772 are clean.Design
Slices.
fanout_graph_edgescuts a sharded tensor per destination rank, and every cut keeps the producer's uuid with its ownoffset,nbytesanddims._take_cached_readreplaces the uuid-only check in all three transports:offset,nbytesanddims: reuse it (+1 ref), as before.(uuid, offset, nbytes, dims)in the request'sslice_uuids; reuse that local uuid while it is still stored, otherwise mint one, re-point the descriptor to it and read into it (+2, as a first read). The per-uuid file and arena reads still use the producer uuid and descriptor placement.Dims are part of the match because a row can hold zero bytes: for a
[4, 0]tensor a TP2 rank's[2, 0]slice and a TP4 rank's[1, 0]slice both have offset 0 and nbytes 0.A re-pointed slice can sit in one edge next to a tensor reused from an earlier edge (an empty tensor's slices are the same for every rank), so
get_ready_tensorsdrops the transit ref and restores the shard layout per tensor read for the first time; only an edge whose every tensor was already emitted passes through as a re-delivery. Gather buffers are keyed by (signal, consumer node), so a slice for one node never counts as a shard of another node's gather on the same worker._collect_and_send_acksmaps a re-pointed uuid back throughslice_origins, so the producer gets one ack per edge for its own uuid. Both maps are per request and dropped at the end ofcleanup_request, after the acks for still-pending edges. No Rust twin: the receive path is Python only (rust/src/shm.rsis the arena allocator). Cost: oneget_infoand a dims compare on a cache hit only, and a dict lookup per acked tensor only for requests that have re-pointed slices.Streams. Serving a gathered stream needs three changes: the worker pre-registers one uuid per source rank, but
get_ready_tensorsemits the consolidated chunk under a new one;BufferedShardsis per (request, signal), so chunk k+1 from one rank arriving before chunk k from another mixes chunks; andlocal_streaming_by_signalcarries only (rid, uuid) from both runtimes, so slicing a local stream needs new columns on the Python and Rust sides. Between equal TP sizescompute_faninis 1 and each rank's shard goes whole to its peer, so those configs are still accepted. The producer's TP size is taken for the walk that emits the stream, since a producer's group can be walk-specific.Tests in detail
test_two_slices_of_one_tensor_on_one_worker_do_not_alias(file, Mooncake):xfrom a TP 1 producer to a TP 2 node and a TP 4 group of two nodes, all with rank 0 on worker_1, for[4, 3]and[4, 0]. Checks each node gets its own rows (and shape), the repeated slice shares one uuid, the producer gets{P: 3}, and the consumer's store is empty after each node releases its input.test_two_slices_of_one_tensor_on_one_worker_do_not_alias(arena): the same for the arena with the Python and the Rust bookkeeper (get_infooffset round trip), both shapes.test_an_edge_mixing_a_reused_and_a_new_slice_is_read_in_full(file, Mooncake, arena):x = [[4, 3], [0, 3]], shard dim 0 and 1; the TP4 node's edge reuses the empty tensor and reads a new slice. Checks rows and layout, acks, and that nothing keeps a ref.test_a_slice_does_not_join_another_nodes_gather(file, Mooncake): TP2 producer to a TP1 node and a TP4 group on one worker, received as gather rank 0, slice, gather rank 1. Checks both nodes' tensors and an empty store after release.test_sharded_stream_across_tp_sizes_is_rejected: realConductor.__init__on a toy stream model, TP 2 to 1 and 1 to 2 refused, 2 to 2 and an unsharded signal accepted.test_a_loop_streaming_a_sharded_signal_across_tp_sizes_is_rejected: the same for a stream inLoop.outputsandLoop.accumulated_outputs.test_sharded_stream_check_uses_the_walk_that_emits_it: a producer group limited to another walk does not count.Question for maintainers
Is refusing sharded streams across TP sizes the intended contract, or should they be served later?
Checklist
ruff check .passesAI assistance: written with AI coding agent