[https://nvbugs/5513423][fix] Correctly respect min_tokens in PyTorch Workflow #7808
Conversation
|
/bot run |
|
PR_Github #19004 [ run ] triggered by Bot |
|
PR_Github #19004 [ run ] completed with state |
stnie
force-pushed
the
bugfix/5513423
branch
from
167aadf to
ad4d9d4
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #19024 [ run ] triggered by Bot |
|
PR_Github #19024 [ run ] completed with state |
stnie
force-pushed
the
bugfix/5513423
branch
from
ad4d9d4 to
16f5f5e
Compare
… workflow using TorchSampler - Added `py_min_length` attribute to `LlmRequest` to store minimum length configuration. - Implemented `_apply_min_length_penalty` method in `TorchSampler` to adjust logits based on minimum length requirements. (Mimics PenaltyLayer) - Updated test case for `min_tokens` to reflect new maximum sequence length and output length constraints from the model Signed-off-by: Stefan Niebler <[email protected]>
…n in TorchSampler to work with speculative decoding - Refactored `_apply_min_length_penalty` method to support step-wise application of minimum length penalties based on the number of steps per request. - Updated the test for `min_tokens` to include a parameterized approach for speculative decoding configurations. Signed-off-by: Stefan Niebler <[email protected]>
stnie
force-pushed
the
bugfix/5513423
branch
from
16f5f5e to
31091b4
Compare
|
/bot run --disable-fail-fast |
|
/bot run |
📝 WalkthroughWalkthroughIntroduces an internal py_min_length field in LlmRequest. Adds and applies a min-length logits penalty in the sampler before decoding step counting. Updates a PyTorch unit test to parameterize speculative decoding, import NGramDecodingConfig, and pass speculative_config into LLM construction, with adjusted token limits. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Req as LlmRequest
participant Sampler as Sampler
participant Model as Model
Note over Req: Request initialized with sampling_config<br/>→ set py_min_length from min_length
Req->>Sampler: generate_step(requests)
Sampler->>Model: forward(context) -> raw_logits
Note over Sampler: Apply min-length penalty to logits<br/>(_apply_min_length_penalty)
Sampler->>Sampler: compute step counts
Sampler->>Sampler: decoding/sampling on modified logits
Sampler-->>Req: tokens
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
PR_Github #19470 [ run ] triggered by Bot |
![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 (8)
tensorrt_llm/_torch/pyexecutor/llm_request.py (2)
1-1: Add NVIDIA Apache-2.0 header.Per repo guidelines, prepend the NVIDIA Apache-2.0 copyright header (year 2025) at the top of this file.
330-330: py_min_length propagation: LGTM; add a tiny guard for robustness.Looks correct and matches how other sampling fields are mirrored. To avoid ambiguous truthiness if bindings ever return a tensor here, coerce to a Python scalar/list once:
- self.py_min_length = self.sampling_config.min_length + ml = self.sampling_config.min_length + # Coerce to a simple Python type (None | int | [int]) to avoid tensor truthiness pitfalls. + if hasattr(ml, "tolist"): + ml = ml.tolist() + self.py_min_length = mltensorrt_llm/_torch/pyexecutor/sampler.py (2)
1-1: Add NVIDIA Apache-2.0 header.Please add the required NVIDIA Apache-2.0 header (2025) at the top of this file.
660-689: Min-length penalty implementation: correct, but harden for input type and semantics.The in-place -inf on EOS until min length is reached mirrors PenaltyLayer behavior. Two small hardening tweaks:
- Coerce r.py_min_length to an int (first element) safely, tolerating None/list/int/tensor.
- Add a brief note that min_length is interpreted as absolute sequence length (prompt + generated). If it is “min_new_tokens” upstream, subtract prompt_len to avoid double counting.
Example minimal tweak:
def _apply_min_length_penalty(logits: torch.Tensor, requests: list[LlmRequest], num_steps: list[int]) -> torch.Tensor: @@ - if any(r.py_min_length and r.max_beam_num_tokens < r.py_min_length[0] + def _ml_val(r: LlmRequest) -> int: + ml = getattr(r, "py_min_length", None) + if ml is None: + return 0 + if isinstance(ml, (list, tuple)): + return int(ml[0]) if ml else 0 + if hasattr(ml, "item"): + return int(ml.item()) + return int(ml) + + if any((_ml_val(r) > 0) and (r.max_beam_num_tokens < _ml_val(r)) for r in requests): current_offset = 0 for index, r in enumerate(requests): - if r.py_min_length: + ml = _ml_val(r) + if ml > 0: for step in range(num_steps[index]): - if r.max_beam_num_tokens + step < r.py_min_length[0]: + # NOTE: Interprets min_length as absolute sequence length (incl. prompt). + if r.max_beam_num_tokens + step < ml: logits[current_offset + step, r.py_end_id] = float('-inf') else: #early exit break current_offset += num_steps[index]tests/unittest/llmapi/test_llm_pytorch.py (4)
1-1: Add NVIDIA Apache-2.0 header.Tests are covered by the same header requirement; please add it.
9-9: Prefer module namespace imports per guidelines.To match import style guidance, import the module and reference members via the namespace:
-from tensorrt_llm.llmapi.llm_args import NGramDecodingConfig, PeftCacheConfig +from tensorrt_llm.llmapi import llm_argsThen use
llm_args.NGramDecodingConfigandllm_args.PeftCacheConfig.
863-887: Min-tokens test: good coverage incl. speculative path; add resource-safety and determinism nits.
- Use
with LLM(...) as llm:to ensure clean shutdown in case of assertion failures.- Consider
temperature=0.0to eliminate randomness and avoid rare flakiness.- if use_speculative: + if use_speculative: spec_config = NGramDecodingConfig( @@ - llm = LLM(**llm_common_config, speculative_config=spec_config) + with LLM(**llm_common_config, speculative_config=spec_config) as llm: + ... - else: - llm = LLM(**llm_common_config) + else: + with LLM(**llm_common_config) as llm: + ...and
- sampling_params = SamplingParams(max_tokens=output_len, - min_tokens=output_len, - temperature=1) + sampling_params = SamplingParams(max_tokens=output_len, + min_tokens=output_len, + temperature=0.0)
887-887: Output length choice: align with model max to exercise the boundary.Optional: set
output_len = llm_common_config["max_num_tokens"] - 48(or derive from tokenizer prompt length) to exercise the near-boundary case and ensure the penalty lifts exactly at the boundary.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/llm_request.py(1 hunks)tensorrt_llm/_torch/pyexecutor/sampler.py(2 hunks)tests/unittest/llmapi/test_llm_pytorch.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Use only spaces, no tabs; indent with 4 spaces.
Files:
tensorrt_llm/_torch/pyexecutor/llm_request.pytests/unittest/llmapi/test_llm_pytorch.pytensorrt_llm/_torch/pyexecutor/sampler.py
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.
Files:
tensorrt_llm/_torch/pyexecutor/llm_request.pytests/unittest/llmapi/test_llm_pytorch.pytensorrt_llm/_torch/pyexecutor/sampler.py
**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).
Files:
tensorrt_llm/_torch/pyexecutor/llm_request.pytests/unittest/llmapi/test_llm_pytorch.pytensorrt_llm/_torch/pyexecutor/sampler.py
🧬 Code graph analysis (2)
tests/unittest/llmapi/test_llm_pytorch.py (2)
tensorrt_llm/llmapi/llm_args.py (2)
NGramDecodingConfig(491-523)PeftCacheConfig(840-906)tensorrt_llm/llmapi/llm.py (1)
LLM(1022-1038)
tensorrt_llm/_torch/pyexecutor/sampler.py (1)
tensorrt_llm/_torch/pyexecutor/llm_request.py (1)
LlmRequest(282-425)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (1)
tensorrt_llm/_torch/pyexecutor/sampler.py (1)
719-721: Apply penalty before sampling: good placement. Verify coverage for context-logged paths.Applying the min-length penalty prior to step counting/sampling is correct. Please confirm that when context logits are requested (filtered indices path above), the sliced
raw_logitsstill aligns withrequests/num_steps, or guard with an assertion:+ assert raw_logits.size(0) >= sum(num_steps), "logits rows < total steps"
|
PR_Github #19470 [ run ] completed with state |
… Workflow (NVIDIA#7808) Signed-off-by: Stefan Niebler <[email protected]> Co-authored-by: Daniel Cámpora <[email protected]>
… Workflow (NVIDIA#7808) Signed-off-by: Stefan Niebler <[email protected]> Co-authored-by: Daniel Cámpora <[email protected]>
4 participants
[https://nvbugs/5513423][fix] Correctly respect min_tokens in PyTorch Workflow
py_min_lengthattribute toLlmRequestto store minimum length configuration._apply_min_length_penaltymethod inTorchSamplerto adjust logits based on minimum length requirements. (Mimics PenaltyLayer)min_tokensto reflect new maximum sequence length and output length constraints from the modelSummary by CodeRabbit
New Features
Tests
Description
min_tokens was not correctly applied in the current PyTorch workflow as the penalty Layer is not executed in the TorchSampler. I added a Python function of the Min_length penalty and added this to the Torch Sampler.
test_min_tokens was adjusted, as the used model does not allow max_seq_len > 2048
Test Coverage
I removed the
xfailannotation fromtest_llm_pytorch.py::test_min_tokensPR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
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.