fix(infer): reduce CUDA OOM retries during evaluation - #6024
Conversation
github-actions
Bot
added
Python
Docs
and removed
Test CUDA
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (5)
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughThe change adds budget-aware batch execution with native CUDA memory calibration, cache release after validation, updated environment documentation, regression coverage, and frame-parameterized atomic statistic test fixtures. ChangesCUDA batch sizing
Frame-aware statistic tests
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~45 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant AutoBatchSize
participant _CudaMemoryBudget
participant CUDAAllocator
participant EvaluationCallable
AutoBatchSize->>_CudaMemoryBudget: request batch limit
_CudaMemoryBudget->>CUDAAllocator: read memory statistics
CUDAAllocator-->>_CudaMemoryBudget: return memory statistics
_CudaMemoryBudget-->>AutoBatchSize: return admissible batch size
AutoBatchSize->>EvaluationCallable: execute batch
EvaluationCallable-->>AutoBatchSize: return result or OOM
AutoBatchSize->>_CudaMemoryBudget: observe completed attempt
Merge Risk: ⚪ Minimal · up to No actionable current-head risk remains from the reviewed changes. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Requesting changes for the cold-start path. The allocator accounting/backoff logic looks well covered, and I did not find a correctness regression in the tail backoff, private-pool handling, or PyTorch 2.1 compatibility. The remaining issue is that the one-frame calibration replaces the original 1024-atom proposal, so a fresh dp test / full-validation evaluation grows geometrically from 1 frame. The PR reports a steady-state benchmark, which does not exercise this cost. Please preserve/jump back toward the pre-calibration proposal after the first observation, or demonstrate with a cold-start benchmark that the extra forward/snapshot calls are negligible and add regression coverage for that path.
CI is still running, so this review is code-level only.
Agent: ChatGPT
Model: GPT-5.6 Sol
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6024 +/- ##
==========================================
- Coverage 77.23% 77.01% -0.23%
==========================================
Files 1153 1153
Lines 139166 139258 +92
Branches 5062 5062
==========================================
- Hits 107481 107243 -238
- Misses 29802 30133 +331
+ Partials 1883 1882 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Addressed the review and CUDA CI failures in c176105.
Local validation: 64 related CPU tests and 69 related GPU tests passed. The three statistics suites also passed all 7 tests with CPU inference forced to one frame per batch. Normal pre-commit hooks passed. CUDA CI is being retriggered for the new commit. |
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
NEEDS HUMAN REVIEW
Re-reviewed the changed head after the previous cold-start concern was addressed. The fix now preserves the pre-calibration atom proposal: the first CUDA calibration call is still one frame, then the runner applies the measured capacity to the original proposal instead of growing geometrically from one frame. The new regression explicitly exercises that path ([1, 7] for an 8-frame input), and the author also provided fresh-process cold-start measurements. I reviewed all 10 changed files, including the shared batch runner, CUDA memory budget/accounting, validation cache release, documentation, and the frame-aware regression fixtures; I did not find a new high-confidence functional or compatibility blocker in the current diff.
I am not approving this head yet because relevant exact-head CI is still incomplete: Test Python, CodeQL, Test C++, one Test CUDA run, Build/upload to PyPI, and Read the Docs are still running/pending. The completed Build C++, Build C library, another Test CUDA run, pre-commit.ci, and CodeRabbit checks are successful. Please treat this as code-level review pending the remaining CI.
Agent: ChatGPT
Model: GPT-5.6 Sol
GitHub account: njzjz-bot
Reviewed head: c176105
Trigger: scheduled review-request monitoring
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
NEEDS HUMAN REVIEW
Re-reviewing this explicitly re-requested head. I re-checked the complete 10-file change, repository guidance, existing review threads/comments, and current-head CI. The previously identified cold-start calibration issue remains fixed: after the one-frame calibration, growth is constrained by the original proposal and the measured capacity, with focused regression coverage. I did not find a new high-confidence functional, compatibility, resource-management, or test-coverage blocker in the retrievable change.
I am not approving yet because relevant CI for this exact head is still incomplete. Build C++, Build C library, Build/upload to PyPI, CodeQL, and one Test CUDA run have completed successfully, while Test Python, Test C++, and another Test CUDA run are still in progress.
Agent: ChatGPT
Model: GPT-5.6 Sol
GitHub account: njzjz-bot
Reviewed head: c176105
Trigger: scheduled review-request monitoring
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
NEEDS HUMAN REVIEW
Re-checked this exact head because the CI context materially changed after the previous pending-CI review. The code-level conclusion is unchanged: the cold-start calibration fix is present, the 10-file diff remains coherent, and I do not see a new high-confidence functional or compatibility blocker in the reviewed changes.
The required CI is now red, however. On exact head c176105445718ba73580557a7ca6f37eab48f066, Test Python, Test C++, Build C++, Build C library, Build/upload to PyPI, CodeQL, one Test CUDA run, pre-commit.ci, CodeRabbit, and docs are successful, but Test CUDA run 34924511513 failed specifically in Test Python on CUDA (job 104239589664); Test C++ on CUDA in the same run passed. I cannot approve while that required CUDA job is failing. The current integration exposes the failed job/step but not the pytest failure log, so the actionable next step is to inspect/re-run that CUDA Python job and determine whether this is a reproducible regression or an infrastructure/flaky failure before approval.
Agent: ChatGPT
Model: GPT-5.6 Sol
GitHub account: njzjz-bot
Reviewed head: c176105
Trigger: scheduled all-PR monitoring
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Re-reviewed this exact head because the CI context materially changed after the previous pending/failed-CUDA review. I re-checked the complete 10-file diff, repository guidance, the existing inline thread and its fix, and the current exact-head workflow results. The cold-start calibration concern remains correctly addressed: the one-frame calibration preserves the original proposal and then jumps to the measured admissible batch, with regression coverage for the [1, 7] path. I did not find a remaining high-confidence functional, compatibility, resource-management, or test-coverage blocker.
The previously failing CUDA workflow has now been rerun successfully. Test Python, both Test CUDA runs, Test C++, Build C++, Build C library, CodeQL, and Build/upload to PyPI are all complete and green for this head, so I am approving.
Agent: ChatGPT
Model: GPT-5.6 Sol
GitHub account: njzjz-bot
Reviewed head: c176105
Trigger: scheduled all-PR monitoring
e8725d0
Summary
Automatic GPU batching currently grows until an allocation fails. Evaluation can recover, but the CUDA allocator may emit OOM warnings and retry allocations while searching for a usable batch size.
This PR reduces those avoidable retries in
dp testand full validation:AutoBatchSizerunner. Back off from the actual attempted batch, including the remaining input length, so the same oversized tail is not retried repeatedly.The memory estimate is a heuristic. Unexpected workload or device-memory changes retain the existing OOM fallback. Explicit
DP_INFER_BATCH_SIZE, CPU, single-frame calls, and non-native CUDA allocators do not opt into memory-based sizing. The existing multiplicative growth/backoff factor is unchanged.Validation
dp testand full-validation tests: 33 passed, plus 9 subtests.expandable_segments=Trueand a 512 MiB process limit: OOM-driven sizing incurred 3 OOMs and 3 allocator retries; memory-based sizing incurred none. Live private-pool and static-buffer CUDA Graph cases also completed without OOM/retries and replayed correctly.Summary by CodeRabbit
New Features
Bug Fixes
Documentation