fix(pt): select device before loading torch ops

[njzjz-bot]

Summary

• Add a shared PyTorch backend helper to choose the rank-local CUDA/HIP device before PyTorch CUDA queries or custom-op loading can create a default context.
• Move the PT C++ model init paths (DeepPotPT, DeepSpinPT, DeepTensorPT, DeepPotPTExpt, DeepSpinPTExpt) to call this helper before deepmd::load_op_library().
• This should avoid each MPI rank leaving a small unused context on GPU 0 while preserving CPU fallback behavior.

Verification

• git diff --check HEAD~1..HEAD
• Static check: verified preselect_torch_device precedes deepmd::load_op_library() in all touched PT init paths.

Fixes #4171

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)

Summary by CodeRabbit

• Bug Fixes
  • Improved device selection so models start on the correct GPU more reliably.
  • Reduced cases where GPU-enabled environments could accidentally fall back to the wrong device.
  • CPU fallback behavior remains intact when no supported GPU is available.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: a4338421-7b9b-450b-867a-37dfe436ba0b

📥 Commits

Reviewing files that changed from the base of the PR and between 73de44b and 050f97e.

📒 Files selected for processing (6)

• source/api_cc/include/commonPT.h
• source/api_cc/src/DeepPotPT.cc
• source/api_cc/src/DeepPotPTExpt.cc
• source/api_cc/src/DeepSpinPT.cc
• source/api_cc/src/DeepSpinPTExpt.cc
• source/api_cc/src/DeepTensorPT.cc

---

📝 Walkthrough

Walkthrough

A new inline helper deepmd::preselect_torch_device is added to commonPT.h. It centralizes rank-local GPU selection (via CUDA/ROCm or Torch APIs) and sets gpu_id/gpu_enabled. All five PyTorch model init functions (DeepPotPT, DeepPotPTExpt, DeepSpinPT, DeepSpinPTExpt, DeepTensorPT) replace their duplicated inline GPU-selection logic with a single call to this helper.

Changes

GPU Preselection Refactor

Layer / File(s)	Summary
preselect_torch_device helper source/api_cc/include/commonPT.h	Adds #include "device.h" and defines deepmd::preselect_torch_device, selecting a rank-local GPU via DPGetDeviceCount/DPSetDevice (under GOOGLE_CUDA/TENSORFLOW_USE_ROCM) or torch::cuda::device_count(), and assigning gpu_enabled from torch::cuda::is_available().
Adopt helper in all init paths source/api_cc/src/DeepPotPT.cc, source/api_cc/src/DeepPotPTExpt.cc, source/api_cc/src/DeepSpinPT.cc, source/api_cc/src/DeepSpinPTExpt.cc, source/api_cc/src/DeepTensorPT.cc	Each ::init replaces its inline torch::cuda::device_count() / torch::cuda::is_available() / DPSetDevice block with preselect_torch_device(gpu_rank, gpu_id, gpu_enabled). DeepTensorPT.cc additionally adds the commonPT.h include. Subsequent CUDA-vs-CPU device construction and logging remain unchanged.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 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 the main change: preselecting the device before loading Torch ops.
Linked Issues check	✅ Passed	The changes match issue #4171 by selecting the rank-local GPU before op-library loading in all touched PT init paths.
Out of Scope Changes check	✅ Passed	No obvious unrelated changes are present beyond the device-preselection refactor needed for the linked bug fix.
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

• X
• Mastodon
• Reddit
• LinkedIn

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

[codecov]

Codecov Report

❌ Patch coverage is 74.28571% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.36%. Comparing base (cc908a8) to head (736b6db).
⚠️ Report is 44 commits behind head on master.

Files with missing lines	Patch %	Lines
source/api_cc/src/DeepPotPTExpt.cc	33.33%	3 Missing and 1 partial ⚠️
source/api_cc/tests/test_neighbor_list_data.cc	60.00%	1 Missing and 1 partial ⚠️
source/api_cc/include/commonPT.h	83.33%	1 Missing ⚠️
...ource/api_cc/tests/test_deeppot_dpa_ptexpt_spin.cc	85.71%	1 Missing ⚠️
source/api_cc/tests/test_deeppot_ptexpt.cc	85.71%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5612      +/-   ##
==========================================
+ Coverage   79.03%   79.36%   +0.33%     
==========================================
  Files        1055     1085      +30     
  Lines      122233   126419    +4186     
  Branches     4401     4597     +196     
==========================================
+ Hits        96607   100334    +3727     
- Misses      24061    24434     +373     
- Partials     1565     1651      +86     

☔ 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.

[njzjz]

Need a real test on a machine with multiple GPUs.

[njzjz]

Possible reviewers based on changed lines, exact file history, and exact-file review history:

• @wanghan-iapcm — 18 commits on changed files; 121 reviews on exact changed files (source/api_cc/include/commonPT.h, source/api_cc/src/DeepPotPT.cc, source/api_cc/src/DeepPotPTExpt.cc, source/api_cc/src/DeepSpinPT.cc, source/api_cc/src/DeepSpinPTExpt.cc, source/api_cc/src/DeepTensorPT.cc).
• @iProzd — 29 reviews on exact changed files (source/api_cc/include/commonPT.h, source/api_cc/src/DeepPotPT.cc, source/api_cc/src/DeepPotPTExpt.cc, source/api_cc/src/DeepSpinPT.cc, source/api_cc/src/DeepSpinPTExpt.cc, source/api_cc/src/DeepTensorPT.cc).

No review request was made automatically.

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

[njzjz-bot]

Reviewed draft head 5532e79536ee314aeea2a90a79f241a16afbe51d with three independent full subagent reviews. Two P2 initialization issues remain inline: raw runtime preselection is keyed to the DeePMD toolkit rather than LibTorch's GPU capability, and the exportable initializers mutate device state before rejecting unsupported in-memory content. The intended selection-before-op-load ordering is otherwise consistent across the five PT initialization paths. Several C++ and wheel checks are currently failing/cancelled and need investigation after the fixes. 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