deepmodeling / deepmodeling/deepmd-kit
[BUG] deepmd/kk can bind Torch to a different GPU than Kokkos
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
### Summary
`deepmd/kk` and `dpa4spin/kk` can run Kokkos on one GPU while loading the Torch model on another. Raw Kokkos pointers are then wrapped as tensors tagged with the independently selected Torch device, causing invalid or illegal device memory access.
This is present on `origin/master` at `8cfd46e37448`.
### Trigger
On a node with at least two visible GPUs, launch one MPI rank while explicitly choosing Kokkos device 1, for example:
```text
lmp -k on g 1 d 1 ...
```
Kokkos allocates its views on GPU 1. DeePMD passes node-local MPI rank 0 into model initialization (`pair_base.cpp:39` and `pair_dpa4spin.cpp:234`), and the PT-experimental APIs independently choose:
```cpp
gpu_id = gpu_rank % torch::cuda::device_count();
```
so the model is loaded on GPU 0 (`DeepPotPTExpt.cc:167-181`, `NativeSpinPTExpt.cc:314-326`).
The device path later passes raw Kokkos pointers from `pair_deepmd_kokkos.cpp` / `pair_dpa4spin_kokkos.cpp`. `DeepPotPTExpt` and `NativeSpinPTExpt` wrap those pointers with `TensorOptions().device(cuda:gpu_id)`, which labels GPU-1 allocations as GPU-0 tensors.
### Expected behavior
The model and all raw-pointer tensor wrappers must use Kokkos's actual device. The pair style should pass the selected Kokkos device to model initialization, or initialization should query and validate the active device and fail clearly on a mismatch.
A two-GPU regression should exercise an explicit nonzero Kokkos device for both ordinary and native-spin `/kk` paths.
---
Coding agent: Codex
Codex version: codex-cli 0.149.0
Model: gpt-5.6-sol
Reasoning effort: xhigh
Contributor guide
Assessment
This issue has not been assessed yet.