[https://nvbugs/5404046][fix] Fix Nemotron-H flaky CUDA graph / overlap scheduler test #6485
Conversation
…add better visibility to failures Signed-off-by: Tomer Asida <[email protected]>
Signed-off-by: Tomer Asida <[email protected]>
Signed-off-by: Tomer Asida <[email protected]>
📝 WalkthroughWalkthroughThe test function Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~7 minutes Suggested reviewers
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. ✨ Finishing Touches
🧪 Generate unit tests
🪧 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
|
![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 (1)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py (1)
289-289: Fix line length violation.The line exceeds the 120-character limit as flagged by static analysis.
- # similar to other unittests comparing with / without CG, compare logits of first generation step (2nd generated token) + # similar to other unittests comparing with / without CG, compare logits of + # first generation step (2nd generated token)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+.
Indent Python code with 4 spaces. Do not use tabs.
Always maintain the namespace when importing in Python, even if only one class or function from a module is used.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case. Prefix k for variable names that start with a number (e.g., k_99th_percentile).
Python global variables should use upper snake_case and prefix G (e.g., G_MY_GLOBAL).
Python constants should use upper snake_case (e.g., MY_CONSTANT).
Avoid shadowing variables declared in an outer scope in Python.
Initialize all externally visible members of a Python class in the constructor.
For interfaces that may be used outside a file, prefer docstrings over comments in Python.
Comments in Python should be reserved for code within a function, or interfaces that are local to a file.
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the docstring for the class.
Avoid using reflection in Python when functionality can be easily achieved without reflection.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.
Files:
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py
**/*.{cpp,h,cu,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.
Files:
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py
🧠 Learnings (1)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py (1)
Learnt from: moraxu
PR: #6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.
🪛 Ruff (0.12.2)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py
289-289: Line too long (127 > 120)
(E501)
🔇 Additional comments (5)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py (5)
242-246: Improved prompt selection for deterministic greedy sampling.The new prompts are well-chosen to ensure significant probability differences between top tokens, addressing the flakiness issue described in the PR objectives. These prompts should produce more deterministic argmax() results during greedy sampling.
248-248: Reduced token generation for focused testing.The reduction from 12 to 10 tokens makes sense for a stability-focused test, reducing the likelihood of encountering problematic probability distributions in longer sequences.
275-287: Enhanced test diagnostics with explicit indexing.The addition of explicit indexing and detailed assertion messages significantly improves test observability. The error messages now clearly identify which prompt failed and the nature of the mismatch.
290-307: Improved assertion diagnostics for logits and logprobs comparison.The addition of descriptive
msglambda functions provides excellent context for test failures, making it much easier to diagnose issues when they occur. The messages clearly identify the prompt index and comparison type.
310-317: Enhanced overlap scheduler assertion with clear diagnostics.The improved error message for overlap scheduler comparison maintains consistency with other assertions and provides clear context for any failures.
There was a problem hiding this comment.
Pull Request Overview
This PR fixes a flaky CUDA graph/overlap scheduler test for Nemotron-H by addressing non-deterministic token selection in greedy sampling. The fix involves replacing prompts that could lead to tied top-2 token probabilities with more deterministic ones, reducing the maximum token count, and improving test failure diagnostics.
- Updated test prompts to use more deterministic patterns that avoid probability ties
- Reduced maximum token generation from 12 to 10 tokens
- Enhanced assertion messages with detailed failure context including prompt indices
Comments suppressed due to low confidence (1)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py:1
- The pytest import is being removed but may still be needed for other tests in this file. Ensure that no other tests in this file use pytest decorators or fixtures.
import torch
|
/bot run |
|
PR_Github #13569 [ run ] triggered by Bot |
|
PR_Github #13569 [ run ] completed with state |
|
/bot run |
|
PR_Github #13579 [ run ] triggered by Bot |
|
PR_Github #13579 [ run ] completed with state |
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #13632 [ run ] triggered by Bot |
|
PR_Github #13632 [ run ] completed with state |
Signed-off-by: Tomer Asida <[email protected]>
|
/bot run |
|
PR_Github #13663 [ run ] triggered by Bot |
|
PR_Github #13663 [ run ] completed with state |
…ap scheduler test (NVIDIA#6485) Signed-off-by: Tomer Asida <[email protected]> Signed-off-by: Lanyu Liao <[email protected]>
…ap scheduler test (NVIDIA#6485) Signed-off-by: Tomer Asida <[email protected]>
Summary by CodeRabbit
Description
This PR improves stability for the Nemotron-H with/without CUDA graph and overlap scheduler unittest.
This test generates greedily from a few prompts, and compares generated text and logits/logprobs when generating with and without CUDA graphs and overlap scheduler.
Flakiness in the test was caused due to generation steps where the top2 most probable tokens had the same probabilities, causing undefined behviour in the outcome of
argmax(), leading to cases where different tokens were selected. Once a single token is different in a generated sequence, all subsequent tokens will be different as well (since they are generated given the previous different token) and comparing logits becomes meaningless.Test stability was improved by choosing prompts for which the difference in probabilities between top2 most probable tokens is significant, so greedy sampling always results in the same outcome. Other than that, test observability was improved by enhancing assertion messages in case of failures.
Test Coverage
test_modeling_nemotron_h.py::test_nemotron_h_cuda_graph_overlap_schedulernow has better stability.