[https://nvbugs/6428092][fix] Forward use_host_stop_criteria alongside host and py_result_diffs in the PP… - #16133
[https://nvbugs/6428092][fix] Forward use_host_stop_criteria alongside host and py_result_diffs in the PP…#16133trtllm-agent wants to merge 2 commits into
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThe pipeline-parallel sample_state broadcast now sends and receives the use_host_stop_criteria flag. The change also removes a waiver for a multi-GPU PyTorch integration test. ChangesSample State Broadcast Update
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to This change forwards the stop-criteria state needed for pipeline-parallel execution while preserving compatibility with sampler variants; no actionable merge-blocking risk remains after normal checks and review. Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
ecfbedf to
a6873af
Compare
| # Propagate use_host_stop_criteria: the last rank's fast host-stop | ||
| # path leaves host.finish_reasons=None, so non-last ranks must | ||
| # know to skip the finish_reasons indexing in update_requests. | ||
| if hasattr(sample_state, "use_host_stop_criteria"): |
| self.send_handles[microbatch_id] = self.dist.isend_object( | ||
| (sample_state.host, py_result_diffs), | ||
| (sample_state.host, py_result_diffs, | ||
| getattr(sample_state, "use_host_stop_criteria", False)), |
a6873af to
ab04caa
Compare
ab04caa to
7eed8b4
Compare
brnguyen2
left a comment
There was a problem hiding this comment.
use_host_stop_criteria doesn't exist in this codebase. grep -rn use_host_stop_criteria tensorrt_llm tests returns only the five lines this PR adds — SampleStateTorch / SampleStateTensorsHostTorch (sampler/sampler.py:1111-1130) have no such field. So getattr(..., False) always sends False, hasattr(...) is always False, and the receive-side assignment never runs. The change is a no-op that only widens the wire tuple, and the waiver removal in waives.txt is unsupported by it.
Also, finish_reasons_host is only None when requests is empty (sampler.py:4137), and update_requests returns early in that case — so the stated root cause needs re-checking against the actual failure.
| # Propagate use_host_stop_criteria: the last rank's fast host-stop | ||
| # path leaves host.finish_reasons=None, so non-last ranks must | ||
| # know to skip the finish_reasons indexing in update_requests. | ||
| if hasattr(sample_state, "use_host_stop_criteria"): |
There was a problem hiding this comment.
hasattr(sample_state, "use_host_stop_criteria") is always False: no sampler state class in the tree defines that attribute, and nothing assigns it. This branch never executes, and the matching getattr(..., False) on line 3115 always sends False. Before revising, confirm the field name against the real sampler state and re-verify the failing test actually reproduces and then passes.
test_tinyllama_logits_processor_2gpu[1-2] (pp_size=2) crashed with IndexError in finish_if_reason because the last PP rank produced a SampleStateTorch with use_host_stop_criteria=True (and host.finish_reasons=None), while only sample_state.host was sent through _ring_broadcast_sample_state. Non-last ranks kept the constructor default False and tried to index into an empty finish_reasons list from process_draft_tokens. Include the flag in the PP send/recv payload and restore it on the receiver. Use getattr/hasattr to stay compatible with sampler flavors that lack this field. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Signed-off-by: handongl <handongl@nvidia.com> Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
7eed8b4 to
526914b
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. |
Summary
Test plan
Links
Summary by CodeRabbit
use_host_stop_criteriathrough the PP ringSAMPLE_STATEpayload.Dev Engineer Review
getattrandhasattrprovide compatibility with sampler variants withoutuse_host_stop_criteria.Falsewhen the field is absent.QA Engineer Review
unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu2"fromtests/integration/test_lists/waives.txt.test-db/orqa/entries were modified.