From 55065c6cd9013618b8c135f08938ab3d30dd4db6 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 25 May 2026 21:00:05 -0700 Subject: [PATCH] [None][fix] py_executor: lift conditional tp_allgather sites out of per-rank-divergent gates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under ADP + disagg on feat/deepseek_v4, GEN workers desync after a small number of successful requests and all 8 ranks crash inside ADPRouter.gather_all_rank_states with either: TypeError: RankState.__init__() takes from 2 to 4 positional arguments but 27 were given (with TLLM_METRICS_ALL_RANKS=1) TypeError: argument of type 'bool' is not iterable (without TLLM_METRICS_ALL_RANKS=1) Root cause is in py_executor.py, not adp_router.py: two conditional tp_allgather collectives live inside per-rank-divergent gates and misalign the MPI collective queue. 1. _handle_responses:4438-4440 (bool payload) — called from _process_previous_batch which is gated by `should_process_previous_batch = can_queue or not can_queue_this_rank` plus `previous_batch is not None`. When one DP rank has an empty batch but its previous_batch was cleared the iter before, that rank skips _process_previous_batch while peers enter it. Peers issue the tp_allgather, the empty-batch rank does not, and the next collective (the very next gather_all_rank_states) receives the bool payload positionally as RankState fields. 2. _append_iter_stats:1436 (dict payload, ~27 keys) — same outer gate via _process_iter_stats. With TLLM_METRICS_ALL_RANKS=1 the same divergent gate fires this gather AND the timeout gather, so the misalignment lands the 27-key dict in the RankState slot. Fix: mirror the existing _flush_pending_transfer_responses pattern. Buffer the timed-out requests and the per-rank iter-stats dict on self, drain at a rank-symmetric position next to the existing flush. Adds two new helpers: - _handle_kv_transfer_timeouts_synced — replaces the inline ADP branch in _handle_responses - _flush_iter_stats_synced — replaces the inline gather in _append_iter_stats Wired in all three executor loops: - _executor_loop_overlap: after the per-rank-divergent `if previous_batch is not None and should_process_previous_batch` block, alongside the existing _flush_pending_transfer_responses. - _executor_loop (non-overlap): right after the `if can_queue` block that calls _handle_responses (already rank-symmetric, but the buffer-drain symmetry is preserved here too). - _executor_loop_pp / _handle_executed_batch: at the bottom of every call (each rank calls it an equal number of times per outer iter via the rank-0-broadcast executed_batch_num), outside the inner `executed_batch is not None` conditional. Bug report: ape-repo/astra-projects/wwfo/artificial-analysis@DSV4-disagg: docs/knowledge/dsv4_disagg_dep8_adp_router_desync.md Repro: DEP8+DEP8 1p1d DSv4-Pro disagg smoke at concurrency 72; failure fires every iter ~900 once at least one GEN rank's batch_size hits 0. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 100 ++++++++++++++---- 1 file changed, 78 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 42acbda93146..363502c6b93b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -424,6 +424,12 @@ def __init__( # responses and flushing them at a synchronised point in the executor # loop avoids the mismatch. self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = [] + # Same buffer-then-synced-flush pattern as _pending_transfer_responses + # above: _handle_responses and _append_iter_stats are reached from + # per-rank-divergent gates, so their tp_allgather collectives are + # lifted to _handle_kv_transfer_timeouts_synced / _flush_iter_stats_synced. + self._pending_timed_out_requests: List[LlmRequest] = [] + self._pending_iter_stats_dict: Optional[Dict] = None # ADP dummy role for _pad_attention_dp_dummy_request. Default is gen; # updated from observed request types. self._adp_dummy_is_gen: bool = True @@ -707,6 +713,52 @@ def _flush_pending_transfer_responses(self): # can complete. self._enqueue_responses(responses) + def _handle_kv_transfer_timeouts_synced(self): + """ADP-safe drain of the KV-transfer-timeout consensus collective. + + Lifted out of _handle_responses, which is reached from per-rank- + divergent gates (_process_previous_batch, _handle_executed_batch). + Non-ADP runs handle timeouts inline; the buffer is empty here. + """ + if not (self.enable_attention_dp and self.dist.world_size != 1): + return + timed_out = self._pending_timed_out_requests + self._pending_timed_out_requests = [] + any_timed_out = any(self.dist.tp_allgather(bool(timed_out))) + if any_timed_out: + self._handle_errors(error_msg="Request timed out (KV transfer)", + requests=timed_out) + + def _flush_iter_stats_synced(self): + """ADP-safe drain of the TLLM_METRICS_ALL_RANKS dict-gather collective. + + Lifted out of _append_iter_stats (same divergent-gate problem as + _handle_kv_transfer_timeouts_synced). Under ADP every rank sends + either the local dict or ``None`` every iter so the gather stays in + lockstep with peer collectives. When the gate doesn't hold, + _append_iter_stats takes its legacy path and never populates the + buffer. + """ + tp_size = self.dist.tp_size + gather_all_ranks = os.environ.get("TLLM_METRICS_ALL_RANKS", "0") == "1" + if not (gather_all_ranks and self.enable_iter_perf_stats and tp_size > 1 + and self.enable_attention_dp): + return + local_dict = self._pending_iter_stats_dict + self._pending_iter_stats_dict = None + gathered = self.dist.tp_allgather(local_dict) + if self.dist.tp_rank != 0: + return + rank_dicts = [d for d in gathered if d is not None] + if not rank_dicts: + return + with self.stats_lock: + for d in rank_dicts: + self.stats.append(("per_rank_dict", d)) + cap = self.max_stats_len * tp_size + if len(self.stats) > cap: + del self.stats[:len(self.stats) - cap] + # Performance metrics methods are in PerfMetricsManager (self.perf_manager) def _event_loop_wrapper(self): @@ -1432,23 +1484,9 @@ def _append_iter_stats(self, append_kv_cache_iteration_stats(local_dict, self._latest_kv_iter_stats) local_dict["rank"] = self.dist.tp_rank - - gathered = self.dist.tp_allgather(local_dict) - - if self.dist.tp_rank == 0: - with self.stats_lock: - # Wrap as ("per_rank_dict", dict) so the serializer can - # distinguish from the legacy (stats, req_stats, kv) tuple. - # Trim before appending so every rank's entry for the - # same iteration lands in the buffer atomically; evicting - # during the append loop would let partial iterations - # survive at the head of self.stats. - cap = self.max_stats_len * tp_size - overflow = max(0, len(self.stats) + len(gathered) - cap) - if overflow: - del self.stats[:overflow] - for d in gathered: - self.stats.append(("per_rank_dict", d)) + # Buffer for _flush_iter_stats_synced; in-place tp_allgather would + # desync (reached from per-rank-divergent gates). + self._pending_iter_stats_dict = local_dict return # Legacy path: rank-0-only (single-rank or iter stats disabled). @@ -2012,6 +2050,13 @@ def _handle_executed_batch(self, executed_batch: Optional[BatchStatePP]): executed_batch.microbatch_id % self.dist.pp_size, ) + # Drain the synced-collective buffers outside the per-rank-divergent + # ``executed_batch is not None`` branch. handle_executed_batches + # invokes this helper the same number of times on every rank + # (broadcast count), so the flushes are rank-symmetric. + self._handle_kv_transfer_timeouts_synced() + self._flush_iter_stats_synced() + @nvtx_range("wait_on_pp_send_handles") def wait_on_pp_send_handles(self, send_handles, microbatch_id): if send_handles[microbatch_id] is not None: @@ -2567,6 +2612,10 @@ def _executor_loop(self): if self.enable_kv_cache_events: self._add_kv_cache_events() + # Drain timeout buffer outside ``if can_queue`` so the synced + # collective fires every iter regardless of future restructuring. + self._handle_kv_transfer_timeouts_synced() + if self.kv_cache_transceiver and self.async_transfer_manager.has_any_inflight_requests( ): self._check_kv_transfer_timeout() @@ -2582,6 +2631,9 @@ def _executor_loop(self): sample_state=sample_state, iter_stats=iter_stats, iter_start_time=iter_start_time)) + # Same lockstep guarantee for iter-stats; no-op when + # TLLM_METRICS_ALL_RANKS=0. + self._flush_iter_stats_synced() self.iter_counter += 1 @@ -2890,6 +2942,12 @@ def _executor_loop_overlap(self): else: self._enqueue_responses([]) + # Drain buffers from the (per-rank-divergent) + # _process_previous_batch above; rank-symmetric companion to + # the _enqueue_responses([]) call in the else branch. + self._handle_kv_transfer_timeouts_synced() + self._flush_iter_stats_synced() + # Call set_exclude_last_generation_logits after _process_previous_batch. # If set before, the response of a request may be incorrect, as it will # use the wrong indices for generation logits when streaming is enabled. @@ -4436,11 +4494,9 @@ def _handle_responses(self): for request in requests_to_terminate: self._terminate_request(request) if self.enable_attention_dp and self.dist.world_size != 1: - any_timed_out = any(self.dist.tp_allgather( - bool(timed_out_requests))) - if any_timed_out: - self._handle_errors(error_msg="Request timed out (KV transfer)", - requests=timed_out_requests) + # Buffer for _handle_kv_transfer_timeouts_synced; in-place + # tp_allgather would desync (reached from per-rank-divergent gates). + self._pending_timed_out_requests.extend(timed_out_requests) else: for req in timed_out_requests: self._handle_errors(