diff --git a/pkgs/development/python-modules/torch/source/default.nix b/pkgs/development/python-modules/torch/source/default.nix index a666b766dd63..cfdb32bd9df5 100644 --- a/pkgs/development/python-modules/torch/source/default.nix +++ b/pkgs/development/python-modules/torch/source/default.nix @@ -3,6 +3,7 @@ lib, fetchFromGitHub, fetchFromGitLab, + fetchpatch, git-unroll, buildPythonPackage, python, @@ -304,6 +305,18 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: { patches = [ ./clang19-template-warning.patch + + # The GCC version upperbounds were wrong for cuda 12.8 and 12.9, which led downstream builds to + # illegitimately fail with: + # RuntimeError: The current installed version of g++ (14.3.0) is greater than the maximum + # required version by CUDA 12.9. Please make sure to use an adequate version of g++ + # (>=6.0.0, <14.0). + # TODO: remove at the next release + (fetchpatch { + name = "allow-gcc-14-with-cuda-12.8-9"; + url = "https://github.com/pytorch/pytorch/commit/39565a7dcf8f93ea22cedeaa20088b24ff6d2634.patch"; + hash = "sha256-Au5fVbs7i33d9c4Xj8koiBP7lGnsTGTaX4VlE2gAfy8="; + }) ] ++ lib.optionals cudaSupport [ ./fix-cmake-cuda-toolkit.patch @@ -741,7 +754,9 @@ buildPythonPackage.override { inherit stdenv; } (finalAttrs: { blasProvider = blas.provider; # To help debug when a package is broken due to CUDA support inherit brokenConditions; - tests = callPackage ../tests { }; + tests = callPackage ../tests { + inherit rocmSupport cudaSupport; + }; }; meta = { diff --git a/pkgs/development/python-modules/torch/tests/default.nix b/pkgs/development/python-modules/torch/tests/default.nix index 90fb21d018d7..28e6081b792a 100644 --- a/pkgs/development/python-modules/torch/tests/default.nix +++ b/pkgs/development/python-modules/torch/tests/default.nix @@ -1,33 +1,58 @@ -{ callPackage }: +{ + lib, + stdenv, + callPackage, + torchvision, + torchcodec, + torchaudio, -rec { - # To perform the runtime check use either - # `nix run .#python3Packages.torch.tests.tester-cudaAvailable` (outside the sandbox), or - # `nix build .#python3Packages.torch.tests.tester-cudaAvailable.gpuCheck` (in a relaxed sandbox) - tester-cudaAvailable = callPackage ./mk-runtime-check.nix { - feature = "cuda"; - versionAttr = "cuda"; - libraries = ps: [ ps.torchWithCuda ]; - }; - tester-rocmAvailable = callPackage ./mk-runtime-check.nix { - feature = "rocm"; - versionAttr = "hip"; - libraries = ps: [ ps.torchWithRocm ]; + rocmSupport, + cudaSupport, +}: + +let + cudaOnly = { + # To perform the runtime check use either + # `nix run .#python3Packages.torch.tests.tester-cudaAvailable` (outside the sandbox), or + # `nix build .#python3Packages.torch.tests.tester-cudaAvailable.gpuCheck` (in a relaxed sandbox) + tester-cudaAvailable = callPackage ./mk-runtime-check.nix { + feature = "cuda"; + versionAttr = "cuda"; + libraries = ps: [ ps.torchWithCuda ]; + }; + + tester-compileCuda = callPackage ./mk-torch-compile-check.nix { + feature = "cuda"; + libraries = ps: [ ps.torchWithCuda ]; + }; }; - compileCpu = tester-compileCpu.gpuCheck; + rocmOnly = { + tester-rocmAvailable = callPackage ./mk-runtime-check.nix { + feature = "rocm"; + versionAttr = "hip"; + libraries = ps: [ ps.torchWithRocm ]; + }; + tester-compileRocm = callPackage ./mk-torch-compile-check.nix { + feature = "rocm"; + libraries = ps: [ ps.torchWithRocm ]; + }; + }; +in +let tester-compileCpu = callPackage ./mk-torch-compile-check.nix { feature = null; libraries = ps: [ ps.torch ]; }; - tester-compileCuda = callPackage ./mk-torch-compile-check.nix { - feature = "cuda"; - libraries = ps: [ ps.torchWithCuda ]; - }; - tester-compileRocm = callPackage ./mk-torch-compile-check.nix { - feature = "rocm"; - libraries = ps: [ ps.torchWithRocm ]; - }; +in +{ + inherit tester-compileCpu; + compileCpu = tester-compileCpu.gpuCheck; mnist-example = callPackage ./mnist-example { }; + + # Core packages from the torch ecosystem + inherit torchvision torchaudio torchcodec; } +// lib.optionalAttrs cudaSupport cudaOnly +// lib.optionalAttrs rocmSupport rocmOnly diff --git a/pkgs/development/python-modules/torchaudio/default.nix b/pkgs/development/python-modules/torchaudio/default.nix index 4e967e98cd50..3ccc529f35a5 100644 --- a/pkgs/development/python-modules/torchaudio/default.nix +++ b/pkgs/development/python-modules/torchaudio/default.nix @@ -1,5 +1,6 @@ { lib, + stdenv, buildPythonPackage, fetchFromGitHub, @@ -53,18 +54,6 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_CUDA = 1; TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_MULTIGPU_CUDA = 1; - # At of 2026-03-27, the default `cudaPackages` version is 12.9.1 - # According to Nvidia, it should support GCC versions up to 14.x: - # -> https://docs.nvidia.com/cuda/archive/12.9.1/cuda-installation-guide-linux/index.html#host-compiler-support-policy - # However, PyTorch's *strict* upper bound is 14.0: - # -> https://github.com/pytorch/pytorch/blob/v2.11.0/torch/utils/cpp_extension.py#L75 - # Hence, the build fails with: - # RuntimeError: The current installed version of g++ (14.3.0) is greater than the maximum - # required version by CUDA 12.9. Please make sure to use an adequate version of g++ - # (>=6.0.0, <14.0). - # Hence, we disable the version check to silence the error: - TORCH_DONT_CHECK_COMPILER_ABI = cudaSupport; - # ROCM USE_ROCM = rocmSupport; PYTORCH_ROCM_ARCH = lib.optionalString rocmSupport torch.gpuTargetString; @@ -80,6 +69,11 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { # RuntimeError: Test is known to fail for Python 3.10, disabling for now # See: https://github.com/pytorch/audio/pull/2224#issuecomment-1048329450 TORCHAUDIO_TEST_ALLOW_SKIP_IF_ON_PYTHON_310 = true; + + # Fails on aarch64-linux with: + # RuntimeError: `fbgemm` is not available + TORCHAUDIO_TEST_ALLOW_SKIP_IF_NO_QUANTIZATION = + stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64; }; build-system = [ @@ -128,9 +122,17 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { # Very long to run "AutogradCPUTest" + ] + ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [ + # AssertionError: Tensor-likes are not close! + "test_batch_inverse_spectrogram" + "test_batch_pitch_shift" + "test_batch_spectrogram" + "test_griffinlim_0_99" ]; passthru.gpuCheck = torchaudio.overridePythonAttrs (old: { + pname = "${finalAttrs.pname}-gpuCheck"; requiredSystemFeatures = [ "cuda" ]; env = (old.env or { }) // { diff --git a/pkgs/development/python-modules/torchvision/default.nix b/pkgs/development/python-modules/torchvision/default.nix index e57eadd789ec..945875ce3220 100644 --- a/pkgs/development/python-modules/torchvision/default.nix +++ b/pkgs/development/python-modules/torchvision/default.nix @@ -63,18 +63,6 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { TORCHVISION_LIBRARY = "${libjpeg_turbo}/lib/"; } // lib.optionalAttrs cudaSupport { - # At of 2026-03-27, the default `cudaPackages` version is 12.9.1 - # According to Nvidia, it should support GCC versions up to 14.x: - # -> https://docs.nvidia.com/cuda/archive/12.9.1/cuda-installation-guide-linux/index.html#host-compiler-support-policy - # However, PyTorch's *strict* upper bound is 14.0: - # -> https://github.com/pytorch/pytorch/blob/v2.11.0/torch/utils/cpp_extension.py#L75 - # Hence, the build fails with: - # RuntimeError: The current installed version of g++ (14.3.0) is greater than the maximum - # required version by CUDA 12.9. Please make sure to use an adequate version of g++ - # (>=6.0.0, <14.0). - # Hence, we disable the version check to silence the error: - TORCH_DONT_CHECK_COMPILER_ABI = 1; - TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" cudaCapabilities}"; FORCE_CUDA = 1; };