fix(pd): handle border scalars and CPU tensors - #5832
Conversation
Store nlocal and nghost in one-element host tensors even when there are no swaps, and select local forward/backward copy primitives from the actual Paddle tensor place. Add direct custom-op regressions because existing Paddle model tests did not exercise nswap == 0 or CPU self-swaps in CUDA-enabled builds. 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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe Paddle border operation now stages scalar control values in single-element tensors and selects local-copy primitives from the destination tensor place. New tests cover zero swaps and CPU self-copy behavior in forward and backward paths. ChangesPaddle border operation fixes
Estimated code review effort: 3 (Moderate) | ~20 minutes 🚥 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.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/tests/pd/test_border_op.py (1)
44-81: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winIncorrect gradient expectation in test.
The backward pass of
border_oproutes the gradient from the ghost slots back to their owner's local slot. In this test, index 2 (the ghost slot) receives its value from index 1 (the owner).
Duringresult.sum().backward(), the gradient for all elements ofresultis1.0. The backward operation accumulates the ghost gradient into the local owner usingindex_add_. Therefore, the gradient at index 1 should be1.0 + 1.0 = 2.0, while indices 0 and 2 retain their original gradients of1.0.The assertion expects
np.ones([3, 2]), which would incorrectly mean all gradients are1.0, causing the test to fail.💚 Proposed fix
- np.testing.assert_array_equal(g1_leaf.grad.numpy(), np.ones([3, 2])) + expected_grad = np.ones([3, 2]) + expected_grad[1] = 2.0 + np.testing.assert_array_equal(g1_leaf.grad.numpy(), expected_grad)🤖 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/tests/pd/test_border_op.py` around lines 44 - 81, Update the gradient assertion in test_border_op_self_copy_uses_cpu_place to expect accumulated gradients: index 1 should contain 2.0 in both columns, while indices 0 and 2 remain 1.0. Keep the forward result assertion and backward invocation 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.
Inline comments:
In `@source/op/pd/comm.cc`:
- Around line 15-37: Wrap the definition of copy_local_tensor_data in `#if`
defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) and the matching `#endif`,
aligning it with its guarded callers so CPU-only builds do not compile
references to gpuMemcpy or gpuMemcpyDeviceToDevice.
---
Outside diff comments:
In `@source/tests/pd/test_border_op.py`:
- Around line 44-81: Update the gradient assertion in
test_border_op_self_copy_uses_cpu_place to expect accumulated gradients: index 1
should contain 2.0 in both columns, while indices 0 and 2 remain 1.0. Keep the
forward result assertion and backward invocation unchanged.
🪄 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
Run ID: 17d433ad-098c-4824-b749-a97e187fecbb
📒 Files selected for processing (2)
source/op/pd/comm.ccsource/tests/pd/test_border_op.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5832 +/- ##
==========================================
+ Coverage 78.57% 79.23% +0.66%
==========================================
Files 1049 1073 +24
Lines 120659 125299 +4640
Branches 4349 4569 +220
==========================================
+ Hits 94807 99286 +4479
- Misses 24288 24371 +83
- Partials 1564 1642 +78 ☔ 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.
Both fixes are correct and this is a well-scoped port — pt/comm.cc does not share either bug (it already reads the scalars with .item() and dispatches the self-copy on is_cuda()), so this is pd catching up rather than an incomplete cross-backend fix. The {1} scalar shape and the phi::is_gpu_place-based copy_local_tensor_data are both right. One coverage note is inline; one out-of-scope parity gap for a follow-up:
pd backward lacks pt's ghost-row gradient zeroing (pre-existing, not this PR's scope). In Border_backward_t, d_local_g1_tensor is initialized as a copy of the incoming gradient (
deepmd-kit/source/op/pd/comm.cc
Lines 262 to 266 in 00eccf4
index_adds received gradient into owner rows — nothing zeroes the ghost range. pt/comm.cc does this explicitly (d_local_g1_tensor.slice(0, nlocal, ntotal).zero_()) because otherwise the ghost rows retain the raw incoming dL/dg_out[ghost], which is spurious whenever an exchanged feature is consumed as a genuine per-node leaf downstream (attention node embeddings in dpa2/dpa3, spin). pd and pt have diverged here — worth a separate follow-up (and while there, double-check the plain-CPU-build writeback path of the computed d_local_g1).
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.
The paddle.set_device("cpu") call is the right change and it fixes the half of my point that was fixable in this file -- on a paddlepaddle-gpu build the control and data tensors now land on CPU, so is_gpu_place() returns false and the branch matches the test's name.
The other half turns out to be worse than I described, and I got a detail wrong myself, so let me set it out properly. Inline below.
Approving anyway: the scalar handling and the CPU copy path read correctly, the change is small and self-contained, and the coverage problem is an infrastructure fact rather than something wrong with this diff. But I would not treat the green checks as evidence that this code works.
| """A CUDA-enabled operator must not use a GPU copy for CPU tensors.""" | ||
| # CUDA Paddle builds otherwise create tensors on the default GPU, which | ||
| # would leave the operator's CPU copy branch untested. | ||
| paddle.set_device("cpu") |
There was a problem hiding this comment.
This line does what I asked, but the job it is written for does not exist.
The reply says "the CUDA Paddle CI job will execute the operator path". There is no such job. .github/workflows/test_cuda.yml sets DP_ENABLE_PADDLE: "0" at workflow level, so both GPU jobs build without Paddle, and .github/workflows/test_python.yml installs Paddle from the cpu nightly index (.../packages/nightly/cpu/paddlepaddle/). Those are the only two places Paddle appears in CI.
Meanwhile copy_local_tensor_data is inside #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM). Putting those together:
- CPU
Test Pythonjob: Paddle present, butGOOGLE_CUDAundefined, so this helper is not compiled and the test exercises the untouched#elsepath. - GPU jobs: helper would compile, but Paddle is disabled, so the custom op is not built and
pytest.importorskipskips the file.
So there is no configuration in which the function this PR changes is both compiled and executed. I checked locally too -- paddle 3.3.1, is_compiled_with_cuda() false, ENABLE_CUSTOMIZED_OP false -- so it is not runnable by hand either.
I owe a correction on my own earlier comment: I wrote "on the CUDA job (the only one that compiles it)", which implied such a job existed and that the only problem was tensor placement. That was wrong. The placement fix was necessary but it cannot be sufficient while no job builds Paddle with CUDA.
Nothing here is yours to fix in this PR, and I am not asking you to. But it does mean the regression cannot fail on unpatched code anywhere today, so it documents the intended behaviour rather than protecting it. Two things that would change that, either of them separate work:
- A CI configuration that builds Paddle with CUDA, which would give this test and the rest of
source/op/pd/real coverage. - Or making the CPU path compile unconditionally -- the
is_gpu_placebranch is only meaningful when a GPU place is possible, but the surrounding function could be built on CPU-only Paddle so the#else memcpyand the placement logic are at least exercised together.
If neither is on the cards soon, a one-line comment in the test saying it is currently unexecuted in CI would stop the next reader from assuming the green tick covers it.
There was a problem hiding this comment.
Agreed on all counts — thank you for the correction about no CUDA Paddle job existing. I added a note to the test declaring that no pipeline currently compiles-and-executes copy_local_tensor_data, so the CPU-branch regression is documented rather than guarded (commit d4843dd). I also filed the CI-coverage gap as #5952, since giving source/op/pd/ real coverage needs a job that builds Paddle with CUDA (out of scope for this PR).
Coding agent: opencode
opencode version: 1.18.9
Model: ustc/deepseek-v4-flash
Reasoning effort: max
No CI job currently builds Paddle with CUDA, so the CPU-branch regression cannot fail in any pipeline. Document that intent in the test. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
2 participants
Fixes #5627
Summary
nlocalandnghostwith one element instead of sizing them bynswapgpuMemcpyonly for actual GPU places and uses hostmemcpyfor CPU or host-pinned placesnswap == 0, CPU self-copy, and the reverse self-swap used by autogradWhy existing tests missed this
The existing Paddle suite did not directly call
border_op. Model-level tests therefore did not construct the two boundary conditions that matter here: a valid no-swap invocation where the atom-count scalars still need storage, and a self-swap using CPU data from an operator compiled with CUDA support. Normal multi-rank runs also tend to use tensors on the configured accelerator, hiding the mismatch between CUDA-awareness and actual tensor place.The new no-swap test passes empty communication arrays with scalar atom counts, while the self-swap test keeps a real LAMMPS-style pointer-valued send list alive and checks both forward data and backward execution on CPU tensors. In a CUDA-enabled CI build, the historical code would route that CPU pointer through device-to-device
gpuMemcpy.Validation
source/op/pd/setup.pypytest source/tests/pd/test_border_op.py -q(2 passed, including backward)GOOGLE_CUDA + USE_MPIbranch with Paddle, CUDA 12.4, and MPI headers usingmpicxx -fsyntax-onlyruff format .ruff check .clang-format --dry-run --Werror source/op/pd/comm.ccgit diff --checkThe local Paddle wheel is CPU-only, so runtime execution of the CUDA-enabled custom op is left to CUDA CI; the CUDA/MPI branch was still compiled locally, and the new CPU-place self-swap test is designed to run unchanged in that build.
Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
Bug Fixes
Tests