fix(pt): select device before loading torch ops

source/api_cc/CMakeLists.txt

@@ -111,14 +111,18 @@ 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_CUDA}>)
 
   add_library(
     deepmd_backend_ptexpt SHARED
     src/DeepPotPTExpt.cc src/DeepPotPTExptPlugin.cc src/DeepSpinPTExpt.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_CUDA}>)
 endif()
 
 if(ENABLE_JAX)

source/api_cc/include/commonPT.h

@@ -14,10 +14,52 @@
 #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;
+#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

@@ -227,31 +227,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

@@ -64,29 +64,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,6 +9,7 @@
 #include <sstream>
 
 #include "common.h"
+#include "commonPT.h"
 #include "device.h"
 #include "errors.h"
 
@@ -74,22 +75,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_CUDA}>)
   # 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

@@ -874,6 +874,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,17 @@ 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);
+}
+
 TEST(TestNeighborListData, CompactCanonicalGraphDropsMaskedGuards) {
   GraphTensorPack graph;
   graph.atype = torch::tensor({0}, torch::kInt64);