deepmodeling / deepmodeling/deepmd-kit

[BUG] deepmd/kk can bind Torch to a different GPU than Kokkos

Open
#5,994 0 comments 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.