diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index 4ccc18542c..c9b388f21e 100644 --- a/source/CMakeLists.txt +++ b/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) diff --git a/source/api_cc/CMakeLists.txt b/source/api_cc/CMakeLists.txt index fe764ab0bf..77bd717241 100644 --- a/source/api_cc/CMakeLists.txt +++ b/source/api_cc/CMakeLists.txt @@ -112,7 +112,9 @@ 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 @@ -120,7 +122,9 @@ if(ENABLE_PYTORCH AND "${OP_CXX_ABI_PT}" EQUAL "${OP_CXX_ABI}") 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=$) endif() if(ENABLE_JAX) diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 873158615c..a87e7c187f 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -2,6 +2,7 @@ #pragma once #ifdef BUILD_PYTORCH +#include #include #include @@ -17,10 +18,61 @@ #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; +#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. * 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 005a18e722..c99ef3aaa1 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/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; 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 2cbb7418e4..aadb03ad57 100644 --- a/source/api_cc/src/DeepSpinPTExpt.cc +++ b/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; diff --git a/source/api_cc/src/DeepTensorPT.cc b/source/api_cc/src/DeepTensorPT.cc index 8ed45d6893..7c95794d8c 100644 --- a/source/api_cc/src/DeepTensorPT.cc +++ b/source/api_cc/src/DeepTensorPT.cc @@ -9,29 +9,12 @@ #include #include "common.h" +#include "commonPT.h" #include "device.h" #include "errors.h" using namespace deepmd; -static torch::Tensor createNlistTensor( - const std::vector>& data) { - size_t total_size = 0; - for (const auto& row : data) { - total_size += row.size(); - } - std::vector 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 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; } diff --git a/source/api_cc/tests/CMakeLists.txt b/source/api_cc/tests/CMakeLists.txt index 1d1b817957..d3994ae75b 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 a8e559d0ee..1add39fe10 100644 --- a/source/api_cc/tests/test_deeppot_ptexpt.cc +++ b/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."; diff --git a/source/api_cc/tests/test_neighbor_list_data.cc b/source/api_cc/tests/test_neighbor_list_data.cc index fa68b99fc2..8470ddac21 100644 --- a/source/api_cc/tests/test_neighbor_list_data.cc +++ b/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) diff --git a/source/cmake/torch_gpu_support.cmake b/source/cmake/torch_gpu_support.cmake new file mode 100644 index 0000000000..6dd3f48b34 --- /dev/null +++ b/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() diff --git a/source/tests/common/test_torch_gpu_support.py b/source/tests/common/test_torch_gpu_support.py new file mode 100644 index 0000000000..37885dca53 --- /dev/null +++ b/source/tests/common/test_torch_gpu_support.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Configure LibTorch capability detection without requiring physical GPUs.""" + +import shutil +import subprocess +from pathlib import ( + Path, +) + +import pytest + + +@pytest.mark.parametrize("toolkit", ["CPU", "CUDA", "ROCM"]) +@pytest.mark.parametrize("torch_runtime", ["CPU", "CUDA", "HIP"]) +def test_torch_gpu_capability(tmp_path: Path, toolkit: str, torch_runtime: str) -> None: + """Only a GPU-enabled LibTorch may enable preselection, including on HIP.""" + cmake = shutil.which("cmake") + if cmake is None: + pytest.skip("CMake is not installed") + module = Path(__file__).resolve().parents[2] / "cmake/torch_gpu_support.cmake" + setup = { + "CPU": "", + "CUDA": "set(TORCH_CUDA_LIBRARIES torch_cuda)", + "HIP": "add_library(torch_hip INTERFACE IMPORTED)", + }[torch_runtime] + (tmp_path / "CMakeLists.txt").write_text( + "cmake_minimum_required(VERSION 3.25)\n" + "project(TorchCapability NONE)\n" + f"set(USE_{toolkit}_TOOLKIT TRUE)\n" + f"{setup}\n" + f'include("{module.as_posix()}")\n' + f'if(NOT "${{DEEPMD_TORCH_HAS_GPU}}" STREQUAL "{str(torch_runtime != "CPU").upper()}")\n' + ' message(FATAL_ERROR "Incorrect generic GPU capability")\n' + "endif()\n" + f'if(NOT "${{DEEPMD_TORCH_HAS_CUDA}}" STREQUAL "{str(torch_runtime == "CUDA").upper()}")\n' + ' message(FATAL_ERROR "Incorrect CUDA-only capability")\n' + "endif()\n" + ) + subprocess.run( + [cmake, "-S", str(tmp_path), "-B", str(tmp_path / "build")], + check=True, + capture_output=True, + text=True, + )