[https://nvbugs/6437410][fix] fix nemotron weight update test - #16712
Conversation
|
/bot run --disable-fail-fast --stage-list "H100_PCIe-PyTorch-Ray-1, DGX_B200-4_GPUs-PyTorch-Ray-1" |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
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 (3)
🚧 Files skipped from review as they are similar to previous changes (3)
WalkthroughThe Ray test setup now installs pinned architecture-specific Mamba dependencies. Update-weights tests pass checkpoint-derived layer configuration directly. The Nemotron-H test runs in-process with explicit dependency checks and an updated logits threshold. Its ChangesUpdate-weights test coverage
Estimated code review effort: 3 (Moderate) | ~25 minutes Mergeability Score: ⚪ Minimal · up to This change is merge-ready after normal checks and review; no actionable merge-blocking risk remains. Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py`:
- Around line 796-847: Update test_llm_update_weights_nemotron_h to launch the
pytest command with subprocess.Popen(start_new_session=True), replacing
subprocess.run while preserving its output capture and timeout behavior. In the
subprocess.TimeoutExpired handler, terminate the entire process group with
os.killpg using the child’s process group ID before failing the test, ensuring
Ray/NCCL workers are also stopped.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: bf9183ba-ae60-4bd9-9a0e-04b6c54031a4
📒 Files selected for processing (3)
tests/integration/test_lists/waives.txttests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.pytests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
|
PR_Github #60897 [ run ] triggered by Bot. Commit: |
|
PR_Github #60897 [ run ] completed with state |
chzblych
left a comment
chzblych
left a comment
There was a problem hiding this comment.
Approved for the waive list change.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: This test-only fix is mechanically mergeable and the core approach (truncating hybrid NemotronH via layers_block_type instead of the silently-ignored num_hidden_layers) looks correct, but the timeout error path leaks GPU-holding worker processes and should be fixed before merge.
Concerns
- [MAJOR]
tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py:831- TimeoutExpired handler leaks the Ray/NCCL process group- What is wrong: The subprocess is started with
subprocess.run(...)with no new session / process group. Onsubprocess.TimeoutExpired, only the directpytestchild is terminated; the Ray/NCCL/CUDA workers it forked are orphaned. - How it fails: The 1800s timeout exists precisely to catch a Ray/NCCL/CUDA deadlock. When that hang occurs ->
TimeoutExpiredfires ->pytest.fail(...)is raised, but the worker grandchildren survive and keep holding multi-GPU memory. Subsequent tests on the same node then fail with OOM / device-busy, turning a single hang into a cascade. - Suggested fix: launch in its own process group and kill the group on timeout, e.g.
proc = subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, start_new_session=True) try: out, err = proc.communicate(timeout=subprocess_timeout_s) except subprocess.TimeoutExpired: os.killpg(os.getpgid(proc.pid), signal.SIGKILL) out, err = proc.communicate() pytest.fail(f"Nemotron-H subprocess hung >{subprocess_timeout_s:.0f}s; killed.\n{out}\n{err}")
- What is wrong: The subprocess is started with
Minor notes (non-blocking)
tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py:772- threshold lowered 0.9 -> 0.8 while observed overlap is ~0.89; the ~0.09 slack could mask a real regression. Consider ~0.85.tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py:44- dropping thenum_hidden_layers=4default means any caller that relied on it now loads the full checkpoint. Please confirm no other instantiations depend on the old default.tests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.py:812--p no:xdistcan error if pytest-xdist is not installed in the subprocess env.
QA view
- Test coverage: adequate for a test-only change - it re-enables the previously waived Nemotron-H test; there is no production code to cover. Runtime pass/fail is unverified (CBTS results unavailable per the PR description) and the timeout/leak path is not itself exercised.
- SM coverage: arch-guarded by
@skip_pre_hopper, so it runs on Hopper+ only; the FP8 Nemotron/Qwen case stays waived, so the fp8 variant of this path remains untested. No new arch introduced without a test. - Test code: see the MAJOR (process-group leak) plus the loosened threshold and
-p no:xdistnotes above. - Test time: significant - removing the
part4waiver re-enables a 4-GPU weight-update test on a (7-layer-truncated) 30B model with an 1800s budget. Exact runtime can't be read from the diff. - Needs
/qa-verify: yes - this un-waives a previously failing bug-fix test whose result is unavailable, changes test infrastructure (subprocess-driven pytest + process lifecycle), and has a leaky error path. A human QA should confirm it passes on Hopper+ and leaves no orphaned GPU workers.
Does this actually fix nvbugs/6437410?
Likely yes. The traced failure is that NemotronHConfig derives num_hidden_layers from layers_block_type and silently ignores the direct override, so the old model_kwargs={'num_hidden_layers':7} kept the full 30B model (OOM / unmatched logits). The diff truncates via layers_block_type[:7] on both the HF reference and the LLM, adds an assert to catch the silent-ignore case, loosens the threshold for mamba/selective-scan drift, and re-runs in a fresh python -m pytest subprocess so the mamba-ssm fast path is picked up. That chain addresses the root cause, but it is unverified at runtime.
Possible new issues
- Leaked Ray/NCCL GPU workers on a hang (MAJOR) can break later tests on the same node.
- Removed
num_hidden_layers=4default may cause unshown callers to load full checkpoints. - The 0.8 threshold reduces sensitivity to future numerical regressions.
What I could not verify
- The other instantiations of
RefHFModelWithIPCHandles(not in the diff) and whether any relied on the old default of 4 layers. - Actual test pass/fail and whether the subprocess leaves orphaned processes at runtime - CBTS coverage was not provided.
- Whether pytest-xdist is guaranteed present in the subprocess environment.
Automated review by NVCortex Lite, run by @fredricz-20070104.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
jenkins/L0_Test.groovy (1)
3811-3813: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winVerify the integrity of the downloaded native wheels.
These commands install native-code wheels from public GitHub during CI. The release assets publish SHA-256 checksums, but this command does not verify them. Add per-architecture checksum validation or mirror the wheels into controlled Artifactory before installation. (github.com)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@jenkins/L0_Test.groovy` around lines 3811 - 3813, Update the native wheel installation block to verify SHA-256 checksums for both causal-conv1d and mamba_ssm assets before pip installation, using the architecture-specific published checksums for mambaArch. Alternatively, source the wheels from controlled Artifactory, but do not install directly from the public URLs without integrity validation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py`:
- Around line 44-51: Update RefHFModelWithIPCHandles.__init__ to annotate its
return as None and replace Optional[int] with int | None and Optional[List[str]]
with list[str] | None, preserving the existing constructor behavior and
parameters.
---
Nitpick comments:
In `@jenkins/L0_Test.groovy`:
- Around line 3811-3813: Update the native wheel installation block to verify
SHA-256 checksums for both causal-conv1d and mamba_ssm assets before pip
installation, using the architecture-specific published checksums for mambaArch.
Alternatively, source the wheels from controlled Artifactory, but do not install
directly from the public URLs without integrity validation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 21608973-137a-4a75-89a8-c7863f851674
📒 Files selected for processing (4)
jenkins/L0_Test.groovyjenkins/scripts/slurm_install.shtests/unittest/_torch/ray_orchestrator/multi_gpu/test_llm_update_weights_multi_gpu.pytests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
|
/bot run --disable-fail-fast --stage-list "H100_PCIe-PyTorch-Ray-1, DGX_B200-4_GPUs-PyTorch-Ray-1" |
|
PR_Github #63636 [ run ] triggered by Bot. Commit: |
|
PR_Github #63636 [ run ] completed with state |
|
/bot run |
|
PR_Github #66447 [ run ] triggered by Bot. Commit: |
|
PR_Github #66447 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66485 [ run ] triggered by Bot. Commit: |
|
PR_Github #66485 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66519 [ run ] triggered by Bot. Commit: |
|
PR_Github #66519 [ run ] completed with state |
|
Hi @zhaoyangwang-nvidia, all review threads have been addressed and the full pre-merge CI has passed. Could you please re-enable auto-merge? Thanks! |
|
Enabled, please check the conflict |
Signed-off-by: Shuyi Xiong <219646547+shuyixiong@users.noreply.github.com>
Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
… drop nested pytest Signed-off-by: shikicloud <shikiw@nvidia.com>
Head branch was pushed to by a user without write access
11e00d9 to
86440e7
Compare
|
/bot run |
|
PR_Github #66620 [ run ] triggered by Bot. Commit: |
|
PR_Github #66620 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66678 [ run ] triggered by Bot. Commit: |
|
PR_Github #66678 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66725 [ run ] triggered by Bot. Commit: |
|
PR_Github #66725 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66728 [ run ] triggered by Bot. Commit: |
|
PR_Github #66728 [ run ] completed with state |
9 participants
Dev Engineer Review
layers_block_typefrom the checkpoint configuration.causal-conv1dandmamba_ssmwheels by architecture.part4waiver was removed.QA Engineer Review
Modified test functions and helpers:
test_llm_update_weights_nemotron_h.mamba_deps._nemotron_h_body._nemotron_h_subprocess_entry.RefHFModelWithIPCHandles.__init__.Coverage:
test_llm_update_weights_nemotron_hwas listed intests/integration/test_lists/waives.txt; its waiver was removed.test-db/orqa/changes were provided.Verdict: needs follow-up — CBTS coverage data is unavailable.
Description
Test Coverage
PR 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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
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
To see a list of available CI bot commands, please comment
/bot help.