Skip to content

[https://nvbugs/5448767][fix] fix mpi4py deadlocks in pp event-loop #6976

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 27, 2025
Merged

[https://nvbugs/5448767][fix] fix mpi4py deadlocks in pp event-loop #6976

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6eb0d0b
finish all send requests before quitting pp event-loop to avoid mpi d…
reasonsolo Aug 18, 2025
883962e
finish all send requests before quitting pp event-loop to avoid mpi d…
reasonsolo Aug 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions tensorrt_llm/_torch/pyexecutor/py_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -627,6 +627,12 @@ def _process_iter_stats(self, finished_requests: list[LlmRequest],
batch_state.sample_state.scheduled_requests), req_stats)

def _executor_loop_cleanup(self):
# Unblock receiving processes. When second-last rank quits before last rank,
# last rank will never return from recv_object.
for req in self.send_handles:
if req is not None:
req.wait()

with self.response_cv:
self.is_shutdown = True
self.response_cv.notify_all()
Expand Down Expand Up @@ -750,6 +756,7 @@ def _executor_loop_pp(self):

sample_state = self._sample_async(
scheduled_batch, batch_outputs)
assert sample_state is not None, "Sampling failed"
sample_state.host.logits = logits_host
self._update_request_states(scheduled_batch)

Expand Down Expand Up @@ -801,6 +808,7 @@ def _executor_loop_pp(self):
if not self.dist.is_second_last_pp_rank:
if self.send_handles[prev_microbatch_id] is not None:
self.send_handles[prev_microbatch_id].wait()
self.send_handles[prev_microbatch_id] = None
needs_logits = (
self._need_return_logits(scheduled_batch)
or (self._need_return_log_probs(scheduled_batch)
Expand Down