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:
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:
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
Summary
deepmd/kkanddpa4spin/kkcan 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/masterat8cfd46e37448.Trigger
On a node with at least two visible GPUs, launch one MPI rank while explicitly choosing Kokkos device 1, for example:
Kokkos allocates its views on GPU 1. DeePMD passes node-local MPI rank 0 into model initialization (
pair_base.cpp:39andpair_dpa4spin.cpp:234), and the PT-experimental APIs independently choose: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.DeepPotPTExptandNativeSpinPTExptwrap those pointers withTensorOptions().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
/kkpaths.Coding agent: Codex
Codex version: codex-cli 0.149.0
Model: gpt-5.6-sol
Reasoning effort: xhigh