Skip to content

Navigation Menu

Sign in
Sign up

fix(tf/pt): validate GPU tabulation sizes - #5846

Draft
njzjz-bot wants to merge 6 commits into
deepmodeling:master from
njzjz-bot:fix/tf-tabulate-gpu-size-validation-5654
Draft

fix(tf/pt): validate GPU tabulation sizes #5846
njzjz-bot wants to merge 6 commits into
deepmodeling:master from
njzjz-bot:fix/tf-tabulate-gpu-size-validation-5654

Conversation

@njzjz-bot

@njzjz-bot njzjz-bot commented Jul 16, 2026 •
edited by coderabbitai Bot
Loading

Copy link
Copy Markdown
Contributor

Summary

  • validate 1 <= last_layer_size <= 1024 before every TensorFlow CUDA/ROCm tabulation wrapper call
  • move the four existing grad-gradient checks ahead of their kernel launches
  • cover all SeA, SeAtten, SeT, and SeR forward/gradient/grad-gradient GPU paths, including zero width

Why 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 -j2
  • constructed all 12 direct-op regression graphs against the fresh build-tree TensorFlow extension
  • pytest 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 .
  • clang-format dry run with --Werror
  • git diff --check

Fixes #5654

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

Summary by CodeRabbit

  • Bug Fixes

    • GPU tabulation operations now reject invalid last_layer_size values outside the supported range of 1–1024.
    • Validation occurs before GPU execution across forward and gradient operations.

  • Tests

    • Added GPU-focused coverage confirming zero, negative, and oversized values are rejected across supported tabulation variants and gradient paths.
 

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
@dosubot dosubot Bot added the bug label Jul 16, 2026

coderabbitai Bot commented Jul 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

i️ Recent review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 8619b4b1-7e25-417c-a771-6c8195ed799b

📥 Commits

Reviewing files that changed from the base of the PR and between 03d6501 and b88dd44.

📒 Files selected for processing (3)
  • source/op/pt/tabulate_multi_device.cc
  • source/op/tf/tabulate_multi_device.cc
  • source/tests/pt/test_tabulate_gpu_size_validation.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • source/op/tf/tabulate_multi_device.cc
📝 Walkthrough

Walkthrough

TensorFlow and PyTorch GPU tabulation wrappers now reject last_layer_size values outside 1–1024 before launching GPU kernels. GPU-gated tests cover invalid sizes across forward, gradient, and grad-gradient paths.

Changes

GPU tabulation validation

Layer / File(s) Summary
TensorFlow GPU validation
source/op/tf/tabulate_multi_device.cc		 Adds shared prelaunch validation for forward, gradient, and grad-gradient GPU paths.
TensorFlow validation tests
source/tests/tf/test_tabulate_gpu_size_validation.py		 Tests standardized InvalidArgumentError responses for oversized and zero sizes.
PyTorch GPU validation
source/op/pt/tabulate_multi_device.cc		 Adds shared validation before descriptor allocation and kernel execution for supported tabulation wrappers.
PyTorch validation tests
source/tests/pt/test_tabulate_gpu_size_validation.py		 Tests oversized, zero, and negative sizes across forward and autograd paths.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

Suggested labels: CUDA, ROCM

