[TRTLLM-6651][feat] Enable Overlap scheduler + Beam Search in TRTLLM Sampler #6223
Conversation
…n TRTLLM Sampler - Move _finalize function to sample_async - store gather_ids and finalize_events in SampleState object - use SampleState buffers in update_requests. Signed-off-by: Stefan Niebler <[email protected]>
WalkthroughThe changes remove a restriction on using the overlap scheduler with beam search, refactor and extend beam search finalization and token gathering logic in the sampler, and add new tests to validate beam search with the overlap scheduler enabled. No public API signatures are removed; some are extended to support new features. Changes
Sequence Diagram(s)sequenceDiagram
participant Test as test_beam_search_output_shapes_overlap
participant LLM as llm_overlap (LLM w/ overlap)
participant Sampler as TRTLLMSampler
participant State as SampleStateTRTLLM
Test->>LLM: Run beam search with overlap scheduler
LLM->>Sampler: sample_async(...)
Sampler->>State: Create SampleStateTRTLLM with finalize_events and gathered_ids
Sampler->>Sampler: update_requests_multiple_beams_or_drafting(...)
Sampler->>Sampler: _post_process_request(request, state)
Sampler->>State: Synchronize on finalize_event, update outputs
LLM->>Test: Return outputs for validation
Estimated code review effort3 (~45 minutes) Suggested labels
Suggested reviewers
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
|
/bot run |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tests/unittest/_torch/test_beam_search.py (1)
128-179: Consider reducing code duplication by parameterizing the fixture.The new test function is identical to
test_beam_search_output_shapesexcept for the fixture used. Consider consolidating both tests by adding a parameter for the overlap scheduler configuration.Here's a suggested refactor to eliminate duplication:
-@pytest.fixture(scope="module") -def llm(fixed_params, input_prompts): - return LLM( - model=os.path.join(llm_models_root(), "llama-models-v2", - "TinyLlama-1.1B-Chat-v1.0"), - kv_cache_config=KvCacheConfig(max_tokens=10000), - max_batch_size=fixed_params["max_beam_width"] * len( - input_prompts - ), # use small batch size to prevent large buffers from possibly hiding wrong data accesses. - max_seq_len=32, - enable_trtllm_sampler=True, - max_beam_width=fixed_params["max_beam_width"], - disable_overlap_scheduler=True, - #TODO: remove this once we have a proper fix for CUDA graph in beam search - cuda_graph_config=None, - ) - - -@pytest.fixture(scope="module") -def llm_overlap(fixed_params, input_prompts): - return LLM( - model=os.path.join(llm_models_root(), "llama-models-v2", - "TinyLlama-1.1B-Chat-v1.0"), - kv_cache_config=KvCacheConfig(max_tokens=10000), - max_batch_size=fixed_params["max_beam_width"] * len( - input_prompts - ), # use small batch size to prevent large buffers from possibly hiding wrong data accesses. - max_seq_len=32, - enable_trtllm_sampler=True, - max_beam_width=fixed_params["max_beam_width"], - disable_overlap_scheduler=False, - #TODO: remove this once we have a proper fix for CUDA graph in beam search - cuda_graph_config=None, - ) +@pytest.fixture(scope="module", params=[True, False], ids=["no_overlap", "with_overlap"]) +def llm(request, fixed_params, input_prompts): + return LLM( + model=os.path.join(llm_models_root(), "llama-models-v2", + "TinyLlama-1.1B-Chat-v1.0"), + kv_cache_config=KvCacheConfig(max_tokens=10000), + max_batch_size=fixed_params["max_beam_width"] * len( + input_prompts + ), # use small batch size to prevent large buffers from possibly hiding wrong data accesses. + max_seq_len=32, + enable_trtllm_sampler=True, + max_beam_width=fixed_params["max_beam_width"], + disable_overlap_scheduler=request.param, + #TODO: remove this once we have a proper fix for CUDA graph in beam search + cuda_graph_config=None, + )Then remove the duplicate test function:
-@force_ampere # Save H100 resource -@pytest.mark.parametrize("return_log_probs", [True, False]) -@pytest.mark.parametrize("gather_generation_logits", [True, False]) -@pytest.mark.parametrize("gather_context_logits", [True, False]) -@pytest.mark.parametrize("num_output_beams", [1, 2]) -@pytest.mark.parametrize("num_prompts", [1, 2]) -@pytest.mark.threadleak(enabled=False) -def test_beam_search_output_shapes_overlap( - gather_context_logits: bool, gather_generation_logits: bool, - return_log_probs: bool, num_output_beams: int, num_prompts: int, - llm_overlap, fixed_params, input_prompts, expected_outputs): - if return_log_probs and num_prompts > 1: - pytest.skip( - "Beam search currently does not support return_log_probs with multiple prompts" - ) - sampling_params = SamplingParams( - max_tokens=fixed_params["max_tokens"], - n=num_output_beams, - best_of=fixed_params["max_beam_width"], - use_beam_search=True, - return_context_logits=gather_context_logits, - return_generation_logits=gather_generation_logits, - logprobs=return_log_probs, - ) - outputs = llm_overlap.generate(input_prompts[:num_prompts], - sampling_params=sampling_params) - assert len(outputs) == num_prompts - for output_idx, output in enumerate(outputs): - if gather_context_logits: - assert output.context_logits is not None - assert len( - output.prompt_token_ids) == output.context_logits.shape[0] - else: - assert output.context_logits is None - assert len(output.outputs) == num_output_beams - for beam_idx, beam in enumerate(output.outputs): - if gather_generation_logits: - gen_logits = beam.generation_logits - assert gen_logits is not None - assert gen_logits.ndim == 2 - assert gen_logits.shape[0] == sampling_params.max_tokens - else: - assert beam.generation_logits is None - - if return_log_probs: - assert len(beam.logprobs) == sampling_params.max_tokens - else: - assert len(beam.logprobs) == 0 - # Check output similarity - assert similar( - beam.text, - expected_outputs[input_prompts[output_idx]][beam_idx])This approach tests both overlap scheduler configurations while avoiding code duplication.
tensorrt_llm/_torch/pyexecutor/sampler.py (1)
911-912: Consider removing redundant synchronization if confirmed unnecessary.The comment suggests this synchronization might be redundant since we already wait for the sampler event in
update_requests. If this is confirmed through testing, consider removing it to improve performance.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/py_executor.py(0 hunks)tensorrt_llm/_torch/pyexecutor/sampler.py(7 hunks)tests/unittest/_torch/test_beam_search.py(2 hunks)
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
🔇 Additional comments (5)
tests/unittest/_torch/test_beam_search.py (1)
54-70: LGTM!The new
llm_overlapfixture properly configures the LLM instance with overlap scheduler enabled to extend test coverage for the new functionality.tensorrt_llm/_torch/pyexecutor/sampler.py (4)
476-476: LGTM!The
gathered_idsfield is appropriately added to store beam search results on the host.
481-481: LGTM!The
finalize_eventsdictionary appropriately tracks CUDA finalization events per request ID.
677-694: Well-implemented beam search finalization logic!The changes properly handle finalization events and gathered token IDs for beam search scenarios, enabling overlap scheduler support.
Also applies to: 721-722, 733-734
822-822: LGTM!The code correctly retrieves finalize events from the state and processes them for each request.
Also applies to: 890-892
|
PR_Github #12455 [ run ] triggered by Bot |
|
PR_Github #12455 [ run ] completed with state |
|
/bot run |
|
PR_Github #12537 [ run ] triggered by Bot |
stnie
changed the title
|
PR_Github #12537 [ run ] completed with state |
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Description
Enable Overlap scheduling together with beam search, which was previously not supported
Test Coverage
test_beam_search.py adjusted to cover overlap scheduling
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.