fix(tf/pt): validate GPU tabulation sizes - #5846
Conversation
Reject zero and oversized tabulation layer widths before launching any TensorFlow CUDA or ROCm wrapper. Cover all forward, gradient, and grad-gradient variants with GPU-only regression cases. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughTensorFlow and PyTorch GPU tabulation wrappers now reject ChangesGPU tabulation validation
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 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 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5846 +/- ##
==========================================
+ Coverage 79.03% 79.14% +0.11%
==========================================
Files 1055 1071 +16
Lines 122233 124875 +2642
Branches 4401 4536 +135
==========================================
+ Hits 96607 98838 +2231
- Misses 24061 24414 +353
- Partials 1565 1623 +58 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The TF-side fix itself is correct: the validation is GPU-scoped (called inside if (device == "GPU"), so CPU runs with large last_layer_size are unaffected even on a CUDA build), the 1..1024 bound matches every kernel's blockDim.x, and moving it before the launch (and extending it to the forward/first-grad ops) is exactly right — it resolves a defect CodeRabbit flagged on #5600 that was never applied. The reject tests will run in the merge-queue CUDA job.
Blocking on one thing: the identical bug is still live in the PyTorch op, source/op/pt/tabulate_multi_device.cc. There the grad-grad TORCH_CHECK(last_layer_size <= 1024, ...) runs after the kernel launch (
deepmd-kit/source/op/pt/tabulate_multi_device.cc
Lines 174 to 182 in e5fdff0
tabulate_fusion_se_a_gpu launch at deepmd-kit/source/op/pt/tabulate_multi_device.cc
Lines 58 to 71 in e5fdff0
…-size-validation-5654
The PyTorch tabulation ops share the CUDA kernels with the TensorFlow ops, so the same last_layer_size > 1024 block-dim overflow reaches users through the PT backend. The grad-grad TORCH_CHECK ran after the launch, and the forward / first-gradient ops had no check at all.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@source/tests/pt/test_tabulate_gpu_size_validation.py`:
- Around line 66-72: Extend test_rejects_zero_width to also validate a negative
width, such as -1, through _assert_rejected, covering the implementation’s
complete last_layer_size <= 0 boundary without changing the existing zero-width
assertion.
- Around line 74-85: The test test_gradient_paths_reject_oversized_width
currently asserts only the forward failure, so it does not exercise gradient
paths. Restructure the tests to use valid forward inputs and explicitly cover
the required tabulate_fusion_se_* _backward and _grad_grad wrappers for PyTorch,
including SeT, SeTTebd, and SeR, while asserting oversized-width rejection in
each relevant gradient operation.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: df27c21c-ccdd-4184-986e-55c3cb1212c7
📒 Files selected for processing (2)
source/op/pt/tabulate_multi_device.ccsource/tests/pt/test_tabulate_gpu_size_validation.py
Address the outstanding requested-change review comments. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Mirroring the validation rather than deferring it to a tracking issue is the better of the two options I offered, and the PT side is more thorough than what I asked for. Verified at 81838508b:
ValidateGpuLastLayerSizerejects<= 0as well as> 1024, and the comment explains why the bound exists (block dimension / dynamic shared memory sizing) rather than just restating it.- Twelve call sites, which matches your count exactly. I checked a forward and three grad-grad wrappers (lines 72, 186, 332, 601): every one sits inside
if (device == "GPU")and immediately precedes the launch, so CPU runs with a largelast_layer_sizekeep working. That scoping was the main thing I wanted to confirm. - All four post-launch
TORCH_CHECK(last_layer_size <= 1024)calls are gone. Removing them was right: sitting after the launch they could never prevent the fault they described, and their presence is a good part of why this survived review on #5600.
On why no test caught this: real configurations use last-layer widths in the tens, so nothing in the suite approaches 1024, and the defect lives entirely inside the device == "GPU" branch that CPU CI never enters. The new test crosses exactly that intersection -- 1025 and 0 across all five forward wrappers plus the descriptor-derived autograd path -- and would fail before the fix, since the forwards had no check and reach the launch instead of raising.
One thing to flag, and it is a correction to my own earlier comment rather than a request. I said the reject tests would run in the merge-queue CUDA job. On this PR Test Python on CUDA is skipping, so neither the new PT tests nor the TF ones from the first commit have actually executed -- Pass testing on CUDA goes green in 8s because it is the aggregate gate over skipped jobs. The gating itself is unavoidable, since a GPU launch guard cannot be tested without a GPU. If the merge queue does not run these either, it would be worth applying the Test CUDA label once before merging so the guards are observed working at least a single time.
Approving; this closes #5654 on both backends.
njzjz
changed the title
|
I'll ask the agent to fix the following errors: |
Reject invalid TensorFlow widths before output allocation on every GPU kernel specialization, and reject negative PyTorch widths before descriptor tensors are created. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Fixed the reported failures in b88dd44:
Validation:
Coding agent: Codex |
3 participants
Summary
1 <= last_layer_size <= 1024before every TensorFlow CUDA/ROCm tabulation wrapper callWhy existing tests missed this
There were no direct TensorFlow tabulation-op GPU tests. Existing PyTorch tabulation coverage uses normal small embedding widths and does not execute these TensorFlow kernels, while real compressed models are also normally well below the device limit. The four existing grad-gradient checks looked like validation but ran only after the GPU wrapper had already launched, and forward/first-gradient paths had no checks at all.
Validation
cmake --build source/build --target deepmd_op -j2pytest source/tests/tf/test_tabulate_gpu_size_validation.py source/tests/tf/test_tabulate.py -q(2 passed, 2 GPU-only tests skipped locally because no usable GPU is exposed)ruff format .ruff check .--Werrorgit diff --checkFixes #5654
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
Bug Fixes
last_layer_sizevalues outside the supported range of 1–1024.Tests