Suggested reviewers: njzjz, wanghan-iapcm

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Out of Scope Changes check ⚠️ Warning The PyTorch implementation and tests extend beyond linked issue [#5654], which specifies TensorFlow GPU tabulation validation. Move the PyTorch changes to a separate pull request or link an issue that explicitly includes PyTorch tabulation size validation.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely describes GPU tabulation size validation for both TensorFlow and PyTorch.
Linked Issues check ✅ Passed The TensorFlow changes validate sizes from 1 through 1024 before GPU execution across forward, gradient, and grad-gradient paths required by issue [#5654].
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with 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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

codecov Bot commented Jul 17, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 50.00000% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.14%. Comparing base (cc908a8) to head (b88dd44).
⚠️ Report is 20 commits behind head on master.

Files with missing lines Patch % Lines
source/op/tf/tabulate_multi_device.cc 31.81% 15 Missing ⚠️
source/op/pt/tabulate_multi_device.cc 75.00% 2 Missing and 2 partials ⚠️
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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@wanghan-iapcm wanghan-iapcm left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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::tabulate_fusion_se_a_grad_grad_gpu(
dz_dy, table, table_info, em_x, em, two_embed, dz_dy_dem_x, dz_dy_dem,
dz_dy_dtwo, nloc, nnei, last_layer_size, is_sorted);
#else
throw std::runtime_error(
"The input tensor is on the GPU, but the GPU support for the "
"customized OP library is not enabled.");
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
TORCH_CHECK(last_layer_size <= 1024,
, and likewise at L328/L472/L596), and the forward / first-gradient PT ops (e.g. the tabulate_fusion_se_a_gpu launch at
const int64_t nnei = em_tensor.size(1);
// compute
if (device == "GPU") {
#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
deepmd::tabulate_fusion_se_a_gpu(descriptor, table, table_info, em_x, em,
two_embed, nloc, nnei, last_layer_size);
#else
throw std::runtime_error(
"The input tensor is on the GPU, but the GPU support for the "
"customized OP library is not enabled.");
#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
} else if (device == "CPU") {
deepmd::tabulate_fusion_se_a_cpu(descriptor, table, table_info, em_x, em,
two_embed, nloc, nnei, last_layer_size);
) have no check at all — the same silent SIGFPE / bad-launch risk #5654 was filed for. Please either apply the same before-launch validation to the PT op in this PR, or open a tracking issue so the vulnerability class is actually closed rather than half-fixed. Details inline.

Comment thread source/op/tf/tabulate_multi_device.cc
njzjz-bot added 2 commits July 27, 2026 11:20 
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 coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
i️ Review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: df27c21c-ccdd-4184-986e-55c3cb1212c7

📥 Commits

Reviewing files that changed from the base of the PR and between 6ce049d and 03d6501.

📒 Files selected for processing (2)
  • source/op/pt/tabulate_multi_device.cc
  • source/tests/pt/test_tabulate_gpu_size_validation.py

Comment thread source/tests/pt/test_tabulate_gpu_size_validation.py
Comment thread source/tests/pt/test_tabulate_gpu_size_validation.py Outdated 
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
njzjz requested a review from wanghan-iapcm July 30, 2026 03:40

@wanghan-iapcm wanghan-iapcm left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  • ValidateGpuLastLayerSize rejects <= 0 as 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 large last_layer_size keep 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 njzjz changed the title (削除) fix(tf): validate GPU tabulation sizes (削除ここまで) (追記) fix(tf/pt): validate GPU tabulation sizes (追記ここまで) Jul 30, 2026
njzjz added this pull request to the merge queue Jul 30, 2026
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jul 30, 2026

njzjz commented Jul 31, 2026

Copy link
Copy Markdown
Member

I'll ask the agent to fix the following errors:

=========================== short test summary info ============================
SUBFAILED(op='se_a', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, 4, -1]"
SUBFAILED(op='se_atten', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, 4, -1]"
SUBFAILED(op='se_t', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, -1]"
SUBFAILED(op='se_t_tebd', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, 1, 1, -1]"
SUBFAILED(op='se_r', last_layer_size=-1) source/tests/pt/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_nonpositive_width - AssertionError: "last_layer_size must be between 1 and 1024" does not match "Trying to create tensor with negative dimension -1: [1, 1, -1]"
SUBFAILED(op='se_a_forward', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_a_grad_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_atten_forward', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_atten_grad_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_t_forward', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_t_grad_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_r_forward', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_r_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_r_grad_grad', last_layer_size=1025) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_oversized_width_for_all_gpu_paths - AssertionError: InvalidArgumentError not raised
SUBFAILED(op='se_a_forward', last_layer_size=0) source/tests/tf/test_tabulate_gpu_size_validation.py::TestTabulateGpuSizeValidation::test_rejects_zero_width - AssertionError: InvalidArgumentError not raised


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

njzjz commented Jul 31, 2026

Copy link
Copy Markdown
Member

Fixed the reported failures in b88dd44:

  • PyTorch now validates GPU last_layer_size before each forward wrapper calls torch::empty, so negative widths raise the tabulation-specific error instead of a generic negative-dimension error.
  • TensorFlow now identifies GPU kernels from the Device specialization and validates all 12 forward/grad/grad-grad paths before allocating any output tensor. CPU widths remain unaffected by the GPU-only 1024 bound.

Validation:

  • TensorFlow and PyTorch custom-op sources compiled and linked successfully in a complete CPU C++ build.
  • Both GPU validation test modules collected successfully (6 tests; skipped because the installed TF/PyTorch builds are CPU-only).
  • clang-format --dry-run --Werror, git diff --check, and targeted Ruff checks passed; ruff format . made no changes.
  • Full ruff check . still reports five unrelated pre-existing findings on this branch: four RUF036 findings in deepmd/jax/jax_md/__init__.py and one duplicate doc_train_input entry in deepmd/tf/entrypoints/__init__.py.

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh

@njzjz njzjz added the Test CUDA Trigger test CUDA workflow label Jul 31, 2026
@github-actions github-actions Bot removed the Test CUDA Trigger test CUDA workflow label Jul 31, 2026
njzjz marked this pull request as draft July 31, 2026 14:16

@njzjz-bot njzjz-bot left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed draft head b88dd447f263a15d3ec4880eb5293e2c3b30999e with three independent full subagent reviews. One P2 API regression remains inline: the shared 1024-width cap is correct for kernels whose launch dimensions or shared memory depend on last_layer_size, but it incorrectly restricts PyTorch SE-T-TEBD, whose CUDA/ROCm kernels use fixed 256-thread blocks and iterate over the width serially. The new test currently encodes the same false restriction. The current CUDA checks are failing and should be examined after this correction. Because this PR was opened by the active njzjz-bot account, this is a comment-only review rather than a self-request-changes event.

The quota is about to reset, so I am concentrating the remaining token budget on these reviews.

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh

Comment thread source/op/pt/tabulate_multi_device.cc Outdated 
Coding-Agent: Codex
Codex-Version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning-Effort: xhigh
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@njzjz njzjz njzjz left review comments
@coderabbitai coderabbitai[bot] coderabbitai[bot] left review comments
@wanghan-iapcm wanghan-iapcm wanghan-iapcm approved these changes

Assignees

No one assigned

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Code scan] Validate TensorFlow tabulation GPU sizes before launching kernels

AltStyle によって変換されたページ (->オリジナル) /