From 1b6efd0fca4d6960e5b1a0a751855efd308b7e7b Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Thu, 25 Jul 2024 08:05:09 +0200 Subject: [PATCH] python312Packages.torch: 2.3.0 -> 2.4.0 Diff: https://github.com/pytorch/pytorch/compare/v2.3.1...v2.4.0 Changelog: https://github.com/pytorch/pytorch/releases/tag/v2.4.0 Co-authored-by: =?UTF-8?q?Danie=CC=88l=20de=20Kok?= Co-authored-by: Connor Baker --- .../python-modules/torch/default.nix | 32 ++++++---- .../torch/fix-cmake-cuda-toolkit.patch | 60 ++++--------------- .../passthrough-python-lib-rel-path.patch | 12 ++++ 3 files changed, 45 insertions(+), 59 deletions(-) create mode 100644 pkgs/development/python-modules/torch/passthrough-python-lib-rel-path.patch diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index afc1fd6e74e1..bd6b1b262837 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -99,7 +99,7 @@ let setBool = v: if v then "1" else "0"; - # https://github.com/pytorch/pytorch/blob/v2.0.1/torch/utils/cpp_extension.py#L1744 + # https://github.com/pytorch/pytorch/blob/v2.4.0/torch/utils/cpp_extension.py#L1953 supportedTorchCudaCapabilities = let real = [ @@ -119,6 +119,7 @@ let "8.7" "8.9" "9.0" + "9.0a" ]; ptx = lists.map (x: "${x}+PTX") real; in @@ -206,13 +207,14 @@ let # In particular, this triggered warnings from cuda's `aliases.nix` "Magma cudaPackages does not match cudaPackages" = cudaSupport && (effectiveMagma.cudaPackages.cudaVersion != cudaPackages.cudaVersion); - "Rocm support is currently broken because `rocmPackages.hipblaslt` is unpackaged. (2024-06-09)" = rocmSupport; + "Rocm support is currently broken because `rocmPackages.hipblaslt` is unpackaged. (2024-06-09)" = + rocmSupport; }; in buildPythonPackage rec { pname = "torch"; # Don't forget to update torch-bin to the same version. - version = "2.3.1"; + version = "2.4.0"; pyproject = true; disabled = pythonOlder "3.8.0"; @@ -230,11 +232,16 @@ buildPythonPackage rec { repo = "pytorch"; rev = "refs/tags/v${version}"; fetchSubmodules = true; - hash = "sha256-vpgtOqzIDKgRuqdT8lB/g6j+oMIH1RPxdbjtlzZFjV8="; + hash = "sha256-s49rtarGNNFpnNG+kfJtZLE8ND53Ma201I0cOjeFSts="; }; patches = - lib.optionals cudaSupport [ ./fix-cmake-cuda-toolkit.patch ] + [ + # Allow setting PYTHON_LIB_REL_PATH with an environment variable. + # https://github.com/pytorch/pytorch/pull/128419 + ./passthrough-python-lib-rel-path.patch + ] + ++ lib.optionals cudaSupport [ ./fix-cmake-cuda-toolkit.patch ] ++ lib.optionals (stdenv.isDarwin && stdenv.isx86_64) [ # pthreadpool added support for Grand Central Dispatch in April # 2020. However, this relies on functionality (DISPATCH_APPLY_AUTO) @@ -279,11 +286,12 @@ buildPythonPackage rec { 'message(FATAL_ERROR "Found NCCL header version and library version' \ 'message(WARNING "Found NCCL header version and library version' '' - # Remove PyTorch's FindCUDAToolkit.cmake and to use CMake's default. - # We do not remove the entirety of cmake/Modules_CUDA_fix because we need FindCUDNN.cmake. + # Remove PyTorch's FindCUDAToolkit.cmake and use CMake's default. + # NOTE: Parts of pytorch rely on unmaintained FindCUDA.cmake with custom patches to support e.g. + # newer architectures (sm_90a). We do want to delete vendored patches, but have to keep them + # until https://github.com/pytorch/pytorch/issues/76082 is addressed + lib.optionalString cudaSupport '' rm cmake/Modules/FindCUDAToolkit.cmake - rm -rf cmake/Modules_CUDA_fix/{upstream,FindCUDA.cmake} '' # error: no member named 'aligned_alloc' in the global namespace; did you mean simply 'aligned_alloc' # This lib overrided aligned_alloc hence the error message. Tltr: his function is linkable but not in header. @@ -375,6 +383,10 @@ buildPythonPackage rec { USE_SYSTEM_NCCL = USE_NCCL; USE_STATIC_NCCL = USE_NCCL; + # Set the correct Python library path, broken since + # https://github.com/pytorch/pytorch/commit/3d617333e + PYTHON_LIB_REL_PATH = "${placeholder "out"}/${python.sitePackages}"; + # Suppress a weird warning in mkl-dnn, part of ideep in pytorch # (upstream seems to have fixed this in the wrong place?) # https://github.com/intel/mkl-dnn/commit/8134d346cdb7fe1695a2aa55771071d455fae0bc @@ -466,9 +478,7 @@ buildPythonPackage rec { libcusolver libcusparse ] - ++ lists.optionals (cudaPackages ? cudnn) [ - cudnn - ] + ++ lists.optionals (cudaPackages ? cudnn) [ cudnn ] ++ lists.optionals useSystemNccl [ # Some platforms do not support NCCL (i.e., Jetson) nccl # Provides nccl.h AND a static copy of NCCL! diff --git a/pkgs/development/python-modules/torch/fix-cmake-cuda-toolkit.patch b/pkgs/development/python-modules/torch/fix-cmake-cuda-toolkit.patch index 45587d3fe7b9..9767d8fee308 100644 --- a/pkgs/development/python-modules/torch/fix-cmake-cuda-toolkit.patch +++ b/pkgs/development/python-modules/torch/fix-cmake-cuda-toolkit.patch @@ -1,51 +1,15 @@ diff --git a/CMakeLists.txt b/CMakeLists.txt -index 9194e520bb0..d05fdcfb6cb 100644 +index c4cd4b2c2a..e983b21353 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt -@@ -1160,10 +1160,6 @@ if(BUILD_SHARED_LIBS) - ${PROJECT_SOURCE_DIR}/cmake/Modules_CUDA_fix - DESTINATION share/cmake/Caffe2/ - COMPONENT dev) -- install(FILES -- ${PROJECT_SOURCE_DIR}/cmake/Modules/FindCUDAToolkit.cmake -- DESTINATION share/cmake/Caffe2/ -- COMPONENT dev) - install(FILES - ${PROJECT_SOURCE_DIR}/cmake/Modules/FindCUSPARSELT.cmake - DESTINATION share/cmake/Caffe2/ -diff --git a/cmake/public/cuda.cmake b/cmake/public/cuda.cmake -index c7595774d81..4fc43771810 100644 ---- a/cmake/public/cuda.cmake -+++ b/cmake/public/cuda.cmake -@@ -61,9 +61,15 @@ find_package(CUDAToolkit REQUIRED) - cmake_policy(POP) - - if(NOT CMAKE_CUDA_COMPILER_VERSION VERSION_EQUAL CUDAToolkit_VERSION) -- message(FATAL_ERROR "Found two conflicting CUDA versions:\n" -- "V${CMAKE_CUDA_COMPILER_VERSION} in '${CUDA_INCLUDE_DIRS}' and\n" -- "V${CUDAToolkit_VERSION} in '${CUDAToolkit_INCLUDE_DIRS}'") -+ if(CUDA_INCLUDE_DIRS IN_LIST CUDAToolkit_INCLUDE_DIR) -+ message(STATUS "CUDA_INCLUDE_DIRS is a substring of CUDAToolkit_INCLUDE_DIR. " -+ "Setting CUDA_INCLUDE_DIRS to CUDAToolkit_INCLUDE_DIR.") -+ set(CUDA_INCLUDE_DIRS "${CUDAToolkit_INCLUDE_DIR}") -+ else() -+ message(FATAL_ERROR "Found two conflicting CUDA installs:\n" -+ "V${CMAKE_CUDA_COMPILER_VERSION} in '${CUDA_INCLUDE_DIRS}' and\n" -+ "V${CUDAToolkit_VERSION} in '${CUDAToolkit_INCLUDE_DIR}'") -+ endif() - endif() - - if(NOT TARGET CUDA::nvToolsExt) -diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py -index fb19b66dfba..3f83bef32fe 100644 ---- a/tools/setup_helpers/cmake.py -+++ b/tools/setup_helpers/cmake.py -@@ -207,6 +207,8 @@ class CMake: - "BUILDING_WITH_TORCH_LIBS", - "CUDA_HOST_COMPILER", - "CUDA_NVCC_EXECUTABLE", -+ "CUDAToolkit_ROOT", -+ "CUDAToolkit_INCLUDE_DIR", - "CUDA_SEPARABLE_COMPILATION", - "CUDNN_LIBRARY", - "CUDNN_INCLUDE_DIR", +@@ -1319,10 +1319,6 @@ if(BUILD_SHARED_LIBS) + DIRECTORY ${PROJECT_SOURCE_DIR}/cmake/Modules_CUDA_fix + DESTINATION share/cmake/Caffe2/ + COMPONENT dev) +- install( +- FILES ${PROJECT_SOURCE_DIR}/cmake/Modules/FindCUDAToolkit.cmake +- DESTINATION share/cmake/Caffe2/ +- COMPONENT dev) + install( + FILES ${PROJECT_SOURCE_DIR}/cmake/Modules/FindCUSPARSELT.cmake + DESTINATION share/cmake/Caffe2/ diff --git a/pkgs/development/python-modules/torch/passthrough-python-lib-rel-path.patch b/pkgs/development/python-modules/torch/passthrough-python-lib-rel-path.patch new file mode 100644 index 000000000000..629a0495c00e --- /dev/null +++ b/pkgs/development/python-modules/torch/passthrough-python-lib-rel-path.patch @@ -0,0 +1,12 @@ +diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py +index 5481ce46031c..d50d9d547399 100644 +--- a/tools/setup_helpers/cmake.py ++++ b/tools/setup_helpers/cmake.py +@@ -231,6 +231,7 @@ def generate( + "SELECTED_OP_LIST", + "TORCH_CUDA_ARCH_LIST", + "TRACING_BASED", ++ "PYTHON_LIB_REL_PATH", + ) + } + )