[https://nvbugs/5458874][fix] Fix Nemotron-H flaky CUDA graph / overlap scheduler test #6996
Conversation
…erated logprob for all tokens in overlap scheduler assertion as well Signed-off-by: Tomer Asida <[email protected]>
Signed-off-by: Tomer Asida <[email protected]>
📝 WalkthroughWalkthroughUnskips and updates a Nemotron-H CUDA graph overlap scheduler unit test to run comparisons between CG no-overlap and CG with-overlap runs. Tightens tolerances to 0.05 for logits and decode logprobs, updates assertion messages, and aligns logprob checks to use extract_decode_logprobs on CG variants. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
Suggested reviewers
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
There was a problem hiding this comment.
Pull Request Overview
This PR fixes a flaky test for Nemotron-H CUDA graph with overlap scheduler by reducing the scope of logits comparison. Instead of comparing all logits for all generation steps, it now compares all logits only for the first generation step and logprobs for all generated tokens, making the test more stable while maintaining adequate validation.
Key changes:
- Removed the
pytest.mark.skipdecorator to re-enable the test - Modified logits comparison to focus on first generation step only
- Added separate logprobs comparison for all generated tokens
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
![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/modeling/test_modeling_nemotron_h.py (2)
319-319: Fix Ruff E501 long line in comment (exceeds 120 chars)Break the comment into shorter lines to satisfy the linter.
Apply this diff:
- # Similar comparison for with / without overlap scheduler, compare logits of first generation step (2nd generated token) + # Compare with vs without overlap scheduler + # Compare logits of first generation step (2nd generated token)
1-1: Missing NVIDIA copyright headerPer repo guidelines, prepend the current-year NVIDIA copyright header.
Add at the top of the file (before imports):
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py
🪛 Ruff (0.12.2)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py
319-319: Line too long (128 > 120)
(E501)
🔇 Additional comments (2)
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py (2)
319-328: Solid focus on first-step logits with tight tolerance—aligns with the stated flakiness fixComparing all logits only for the first generation step and tightening atol/rtol to 0.05 for the overlap scheduler case is appropriate and mirrors the CUDA-graphs pattern. This should materially reduce flakiness while still catching regressions (scheduler should not affect logits).
330-338: Good choice to compare only generated-token logprobs for subsequent stepsUsing extract_decode_logprobs to match only the selected tokens across the decode steps is the right granularity and should further minimize spurious diffs. The tighter tolerances are reasonable here as well.
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #15615 [ run ] triggered by Bot |
|
PR_Github #15615 [ run ] completed with state |
tomeras91
changed the title
Summary by CodeRabbit
Description
Make test less flaky by not comparing all logits for all generated steps, but just all logits for first generation step and just generated logits for all generation steps (similar to what is tested for CUDA graphs in this test).
Still use a small tolerance since overlap scheduler should not affect logits.
Test Coverage
tests/unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_cuda_graph_overlap_schedulernow passes and should not be flaky