Skip to content

WIP: python3Packages.torch: migrate to setupCudaHook #245516

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
SomeoneSerge wants to merge 3 commits into from
Closed

WIP: python3Packages.torch: migrate to setupCudaHook #245516

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4d2272c
python3Packages.torch: asserts -> meta.broken
SomeoneSerge May 23, 2023
4cd492b
python3Packages.torch: symlinkJoin -> CUDAToolkit_ROOT
SomeoneSerge May 23, 2023
667358d
python3Packages.pynvml: drop the unused cudatoolkit input
SomeoneSerge May 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions pkgs/development/python-modules/pynvml/default.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
, fetchPypi
, substituteAll
, pythonOlder
, cudatoolkit
, addOpenGLRunpath
}:

Expand All @@ -24,7 +23,6 @@ buildPythonPackage rec {
})
];

propagatedBuildInputs = [ cudatoolkit ];

doCheck = false; # no tests in PyPi dist
pythonImportsCheck = [ "pynvml" "pynvml.smi" ];
Expand Down
36 changes: 14 additions & 22 deletions pkgs/development/python-modules/torch/default.nix
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,15 +52,9 @@

let
inherit (lib) lists strings trivial;
inherit (cudaPackages) cudatoolkit cudaFlags cudnn nccl;
in

assert cudaSupport -> stdenv.isLinux;
assert cudaSupport -> (cudaPackages.cudaMajorVersion == "11");

# confirm that cudatoolkits are sync'd across dependencies
assert !(MPISupport && cudaSupport) || mpi.cudatoolkit == cudatoolkit;
assert !cudaSupport || magma.cudaPackages.cudatoolkit == cudatoolkit;
with cudaPackages;
Copy link
Contributor Author

@SomeoneSerge SomeoneSerge Jul 26, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I probably shouldn't use with

Copy link
Member

@mweinelt mweinelt Aug 2, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At least not with that big of a scope.


let
setBool = v: if v then "1" else "0";
Expand Down Expand Up @@ -103,22 +97,11 @@ let
throw "No GPU targets specified"
);

cudatoolkit_joined = symlinkJoin {
name = "${cudatoolkit.name}-unsplit";
# nccl is here purely for semantic grouping it could be moved to nativeBuildInputs
paths = [ cudatoolkit.out cudatoolkit.lib nccl.dev nccl.out ];
};

# Normally libcuda.so.1 is provided at runtime by nvidia-x11 via
# LD_LIBRARY_PATH=/run/opengl-driver/lib. We only use the stub
# libcuda.so from cudatoolkit for running tests, so that we don’t have
# to recompile pytorch on every update to nvidia-x11 or the kernel.
cudaStub = linkFarm "cuda-stub" [{
name = "libcuda.so.1";
path = "${cudatoolkit}/lib/stubs/libcuda.so";
}];
cudaStubEnv = lib.optionalString cudaSupport
"LD_LIBRARY_PATH=${cudaStub}\${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH ";
cudaStubEnv = lib.optionalString cudaSupport ''addToSearchPath LD_LIBRARY_PATH "${cuda_cudart}/lib/stubs" ;'';

rocmtoolkit_joined = symlinkJoin {
name = "rocm-merged";
Expand Down Expand Up @@ -194,7 +177,6 @@ in buildPythonPackage rec {

preConfigure = lib.optionalString cudaSupport ''
export TORCH_CUDA_ARCH_LIST="${gpuTargetString}"
export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++
'' + lib.optionalString (cudaSupport && cudnn != null) ''
export CUDNN_INCLUDE_DIR=${cudnn}/include
'' + lib.optionalString rocmSupport ''
Expand Down Expand Up @@ -285,7 +267,7 @@ in buildPythonPackage rec {
pybind11
pythonRelaxDepsHook
removeReferencesTo
] ++ lib.optionals cudaSupport [ cudatoolkit_joined ]
] ++ lib.optionals cudaSupport [ cuda_nvcc ]
++ lib.optionals rocmSupport [ rocmtoolkit_joined ];

buildInputs = [ blas blas.provider pybind11 ]
Expand Down Expand Up @@ -418,6 +400,16 @@ in buildPythonPackage rec {
license = licenses.bsd3;
maintainers = with maintainers; [ teh thoughtpolice tscholak ]; # tscholak esp. for darwin-related builds
platforms = with platforms; linux ++ lib.optionals (!cudaSupport && !rocmSupport) darwin;
broken = rocmSupport && cudaSupport; # CUDA and ROCm are mutually exclusive
broken =
let
mutuallyExclusiveAccelerators = rocmSupport && cudaSupport;
unsupportedCuda = cudaSupport && (lib.versionOlder cudaPackages.cudaVersion "11");
cudaUnsupportedPlatform = cudaSupport && !stdenv.isLinux;

mpiCudaSynchronized = mpi.cudatoolkit == cudaPackages.cudatoolkit;
magmaCudaSynchronized = magma.cudaPackages.cudatoolkit == cudaPackages.cudatoolkit;
cudaUnsynchronized = cudaSupport && (!mpiCudaSynchronized || !magmaCudaSynchronized);
in
mutuallyExclusiveAccelerators || unsupportedCuda || cudaUnsynchronized;
};
}