diff --git a/source/api_cc/CMakeLists.txt b/source/api_cc/CMakeLists.txt index fde1a99c42..c89e399c9e 100644 --- a/source/api_cc/CMakeLists.txt +++ b/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=$) 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=$) endif() if(ENABLE_JAX) diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 806caab494..3a5aae3d20 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -14,10 +14,52 @@ #include #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. * diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index 64ea871d97..1b1c9c6e29 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/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; } diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index c1049cdfea..69fdcaf682 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/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; diff --git a/source/api_cc/src/DeepSpinPT.cc b/source/api_cc/src/DeepSpinPT.cc index 7537ca2588..f7a9be1c5d 100644 --- a/source/api_cc/src/DeepSpinPT.cc +++ b/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; } diff --git a/source/api_cc/src/DeepSpinPTExpt.cc b/source/api_cc/src/DeepSpinPTExpt.cc index c908be18c7..ff594e73b9 100644 --- a/source/api_cc/src/DeepSpinPTExpt.cc +++ b/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; diff --git a/source/api_cc/src/DeepTensorPT.cc b/source/api_cc/src/DeepTensorPT.cc index 8ed45d6893..a795cb665b 100644 --- a/source/api_cc/src/DeepTensorPT.cc +++ b/source/api_cc/src/DeepTensorPT.cc @@ -9,6 +9,7 @@ #include #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; } diff --git a/source/api_cc/tests/CMakeLists.txt b/source/api_cc/tests/CMakeLists.txt index 1d1b817957..06ee8c4107 100644 --- a/source/api_cc/tests/CMakeLists.txt +++ b/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=$) # Link torch so __has_include() succeeds and # BUILD_PT_EXPT is set for the test binary; otherwise pt_expt tests all # GTEST_SKIP() with "PyTorch support is not enabled". diff --git a/source/api_cc/tests/test_deeppot_dpa_ptexpt_spin.cc b/source/api_cc/tests/test_deeppot_dpa_ptexpt_spin.cc index dc43cb3719..22c3f4fddf 100644 --- a/source/api_cc/tests/test_deeppot_dpa_ptexpt_spin.cc +++ b/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 // ============================================================================ diff --git a/source/api_cc/tests/test_deeppot_ptexpt.cc b/source/api_cc/tests/test_deeppot_ptexpt.cc index b4f12d1453..71a57be648 100644 --- a/source/api_cc/tests/test_deeppot_ptexpt.cc +++ b/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."; diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index ef83eae28b..a2fd92c4e7 100644 --- a/source/api_cc/tests/test_neighbor_list_data.cc +++ b/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);