diff --git a/pkgs/development/compilers/cudatoolkit/default.nix b/pkgs/development/compilers/cudatoolkit/default.nix index 65e9eb827605227..4fa99f80c8e363c 100644 --- a/pkgs/development/compilers/cudatoolkit/default.nix +++ b/pkgs/development/compilers/cudatoolkit/default.nix @@ -86,5 +86,9 @@ rec { gcc = gcc10; # can bump to 11 along with stdenv.cc }; + # Make sure to only ever update this to a version that is compatible with the + # latest cudnn, nccl, cutensor, etc! It sometimes happens that CUDA versions + # are released prior to compatibility with the rest of the ecosystem. And + # don't forget to request a review from @NixOS/cuda-maintainers! cudatoolkit_11 = cudatoolkit_11_5; } diff --git a/pkgs/development/libraries/science/math/cudnn/default.nix b/pkgs/development/libraries/science/math/cudnn/default.nix index d05cc3f90ec08b6..a3f785c5b38f6ec 100644 --- a/pkgs/development/libraries/science/math/cudnn/default.nix +++ b/pkgs/development/libraries/science/math/cudnn/default.nix @@ -81,31 +81,27 @@ rec { cudnn_8_1_cudatoolkit_11_2 = cudnn_8_1_cudatoolkit_10_2.override { cudatoolkit = cudatoolkit_11_2; }; cudnn_8_1_cudatoolkit_10 = cudnn_8_1_cudatoolkit_10_2.override { cudatoolkit = cudatoolkit_10; }; - cudnn_8_1_cudatoolkit_11 = cudnn_8_1_cudatoolkit_10_2.override { cudatoolkit = cudatoolkit_11; }; # cuDNN 8.3 is necessary for the latest jaxlib, esp. jaxlib-bin. See # https://github.com/google/jax/discussions/9455 for more info. - cudnn_8_3_cudatoolkit_10_2 = - generic - rec { - version = "8.3.2"; - cudatoolkit = cudatoolkit_10_2; - # See https://docs.nvidia.com/deeplearning/cudnn/archives/cudnn-832/support-matrix/index.html#cudnn-cuda-hardware-versions. - minCudaVersion = "10.2.00000"; - maxCudaVersion = "11.5.99999"; - mkSrc = cudatoolkit: - let v = if lib.versions.majorMinor cudatoolkit.version == "10.2" then "10.2" else "11.5"; in - fetchurl { - # Starting at version 8.3.1 there's a new directory layout including - # a subdirectory `local_installers`. - url = "https://developer.download.nvidia.com/compute/redist/cudnn/v${version}/local_installers/${v}/cudnn-linux-x86_64-8.3.2.44_cuda${v}-archive.tar.xz"; - hash = { - "10.2" = "sha256-1vVu+cqM+PketzIQumw9ykm6REbBZhv6/lXB7EC2aaw="; - "11.5" = "sha256-VQCVPAjF5dHd3P2iNPnvvdzb5DpTsm3AqCxyP6FwxFc="; - }."${v}"; - }; - } - ; + cudnn_8_3_cudatoolkit_10_2 = generic rec { + version = "8.3.2"; + cudatoolkit = cudatoolkit_10_2; + # See https://docs.nvidia.com/deeplearning/cudnn/archives/cudnn-832/support-matrix/index.html#cudnn-cuda-hardware-versions. + minCudaVersion = "10.2.00000"; + maxCudaVersion = "11.5.99999"; + mkSrc = cudatoolkit: + let v = if lib.versions.majorMinor cudatoolkit.version == "10.2" then "10.2" else "11.5"; in + fetchurl { + # Starting at version 8.3.1 there's a new directory layout including + # a subdirectory `local_installers`. + url = "https://developer.download.nvidia.com/compute/redist/cudnn/v${version}/local_installers/${v}/cudnn-linux-x86_64-8.3.2.44_cuda${v}-archive.tar.xz"; + hash = { + "10.2" = "sha256-1vVu+cqM+PketzIQumw9ykm6REbBZhv6/lXB7EC2aaw="; + "11.5" = "sha256-VQCVPAjF5dHd3P2iNPnvvdzb5DpTsm3AqCxyP6FwxFc="; + }."${v}"; + }; + }; cudnn_8_3_cudatoolkit_11_0 = cudnn_8_3_cudatoolkit_10_2.override { cudatoolkit = cudatoolkit_11_0; }; cudnn_8_3_cudatoolkit_11_1 = cudnn_8_3_cudatoolkit_10_2.override { cudatoolkit = cudatoolkit_11_1; }; cudnn_8_3_cudatoolkit_11_2 = cudnn_8_3_cudatoolkit_10_2.override { cudatoolkit = cudatoolkit_11_2; }; diff --git a/pkgs/development/libraries/science/math/cutensor/default.nix b/pkgs/development/libraries/science/math/cutensor/default.nix index 25633ca63e6b1be..d0d24bd1adbe86d 100644 --- a/pkgs/development/libraries/science/math/cutensor/default.nix +++ b/pkgs/development/libraries/science/math/cutensor/default.nix @@ -1,6 +1,14 @@ { callPackage -, cudatoolkit_10_1, cudatoolkit_10_2 -, cudatoolkit_11_0, cudatoolkit_11_1, cudatoolkit_11_2, cudatoolkit_11_3, cudatoolkit_11_4 +, cudatoolkit_10_1 +, cudatoolkit_10_2 +, cudatoolkit_11 +, cudatoolkit_11_0 +, cudatoolkit_11_1 +, cudatoolkit_11_2 +, cudatoolkit_11_3 +, cudatoolkit_11_4 +, cudatoolkit_11_5 +, cudatoolkit_11_6 }: rec { @@ -8,18 +16,18 @@ rec { version = "1.2.2.5"; libPath = "lib/10.1"; cudatoolkit = cudatoolkit_10_1; - # 1.2.2 is compatible with CUDA 11.0, 11.1, and 11.2: - # ephemeral doc at https://developer.nvidia.com/cutensor/downloads - sha256 = "1dl9bd71frhac9cb8lvnh71zfsnqxbxbfhndvva2zf6nh0my4klm"; + # 1.2.2 is compatible with CUDA 10.1, 10.2, and 11.x. + # See https://docs.nvidia.com/cuda/cutensor/release_notes.html#cutensor-v1-2-2. + hash = "sha256-lU7iK4DWuC/U3s1Ct/rq2Gr3w4F2U7RYYgpmF05bibY="; }; cutensor_cudatoolkit_10_2 = cutensor_cudatoolkit_10_1.override { version = "1.3.1.3"; libPath = "lib/10.2"; cudatoolkit = cudatoolkit_10_2; - # 1.3.1 is compatible with CUDA 11.0, 11.1, and 11.2: - # ephemeral doc at https://developer.nvidia.com/cutensor/downloads - sha256 = "sha256-mNlVnabB2IC3HnYY0mb06RLqQzDxN9ePGVeBy3hkBC8="; + # 1.3.1 is compatible with CUDA 10.2 and 11.x. + # See https://docs.nvidia.com/cuda/cutensor/release_notes.html#cutensor-v1-3-1. + hash = "sha256-mNlVnabB2IC3HnYY0mb06RLqQzDxN9ePGVeBy3hkBC8="; }; cutensor_cudatoolkit_10 = cutensor_cudatoolkit_10_2; @@ -29,21 +37,12 @@ rec { cudatoolkit = cudatoolkit_11_0; }; - cutensor_cudatoolkit_11_1 = cutensor_cudatoolkit_11_0.override { - cudatoolkit = cudatoolkit_11_1; - }; - - cutensor_cudatoolkit_11_2 = cutensor_cudatoolkit_11_0.override { - cudatoolkit = cudatoolkit_11_2; - }; - - cutensor_cudatoolkit_11_3 = cutensor_cudatoolkit_11_0.override { - cudatoolkit = cudatoolkit_11_3; - }; - - cutensor_cudatoolkit_11_4 = cutensor_cudatoolkit_11_0.override { - cudatoolkit = cudatoolkit_11_4; - }; + cutensor_cudatoolkit_11_1 = cutensor_cudatoolkit_11_0.override { cudatoolkit = cudatoolkit_11_1; }; + cutensor_cudatoolkit_11_2 = cutensor_cudatoolkit_11_0.override { cudatoolkit = cudatoolkit_11_2; }; + cutensor_cudatoolkit_11_3 = cutensor_cudatoolkit_11_0.override { cudatoolkit = cudatoolkit_11_3; }; + cutensor_cudatoolkit_11_4 = cutensor_cudatoolkit_11_0.override { cudatoolkit = cudatoolkit_11_4; }; + cutensor_cudatoolkit_11_5 = cutensor_cudatoolkit_11_0.override { cudatoolkit = cudatoolkit_11_5; }; + cutensor_cudatoolkit_11_6 = cutensor_cudatoolkit_11_0.override { cudatoolkit = cudatoolkit_11_6; }; - cutensor_cudatoolkit_11 = cutensor_cudatoolkit_11_4; + cutensor_cudatoolkit_11 = cutensor_cudatoolkit_11_0.override { cudatoolkit = cudatoolkit_11; }; } diff --git a/pkgs/development/libraries/science/math/cutensor/generic.nix b/pkgs/development/libraries/science/math/cutensor/generic.nix index 0b2bd31b2c3034c..e7c5e0b298c07e1 100644 --- a/pkgs/development/libraries/science/math/cutensor/generic.nix +++ b/pkgs/development/libraries/science/math/cutensor/generic.nix @@ -7,7 +7,7 @@ , addOpenGLRunpath , version -, sha256 +, hash }: let @@ -21,7 +21,7 @@ stdenv.mkDerivation { src = fetchurl { url = "https://developer.download.nvidia.com/compute/cutensor/${mostOfVersion}/local_installers/libcutensor-${stdenv.hostPlatform.parsed.kernel.name}-${stdenv.hostPlatform.parsed.cpu.name}-${version}.tar.gz"; - inherit sha256; + inherit hash; }; outputs = [ "out" "dev" ]; diff --git a/pkgs/development/libraries/science/math/nccl/default.nix b/pkgs/development/libraries/science/math/nccl/default.nix index 5036ad890799eb5..c08b5aa5851dc9a 100644 --- a/pkgs/development/libraries/science/math/nccl/default.nix +++ b/pkgs/development/libraries/science/math/nccl/default.nix @@ -38,6 +38,10 @@ stdenv.mkDerivation rec { enableParallelBuilding = true; + passthru = { + inherit cudatoolkit; + }; + meta = with lib; { description = "Multi-GPU and multi-node collective communication primitives for NVIDIA GPUs"; homepage = "https://developer.nvidia.com/nccl"; diff --git a/pkgs/development/python-modules/cupy/default.nix b/pkgs/development/python-modules/cupy/default.nix index 2ae38e50513aab6..a798920ef7a87ba 100644 --- a/pkgs/development/python-modules/cupy/default.nix +++ b/pkgs/development/python-modules/cupy/default.nix @@ -5,6 +5,10 @@ , addOpenGLRunpath }: +assert cudnn.cudatoolkit == cudatoolkit; +assert cutensor.cudatoolkit == cudatoolkit; +assert nccl.cudatoolkit == cudatoolkit; + buildPythonPackage rec { pname = "cupy"; version = "10.2.0"; @@ -15,8 +19,15 @@ buildPythonPackage rec { sha256 = "sha256-5ovvA76QGOsOnVztMfDgLerks5nJrKR08rLc+ArmWA8="; }; + # See https://docs.cupy.dev/en/v10.2.0/reference/environment.html. Seting both + # CUPY_NUM_BUILD_JOBS and CUPY_NUM_NVCC_THREADS to NIX_BUILD_CORES results in + # a small amount of thrashing but it turns out there are a large number of + # very short builds and a few extremely long ones, so setting both ends up + # working nicely in practice. preConfigure = '' export CUDA_PATH=${cudatoolkit} + export CUPY_NUM_BUILD_JOBS="$NIX_BUILD_CORES" + export CUPY_NUM_NVCC_THREADS="$NIX_BUILD_CORES" ''; nativeBuildInputs = [ diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix index 3a88b23b73b5017..59442957931a78c 100644 --- a/pkgs/top-level/all-packages.nix +++ b/pkgs/top-level/all-packages.nix @@ -4647,7 +4647,6 @@ with pkgs; cudnn_8_1_cudatoolkit_11_1 cudnn_8_1_cudatoolkit_11_2 cudnn_8_1_cudatoolkit_10 - cudnn_8_1_cudatoolkit_11 cudnn_8_3_cudatoolkit_10_2 cudnn_8_3_cudatoolkit_11_0 cudnn_8_3_cudatoolkit_11_1 @@ -4658,8 +4657,8 @@ with pkgs; cudnn_8_3_cudatoolkit_10 cudnn_8_3_cudatoolkit_11; - # TODO(samuela): This is old and should be upgraded to 8.3 at some point. - cudnn = cudnn_7_6_cudatoolkit_10_1; + # Make sure to keep this in sync with the `cudatoolkit` version! + cudnn = cudnn_8_3_cudatoolkit_10; cutensorPackages = callPackages ../development/libraries/science/math/cutensor { }; inherit (cutensorPackages) @@ -33022,11 +33021,16 @@ with pkgs; ### SCIENCE / MATH caffe = callPackage ../applications/science/math/caffe ({ + cudaSupport = config.cudaSupport or false; + cudatoolkit = cudatoolkit_10_1; + cudnn = cudnn_7_6_cudatoolkit_10_1; opencv3 = opencv3WithoutCuda; # Used only for image loading. blas = openblas; inherit (darwin.apple_sdk.frameworks) Accelerate CoreGraphics CoreVideo; } // (config.caffe or {})); + caffeWithCuda = caffe.override { cudaSupport = true; }; + caffe2 = callPackage ../development/libraries/science/math/caffe2 (rec { inherit (python3Packages) python future six numpy pydot; protobuf = protobuf3_1; diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index a8fa9c95307c3b1..8acdbdd14b8d3f4 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -1459,6 +1459,11 @@ in { inherit (self) python numpy boost; }); + caffeWithCuda = toPythonModule (pkgs.caffeWithCuda.override { + pythonSupport = true; + inherit (self) python numpy boost; + }); + cairocffi = callPackage ../development/python-modules/cairocffi { }; cairosvg = callPackage ../development/python-modules/cairosvg { }; @@ -1961,7 +1966,7 @@ in { cupy = callPackage ../development/python-modules/cupy { cudatoolkit = pkgs.cudatoolkit_11; - cudnn = pkgs.cudnn_8_1_cudatoolkit_11; + cudnn = pkgs.cudnn_8_3_cudatoolkit_11; nccl = pkgs.nccl_cudatoolkit_11; cutensor = pkgs.cutensor_cudatoolkit_11; }; @@ -8362,6 +8367,24 @@ in { pytorch = callPackage ../development/python-modules/pytorch { cudaSupport = pkgs.config.cudaSupport or false; + + # TODO: next time pytorch is updated (to 1.11.0, currently in staging as of + # 2022-03-31), make the following changes: + + # -> cudatoolk_11 + cudatoolkit = pkgs.cudatoolkit_10; + + # -> cudnn_8_3_cudatoolkit_11 + cudnn = pkgs.cudnn_8_1_cudatoolkit_10; + + # -> cutensor_cudatoolkit_11 (cutensor is a new dependency in v1.11.0) + # cutensor = pkgs.cutensor_cudatoolkit_11; + + # -> setting a custom magma should be unnecessary with v1.11.0 + magma = pkgs.magma.override { cudatoolkit = pkgs.cudatoolkit_10; }; + + # -> nccl_cudatoolkit_11 + nccl = pkgs.nccl.override { cudatoolkit = pkgs.cudatoolkit_10; }; }; pytorch-bin = callPackage ../development/python-modules/pytorch/bin.nix { };