gador/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

python3Packages.torch: only build with NCCL when targeting CUDA on a supported platform

Browse Source
This commit is contained in:
Connor Baker
2024-01-10 01:15:01 +00:00
parent 5c260fa532
commit 39cab2b768
1 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkgs/development/python-modules/torch/default.nix
+8 -8
View File
@@ -7,8 +7,8 @@
magma,
magma-hip,
magma-cuda-static,
# Use the system NCCL as long as it is supported.
useSystemNccl ? !cudaPackages.nccl.meta.unsupported,
# Use the system NCCL as long as we're targeting CUDA on a supported platform.
useSystemNccl ? (cudaSupport && !cudaPackages.nccl.meta.unsupported),
MPISupport ? false, mpi,
buildDocs ? false,
@@ -58,7 +58,6 @@
let
inherit (lib) attrsets lists strings trivial;
inherit (cudaPackages) cudaFlags cudnn nccl;
ncclSupported = cudaSupport && !cudaPackages.nccl.meta.unsupported;
setBool = v: if v then "1" else "0";
@@ -123,7 +122,6 @@ let
"Unsupported CUDA version" = cudaSupport && !(builtins.elem cudaPackages.cudaMajorVersion [ "11" "12" ]);
"MPI cudatoolkit does not match cudaPackages.cudatoolkit" = MPISupport && cudaSupport && (mpi.cudatoolkit != cudaPackages.cudatoolkit);
"Magma cudaPackages does not match cudaPackages" = cudaSupport && (effectiveMagma.cudaPackages != cudaPackages);
"Requested system NCCL, but cudaPackages.nccl is not supported" = useSystemNccl && !ncclSupported;
};
in buildPythonPackage rec {
pname = "torch";
@@ -276,9 +274,11 @@ in buildPythonPackage rec {
PYTORCH_BUILD_VERSION = version;
PYTORCH_BUILD_NUMBER = 0;
USE_NCCL = setBool (cudaSupport && ncclSupported);
USE_SYSTEM_NCCL = setBool (cudaSupport && useSystemNccl); # don't build pytorch's third_party NCCL
USE_STATIC_NCCL = setBool (cudaSupport && useSystemNccl);
# In-tree builds of NCCL are not supported.
# Use NCCL when cudaSupport is enabled and nccl is available.
USE_NCCL = setBool useSystemNccl;
USE_SYSTEM_NCCL = USE_NCCL;
USE_STATIC_NCCL = USE_NCCL;
# Suppress a weird warning in mkl-dnn, part of ideep in pytorch
# (upstream seems to have fixed this in the wrong place?)
@@ -366,7 +366,7 @@ in buildPythonPackage rec {
] ++ lists.optionals (cudaPackages ? cudnn) [
cudnn.dev
cudnn.lib
] ++ lists.optionals (useSystemNccl && ncclSupported) [
] ++ lists.optionals useSystemNccl [
# Some platforms do not support NCCL (i.e., Jetson)
nccl.dev # Provides nccl.h AND a static copy of NCCL!
] ++ lists.optionals (strings.versionOlder cudaVersion "11.8") [