[https://nvbugs/6388153][fix] Modify the broadcast of sample state in pp to default to synchronous mode. - #16170
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughPyExecutor initializes PP sample-state relay before microbatch allocation and defaults to synchronous relay. Synchronous execution drains pending sends and broadcasts inline. The waiver list removes obsolete DeepSeekV3Lite entries. ChangesPP Synchronous Sample-State Relay
DeepSeekV3Lite Test Waiver Updates
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant ExecutedBatchQueue
participant RingBroadcast
alt synchronous relay
PyExecutor->>PyExecutor: drain pending sample-state sends
PyExecutor->>RingBroadcast: broadcast executed batch inline
else asynchronous relay
PyExecutor->>ExecutedBatchQueue: enqueue executed batch
ExecutedBatchQueue->>RingBroadcast: relay sample state in background
end
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
354e901 to
5c7af1d
Compare
|
LGTM. Can merge after @yuxianq take a look. |
|
/bot run --disable-fail-fast |
|
PR_Github #58597 [ run ] triggered by Bot. Commit: |
|
PR_Github #58597 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58670 [ run ] triggered by Bot. Commit: |
|
PR_Github #58670 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58757 [ run ] triggered by Bot. Commit: |
|
PR_Github #58757 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58923 [ run ] triggered by Bot. Commit: |
|
PR_Github #58923 [ run ] completed with state
|
42e4067 to
7255092
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58956 [ run ] triggered by Bot. Commit: |
|
PR_Github #58956 [ run ] completed with state
|
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #63986 [ run ] completed with state
|
|
The deadlock analysis reads right, and relaying inline removes the GIL-starvation window cleanly. Two things I'd want said out loud before this merges, since it's already at five approvals. 1. There's an unrelated new waiver in here. The diff removes the two Different bug, and 2. The title says "relay synchronously"; the change is a default flip. Related: Has anyone measured PP throughput before/after? The relay now serializes with forward instead of overlapping it, plus there's a new pre-forward Not blocking — deferring to the PP/runtime owners on the tradeoff. But I'd split the waiver out. |
|
/bot run |
|
PR_Github #64151 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #64158 [ run ] triggered by Bot. Commit: |
|
PR_Github #64151 [ run ] completed with state |
|
PR_Github #64158 [ run ] completed with state
|
|
/bot run |
|
PR_Github #64272 [ run ] triggered by Bot. Commit: |
|
PR_Github #64272 [ run ] completed with state
|
c068f0a to
9559a9e
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
740-754: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCreate relay-only resources only in async mode.
Line 744 selects synchronous relay by default. However,
start_worker()still creates the relay queue, duplicates the MPI communicator, and startsbroadcast_sample_state_handlerfor every PP executor. This adds an unnecessary MPI collective, daemon thread, and CUDA thread initialization on the default path.Keep
executed_batch_response_queuefor both modes. Guardexecuted_batch_queue, communicator duplication, thread startup, and the matching shutdown sentinel and join withself.pp_async_broadcast_sample_state.Proposed fix
if self.dist.pp_size > 1: - self.executed_batch_queue = Queue(maxsize=self.num_micro_batches) self.executed_batch_response_queue = Queue(maxsize=-1) + if self.pp_async_broadcast_sample_state: + self.executed_batch_queue = Queue(maxsize=self.num_micro_batches) + self._broadcast_mpi_comm = mpi_comm().Dup() + self.broadcast_sample_state_handler = threading.Thread(...) + self.broadcast_sample_state_handler.start() -if self.dist.pp_size > 1: +if self.dist.pp_size > 1 and self.pp_async_broadcast_sample_state: self.executed_batch_queue.put(None) self.broadcast_sample_state_handler.join()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 740 - 754, Update start_worker() so relay-only resources are created only when self.pp_async_broadcast_sample_state is true: conditionally initialize executed_batch_queue, duplicate the MPI communicator, and start broadcast_sample_state_handler, while retaining executed_batch_response_queue for both modes. Apply the same condition to the matching shutdown sentinel and thread join, leaving synchronous relay behavior free of those resources.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 740-754: Update start_worker() so relay-only resources are created
only when self.pp_async_broadcast_sample_state is true: conditionally initialize
executed_batch_queue, duplicate the MPI communicator, and start
broadcast_sample_state_handler, while retaining executed_batch_response_queue
for both modes. Apply the same condition to the matching shutdown sentinel and
thread join, leaving synchronous relay behavior free of those resources.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 373f7743-18cd-4535-846f-b3e0f62cfe53
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/test_lists/waives.txt
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
9559a9e to
62e9a01
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
1 similar comment
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot skip --comment "failed tests are unrelated to this PR, skip CI" |
|
PR_Github #64298 [ skip ] triggered by Bot. Commit: |
Signed-off-by: WeiHaocheng <20514172+WeiHaocheng@users.noreply.github.com>
33861f1 to
3de1064
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Signed-off-by: WeiHaocheng <469657736@qq.com>
|
/bot skip --comment "failed tests are unrelated to this PR, skip CI" |
|
PR_Github #64298 [ skip ] completed with state |
|
PR_Github #64302 [ skip ] triggered by Bot. Commit: |
|
PR_Github #64302 [ skip ] completed with state |
Under pipeline parallelism, sample states are relayed between PP ranks
(last rank -> 0 -> 1 -> ...) by a dedicated background thread. The relay
thread needs the GIL to run: when the executor thread blocks inside a
GIL-holding native call that waits on GPU progress (e.g. a DeepGEMM JIT
kernel cold load whose lazy loading requires context synchronization),
the relay starves, the downstream rank stops advancing, its unlaunched
forward never pairs the in-flight NCCL p2p kernel the blocked rank's GPU
is waiting on, and the ranks deadlock.
Summary
TLLM_PP_ASYNC_BROADCAST_SAMPLE_STATE=1.tests/integration/test_lists/waives.txt.Dev Engineer Review
num_micro_batchesfrommax(pp_size, 1024)topp_size.TestGPTOSS::test_w4_chunked_prefill[trtllm-fp8]waiver for nvbug 6529692 should remain in a separate PR.QA Engineer Review
tests/integration/test_lists/waives.txtchanged.test-db/orqa/files were modified.