fix(tabulate): correct sorted padding two-embed gradients - #5904
Conversation
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 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 (2)
📝 WalkthroughWalkthroughChangesSorted-padding gradient corrections
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
source/lib/tests/test_tabulate_se_a.cc (2)
881-888: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse an unsigned index and assert matching extents in
dot.
iiisintwhilelhs.size()issize_type, so-Wsign-comparefires here and in the test loops at Lines 1008, 1019, 1031, 1049, 1060, and 1070 (a-Werrortest target would fail to build). Indexingrhswithlhs's extent is also unguarded.♻️ Proposed fix
static double dot(const std::vector<double>& lhs, const std::vector<double>& rhs) { + EXPECT_EQ(lhs.size(), rhs.size()); double result = 0.0; - for (int ii = 0; ii < lhs.size(); ++ii) { + for (std::size_t ii = 0; ii < lhs.size(); ++ii) { result += lhs[ii] * rhs[ii]; } return result; }The test loops can likewise use
std::size_t ii.🤖 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 `@source/lib/tests/test_tabulate_se_a.cc` around lines 881 - 888, Update dot to use std::size_t for its loop index and assert that lhs and rhs have matching extents before indexing. Apply the same unsigned index type to the test loops identified at lines 1008, 1019, 1031, 1049, 1060, and 1070, preserving their existing loop behavior.
1045-1058: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd GPU forward parity to this fixture.
two_embed_gradient_matches_forward_gpureads GPU gradients while the finite-difference reference comes fromforward_projection_cpu, so a CPU/GPU forward mismatch reports as a backward failure. Add an assertion comparingtabulate_fusion_se_a_gpuoutput with the CPU forward for thislast_layer_size == 1configuration before using the GPU gradient in other wrapper coverage.🤖 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 `@source/lib/tests/test_tabulate_se_a.cc` around lines 1045 - 1058, Add a CPU/GPU forward-output parity assertion in TestTabulateSeASortedPaddingTwoEmbed’s two_embed_gradient_matches_forward_gpu test, comparing tabulate_fusion_se_a_gpu with forward_projection_cpu for the last_layer_size == 1 configuration before validating gradients. Keep the existing finite-difference gradient checks unchanged.
🤖 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.
Nitpick comments:
In `@source/lib/tests/test_tabulate_se_a.cc`:
- Around line 881-888: Update dot to use std::size_t for its loop index and
assert that lhs and rhs have matching extents before indexing. Apply the same
unsigned index type to the test loops identified at lines 1008, 1019, 1031,
1049, 1060, and 1070, preserving their existing loop behavior.
- Around line 1045-1058: Add a CPU/GPU forward-output parity assertion in
TestTabulateSeASortedPaddingTwoEmbed’s two_embed_gradient_matches_forward_gpu
test, comparing tabulate_fusion_se_a_gpu with forward_projection_cpu for the
last_layer_size == 1 configuration before validating gradients. Keep the
existing finite-difference gradient checks unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 4544f118-1fa2-4780-8a29-b5c4895c7753
📒 Files selected for processing (3)
source/lib/src/gpu/tabulate.cusource/lib/src/tabulate.ccsource/lib/tests/test_tabulate_se_a.cc
There was a problem hiding this comment.
Pull request overview
Fixes the SE-A tabulation attention (two_embed) gradient contract for sorted padding tails so that only the first padding sentinel contributes to the folded forward output, with its gradient scaled by the tail length, and ensures the GPU path clears unused tail gradients. Adds targeted CPU/GPU finite-difference and grad-grad regression tests for the corrected behavior.
Changes:
- Update CPU backward to write a repeat-scaled
two_embedgradient only at the first sorted-padding sentinel (tail entries remain zero). - Update GPU backward to mirror the same folded-tail gradient ownership and explicitly zero
dy_dtwooutputs to avoid leaving uninitialized tail values. - Add CPU/GPU finite-difference and double-backward regression tests for nonuniform tail cotangents and a tail longer than the 4-warp tile.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| source/lib/tests/test_tabulate_se_a.cc | Adds CPU/GPU finite-difference + grad-grad regression tests for sorted-padding two_embed gradients. |
| source/lib/src/tabulate.cc | Fixes CPU two_embed backward to match folded sorted-padding forward contract (only first sentinel gets scaled grad). |
| source/lib/src/gpu/tabulate.cu | Fixes GPU two_embed backward for sorted padding and clears dy_dtwo to keep unused tail gradients at zero. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5904 +/- ##
==========================================
- Coverage 79.03% 78.96% -0.07%
==========================================
Files 1055 1069 +14
Lines 122233 124120 +1887
Branches 4401 4527 +126
==========================================
+ Hits 96607 98017 +1410
- Misses 24061 24483 +422
- Partials 1565 1620 +55 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/lib/tests/test_tabulate_se_a.cc (1)
986-989: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winKeep the
emfixture length consistent withnnei.
nnei * 4is 24, but this initializer contains 28 values. The final four values are ignored, leaving an accidental seventh neighbor in the test data. Remove them or updatenneito the intended size.🤖 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 `@source/lib/tests/test_tabulate_se_a.cc` around lines 986 - 989, Make the em fixture in the test consistent with nnei: nnei * 4 is 24, so remove the final four initializer values, unless the test intentionally requires seven neighbors and nnei should be updated accordingly.
🤖 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.
Outside diff comments:
In `@source/lib/tests/test_tabulate_se_a.cc`:
- Around line 986-989: Make the em fixture in the test consistent with nnei:
nnei * 4 is 24, so remove the final four initializer values, unless the test
intentionally requires seven neighbors and nnei should be updated accordingly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: edd6ee5c-ad93-4a5e-8d85-7201952a53e7
📒 Files selected for processing (2)
source/lib/src/gpu/tabulate.cusource/lib/tests/test_tabulate_se_a.cc
🚧 Files skipped from review as they are similar to previous changes (1)
- source/lib/src/gpu/tabulate.cu
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
source/lib/src/gpu/tabulate.cu:1101
dy_dtwois memset to 0 whenevertwo_embed != nullptr, but the comment and the rationale only apply to sorted-padding mode. Whenis_sorted == false, the kernel writes everydy_dtwoentry, so this extra memset is unnecessary work on a hot path (backward) and can become a noticeable overhead for largenloc * nnei * last_layer_size. Consider guarding the memset withis_sortedso it only runs when there may be an unwritten padding tail.
if (two_embed != nullptr) {
// The sorted-padding fast path writes only the first sentinel. Explicitly
// clear the unused tail because framework output buffers are uninitialized.
DPErrcheck(
gpuMemset(dy_dtwo, 0, sizeof(FPTYPE) * nloc * nnei * last_layer_size));
}
source/lib/tests/test_tabulate_se_a.cc:1158
- This test is named
two_embed_gradient_matches_forward_gpu, but its finite-difference reference usesforward_projection_cpu(...). That makes the intent ambiguous and it won’t catch GPU-only forward/backward inconsistencies. Either rename the test to reflect the CPU-forward reference, or add a GPU forward-projection helper and use it here.
TEST_F(TestTabulateSeASortedPaddingTwoEmbed,
two_embed_gradient_matches_forward_gpu) {
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Verified this by building tabulate.cc at both the base and the head and running finite-difference drivers against the real deepmd:: entry points, rather than by reading the diff.
The new contract is exactly the derivative of the folded forward. Since the fold multiplies the sentinel row's contribution by (nnei - jj) and uses only that row's two-embedding value before breaking, the output depends on two_embed[jj] with that multiplicity and does not depend on the later padding rows at all. Analytic and central-difference agree at the sentinel (1.0) and on the tail (0.0), where the old code produced 0.2 in every slot -- wrong in magnitude at the sentinel and wrong in support on four rows the forward never reads. The old tail sum happened to come out right, which is presumably why this survived so long. dy_dem and dy_dem_x are byte-identical to base.
The part I found most convincing is that grad_grad needed no functional change, because it was already the JVP of the new contract. That means grad and grad_grad were implementing mutually inconsistent contracts before this PR, and second derivatives of compressed se_atten were wrong accordingly. With a non-uniform cotangent, the finite difference of the backward gives 6.2 at base against the grad-grad kernel's 1.0; at head both are 1.0. So this is not only a gradient fix, it restores self-consistency between the two stages -- and the two comment-only hunks there correctly describe what the surrounding code already did.
The added gpuMemset is genuinely required rather than defensive: the tail used to be fully overwritten by the copy loop, so a partial write now needs it, and the TF op hands in allocate_output memory that is uninitialized. The CPU side already had the equivalent memset, so the "retains the zero initialized above" comment is accurate. grad_grad has a single output that is fully zeroed on both sides, so there is no missing counterpart there.
On the tests: these fail pre-fix, which I checked by running them rather than inferring it. Eleven assertions fail at the base and none at the head, with discrepancies of 0.2 to 5.2 against a 1e-9 tolerance. The oracles are independent of the implementation -- central difference of the forward, plus a grad-to-grad-grad consistency relation with a deliberately non-uniform cotangent so that any tail leakage shows up -- rather than restated implementation constants. Seeding the GPU output buffer with 7.0 to make the memset load-bearing, and choosing nnei = 6 so the tail exceeds the four-warp tile and covers entries no warp visits, are both good touches.
Worth recording why nothing caught this earlier: dy_dtwo has been produced since 2023 and no test asserted its values at all until three days ago. The one that did asserted CPU/GPU parity, which cannot catch a bug both sides share. This PR is the first time the output has been checked against an oracle independent of the kernels.
Avoid clearing fully overwritten unsorted GPU gradients, correct the test fixture extent, and add explicit CPU/GPU forward parity coverage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Addressed the remaining review findings in 82ab1e5:
Validation: Ruff and formatting checks passed; the CUDA sources compiled; all five targeted CPU/GPU tabulation tests passed. Coding agent: Codex |
Summary
two_embedValidation
source/build/lib/tests/runUnitTests_lib --gtest_filter='TestTabulateSeA.*:TestTabulateSeASortedPaddingTwoEmbed.*'deepmd_op_cudaandrunUnitTests_libruff format .ruff check .clang-format --dry-run --Werroron changed C++/CUDA filesCloses #5891
Note: #5844 modifies the same GPU kernel for the shared-breakpoint fix, so the later-merging branch may need a small conflict resolution; the gradient contract fixed here is independent.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit