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
📝 WalkthroughWalkthroughChangesSorted-padding gradient corrections
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related issues
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.
| for (int ii = 0; ii < two_embed.size(); ++ii) { | ||
| std::vector<double> plus = two_embed; | ||
| std::vector<double> minus = two_embed; | ||
| plus[ii] += step; | ||
| minus[ii] -= step; | ||
| const double expected = | ||
| (forward_projection_cpu(plus) - forward_projection_cpu(minus)) / | ||
| (2.0 * step); | ||
| EXPECT_NEAR(actual[ii], expected, 1e-9); |
| 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)); | ||
| } |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5904 +/- ##
==========================================
- Coverage 78.89% 78.64% -0.25%
==========================================
Files 1054 1054
Lines 121774 121820 +46
Branches 4408 4412 +4
==========================================
- Hits 96076 95810 -266
- Misses 24121 24429 +308
- Partials 1577 1581 +4 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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
Bug Fixes
Tests