fix(pt): select device before loading torch ops

source/CMakeLists.txt

@@ -394,16 +394,7 @@ if(ENABLE_PYTORCH AND NOT DEEPMD_C_ROOT)
     endif()
   endif()
   find_package(Torch REQUIRED)
-  # TorchConfig populates TORCH_CUDA_LIBRARIES only when PyTorch itself was
-  # built with CUDA. The fused graph-lower operators in op/pt include ATen CUDA
-  # headers and link libtorch_cuda, so they can be compiled only against a
-  # CUDA-enabled PyTorch: a CPU-only wheel ships the ATen/c10 CUDA headers but
-  # neither the generated c10/cuda/impl/cuda_cmake_macros.h nor libtorch_cuda.
-  if(TORCH_CUDA_LIBRARIES)
-    set(DEEPMD_TORCH_HAS_CUDA TRUE)
-  else()
-    set(DEEPMD_TORCH_HAS_CUDA FALSE)
-  endif()
+  include(cmake/torch_gpu_support.cmake)
   # PyTorch 2.14 raised the minimum C++ standard from C++17 to C++20.
   if(NOT Torch_VERSION VERSION_LESS "2.14.0")
     set_if_higher(CMAKE_CXX_STANDARD 20)

source/api_cc/CMakeLists.txt

@@ -112,15 +112,19 @@ if(ENABLE_PYTORCH AND "${OP_CXX_ABI_PT}" EQUAL "${OP_CXX_ABI}")
                                        src/DeepSpinPT.cc src/DeepTensorPT.cc)
   deepmd_configure_backend_plugin(deepmd_backend_pt)
   target_link_libraries(deepmd_backend_pt PRIVATE "${TORCH_LIBRARIES}")
-  target_compile_definitions(deepmd_backend_pt PRIVATE BUILD_PYTORCH)
+  target_compile_definitions(
+    deepmd_backend_pt
+    PRIVATE BUILD_PYTORCH DEEPMD_TORCH_HAS_GPU=$<BOOL:${DEEPMD_TORCH_HAS_GPU}>)
 
   add_library(
     deepmd_backend_ptexpt SHARED
     src/DeepPotPTExpt.cc src/DeepPotPTExptPlugin.cc src/DeepSpinPTExpt.cc
     src/NativeSpinPTExpt.cc)
   deepmd_configure_backend_plugin(deepmd_backend_ptexpt)
   target_link_libraries(deepmd_backend_ptexpt PRIVATE "${TORCH_LIBRARIES}")
-  target_compile_definitions(deepmd_backend_ptexpt PRIVATE BUILD_PYTORCH)
+  target_compile_definitions(
+    deepmd_backend_ptexpt
+    PRIVATE BUILD_PYTORCH DEEPMD_TORCH_HAS_GPU=$<BOOL:${DEEPMD_TORCH_HAS_GPU}>)
 endif()
 
 if(ENABLE_JAX)

source/api_cc/include/commonPT.h

@@ -2,6 +2,7 @@
 #pragma once
 
 #ifdef BUILD_PYTORCH
+#include <c10/core/impl/VirtualGuardImpl.h>
 #include <torch/torch.h>
 #include <torch/version.h>
 
@@ -17,10 +18,61 @@
 #include <vector>
 
 #include "common.h"
+#include "device.h"
 #include "neighbor_list.h"
 
 namespace deepmd {
 
+// This definition is supplied by CMake for DeePMD's PyTorch backend targets.
+// Default to false for external consumers that include this internal helper.
+#ifndef DEEPMD_TORCH_HAS_GPU
+#define DEEPMD_TORCH_HAS_GPU 0
+#endif
+
+/** @brief Whether the linked LibTorch build provides a GPU runtime. */
+inline constexpr bool torch_has_gpu_support() {
+  return DEEPMD_TORCH_HAS_GPU != 0;
+}
+
+/**
+ * @brief Select the per-rank GPU before PyTorch can create a default context.
+ *
+ * Some PyTorch/CUDA queries and the torch custom-op library loader may create a
+ * CUDA/HIP context on the current runtime device. In MPI jobs the runtime
+ * default is usually GPU 0, so selecting the rank-local GPU first avoids every
+ * rank leaving a small, unused context on GPU 0.
+ *
+ * @param[in] gpu_rank Rank-local GPU index passed by the caller.
+ * @param[out] gpu_id Visible GPU selected for this rank.
+ * @param[out] gpu_enabled Whether PyTorch reports CUDA/HIP availability.
+ */
+inline void preselect_torch_device(const int& gpu_rank,
+                                   int& gpu_id,
+                                   bool& gpu_enabled) {
+#if (defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)) && \
+    DEEPMD_TORCH_HAS_GPU
+  int gpu_num = 0;
+  DPGetDeviceCount(gpu_num);
+  gpu_id = (gpu_num > 0) ? (gpu_rank % gpu_num) : 0;
+  if (gpu_num > 0) {
+    DPErrcheck(DPSetDevice(gpu_id));
+  }
+#else
+  int gpu_num = torch::cuda::device_count();
+  gpu_id = (gpu_num > 0) ? (gpu_rank % gpu_num) : 0;
+#if DEEPMD_TORCH_HAS_GPU
+  if (gpu_num > 0) {
+    // LibTorch can provide CUDA/HIP even when DeePMD's toolkit is disabled.
+    // Dispatch through its device registry to avoid toolkit headers and keep
+    // the selected device current after returning (unlike an RAII guard).
+    c10::impl::VirtualGuardImpl(c10::DeviceType::CUDA)
+        .setDevice(c10::Device(c10::DeviceType::CUDA, gpu_id));
+  }
+#endif
+#endif  // DeePMD toolkit and LibTorch GPU support
+  gpu_enabled = torch::cuda::is_available();
+}
+
 /**
  * @brief Build comm_dict tensors from sendlist/sendnum/recvnum buffers.
  *

source/api_cc/src/DeepPotPT.cc

@@ -53,18 +53,13 @@ void DeepPotPT::init(const std::string& model,
               << std::endl;
     return;
   }
+  preselect_torch_device(gpu_rank, gpu_id, gpu_enabled);
   deepmd::load_op_library(deepmd::DPBackend::PyTorch);
-  int gpu_num = torch::cuda::device_count();
-  gpu_id = (gpu_num > 0) ? (gpu_rank % gpu_num) : 0;
-  gpu_enabled = torch::cuda::is_available();
   torch::Device device(torch::kCUDA, gpu_id);
   if (!gpu_enabled) {
     device = torch::Device(torch::kCPU);
     std::cout << "load model from: " << model << " to cpu " << std::endl;
   } else {
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    DPErrcheck(DPSetDevice(gpu_id));
-#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     std::cout << "load model from: " << model << " to gpu " << gpu_id
               << std::endl;
   }

source/api_cc/src/DeepPotPTExpt.cc

@@ -151,31 +151,26 @@ void DeepPotPTExpt::init(const std::string& model,
     return;
   }
 
-  // Load libdeepmd_op_pt.so so its TORCH_LIBRARY_FRAGMENT entries
-  // (deepmd::*, deepmd_export::*) are visible to torch's dispatcher
-  // before the AOTI module loads.  Without this, multi-rank message-passing
-  // .pt2 archives fail at pair_style time with
-  // ``Could not find schema for deepmd_export::border_op``.
-  deepmd::load_op_library(deepmd::DPBackend::PyTorchExportable);
-
   if (!file_content.empty()) {
     throw deepmd::deepmd_exception(
         "In-memory file_content loading is not supported for .pt2 models. "
         "Please provide a file path instead.");
   }
 
-  int gpu_num = torch::cuda::device_count();
-  gpu_id = (gpu_num > 0) ? (gpu_rank % gpu_num) : 0;
-  gpu_enabled = torch::cuda::is_available();
+  preselect_torch_device(gpu_rank, gpu_id, gpu_enabled);
+
+  // Load libdeepmd_op_pt.so so its TORCH_LIBRARY_FRAGMENT entries
+  // (deepmd::*, deepmd_export::*) are visible to torch's dispatcher
+  // before the AOTI module loads.  Without this, multi-rank message-passing
+  // .pt2 archives fail at pair_style time with
+  // ``Could not find schema for deepmd_export::border_op``.
+  deepmd::load_op_library(deepmd::DPBackend::PyTorchExportable);
 
   std::string device_str;
   if (!gpu_enabled) {
     device_str = "cpu";
     std::cout << "load model from: " << model << " to cpu" << std::endl;
   } else {
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    DPErrcheck(DPSetDevice(gpu_id));
-#endif
     device_str = "cuda:" + std::to_string(gpu_id);
     std::cout << "load model from: " << model << " to gpu " << gpu_id
               << std::endl;

source/api_cc/src/DeepSpinPT.cc

@@ -52,22 +52,13 @@ void DeepSpinPT::init(const std::string& model,
               << std::endl;
     return;
   }
+  preselect_torch_device(gpu_rank, gpu_id, gpu_enabled);
   deepmd::load_op_library(deepmd::DPBackend::PyTorch);
-  int gpu_num = torch::cuda::device_count();
-  if (gpu_num > 0) {
-    gpu_id = gpu_rank % gpu_num;
-  } else {
-    gpu_id = 0;
-  }
   torch::Device device(torch::kCUDA, gpu_id);
-  gpu_enabled = torch::cuda::is_available();
   if (!gpu_enabled) {
     device = torch::Device(torch::kCPU);
     std::cout << "load model from: " << model << " to cpu " << std::endl;
   } else {
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    DPErrcheck(DPSetDevice(gpu_id));
-#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     std::cout << "load model from: " << model << " to gpu " << gpu_id
               << std::endl;
   }

source/api_cc/src/DeepSpinPTExpt.cc

@@ -68,29 +68,24 @@ void DeepSpinPTExpt::init(const std::string& model,
     return;
   }
 
-  // Load libdeepmd_op_pt.so so deepmd_export::* schemas are visible
-  // to torch's dispatcher before the AOTI module loads.  See
-  // DeepPotPTExpt::init for the full rationale.
-  deepmd::load_op_library(deepmd::DPBackend::PyTorchExportable);
-
   if (!file_content.empty()) {
     throw deepmd::deepmd_exception(
         "In-memory file_content loading is not supported for .pt2 models. "
         "Please provide a file path instead.");
   }
 
-  int gpu_num = torch::cuda::device_count();
-  gpu_id = (gpu_num > 0) ? (gpu_rank % gpu_num) : 0;
-  gpu_enabled = torch::cuda::is_available();
+  preselect_torch_device(gpu_rank, gpu_id, gpu_enabled);
+
+  // Load libdeepmd_op_pt.so so deepmd_export::* schemas are visible
+  // to torch's dispatcher before the AOTI module loads. See
+  // DeepPotPTExpt::init for the full rationale.
+  deepmd::load_op_library(deepmd::DPBackend::PyTorchExportable);
 
   std::string device_str;
   if (!gpu_enabled) {
     device_str = "cpu";
     std::cout << "load model from: " << model << " to cpu" << std::endl;
   } else {
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    DPErrcheck(DPSetDevice(gpu_id));
-#endif
     device_str = "cuda:" + std::to_string(gpu_id);
     std::cout << "load model from: " << model << " to gpu " << gpu_id
               << std::endl;

source/api_cc/src/DeepTensorPT.cc

@@ -9,29 +9,12 @@
 #include <sstream>
 
 #include "common.h"
+#include "commonPT.h"
 #include "device.h"
 #include "errors.h"
 
 using namespace deepmd;
 
-static torch::Tensor createNlistTensor(
-    const std::vector<std::vector<int>>& data) {
-  size_t total_size = 0;
-  for (const auto& row : data) {
-    total_size += row.size();
-  }
-  std::vector<int> flat_data;
-  flat_data.reserve(total_size);
-  for (const auto& row : data) {
-    flat_data.insert(flat_data.end(), row.begin(), row.end());
-  }
-
-  torch::Tensor flat_tensor = torch::tensor(flat_data, torch::kInt32);
-  int nloc = data.size();
-  int nnei = nloc > 0 ? total_size / nloc : 0;
-  return flat_tensor.view({1, nloc, nnei});
-}
-
 void DeepTensorPT::translate_error(std::function<void()> f) {
   try {
     f();
@@ -74,22 +57,13 @@ void DeepTensorPT::init(const std::string& model,
     return;
   }
   name_scope = name_scope_;
+  preselect_torch_device(gpu_rank, gpu_id, gpu_enabled);
   deepmd::load_op_library(deepmd::DPBackend::PyTorch);
-  int gpu_num = torch::cuda::device_count();
-  if (gpu_num > 0) {
-    gpu_id = gpu_rank % gpu_num;
-  } else {
-    gpu_id = 0;
-  }
   torch::Device device(torch::kCUDA, gpu_id);
-  gpu_enabled = torch::cuda::is_available();
   if (!gpu_enabled) {
     device = torch::Device(torch::kCPU);
     std::cout << "load model from: " << model << " to cpu " << std::endl;
   } else {
-#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    DPErrcheck(DPSetDevice(gpu_id));
-#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
     std::cout << "load model from: " << model << " to gpu " << gpu_id
               << std::endl;
   }

source/api_cc/tests/CMakeLists.txt

@@ -15,7 +15,9 @@ if(ENABLE_TENSORFLOW)
     TensorFlow::tensorflow_framework)
 endif()
 if(ENABLE_PYTORCH)
-  target_compile_definitions(runUnitTests_cc PRIVATE BUILD_PYTORCH)
+  target_compile_definitions(
+    runUnitTests_cc
+    PRIVATE BUILD_PYTORCH DEEPMD_TORCH_HAS_GPU=$<BOOL:${DEEPMD_TORCH_HAS_GPU}>)
   # Link torch so __has_include(<torch/csrc/inductor/...>) succeeds and
   # BUILD_PT_EXPT is set for the test binary; otherwise pt_expt tests all
   # GTEST_SKIP() with "PyTorch support is not enabled".

source/api_cc/tests/test_deeppot_dpa_ptexpt_spin.cc

@@ -23,6 +23,21 @@ constexpr const char* kRefPath = "../../tests/infer/deeppot_dpa_spin.expected";
 constexpr const char* kModelPath = "../../tests/infer/deeppot_dpa_spin.pt2";
 }  // namespace
 
+TEST(TestDeepSpinPTExptInit, rejects_in_memory_file_content) {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT_SPIN
+  GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+  deepmd::DeepSpin dp;
+  try {
+    dp.init("unused.pt2", 0, "unsupported in-memory model");
+    FAIL() << "Expected in-memory .pt2 loading to be rejected.";
+  } catch (const deepmd::deepmd_exception& error) {
+    EXPECT_NE(std::string(error.what())
+                  .find("In-memory file_content loading is not supported"),
+              std::string::npos);
+  }
+}
+
 // ============================================================================
 // PBC test fixture
 // ============================================================================

source/api_cc/tests/test_deeppot_ptexpt.cc

@@ -913,6 +913,21 @@ TYPED_TEST(TestInferDeepPotAPtExptNoPbc, cpu_build_nlist_nframes) {
 
 // ========== Parser / metadata coverage tests ==========
 
+TEST(TestDeepPotPTExptParser, rejects_in_memory_file_content) {
+#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
+  GTEST_SKIP() << "Skip because PyTorch support is not enabled.";
+#endif
+  deepmd::DeepPot dp;
+  try {
+    dp.init("unused.pt2", 0, "unsupported in-memory model");
+    FAIL() << "Expected in-memory .pt2 loading to be rejected.";
+  } catch (const deepmd::deepmd_exception& error) {
+    EXPECT_NE(std::string(error.what())
+                  .find("In-memory file_content loading is not supported"),
+              std::string::npos);
+  }
+}
+
 TEST(TestDeepPotPTExptParser, load_nonexistent_file) {
 #if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT
   GTEST_SKIP() << "Skip because PyTorch support is not enabled.";

source/api_cc/tests/test_neighbor_list_data.cc

@@ -139,6 +139,41 @@ TEST(TestNeighborListData, RoundTripWithEmptyRows) {
 }
 
 #ifdef BUILD_PYTORCH
+TEST(TestTorchDeviceSelection, CpuLibTorchUsesTorchFallback) {
+  if (torch_has_gpu_support()) {
+    GTEST_SKIP() << "This regression requires a CPU-only LibTorch build.";
+  }
+  int gpu_id = -1;
+  bool gpu_enabled = true;
+  ASSERT_NO_THROW(preselect_torch_device(/*gpu_rank=*/7, gpu_id, gpu_enabled));
+  EXPECT_EQ(gpu_id, 0);
+  EXPECT_FALSE(gpu_enabled);
+}
+
+#if DEEPMD_TORCH_HAS_GPU
+TEST(TestTorchDeviceSelection, SelectsRankLocalRuntimeDevice) {
+  const int count = torch::cuda::device_count();
+  if (count < 2) {
+    GTEST_SKIP() << "This regression requires two visible GPUs.";
+  }
+  // Exercise the actual current device, including toolkit-disabled builds
+  // where computing gpu_id alone would leave allocations on GPU 0.
+  c10::DeviceGuard restore(c10::Device(c10::DeviceType::CUDA, 0));
+  int gpu_id = -1;
+  bool gpu_enabled = false;
+  preselect_torch_device(count + 1, gpu_id, gpu_enabled);
+  EXPECT_TRUE(gpu_enabled);
+  EXPECT_EQ(gpu_id, 1);
+  EXPECT_EQ(
+      c10::impl::VirtualGuardImpl(c10::DeviceType::CUDA).getDevice().index(),
+      1);
+  EXPECT_EQ(torch::ones({1}, torch::TensorOptions().device(torch::kCUDA))
+                .device()
+                .index(),
+            1);
+}
+#endif
+
 TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) {
 #if TORCH_VERSION_MAJOR < 2 || \
     (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR < 3)

source/cmake/torch_gpu_support.cmake

@@ -0,0 +1,14 @@
+# Query the linked LibTorch build, not DeePMD's selected toolkit: a CPU-only
+# LibTorch may be used in a CUDA/ROCm-enabled DeePMD build. CUDA-only fused ops
+# still require their narrower capability, while device preselection supports
+# both runtimes. ROCm TorchConfig exports torch_hip through Caffe2Targets.
+if(TORCH_CUDA_LIBRARIES)
+  set(DEEPMD_TORCH_HAS_CUDA TRUE)
+else()
+  set(DEEPMD_TORCH_HAS_CUDA FALSE)
+endif()
+if(DEEPMD_TORCH_HAS_CUDA OR TARGET torch_hip)
+  set(DEEPMD_TORCH_HAS_GPU TRUE)
+else()
+  set(DEEPMD_TORCH_HAS_GPU FALSE)
+endif()